updates from feedback

Author: eavanvalkenburg

python/packages/azure-ai/agent_framework_azure_ai/_embedding_client.py

@@ -5,6 +5,7 @@
 import logging
 import sys
 from collections.abc import Sequence
+from contextlib import suppress
 from typing import Any, ClassVar, Generic, TypedDict
 
 from agent_framework import (
@@ -13,8 +14,9 @@
     Embedding,
     EmbeddingGenerationOptions,
     GeneratedEmbeddings,
+    UsageDetails,
+    load_settings,
 )
-from agent_framework._settings import load_settings
 from agent_framework.observability import EmbeddingTelemetryLayer
 from azure.ai.inference.aio import EmbeddingsClient, ImageEmbeddingsClient
 from azure.ai.inference.models import ImageEmbeddingInput
@@ -162,8 +164,10 @@ def __init__(
 
     async def close(self) -> None:
         """Close the underlying SDK clients and release resources."""
-        await self._text_client.close()
-        await self._image_client.close()
+        with suppress(Exception):
+            await self._text_client.close()
+        with suppress(Exception):
+            await self._image_client.close()
 
     async def __aenter__(self) -> RawAzureAIInferenceEmbeddingClient[AzureAIInferenceEmbeddingOptionsT]:
         """Enter the async context manager."""
@@ -204,10 +208,6 @@ async def get_embeddings(
             return GeneratedEmbeddings([], options=options)  # type: ignore[reportReturnType]
 
         opts: dict[str, Any] = dict(options) if options else {}
-        text_model = opts.get("model_id") or self.model_id
-        image_model = opts.get("image_model_id") or self.image_model_id
-        if not text_model:
-            raise ValueError("model_id is required")
 
         # Separate text and image inputs, tracking original indices.
         text_items: list[tuple[int, str]] = []
@@ -249,12 +249,13 @@ async def get_embeddings(
             common_kwargs["model_extras"] = extra_parameters
 
         # Allocate results array.
-        results: list[Embedding[list[float]] | None] = [None] * len(values)
-        total_prompt_tokens = 0
-        total_completion_tokens = 0
+        embeddings: list[Embedding[list[float]] | None] = [None] * len(values)
+        usage_details: UsageDetails = {"input_token_count": 0, "output_token_count": 0}
 
         # Embed text inputs.
         if text_items:
+            if not (text_model := opts.get("model_id") or self.model_id):
+                raise ValueError("An model_id is required, either in the client or options, for text inputs.")
             text_inputs = [t for _, t in text_items]
             response = await self._text_client.embed(
                 input=text_inputs,
@@ -263,18 +264,19 @@ async def get_embeddings(
             )
             for i, item in enumerate(response.data):
                 original_idx = text_items[i][0]
-                vector: list[float] = [float(v) for v in item.embedding]
-                results[original_idx] = Embedding(
-                    vector=vector,
-                    dimensions=len(vector),
+                embeddings[original_idx] = Embedding(
+                    vector=item.embedding,
+                    dimensions=len(item.embedding),
                     model_id=response.model or text_model,
                 )
             if response.usage:
-                total_prompt_tokens += response.usage.prompt_tokens
-                total_completion_tokens += getattr(response.usage, "completion_tokens", 0) or 0
+                usage_details["input_token_count"] += response.usage.prompt_tokens
+                usage_details["output_token_count"] += getattr(response.usage, "completion_tokens", 0) or 0
 
         # Embed image inputs.
         if image_items:
+            if not (image_model := opts.get("image_model_id") or self.image_model_id):
+                raise ValueError("An image_model_id is required, either in the client or options, for image inputs.")
             image_inputs = [img for _, img in image_items]
             response = await self._image_client.embed(
                 input=image_inputs,
@@ -283,25 +285,16 @@ async def get_embeddings(
             )
             for i, item in enumerate(response.data):
                 original_idx = image_items[i][0]
-                img_vector: list[float] = [float(v) for v in item.embedding]
-                results[original_idx] = Embedding(
-                    vector=img_vector,
-                    dimensions=len(img_vector),
+                embeddings[original_idx] = Embedding(
+                    vector=item.embedding,
+                    dimensions=len(item.embedding),
                     model_id=response.model or image_model,
                 )
             if response.usage:
-                total_prompt_tokens += response.usage.prompt_tokens
-                total_completion_tokens += getattr(response.usage, "completion_tokens", 0) or 0
+                usage_details["input_token_count"] += response.usage.prompt_tokens
+                usage_details["output_token_count"] += getattr(response.usage, "completion_tokens", 0) or 0
 
-        embeddings = [r for r in results if r is not None]
-
-        usage_dict: dict[str, Any] | None = None
-        if total_prompt_tokens > 0 or total_completion_tokens > 0:
-            usage_dict = {"prompt_tokens": total_prompt_tokens}
-            if total_completion_tokens > 0:
-                usage_dict["completion_tokens"] = total_completion_tokens
-
-        return GeneratedEmbeddings(embeddings, options=options, usage=usage_dict)  # type: ignore[reportReturnType]
+        return GeneratedEmbeddings(embeddings, options=options, usage=usage_details)  # type: ignore[reportReturnType]
 
 
 class AzureAIInferenceEmbeddingClient(

python/packages/azure-ai/tests/azure_ai/test_azure_ai_inference_embedding_client.py

@@ -169,9 +169,7 @@ async def test_model_override_in_options(
         call_kwargs = mock_text_client.embed.call_args
         assert call_kwargs.kwargs["model"] == "custom-model"
 
-    async def test_unsupported_content_type_raises(
-        self, raw_client: RawAzureAIInferenceEmbeddingClient[Any]
-    ) -> None:
+    async def test_unsupported_content_type_raises(self, raw_client: RawAzureAIInferenceEmbeddingClient[Any]) -> None:
         """Non-text, non-image Content raises ValueError."""
         error_content = Content("error", message="fail")
         with pytest.raises(ValueError, match="Unsupported Content type"):
@@ -181,12 +179,10 @@ async def test_usage_metadata(
         self, raw_client: RawAzureAIInferenceEmbeddingClient[Any], mock_text_client: AsyncMock
     ) -> None:
         """Usage metadata is populated from the response."""
-        mock_text_client.embed.return_value = _make_embed_response(
-            [[0.1, 0.2]], prompt_tokens=42
-        )
+        mock_text_client.embed.return_value = _make_embed_response([[0.1, 0.2]], prompt_tokens=42)
         result = await raw_client.get_embeddings(["hello"])
         assert result.usage is not None
-        assert result.usage["prompt_tokens"] == 42
+        assert result.usage["input_token_count"] == 42
 
     def test_service_url(self, raw_client: RawAzureAIInferenceEmbeddingClient[Any]) -> None:
         """service_url returns the configured endpoint."""
@@ -229,9 +225,7 @@ def test_image_model_id_from_env(self) -> None:
             assert client.model_id == "text-model"
             assert client.image_model_id == "image-model"
 
-    def test_image_model_id_explicit(
-        self, mock_text_client: AsyncMock, mock_image_client: AsyncMock
-    ) -> None:
+    def test_image_model_id_explicit(self, mock_text_client: AsyncMock, mock_image_client: AsyncMock) -> None:
         """image_model_id can be set explicitly."""
         client = RawAzureAIInferenceEmbeddingClient(
             model_id="text-model",
@@ -277,9 +271,7 @@ async def test_otel_provider_name_default(self) -> None:
         """Default OTEL provider name is azure.ai.inference."""
         assert AzureAIInferenceEmbeddingClient.OTEL_PROVIDER_NAME == "azure.ai.inference"
 
-    async def test_otel_provider_name_override(
-        self, mock_text_client: AsyncMock, mock_image_client: AsyncMock
-    ) -> None:
+    async def test_otel_provider_name_override(self, mock_text_client: AsyncMock, mock_image_client: AsyncMock) -> None:
         """OTEL provider name can be overridden."""
         client = AzureAIInferenceEmbeddingClient(
             model_id="test-model",

python/packages/bedrock/agent_framework_bedrock/_embedding_client.py

@@ -15,8 +15,10 @@
     Embedding,
     EmbeddingGenerationOptions,
     GeneratedEmbeddings,
+    SecretString,
+    UsageDetails,
+    load_settings,
 )
-from agent_framework._settings import SecretString, load_settings
 from agent_framework.observability import EmbeddingTelemetryLayer
 from boto3.session import Session as Boto3Session
 from botocore.client import BaseClient
@@ -29,10 +31,19 @@
 
 
 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.
 
@@ -61,16 +72,6 @@ class BedrockEmbeddingOptions(EmbeddingGenerationOptions, total=False):
 )
 
 
-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 RawBedrockEmbeddingClient(
     BaseEmbeddingClient[str, list[float], BedrockEmbeddingOptionsT],
     Generic[BedrockEmbeddingOptionsT],
@@ -80,8 +81,9 @@ class RawBedrockEmbeddingClient(
     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.
+        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.
@@ -118,39 +120,33 @@ def __init__(
             env_file_path=env_file_path,
             env_file_encoding=env_file_encoding,
         )
-        if not settings.get("region"):
-            settings["region"] = DEFAULT_REGION
+        resolved_region = settings.get("region") or DEFAULT_REGION
 
         if client is None:
-            session = boto3_session or self._create_session(settings)
-            client = session.client(
+            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=settings["region"],
+                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"]
-        self.region = settings["region"]
+        self.model_id = settings["embedding_model_id"]  # type: ignore[assignment]
+        self.region = resolved_region
         super().__init__(**kwargs)
 
-    @staticmethod
-    def _create_session(settings: BedrockEmbeddingSettings) -> Boto3Session:
-        """Create a boto3 session from settings."""
-        session_kwargs: dict[str, Any] = {"region_name": settings.get("region") or DEFAULT_REGION}
-        if settings.get("access_key") and settings.get("secret_key"):
-            session_kwargs["aws_access_key_id"] = settings["access_key"].get_secret_value()  # type: ignore[union-attr]
-            session_kwargs["aws_secret_access_key"] = settings["secret_key"].get_secret_value()  # type: ignore[union-attr]
-        if settings.get("session_token"):
-            session_kwargs["aws_session_token"] = settings["session_token"].get_secret_value()  # type: ignore[union-attr]
-        return Boto3Session(**session_kwargs)
-
     def service_url(self) -> str:
         """Get the URL of the service."""
-        meta = getattr(self._bedrock_client, "meta", None)
-        if meta and hasattr(meta, "endpoint_url"):
-            return str(meta.endpoint_url)
-        return f"https://bedrock-runtime.{self.region}.amazonaws.com"
+        return str(self._bedrock_client.meta.endpoint_url)
 
     async def get_embeddings(
         self,
@@ -181,41 +177,50 @@ async def get_embeddings(
         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
 
-        for text in values:
-            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())
-            vector = response_body["embedding"]
-            embeddings.append(
-                Embedding(
-                    vector=vector,
-                    dimensions=len(vector),
-                    model_id=model,
-                )
-            )
-            total_input_tokens += response_body.get("inputTextTokenCount", 0)
-
-        usage_dict: dict[str, Any] | None = None
+        usage_dict: UsageDetails | None = None
         if total_input_tokens > 0:
-            usage_dict = {"prompt_tokens": total_input_tokens}
+            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],

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

@@ -26,12 +26,10 @@ def invoke_model(self, **kwargs: Any) -> dict[str, Any]:
         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()
+                read=lambda: json.dumps({
+                    "embedding": [0.1 * (i + 1) for i in range(dimensions)],
+                    "inputTextTokenCount": 5,
+                }).encode()
             ),
         }
 
@@ -73,14 +71,12 @@ async def test_bedrock_embedding_get_embeddings() -> None:
     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 == {"prompt_tokens": 10}
+    assert result.usage == {"input_token_count": 10}
 
     # Two calls since Titan processes one input at a time
     assert len(stub.calls) == 2
-    body0 = json.loads(stub.calls[0]["body"])
-    assert body0["inputText"] == "hello"
-    body1 = json.loads(stub.calls[1]["body"])
-    assert body1["inputText"] == "world"
+    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: