Python: Phase 2: Embedding clients for Ollama, Bedrock, and Azure AI Inference (#4207)

* Phase 2: Embedding clients for Ollama, Bedrock, and Azure AI Inference

Add embedding client implementations to existing provider packages:

- OllamaEmbeddingClient: Text embeddings via Ollama's embed API
- BedrockEmbeddingClient: Text embeddings via Amazon Titan on Bedrock
- AzureAIInferenceEmbeddingClient: Text and image embeddings via Azure AI
  Inference, supporting Content | str input with separate model IDs for
  text (AZURE_AI_INFERENCE_EMBEDDING_MODEL_ID) and image
  (AZURE_AI_INFERENCE_IMAGE_EMBEDDING_MODEL_ID) endpoints

Additional changes:
- Rename EmbeddingCoT -> EmbeddingT, EmbeddingOptionsCoT -> EmbeddingOptionsT
- Add otel_provider_name passthrough to all embedding clients
- Register integration pytest marker in all packages
- Add lazy-loading namespace exports for Ollama and Bedrock embeddings
- Add image embedding sample using Cohere-embed-v3-english
- Add azure-ai-inference dependency to azure-ai package

Part of #1188

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix mypy duplicate name and ruff lint issues

- Rename second 'vector' variable to 'img_vector' in image embedding loop
- Combine nested with statements in tests
- Remove unused result assignments in tests

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* updates from feedback

* Fix CI failures in embedding usage handling

- Fix Azure AI embedding mypy issues by normalizing vectors to list[float],
  safely accumulating optional usage token fields, and filtering None entries
  before constructing GeneratedEmbeddings
- Avoid Bandit false positive by initializing usage details as an empty dict
- Update OpenAI embedding tests to assert canonical usage keys
  (input_token_count/total_token_count)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

python/packages/core/agent_framework/_clients.py

@@ -667,16 +667,12 @@ class SupportsFileSearchTool(Protocol):
 
 # region SupportsGetEmbeddings Protocol
 
-# Contravariant/covariant TypeVars for the Protocol
+# Contravariant TypeVars for the Protocol
 EmbeddingInputContraT = TypeVar(
     "EmbeddingInputContraT",
     default="str",
     contravariant=True,
 )
-EmbeddingCoT = TypeVar(
-    "EmbeddingCoT",
-    default="list[float]",
-)
 EmbeddingOptionsContraT = TypeVar(
     "EmbeddingOptionsContraT",
     bound=TypedDict,  # type: ignore[valid-type]
@@ -686,7 +682,7 @@ EmbeddingOptionsContraT = TypeVar(
 
 
 @runtime_checkable
-class SupportsGetEmbeddings(Protocol[EmbeddingInputContraT, EmbeddingCoT, EmbeddingOptionsContraT]):
+class SupportsGetEmbeddings(Protocol[EmbeddingInputContraT, EmbeddingT, EmbeddingOptionsContraT]):
     """Protocol for an embedding client that can generate embeddings.
 
     This protocol enables duck-typing for embedding generation. Any class that
@@ -714,7 +710,7 @@ class SupportsGetEmbeddings(Protocol[EmbeddingInputContraT, EmbeddingCoT, Embedd
         values: Sequence[EmbeddingInputContraT],
         *,
         options: EmbeddingOptionsContraT | None = None,
-    ) -> Awaitable[GeneratedEmbeddings[EmbeddingCoT]]:
+    ) -> Awaitable[GeneratedEmbeddings[EmbeddingT]]:
         """Generate embeddings for the given values.
 
         Args:
@@ -733,15 +729,15 @@ class SupportsGetEmbeddings(Protocol[EmbeddingInputContraT, EmbeddingCoT, Embedd
 # region BaseEmbeddingClient
 
 # Covariant for the BaseEmbeddingClient
-EmbeddingOptionsCoT = TypeVar(
-    "EmbeddingOptionsCoT",
+EmbeddingOptionsT = TypeVar(
+    "EmbeddingOptionsT",
     bound=TypedDict,  # type: ignore[valid-type]
     default="EmbeddingGenerationOptions",
     covariant=True,
 )
 
 
-class BaseEmbeddingClient(SerializationMixin, ABC, Generic[EmbeddingInputT, EmbeddingT, EmbeddingOptionsCoT]):
+class BaseEmbeddingClient(SerializationMixin, ABC, Generic[EmbeddingInputT, EmbeddingT, EmbeddingOptionsT]):
     """Abstract base class for embedding clients.
 
     Subclasses implement ``get_embeddings`` to provide the actual
@@ -785,7 +781,7 @@ class BaseEmbeddingClient(SerializationMixin, ABC, Generic[EmbeddingInputT, Embe
         self,
         values: Sequence[EmbeddingInputT],
         *,
-        options: EmbeddingOptionsCoT | None = None,
+        options: EmbeddingOptionsT | None = None,
     ) -> GeneratedEmbeddings[EmbeddingT]:
         """Generate embeddings for the given values.
 

python/packages/core/agent_framework/_types.py

@@ -377,6 +377,12 @@ class UsageDetails(TypedDict, total=False):
 
     This is a non-closed dictionary, so any specific provider fields can be added as needed.
     Whenever they can be mapped to standard fields, they will be.
+
+    Keys:
+        input_token_count: The number of input tokens used.
+        output_token_count: The number of output tokens generated.
+        total_token_count: The total number of tokens (input + output).
+
     """
 
     input_token_count: int | None
@@ -3289,7 +3295,7 @@ class GeneratedEmbeddings(list[Embedding[EmbeddingT]], Generic[EmbeddingT, Embed
         embeddings: Iterable[Embedding[EmbeddingT]] | None = None,
         *,
         options: EmbeddingOptionsT | None = None,
-        usage: dict[str, Any] | None = None,
+        usage: UsageDetails | None = None,
         additional_properties: dict[str, Any] | None = None,
     ) -> None:
         super().__init__(embeddings or [])

python/packages/core/agent_framework/amazon/__init__.py

@@ -8,6 +8,9 @@ This module lazily re-exports objects from:
 Supported classes:
 - BedrockChatClient
 - BedrockChatOptions
+- BedrockEmbeddingClient
+- BedrockEmbeddingOptions
+- BedrockEmbeddingSettings
 - BedrockGuardrailConfig
 - BedrockSettings
 """
@@ -17,7 +20,15 @@ from typing import Any
 
 IMPORT_PATH = "agent_framework_bedrock"
 PACKAGE_NAME = "agent-framework-bedrock"
-_IMPORTS = ["BedrockChatClient", "BedrockChatOptions", "BedrockGuardrailConfig", "BedrockSettings"]
+_IMPORTS = [
+    "BedrockChatClient",
+    "BedrockChatOptions",
+    "BedrockEmbeddingClient",
+    "BedrockEmbeddingOptions",
+    "BedrockEmbeddingSettings",
+    "BedrockGuardrailConfig",
+    "BedrockSettings",
+]
 
 
 def __getattr__(name: str) -> Any:

python/packages/core/agent_framework/amazon/__init__.pyi

@@ -3,6 +3,9 @@
 from agent_framework_bedrock import (
     BedrockChatClient,
     BedrockChatOptions,
+    BedrockEmbeddingClient,
+    BedrockEmbeddingOptions,
+    BedrockEmbeddingSettings,
     BedrockGuardrailConfig,
     BedrockSettings,
 )
@@ -10,6 +13,9 @@ from agent_framework_bedrock import (
 __all__ = [
     "BedrockChatClient",
     "BedrockChatOptions",
+    "BedrockEmbeddingClient",
+    "BedrockEmbeddingOptions",
+    "BedrockEmbeddingSettings",
     "BedrockGuardrailConfig",
     "BedrockSettings",
 ]

python/packages/core/agent_framework/azure/_embedding_client.py

@@ -99,6 +99,7 @@ class AzureOpenAIEmbeddingClient(
         credential: AzureCredentialTypes | AzureTokenProvider | None = None,
         default_headers: Mapping[str, str] | None = None,
         async_client: AsyncAzureOpenAI | None = None,
+        otel_provider_name: str | None = None,
         env_file_path: str | None = None,
         env_file_encoding: str | None = None,
     ) -> None:
@@ -133,4 +134,5 @@ class AzureOpenAIEmbeddingClient(
             credential=credential,
             default_headers=default_headers,
             client=async_client,
+            otel_provider_name=otel_provider_name,
         )

python/packages/core/agent_framework/observability.py

@@ -1279,15 +1279,15 @@ class ChatTelemetryLayer(Generic[OptionsCoT]):
         return _get_response()
 
 
-EmbeddingOptionsCoT = TypeVar(
-    "EmbeddingOptionsCoT",
+EmbeddingOptionsT = TypeVar(
+    "EmbeddingOptionsT",
     bound=TypedDict,  # type: ignore[valid-type]
     default="EmbeddingGenerationOptions",
     covariant=True,
 )
 
 
-class EmbeddingTelemetryLayer(Generic[EmbeddingInputT, EmbeddingT, EmbeddingOptionsCoT]):
+class EmbeddingTelemetryLayer(Generic[EmbeddingInputT, EmbeddingT, EmbeddingOptionsT]):
     """Layer that wraps embedding client get_embeddings with OpenTelemetry tracing."""
 
     def __init__(self, *args: Any, otel_provider_name: str | None = None, **kwargs: Any) -> None:
@@ -1301,7 +1301,7 @@ class EmbeddingTelemetryLayer(Generic[EmbeddingInputT, EmbeddingT, EmbeddingOpti
         self,
         values: Sequence[EmbeddingInputT],
         *,
-        options: EmbeddingOptionsCoT | None = None,
+        options: EmbeddingOptionsT | None = None,
     ) -> GeneratedEmbeddings[EmbeddingT]:
         """Trace embedding generation with OpenTelemetry spans and metrics."""
         global OBSERVABILITY_SETTINGS

python/packages/core/agent_framework/ollama/__init__.py

@@ -7,6 +7,10 @@ This module lazily re-exports objects from:
 
 Supported classes:
 - OllamaChatClient
+- OllamaChatOptions
+- OllamaEmbeddingClient
+- OllamaEmbeddingOptions
+- OllamaEmbeddingSettings
 - OllamaSettings
 """
 
@@ -15,7 +19,14 @@ from typing import Any
 
 IMPORT_PATH = "agent_framework_ollama"
 PACKAGE_NAME = "agent-framework-ollama"
-_IMPORTS = ["OllamaChatClient", "OllamaSettings"]
+_IMPORTS = [
+    "OllamaChatClient",
+    "OllamaChatOptions",
+    "OllamaEmbeddingClient",
+    "OllamaEmbeddingOptions",
+    "OllamaEmbeddingSettings",
+    "OllamaSettings",
+]
 
 
 def __getattr__(name: str) -> Any:

python/packages/core/agent_framework/ollama/__init__.pyi

@@ -2,10 +2,18 @@
 
 from agent_framework_ollama import (
     OllamaChatClient,
+    OllamaChatOptions,
+    OllamaEmbeddingClient,
+    OllamaEmbeddingOptions,
+    OllamaEmbeddingSettings,
     OllamaSettings,
 )
 
 __all__ = [
     "OllamaChatClient",
+    "OllamaChatOptions",
+    "OllamaEmbeddingClient",
+    "OllamaEmbeddingOptions",
+    "OllamaEmbeddingSettings",
     "OllamaSettings",
 ]

python/packages/core/agent_framework/openai/_embedding_client.py

@@ -12,7 +12,7 @@ from openai import AsyncOpenAI
 
 from .._clients import BaseEmbeddingClient
 from .._settings import load_settings
-from .._types import Embedding, EmbeddingGenerationOptions, GeneratedEmbeddings
+from .._types import Embedding, EmbeddingGenerationOptions, GeneratedEmbeddings, UsageDetails
 from ..observability import EmbeddingTelemetryLayer
 from ._shared import OpenAIBase, OpenAIConfigMixin, OpenAISettings
 
@@ -116,11 +116,11 @@ class RawOpenAIEmbeddingClient(
                 )
             )
 
-        usage_dict: dict[str, Any] | None = None
+        usage_dict: UsageDetails | None = None
         if response.usage:
             usage_dict = {
-                "prompt_tokens": response.usage.prompt_tokens,
-                "total_tokens": response.usage.total_tokens,
+                "input_token_count": response.usage.prompt_tokens,
+                "total_token_count": response.usage.total_tokens,
             }
 
         return GeneratedEmbeddings(embeddings, options=options, usage=usage_dict)
@@ -143,6 +143,7 @@ class OpenAIEmbeddingClient(
         default_headers: Additional HTTP headers.
         async_client: Pre-configured AsyncOpenAI client.
         base_url: Custom API base URL.
+        otel_provider_name: Override the OpenTelemetry provider name for telemetry.
         env_file_path: Path to .env file for settings.
         env_file_encoding: Encoding for .env file.
 
@@ -176,6 +177,7 @@ class OpenAIEmbeddingClient(
         default_headers: Mapping[str, str] | None = None,
         async_client: AsyncOpenAI | None = None,
         base_url: str | None = None,
+        otel_provider_name: str | None = None,
         env_file_path: str | None = None,
         env_file_encoding: str | None = None,
     ) -> None:
@@ -208,4 +210,5 @@ class OpenAIEmbeddingClient(
             org_id=openai_settings["org_id"],
             default_headers=default_headers,
             client=async_client,
+            otel_provider_name=otel_provider_name,
         )

python/packages/core/pyproject.toml

@@ -91,6 +91,9 @@ asyncio_mode = "auto"
 asyncio_default_fixture_loop_scope = "function"
 filterwarnings = []
 timeout = 120
+markers = [
+    "integration: marks tests as integration tests that require external services",
+]
 
 [tool.coverage.run]
 omit = [
@@ -124,6 +127,7 @@ exclude_dirs = ["tests"]
 [tool.poe]
 executor.type = "uv"
 include = "../../shared_tasks.toml"
+
 [tool.poe.tasks]
 mypy = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework"
 test = "pytest --cov=agent_framework --cov-report=term-missing:skip-covered -n auto --dist worksteal tests"

python/packages/core/tests/openai/test_openai_embedding_client.py

@@ -100,8 +100,8 @@ async def test_openai_get_embeddings_usage(openai_unit_test_env: None) -> None:
     result = await client.get_embeddings(["test"])
 
     assert result.usage is not None
-    assert result.usage["prompt_tokens"] == 10
-    assert result.usage["total_tokens"] == 10
+    assert result.usage["input_token_count"] == 10
+    assert result.usage["total_token_count"] == 10
 
 
 async def test_openai_options_passthrough_dimensions(openai_unit_test_env: None) -> None:
@@ -284,7 +284,7 @@ async def test_integration_openai_get_embeddings() -> None:
     assert all(isinstance(v, float) for v in result[0].vector)
     assert result[0].model_id is not None
     assert result.usage is not None
-    assert result.usage["prompt_tokens"] > 0
+    assert result.usage["input_token_count"] > 0
 
 
 @skip_if_openai_integration_tests_disabled
@@ -327,7 +327,7 @@ async def test_integration_azure_openai_get_embeddings() -> None:
     assert all(isinstance(v, float) for v in result[0].vector)
     assert result[0].model_id is not None
     assert result.usage is not None
-    assert result.usage["prompt_tokens"] > 0
+    assert result.usage["input_token_count"] > 0
 
 
 @skip_if_azure_openai_integration_tests_disabled