Python: Add async context managers and ownership close to embedding clients

python/packages/bedrock/agent_framework_bedrock/_embedding_client.py

@@ -31,6 +31,10 @@
     from typing import TypeVar  # pragma: no cover
 else:
     from typing_extensions import TypeVar  # pragma: no cover
+if sys.version_info >= (3, 12):
+    from typing import override  # pragma: no cover
+else:
+    from typing_extensions import override  # pragma: no cover
 
 
 logger = logging.getLogger("agent_framework.bedrock")
@@ -126,6 +130,8 @@ def __init__(
         )
         resolved_region = settings.get("region") or DEFAULT_REGION
 
+        self._owns_client = client is None
+        self._client_closed = False
         if client:
             self._bedrock_client = client
         else:
@@ -150,6 +156,13 @@ def __init__(
         self.region = resolved_region
         super().__init__(additional_properties=additional_properties)
 
+    @override
+    async def close(self) -> None:
+        """Close the owned runtime connection without closing caller-owned sessions."""
+        if self._owns_client and not self._client_closed:
+            await asyncio.to_thread(self._bedrock_client.close)
+            self._client_closed = True
+
     def service_url(self) -> str:
         """Get the URL of the service."""
         return str(self._bedrock_client.meta.endpoint_url)

python/packages/bedrock/tests/test_bedrock_client.py

@@ -1338,3 +1338,32 @@ def test_prepare_options_prompt_cleanup_uses_overridden_model_id() -> None:
     assert "inferenceConfig" not in request, "prompt ARN requests must not carry inferenceConfig"
     assert "system" not in request
     assert "toolConfig" not in request
+
+
+@pytest.mark.parametrize("injected", [False, True])
+async def test_bedrock_embedding_client_context_ownership(injected: bool) -> None:
+    """Context exit and close() release only the connector-owned runtime client."""
+    runtime = MagicMock()
+    kwargs: dict[str, Any] = (
+        {"client": cast(BaseClient, runtime)}
+        if injected
+        else {
+            "region": "us-east-1",
+            "access_key": "access",
+            "secret_key": "secret",
+        }
+    )
+
+    if injected:
+        client = BedrockEmbeddingClient(model="amazon.titan-embed-text-v2:0", **kwargs)
+    else:
+        with patch("agent_framework_bedrock._embedding_client.Boto3Session") as session_cls:
+            session_cls.return_value.client.return_value = runtime
+            client = BedrockEmbeddingClient(model="amazon.titan-embed-text-v2:0", **kwargs)
+
+    async with client as entered:
+        assert entered is client
+    await client.close()
+    assert runtime.close.call_count == (0 if injected else 1)
+    assert "_owns_client" not in client.to_dict()
+    assert "_client_closed" not in client.to_dict()

python/packages/core/agent_framework/_clients.py

@@ -1032,6 +1032,32 @@ def __init__(
         self.additional_properties = additional_properties or {}
         super().__init__()
 
+    async def __aenter__(self) -> Self:
+        """Enter the client context without requiring it for ordinary requests."""
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+        /,
+    ) -> None:
+        """Release resources owned by the client without suppressing exceptions."""
+        result = self.close()
+        if isawaitable(result):
+            await result
+
+    async def close(self) -> None:
+        """Release client-owned connections and credentials.
+
+        The default implementation does nothing. Resource-owning implementations
+        override this method and must leave caller-supplied resources open.
+        Context exit also accepts legacy synchronous ``close()`` overrides.
+        Clients remain usable without entering an async context; callers can use
+        ``await client.close()`` to release their resources explicitly.
+        """
+
     @abstractmethod
     async def get_embeddings(
         self,

python/packages/core/tests/core/test_clients.py

@@ -18,6 +18,7 @@
     SUMMARY_OF_MESSAGE_IDS_KEY,
     AgentTriggeredCompactionStrategy,
     BaseChatClient,
+    BaseEmbeddingClient,
     ChatMiddleware,
     ChatOptions,
     ChatResponse,
@@ -36,6 +37,7 @@
     tool,
 )
 from agent_framework._sessions import LOCAL_HISTORY_CONVERSATION_ID
+from agent_framework._types import Embedding, GeneratedEmbeddings
 from agent_framework.exceptions import ChatClientInvalidRequestException
 
 
@@ -1564,3 +1566,53 @@ async def fake_inner_get_response(**kwargs):
         assert appended_messages[0].text == "You are a helpful assistant."
         assert appended_messages[1].role == "user"
         assert appended_messages[1].text == "hello"
+
+
+class _MockEmbeddingClient(BaseEmbeddingClient[str, list[float], None]):
+    """Minimal embedding client used to exercise the base lifecycle methods."""
+
+    async def get_embeddings(self, values, *, options=None):  # type: ignore[override,no-untyped-def]
+        return GeneratedEmbeddings([Embedding(vector=[0.1, 0.2, 0.3]) for _ in values], options=options)
+
+
+async def test_base_embedding_client_default_lifecycle() -> None:
+    client = _MockEmbeddingClient()
+    await client.close()
+    async with client as entered:
+        assert entered is client
+        embeddings = await entered.get_embeddings(["hello"])
+        assert len(embeddings) == 1
+    await client.close()
+    assert len(await client.get_embeddings(["again"])) == 1
+
+
+@pytest.mark.parametrize("error", [None, RuntimeError("body failed")])
+async def test_base_embedding_client_context_calls_close(error: BaseException | None) -> None:
+    client = _MockEmbeddingClient()
+    with patch.object(client, "close", new_callable=AsyncMock) as close:
+        if error is None:
+            async with client:
+                close.assert_not_awaited()
+        else:
+            with pytest.raises(type(error)) as raised:
+                async with client:
+                    raise error
+            assert raised.value is error
+        close.assert_awaited_once_with()
+
+
+@pytest.mark.parametrize("error", [None, RuntimeError("body failed")])
+async def test_base_embedding_client_context_accepts_sync_close(
+    error: RuntimeError | None,
+) -> None:
+    client = _MockEmbeddingClient()
+    with patch.object(client, "close", new_callable=MagicMock, return_value=None) as close:
+        if error is None:
+            async with client:
+                close.assert_not_called()
+        else:
+            with pytest.raises(type(error)) as raised:
+                async with client:
+                    raise error
+            assert raised.value is error
+        close.assert_called_once_with()

python/packages/ollama/agent_framework_ollama/_embedding_client.py

@@ -25,6 +25,10 @@
     from typing import TypeVar  # pragma: no cover
 else:
     from typing_extensions import TypeVar  # pragma: no cover
+if sys.version_info >= (3, 12):
+    from typing import override  # pragma: no cover
+else:
+    from typing_extensions import override  # pragma: no cover
 
 
 logger = logging.getLogger("agent_framework.ollama")
@@ -112,9 +116,22 @@ def __init__(
 
         self.model = ollama_settings["embedding_model"]  # type: ignore[assignment,reportTypedDictNotRequiredAccess]
         self.client = client or AsyncClient(host=ollama_settings.get("host"))
+        self._owns_client = client is None
+        self._client_closed = False
         self.host = str(self.client._client.base_url)  # type: ignore[reportUnknownMemberType,reportPrivateUsage,reportUnknownArgumentType]
         super().__init__(additional_properties=additional_properties)
 
+    @override
+    async def close(self) -> None:
+        """Close the owned HTTP transport, leaving caller-supplied clients open.
+
+        ``ollama<0.6`` has no ``AsyncClient.aclose``; the underlying httpx
+        transport is closed directly, matching the constructor's own access.
+        """
+        if self._owns_client and not self._client_closed:
+            await self.client._client.aclose()  # type: ignore[reportPrivateUsage]
+            self._client_closed = True
+
     def service_url(self) -> str:
         """Get the URL of the service."""
         return self.host

python/packages/ollama/tests/ollama/test_ollama_embedding_client.py

@@ -180,3 +180,28 @@ async def test_ollama_embedding_integration() -> None:
         assert isinstance(embedding.vector, list)
         assert len(embedding.vector) > 0
         assert all(isinstance(v, float) for v in embedding.vector)
+
+
+@pytest.mark.parametrize("injected", [False, True])
+async def test_ollama_embedding_client_context_ownership(injected: bool) -> None:
+    """Context exit and close() release only the connector-owned HTTP transport."""
+    sdk_client = AsyncClient(host="http://localhost:11434") if injected else None
+    kwargs: dict[str, Any] = {"client": sdk_client} if injected else {"host": "http://localhost:11434"}
+
+    with patch("agent_framework_ollama._embedding_client.AsyncClient", wraps=AsyncClient) as client_cls:
+        client = OllamaEmbeddingClient(model="nomic-embed-text", **kwargs)
+        if injected:
+            client_cls.assert_not_called()
+        else:
+            client_cls.assert_called_once()
+
+    with patch.object(client.client._client, "aclose", wraps=client.client._client.aclose) as aclose:
+        async with client as entered:
+            assert entered is client
+        await client.close()
+    assert aclose.await_count == (0 if injected else 1)
+    assert client.client._client.is_closed is not injected
+    assert "_owns_client" not in client.to_dict()
+    assert "_client_closed" not in client.to_dict()
+    if sdk_client is not None:
+        await sdk_client._client.aclose()

python/packages/openai/agent_framework_openai/_embedding_client.py

@@ -23,6 +23,10 @@
     from typing import TypeVar  # pragma: no cover
 else:
     from typing_extensions import TypeVar  # pragma: no cover
+if sys.version_info >= (3, 12):
+    from typing import override  # pragma: no cover
+else:
+    from typing_extensions import override  # pragma: no cover
 
 if TYPE_CHECKING:
     from azure.core.credentials import TokenCredential
@@ -231,6 +235,8 @@ def __init__(
         )
 
         self.client = client
+        self._owns_client = async_client is None
+        self._client_closed = False
         resolved_model = settings.get("model")
         self.model: str | None = resolved_model.strip() if isinstance(resolved_model, str) and resolved_model else None
 
@@ -250,6 +256,13 @@ def __init__(
 
         super().__init__(additional_properties=additional_properties)
 
+    @override
+    async def close(self) -> None:
+        """Close the connector-owned SDK connection, not caller-supplied credentials."""
+        if self._owns_client and not self._client_closed:
+            await self.client.close()
+            self._client_closed = True
+
     def service_url(self) -> str:
         """Get the URL of the service."""
         return str(self.client.base_url) if self.client else "Unknown"

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

@@ -5,7 +5,7 @@
 import inspect
 import os
 from typing import Any, cast
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import agent_framework._telemetry as telemetry
 import pytest
@@ -432,3 +432,25 @@ async def test_integration_openai_get_embeddings_with_dimensions() -> None:
 
     assert len(result) == 1
     assert len(result[0].vector) == 256
+
+
+@pytest.mark.parametrize("injected", [False, True])
+async def test_openai_embedding_client_context_ownership(injected: bool) -> None:
+    """Context exit and close() release only the connector-owned SDK client."""
+    from openai import AsyncOpenAI
+
+    sdk_client = AsyncOpenAI(api_key="test-key") if injected else None
+    kwargs: dict[str, Any] = {"async_client": sdk_client} if injected else {"api_key": "test-key"}
+
+    client = OpenAIEmbeddingClient(model="text-embedding-3-small", **kwargs)
+    with patch.object(client.client, "close", wraps=client.client.close) as close:
+        async with client as entered:
+            assert entered is client
+            close.assert_not_awaited()
+        await client.close()
+        assert close.await_count == (0 if injected else 1)
+    assert client.client.is_closed() is not injected
+    assert "_owns_client" not in client.to_dict()
+    assert "_client_closed" not in client.to_dict()
+    if sdk_client is not None:
+        await sdk_client.close()