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>
2026-06-16 21:04:09 +08:00 · 2026-02-25 17:45:08 +00:00
parent e3a5b915a6
commit 6138487888
44 changed files with 1836 additions and 34 deletions
@@ -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)