diff --git a/python/packages/bedrock/agent_framework_bedrock/_embedding_client.py b/python/packages/bedrock/agent_framework_bedrock/_embedding_client.py index 562b4229f5a..63e69445121 100644 --- a/python/packages/bedrock/agent_framework_bedrock/_embedding_client.py +++ b/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) diff --git a/python/packages/bedrock/tests/test_bedrock_client.py b/python/packages/bedrock/tests/test_bedrock_client.py index 9f753e8835d..a8c983e4e49 100644 --- a/python/packages/bedrock/tests/test_bedrock_client.py +++ b/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() diff --git a/python/packages/core/agent_framework/_clients.py b/python/packages/core/agent_framework/_clients.py index ed49727411c..a5e508bafa1 100644 --- a/python/packages/core/agent_framework/_clients.py +++ b/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, diff --git a/python/packages/core/tests/core/test_clients.py b/python/packages/core/tests/core/test_clients.py index 1bdc4ce5c82..692ae0a16fe 100644 --- a/python/packages/core/tests/core/test_clients.py +++ b/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() diff --git a/python/packages/ollama/agent_framework_ollama/_embedding_client.py b/python/packages/ollama/agent_framework_ollama/_embedding_client.py index 6fb4125a123..d2b66be9ee2 100644 --- a/python/packages/ollama/agent_framework_ollama/_embedding_client.py +++ b/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 diff --git a/python/packages/ollama/tests/ollama/test_ollama_embedding_client.py b/python/packages/ollama/tests/ollama/test_ollama_embedding_client.py index 6cf2b67405d..8697139dd64 100644 --- a/python/packages/ollama/tests/ollama/test_ollama_embedding_client.py +++ b/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() diff --git a/python/packages/openai/agent_framework_openai/_embedding_client.py b/python/packages/openai/agent_framework_openai/_embedding_client.py index 20823c0b957..615548c8981 100644 --- a/python/packages/openai/agent_framework_openai/_embedding_client.py +++ b/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" 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..778481268a1 100644 --- a/python/packages/openai/tests/openai/test_openai_embedding_client.py +++ b/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()