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/bedrock/agent_framework_bedrock/__init__.py

@@ -3,6 +3,7 @@
 import importlib.metadata
 
 from ._chat_client import BedrockChatClient, BedrockChatOptions, BedrockGuardrailConfig, BedrockSettings
+from ._embedding_client import BedrockEmbeddingClient, BedrockEmbeddingOptions, BedrockEmbeddingSettings
 
 try:
     __version__ = importlib.metadata.version(__name__)
@@ -12,6 +13,9 @@ except importlib.metadata.PackageNotFoundError:
 __all__ = [
     "BedrockChatClient",
     "BedrockChatOptions",
+    "BedrockEmbeddingClient",
+    "BedrockEmbeddingOptions",
+    "BedrockEmbeddingSettings",
     "BedrockGuardrailConfig",
     "BedrockSettings",
     "__version__",

python/packages/bedrock/agent_framework_bedrock/_embedding_client.py

@@ -0,0 +1,292 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import sys
+from collections.abc import Sequence
+from typing import Any, ClassVar, Generic, TypedDict
+
+from agent_framework import (
+    AGENT_FRAMEWORK_USER_AGENT,
+    BaseEmbeddingClient,
+    Embedding,
+    EmbeddingGenerationOptions,
+    GeneratedEmbeddings,
+    SecretString,
+    UsageDetails,
+    load_settings,
+)
+from agent_framework.observability import EmbeddingTelemetryLayer
+from boto3.session import Session as Boto3Session
+from botocore.client import BaseClient
+from botocore.config import Config as BotoConfig
+
+if sys.version_info >= (3, 13):
+    from typing import TypeVar  # type: ignore # pragma: no cover
+else:
+    from typing_extensions import TypeVar  # type: ignore # pragma: no cover
+
+
+logger = logging.getLogger("agent_framework.bedrock")
+DEFAULT_REGION = "us-east-1"
+
+
+class BedrockEmbeddingSettings(TypedDict, total=False):
+    """Bedrock embedding settings."""
+
+    region: str | None
+    embedding_model_id: str | None
+    access_key: SecretString | None
+    secret_key: SecretString | None
+    session_token: SecretString | None
+
+
+class BedrockEmbeddingOptions(EmbeddingGenerationOptions, total=False):
+    """Bedrock-specific embedding options.
+
+    Extends EmbeddingGenerationOptions with Bedrock-specific fields.
+
+    Examples:
+        .. code-block:: python
+
+            from agent_framework_bedrock import BedrockEmbeddingOptions
+
+            options: BedrockEmbeddingOptions = {
+                "model_id": "amazon.titan-embed-text-v2:0",
+                "dimensions": 1024,
+                "normalize": True,
+            }
+    """
+
+    normalize: bool
+
+
+BedrockEmbeddingOptionsT = TypeVar(
+    "BedrockEmbeddingOptionsT",
+    bound=TypedDict,  # type: ignore[valid-type]
+    default="BedrockEmbeddingOptions",
+    covariant=True,
+)
+
+
+class RawBedrockEmbeddingClient(
+    BaseEmbeddingClient[str, list[float], BedrockEmbeddingOptionsT],
+    Generic[BedrockEmbeddingOptionsT],
+):
+    """Raw Bedrock embedding client without telemetry.
+
+    Keyword Args:
+        model_id: The Bedrock embedding model ID (e.g. "amazon.titan-embed-text-v2:0").
+            Can also be set via environment variable BEDROCK_EMBEDDING_MODEL_ID.
+        region: AWS region. Will try to load from BEDROCK_REGION env var,
+            if not set, the regular Boto3 configuration/loading applies
+            (which may include other env vars, config files, or instance metadata).
+        access_key: AWS access key for manual credential injection.
+        secret_key: AWS secret key paired with access_key.
+        session_token: AWS session token for temporary credentials.
+        client: Preconfigured Bedrock runtime client.
+        boto3_session: Custom boto3 session used to build the runtime client.
+        env_file_path: Path to .env file for settings.
+        env_file_encoding: Encoding for .env file.
+    """
+
+    def __init__(
+        self,
+        *,
+        region: str | None = None,
+        model_id: str | None = None,
+        access_key: str | None = None,
+        secret_key: str | None = None,
+        session_token: str | None = None,
+        client: BaseClient | None = None,
+        boto3_session: Boto3Session | None = None,
+        env_file_path: str | None = None,
+        env_file_encoding: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize a raw Bedrock embedding client."""
+        settings = load_settings(
+            BedrockEmbeddingSettings,
+            env_prefix="BEDROCK_",
+            required_fields=["embedding_model_id"],
+            region=region,
+            embedding_model_id=model_id,
+            access_key=access_key,
+            secret_key=secret_key,
+            session_token=session_token,
+            env_file_path=env_file_path,
+            env_file_encoding=env_file_encoding,
+        )
+        resolved_region = settings.get("region") or DEFAULT_REGION
+
+        if client is None:
+            if not boto3_session:
+                session_kwargs: dict[str, Any] = {}
+                if region := settings.get("region"):
+                    session_kwargs["region_name"] = region
+                if (access_key := settings.get("access_key")) and (secret_key := settings.get("secret_key")):
+                    session_kwargs["aws_access_key_id"] = access_key.get_secret_value()  # type: ignore[union-attr]
+                    session_kwargs["aws_secret_access_key"] = secret_key.get_secret_value()  # type: ignore[union-attr]
+                if session_token := settings.get("session_token"):
+                    session_kwargs["aws_session_token"] = session_token.get_secret_value()  # type: ignore[union-attr]
+                boto3_session = Boto3Session(**session_kwargs)
+            client = boto3_session.client(
+                "bedrock-runtime",
+                region_name=boto3_session.region_name or resolved_region,
+                config=BotoConfig(user_agent_extra=AGENT_FRAMEWORK_USER_AGENT),
+            )
+
+        self._bedrock_client = client
+        self.model_id = settings["embedding_model_id"]  # type: ignore[assignment]
+        self.region = resolved_region
+        super().__init__(**kwargs)
+
+    def service_url(self) -> str:
+        """Get the URL of the service."""
+        return str(self._bedrock_client.meta.endpoint_url)
+
+    async def get_embeddings(
+        self,
+        values: Sequence[str],
+        *,
+        options: BedrockEmbeddingOptionsT | None = None,
+    ) -> GeneratedEmbeddings[list[float]]:
+        """Call the Bedrock invoke_model API for embeddings.
+
+        Uses the Amazon Titan Embeddings model format. Each value is embedded
+        individually since Titan's invoke_model API accepts one input at a time.
+
+        Args:
+            values: The text values to generate embeddings for.
+            options: Optional embedding generation options.
+
+        Returns:
+            Generated embeddings with usage metadata.
+
+        Raises:
+            ValueError: If model_id is not provided or values is empty.
+        """
+        if not values:
+            return GeneratedEmbeddings([], options=options)
+
+        opts: dict[str, Any] = dict(options) if options else {}
+        model = opts.get("model_id") or self.model_id
+        if not model:
+            raise ValueError("model_id is required")
+
+        embedding_results = await asyncio.gather(
+            *(self._generate_embedding_for_text(opts, model, text) for text in values)
+        )
+        embeddings: list[Embedding[list[float]]] = []
+        total_input_tokens = 0
+        for embedding, input_tokens in embedding_results:
+            embeddings.append(embedding)
+            total_input_tokens += input_tokens
+
+        usage_dict: UsageDetails | None = None
+        if total_input_tokens > 0:
+            usage_dict = {"input_token_count": total_input_tokens}
+
+        return GeneratedEmbeddings(embeddings, options=options, usage=usage_dict)
+
+    async def _generate_embedding_for_text(
+        self,
+        opts: dict[str, Any],
+        model: str,
+        text: str,
+    ) -> tuple[Embedding[list[float]], int]:
+        body: dict[str, Any] = {"inputText": text}
+        if dimensions := opts.get("dimensions"):
+            body["dimensions"] = dimensions
+        if (normalize := opts.get("normalize")) is not None:
+            body["normalize"] = normalize
+
+        response = await asyncio.to_thread(
+            self._bedrock_client.invoke_model,
+            modelId=model,
+            contentType="application/json",
+            accept="application/json",
+            body=json.dumps(body),
+        )
+
+        response_body = json.loads(response["body"].read())
+        embedding = Embedding(
+            vector=response_body["embedding"],
+            dimensions=len(response_body["embedding"]),
+            model_id=model,
+        )
+        input_tokens = int(response_body.get("inputTextTokenCount", 0))
+        return embedding, input_tokens
+
+
+class BedrockEmbeddingClient(
+    EmbeddingTelemetryLayer[str, list[float], BedrockEmbeddingOptionsT],
+    RawBedrockEmbeddingClient[BedrockEmbeddingOptionsT],
+    Generic[BedrockEmbeddingOptionsT],
+):
+    """Bedrock embedding client with telemetry support.
+
+    Uses the Amazon Titan Embeddings model via Bedrock's invoke_model API.
+
+    Keyword Args:
+        model_id: The Bedrock embedding model ID (e.g. "amazon.titan-embed-text-v2:0").
+            Can also be set via environment variable BEDROCK_EMBEDDING_MODEL_ID.
+        region: AWS region. Defaults to "us-east-1".
+            Can also be set via environment variable BEDROCK_REGION.
+        access_key: AWS access key for manual credential injection.
+        secret_key: AWS secret key paired with access_key.
+        session_token: AWS session token for temporary credentials.
+        client: Preconfigured Bedrock runtime client.
+        boto3_session: Custom boto3 session used to build the runtime client.
+        env_file_path: Path to .env file for settings.
+        env_file_encoding: Encoding for .env file.
+
+    Examples:
+        .. code-block:: python
+
+            from agent_framework_bedrock import BedrockEmbeddingClient
+
+            # Using default AWS credentials
+            client = BedrockEmbeddingClient(
+                model_id="amazon.titan-embed-text-v2:0",
+            )
+
+            # Generate embeddings
+            result = await client.get_embeddings(["Hello, world!"])
+            print(result[0].vector)
+    """
+
+    OTEL_PROVIDER_NAME: ClassVar[str] = "aws.bedrock"  # type: ignore[reportIncompatibleVariableOverride, misc]
+
+    def __init__(
+        self,
+        *,
+        region: str | None = None,
+        model_id: str | None = None,
+        access_key: str | None = None,
+        secret_key: str | None = None,
+        session_token: str | None = None,
+        client: BaseClient | None = None,
+        boto3_session: Boto3Session | None = None,
+        otel_provider_name: str | None = None,
+        env_file_path: str | None = None,
+        env_file_encoding: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize a Bedrock embedding client."""
+        super().__init__(
+            region=region,
+            model_id=model_id,
+            access_key=access_key,
+            secret_key=secret_key,
+            session_token=session_token,
+            client=client,
+            boto3_session=boto3_session,
+            otel_provider_name=otel_provider_name,
+            env_file_path=env_file_path,
+            env_file_encoding=env_file_encoding,
+            **kwargs,
+        )

python/packages/bedrock/pyproject.toml

@@ -28,7 +28,6 @@ dependencies = [
     "botocore>=1.35.0,<2.0.0",
 ]
 
-
 [tool.uv]
 prerelease = "if-necessary-or-explicit"
 environments = [
@@ -46,6 +45,9 @@ addopts = "-ra -q -r fEX"
 asyncio_mode = "auto"
 asyncio_default_fixture_loop_scope = "function"
 filterwarnings = []
+markers = [
+    "integration: marks tests as integration tests that require external services",
+]
 timeout = 120
 
 [tool.ruff]

python/packages/bedrock/tests/bedrock/test_bedrock_embedding_client.py

@@ -0,0 +1,168 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import json
+import os
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+from agent_framework import Embedding, GeneratedEmbeddings
+
+from agent_framework_bedrock import BedrockEmbeddingClient, BedrockEmbeddingOptions
+
+
+class _StubBedrockEmbeddingRuntime:
+    """Stub for the Bedrock runtime client that handles invoke_model for embeddings."""
+
+    def __init__(self) -> None:
+        self.calls: list[dict[str, Any]] = []
+
+    def invoke_model(self, **kwargs: Any) -> dict[str, Any]:
+        self.calls.append(kwargs)
+        body = json.loads(kwargs.get("body", "{}"))
+        # Simulate Titan embedding response
+        dimensions = body.get("dimensions", 3)
+        return {
+            "body": MagicMock(
+                read=lambda: json.dumps({
+                    "embedding": [0.1 * (i + 1) for i in range(dimensions)],
+                    "inputTextTokenCount": 5,
+                }).encode()
+            ),
+        }
+
+
+async def test_bedrock_embedding_construction() -> None:
+    """Test construction with explicit parameters."""
+    stub = _StubBedrockEmbeddingRuntime()
+    client = BedrockEmbeddingClient(
+        model_id="amazon.titan-embed-text-v2:0",
+        region="us-west-2",
+        client=stub,
+    )
+    assert client.model_id == "amazon.titan-embed-text-v2:0"
+    assert client.region == "us-west-2"
+
+
+async def test_bedrock_embedding_construction_missing_model_raises(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Test that missing model_id raises an error."""
+    monkeypatch.delenv("BEDROCK_EMBEDDING_MODEL_ID", raising=False)
+    from agent_framework.exceptions import SettingNotFoundError
+
+    with pytest.raises(SettingNotFoundError):
+        BedrockEmbeddingClient(region="us-west-2")
+
+
+async def test_bedrock_embedding_get_embeddings() -> None:
+    """Test generating embeddings via the Bedrock invoke_model API."""
+    stub = _StubBedrockEmbeddingRuntime()
+    client = BedrockEmbeddingClient(
+        model_id="amazon.titan-embed-text-v2:0",
+        region="us-west-2",
+        client=stub,
+    )
+
+    result = await client.get_embeddings(["hello", "world"])
+
+    assert isinstance(result, GeneratedEmbeddings)
+    assert len(result) == 2
+    assert len(result[0].vector) == 3
+    assert len(result[1].vector) == 3
+    assert result[0].model_id == "amazon.titan-embed-text-v2:0"
+    assert result.usage == {"input_token_count": 10}
+
+    # Two calls since Titan processes one input at a time
+    assert len(stub.calls) == 2
+    call_texts = {json.loads(call["body"])["inputText"] for call in stub.calls}
+    assert call_texts == {"hello", "world"}
+
+
+async def test_bedrock_embedding_get_embeddings_empty_input() -> None:
+    """Test generating embeddings with empty input."""
+    stub = _StubBedrockEmbeddingRuntime()
+    client = BedrockEmbeddingClient(
+        model_id="amazon.titan-embed-text-v2:0",
+        region="us-west-2",
+        client=stub,
+    )
+
+    result = await client.get_embeddings([])
+
+    assert isinstance(result, GeneratedEmbeddings)
+    assert len(result) == 0
+    assert len(stub.calls) == 0
+
+
+async def test_bedrock_embedding_get_embeddings_with_options() -> None:
+    """Test generating embeddings with custom options."""
+    stub = _StubBedrockEmbeddingRuntime()
+    client = BedrockEmbeddingClient(
+        model_id="amazon.titan-embed-text-v2:0",
+        region="us-west-2",
+        client=stub,
+    )
+
+    options: BedrockEmbeddingOptions = {
+        "dimensions": 5,
+        "normalize": True,
+    }
+    result = await client.get_embeddings(["hello"], options=options)
+
+    assert len(result) == 1
+    assert len(result[0].vector) == 5
+
+    body = json.loads(stub.calls[0]["body"])
+    assert body["dimensions"] == 5
+    assert body["normalize"] is True
+
+
+async def test_bedrock_embedding_get_embeddings_no_model_raises() -> None:
+    """Test that missing model_id at call time raises ValueError."""
+    stub = _StubBedrockEmbeddingRuntime()
+    client = BedrockEmbeddingClient(
+        model_id="amazon.titan-embed-text-v2:0",
+        region="us-west-2",
+        client=stub,
+    )
+    client.model_id = None  # type: ignore[assignment]
+
+    with pytest.raises(ValueError, match="model_id is required"):
+        await client.get_embeddings(["hello"])
+
+
+async def test_bedrock_embedding_default_region() -> None:
+    """Test that default region is us-east-1."""
+    stub = _StubBedrockEmbeddingRuntime()
+    client = BedrockEmbeddingClient(
+        model_id="amazon.titan-embed-text-v2:0",
+        client=stub,
+    )
+    assert client.region == "us-east-1"
+
+
+# region: Integration Tests
+
+skip_if_bedrock_embedding_integration_tests_disabled = pytest.mark.skipif(
+    os.getenv("BEDROCK_EMBEDDING_MODEL_ID", "") in ("", "test-model")
+    or not (os.getenv("AWS_ACCESS_KEY_ID") or os.getenv("BEDROCK_ACCESS_KEY")),
+    reason="No real Bedrock embedding model or AWS credentials provided; skipping integration tests.",
+)
+
+
+@pytest.mark.flaky
+@pytest.mark.integration
+@skip_if_bedrock_embedding_integration_tests_disabled
+async def test_bedrock_embedding_integration() -> None:
+    """Integration test for Bedrock embedding client."""
+    client = BedrockEmbeddingClient()
+    result = await client.get_embeddings(["Hello, world!", "How are you?"])
+
+    assert isinstance(result, GeneratedEmbeddings)
+    assert len(result) == 2
+    for embedding in result:
+        assert isinstance(embedding, Embedding)
+        assert isinstance(embedding.vector, list)
+        assert len(embedding.vector) > 0
+        assert all(isinstance(v, float) for v in embedding.vector)