diff --git a/python/packages/azure-cosmos-memory/README.md b/python/packages/azure-cosmos-memory/README.md index 5a99b4f30ed..87d0857ce24 100644 --- a/python/packages/azure-cosmos-memory/README.md +++ b/python/packages/azure-cosmos-memory/README.md @@ -60,6 +60,14 @@ The provider supports the same authentication modes as other Azure integrations: Existing configurations can continue using `FOUNDRY_ENDPOINT`; the provider treats it as a legacy fallback. +### User-Agent Identification + +When the provider creates its memory client, Cosmos DB requests include +`agent-framework-azure-cosmos-memory/` alongside the +Agent Framework and Agent Memory Toolkit identifiers. This dedicated token +distinguishes the provider from direct toolkit usage. Caller-supplied +`memory_client` instances retain their existing User-Agent configuration. + ### Development Setup To avoid dependency conflicts with your system Python, it's recommended to use a virtual environment: diff --git a/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py b/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py index d316455b43d..b8452fd4e39 100644 --- a/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py +++ b/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py @@ -19,7 +19,7 @@ from agent_framework import AgentSession, ContextProvider, Message, SessionContext from agent_framework._settings import load_settings -from agent_framework._telemetry import mark_feature_used +from agent_framework._telemetry import get_user_agent, mark_feature_used from ._feature_usage import FeatureIndex @@ -211,6 +211,11 @@ def __init__( embedding_model = cast("str", settings.get("embedding_model")) chat_model = cast("str", settings.get("chat_model")) + # Defer the import because the package exports this provider before setting its version. + from . import __version__ + + user_agent = f"{get_user_agent()} agent-framework-azure-cosmos-memory/{__version__}" + # Authentication: if the caller supplies a credential, wire it into both the Cosmos # and AI Foundry clients and disable the toolkit's default-credential creation. # Otherwise let the toolkit build a DefaultAzureCredential (EnvironmentCredential → @@ -226,6 +231,7 @@ def __init__( cosmos_credential=credential, ai_foundry_credential=credential, use_default_credential=False, + user_agent=user_agent, cadence_thresholds=cadence_thresholds or None, ) else: @@ -236,6 +242,7 @@ def __init__( embedding_deployment_name=embedding_model, chat_deployment_name=chat_model, use_default_credential=True, + user_agent=user_agent, cadence_thresholds=cadence_thresholds or None, ) self._should_close_client = True diff --git a/python/packages/azure-cosmos-memory/tests/test_context_provider.py b/python/packages/azure-cosmos-memory/tests/test_context_provider.py index f48c7602102..b12296c3331 100644 --- a/python/packages/azure-cosmos-memory/tests/test_context_provider.py +++ b/python/packages/azure-cosmos-memory/tests/test_context_provider.py @@ -19,6 +19,7 @@ from agent_framework._sessions import AgentSession, SessionContext from agent_framework.exceptions import SettingNotFoundError +from agent_framework_azure_cosmos_memory import __version__ from agent_framework_azure_cosmos_memory._context_provider import ( DEFAULT_CONTEXT_PROMPT, CosmosMemoryContextProvider, @@ -99,15 +100,16 @@ class TestInit: def test_init_with_all_params(self, mock_memory_client: AsyncMock) -> None: """Initialize with all parameters provided.""" - provider = CosmosMemoryContextProvider( - source_id="test_memory", - memory_client=mock_memory_client, - top_k=10, - min_confidence=0.8, - memory_types=["fact", "episodic"], - context_prompt="Custom prompt:", - auto_extract=True, - ) + with patch("agent_framework_azure_cosmos_memory._context_provider.get_user_agent") as mock_get_user_agent: + provider = CosmosMemoryContextProvider( + source_id="test_memory", + memory_client=mock_memory_client, + top_k=10, + min_confidence=0.8, + memory_types=["fact", "episodic"], + context_prompt="Custom prompt:", + auto_extract=True, + ) assert provider.source_id == "test_memory" assert provider.top_k == 10 @@ -117,6 +119,7 @@ def test_init_with_all_params(self, mock_memory_client: AsyncMock) -> None: assert provider.auto_extract is True assert provider.memory_client is mock_memory_client assert provider._should_close_client is False + mock_get_user_agent.assert_not_called() def test_init_default_values(self, mock_memory_client: AsyncMock) -> None: """Initialize with default values.""" @@ -131,9 +134,13 @@ def test_init_default_values(self, mock_memory_client: AsyncMock) -> None: def test_init_creates_client_when_none(self) -> None: """When no client provided, creates AsyncCosmosMemoryClient with default credential.""" - with patch( - "agent_framework_azure_cosmos_memory._context_provider.AsyncCosmosMemoryClient" - ) as mock_client_class: + with ( + patch("agent_framework_azure_cosmos_memory._context_provider.AsyncCosmosMemoryClient") as mock_client_class, + patch( + "agent_framework_azure_cosmos_memory._context_provider.get_user_agent", + return_value="agent-framework-python/test", + ) as mock_get_user_agent, + ): mock_client_class.return_value = AsyncMock() provider = CosmosMemoryContextProvider( @@ -149,6 +156,10 @@ def test_init_creates_client_when_none(self) -> None: _, kwargs = mock_client_class.call_args assert kwargs["use_default_credential"] is True assert "cosmos_credential" not in kwargs + assert kwargs["user_agent"] == ( + f"agent-framework-python/test agent-framework-azure-cosmos-memory/{__version__}" + ) + mock_get_user_agent.assert_called_once_with() # The explicitly provided models are forwarded to the toolkit client. assert kwargs["embedding_deployment_name"] == "text-embedding-3-large" assert kwargs["chat_deployment_name"] == "gpt-4o-mini" @@ -156,9 +167,13 @@ def test_init_creates_client_when_none(self) -> None: def test_init_wires_explicit_credential(self) -> None: """An explicit credential is passed to both Cosmos and AI Foundry, disabling default.""" - with patch( - "agent_framework_azure_cosmos_memory._context_provider.AsyncCosmosMemoryClient" - ) as mock_client_class: + with ( + patch("agent_framework_azure_cosmos_memory._context_provider.AsyncCosmosMemoryClient") as mock_client_class, + patch( + "agent_framework_azure_cosmos_memory._context_provider.get_user_agent", + return_value="agent-framework-python/test", + ) as mock_get_user_agent, + ): mock_client_class.return_value = AsyncMock() sentinel = MagicMock() @@ -174,6 +189,10 @@ def test_init_wires_explicit_credential(self) -> None: assert kwargs["cosmos_credential"] is sentinel assert kwargs["ai_foundry_credential"] is sentinel assert kwargs["use_default_credential"] is False + assert kwargs["user_agent"] == ( + f"agent-framework-python/test agent-framework-azure-cosmos-memory/{__version__}" + ) + mock_get_user_agent.assert_called_once_with() def test_init_raises_without_endpoints(self, monkeypatch: pytest.MonkeyPatch) -> None: """Raises SettingNotFoundError when the Cosmos endpoint is not provided."""