Python: feat(python): add close() and async context manager to embedding clients

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
 

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

python/packages/ollama/agent_framework_ollama/_embedding_client.py

@@ -112,13 +112,23 @@ 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)
 
     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],

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.
 

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()