diff --git a/python/packages/core/agent_framework/_clients.py b/python/packages/core/agent_framework/_clients.py index d353b0a552..b0503e079b 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 2de74b886a..e40f23b9be 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 6fb4125a12..1f08aff015 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 9560f99cc9..6587d183c1 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 cf2cb071a0..79f4fa68d9 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()