From 8664a6b1406a7aeb2f3618d3f20e6a76b4e034f1 Mon Sep 17 00:00:00 2001 From: Lushiyao5 <321342421+Lushiyao5@users.noreply.github.com> Date: Fri, 30 Oct 2026 12:07:53 -0500 Subject: [PATCH] feat(python): add close() and async context manager to embedding clients Embedding clients were left out of the deterministic lifecycle that #9748 gave chat clients: BaseEmbeddingClient had no __aenter__/__aexit__/close(), so 'async with' did not work and internally created SDK clients (OpenAI httpx pools, Ollama httpx transports) were only released by GC. This aligns the embedding family with the chat clients and with the existing GeminiEmbeddingClient.close()/_owns_client pattern. - BaseEmbeddingClient: async context manager and a no-op close() default, mirroring BaseChatClient. - OpenAI embedding client: track _owns_client; close() awaits client.close() only when the SDK client was created internally. - Ollama embedding client: same ownership pattern via AsyncClient.aclose(). Caller-supplied clients always stay open: their lifecycle belongs to the caller. Bedrock's boto3 client has no close() to call and keeps its default no-op. Closes #9814 --- .../packages/core/agent_framework/_clients.py | 23 ++++++++++++++++++ .../core/tests/core/test_embedding_client.py | 13 ++++++++++ .../_embedding_client.py | 10 ++++++++ .../_embedding_client.py | 11 +++++++++ .../openai/test_openai_embedding_client.py | 24 +++++++++++++++++++ 5 files changed, 81 insertions(+) diff --git a/python/packages/core/agent_framework/_clients.py b/python/packages/core/agent_framework/_clients.py index d353b0a5521..b0503e079b1 100644 --- a/python/packages/core/agent_framework/_clients.py +++ b/python/packages/core/agent_framework/_clients.py @@ -1049,6 +1049,29 @@ async def get_embeddings( """ ... + async def __aenter__(self) -> Self: + """Enter the embedding 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 embedding client without suppressing exceptions.""" + result = self.close() + if isawaitable(result): + await result + + async def close(self) -> None: + """Release embedding-client-owned connections and credentials. + + The default implementation does nothing. Resource-owning implementations + override this method and must leave caller-supplied resources open. + """ + # endregion diff --git a/python/packages/core/tests/core/test_embedding_client.py b/python/packages/core/tests/core/test_embedding_client.py index 2de74b886a1..e40f23b9bed 100644 --- a/python/packages/core/tests/core/test_embedding_client.py +++ b/python/packages/core/tests/core/test_embedding_client.py @@ -102,3 +102,16 @@ async def generate(self, values): client = NotAnEmbeddingClient() assert not isinstance(client, SupportsGetEmbeddings) + + +async def test_base_embedding_client_async_context_manager() -> None: + """The base close() is a no-op, but the async context manager protocol works.""" + client = MockEmbeddingClient() + + async with client as entered: + assert entered is client + + +async def test_base_embedding_client_close_default_is_noop() -> None: + client = MockEmbeddingClient() + await client.close() # must not raise diff --git a/python/packages/ollama/agent_framework_ollama/_embedding_client.py b/python/packages/ollama/agent_framework_ollama/_embedding_client.py index 6fb4125a123..1f08aff015a 100644 --- a/python/packages/ollama/agent_framework_ollama/_embedding_client.py +++ b/python/packages/ollama/agent_framework_ollama/_embedding_client.py @@ -112,6 +112,8 @@ def __init__( self.model = ollama_settings["embedding_model"] # type: ignore[assignment,reportTypedDictNotRequiredAccess] self.client = client or AsyncClient(host=ollama_settings.get("host")) + # Track ownership so close() only releases clients this instance created. + self._owns_client = client is None self.host = str(self.client._client.base_url) # type: ignore[reportUnknownMemberType,reportPrivateUsage,reportUnknownArgumentType] super().__init__(additional_properties=additional_properties) @@ -119,6 +121,14 @@ def service_url(self) -> str: """Get the URL of the service.""" return self.host + async def close(self) -> None: + """Close the internally created Ollama client when this instance owns it. + + A caller-supplied ``client`` stays open: its lifecycle belongs to the caller. + """ + if self._owns_client: + await self.client.close() + async def get_embeddings( self, values: Sequence[str], diff --git a/python/packages/openai/agent_framework_openai/_embedding_client.py b/python/packages/openai/agent_framework_openai/_embedding_client.py index 9560f99cc91..6587d183c11 100644 --- a/python/packages/openai/agent_framework_openai/_embedding_client.py +++ b/python/packages/openai/agent_framework_openai/_embedding_client.py @@ -229,6 +229,8 @@ def __init__( ) self.client = client + # Track ownership so close() only releases clients this instance created. + self._owns_client = async_client is None resolved_model = settings.get("model") self.model: str | None = resolved_model.strip() if isinstance(resolved_model, str) and resolved_model else None @@ -252,6 +254,15 @@ def service_url(self) -> str: """Get the URL of the service.""" return str(self.client.base_url) if self.client else "Unknown" + async def close(self) -> None: + """Close the internally created OpenAI client when this instance owns it. + + A caller-supplied ``async_client`` stays open: its lifecycle belongs to + the caller. + """ + if self._owns_client and self.client is not None: + await self.client.close() + def _prepare_extra_request_options(self, options: Mapping[str, Any]) -> dict[str, Any]: """Return extra keyword arguments for the ``embeddings.create`` request. diff --git a/python/packages/openai/tests/openai/test_openai_embedding_client.py b/python/packages/openai/tests/openai/test_openai_embedding_client.py index cf2cb071a0d..79f4fa68d95 100644 --- a/python/packages/openai/tests/openai/test_openai_embedding_client.py +++ b/python/packages/openai/tests/openai/test_openai_embedding_client.py @@ -432,3 +432,27 @@ async def test_integration_openai_get_embeddings_with_dimensions() -> None: assert len(result) == 1 assert len(result[0].vector) == 256 + + +async def test_openai_embedding_client_close_ownership() -> None: + """close() releases an internally created client and leaves a caller-supplied one open.""" + from unittest.mock import AsyncMock, MagicMock + + from agent_framework_openai import OpenAIEmbeddingClient + + # Internally created client: close() must release it. + owned = OpenAIEmbeddingClient(model="text-embedding-3-small", api_key="test-key") + assert owned._owns_client is True + owned.client.close = AsyncMock() + await owned.close() + owned.client.close.assert_awaited_once() + + # Caller-supplied client: close() must leave it open. + supplied = MagicMock() + supplied.close = AsyncMock() + caller_owned = OpenAIEmbeddingClient( + model="text-embedding-3-small", api_key="test-key", async_client=supplied + ) + assert caller_owned._owns_client is False + await caller_owned.close() + supplied.close.assert_not_awaited()