From 9239573861e07f56e3fc53377be8b4e18f11c55b Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Wed, 8 Jul 2026 15:08:04 -0700 Subject: [PATCH 1/3] Adding user agent --- CHANGELOG.md | 11 ++++-- .../cosmos/agent_memory/_base/base_client.py | 3 ++ azure/cosmos/agent_memory/_utils.py | 36 +++++++++++++++++++ .../agent_memory/aio/cosmos_memory_client.py | 10 ++++-- .../agent_memory/cosmos_memory_client.py | 10 ++++-- pyproject.toml | 2 +- tests/unit/test_cosmos_memory_client.py | 32 +++++++++++++++++ tests/unit/test_utils.py | 28 +++++++++++++++ uv.lock | 2 +- 9 files changed, 125 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7bc3f45..917cf41 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,16 +2,21 @@ ## [Unreleased] +## [0.2.0b3] (2026-07-08) + #### Features Added +* A custom user-agent can now be supplied via the new `user_agent` constructor + argument on `CosmosMemoryClient` and `AsyncCosmosMemoryClient`. The toolkit's + own user-agent (`azsdk-python-cosmos-agent-memory/`) is always sent to + Azure Cosmos DB so usage can be tracked; when a custom value is provided it is prefixed and + the toolkit's user-agent is suffixed behind it (`" "`). * Per-turn processing cadence can now be set in-process via the new `cadence_thresholds` constructor argument on `CosmosMemoryClient` and `AsyncCosmosMemoryClient`, instead of only through environment variables. Pass a mapping keyed by the same names as the env vars (e.g. `FACT_EXTRACTION_EVERY_N`, `DEDUP_EVERY_N`, `THREAD_SUMMARY_EVERY_N`, `USER_SUMMARY_EVERY_N`); any key not present falls back to the environment/defaults, and `None` preserves today's - env-only behavior. The mapping is validated and defensively copied at - construction: values are coerced to `int`, negatives are rejected (`0` disables - a step), and later mutation of the caller's mapping cannot change client behavior. + env-only behavior. See [PR:#29](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/29) ## [0.2.0b2] (2026-07-01) diff --git a/azure/cosmos/agent_memory/_base/base_client.py b/azure/cosmos/agent_memory/_base/base_client.py index 73d662b..ba154e8 100644 --- a/azure/cosmos/agent_memory/_base/base_client.py +++ b/azure/cosmos/agent_memory/_base/base_client.py @@ -14,6 +14,7 @@ _resolve_cosmos_provisioning_autoscale_max_ru, _resolve_cosmos_throughput_mode, _resolve_embedding_dimensions, + build_cosmos_user_agent, normalize_ai_foundry_endpoint, ) from azure.cosmos.agent_memory.exceptions import CosmosNotConnectedError, MemoryNotFoundError, ValidationError @@ -48,6 +49,7 @@ def _init_base_config( chat_deployment_name: str, use_default_credential: bool, enable_turn_embeddings: Optional[bool] = None, + user_agent: Optional[str] = None, default_credential_module: str = "azure.identity", ) -> None: """Initialize shared local state, config values, and default credentials.""" @@ -81,6 +83,7 @@ def _init_base_config( self._enable_turn_embeddings = ( enable_turn_embeddings if enable_turn_embeddings is not None else get_enable_turn_embeddings() ) + self._cosmos_user_agent = build_cosmos_user_agent(user_agent) self._owns_cosmos_credential = False self._owns_ai_foundry_credential = False diff --git a/azure/cosmos/agent_memory/_utils.py b/azure/cosmos/agent_memory/_utils.py index 849ef48..3cae259 100644 --- a/azure/cosmos/agent_memory/_utils.py +++ b/azure/cosmos/agent_memory/_utils.py @@ -24,6 +24,42 @@ VALID_TYPES = {"turn", "thread_summary", "fact", "user_summary", "procedural", "episodic"} +def _sdk_user_agent() -> str: + """Return the SDK user-agent suffix used to tag Cosmos DB telemetry. + + The suffix is appended to the azure-cosmos SDK's default user-agent and + surfaces in Cosmos DB backend telemetry (e.g. the ``userAgent`` column in + Kusto), enabling usage tracking of the Agent Memory Toolkit. + """ + try: + from importlib.metadata import PackageNotFoundError, version + + try: + pkg_version = version("azure-cosmos-agent-memory") + except PackageNotFoundError: + pkg_version = "unknown" + except Exception: + pkg_version = "unknown" + return f"azsdk-python-cosmos-agent-memory/{pkg_version}" + + +COSMOS_USER_AGENT = _sdk_user_agent() + + +def build_cosmos_user_agent(custom_user_agent: Optional[str] = None) -> str: + """Return the user-agent to pass to the Cosmos client. + + When ``custom_user_agent`` is provided, the toolkit's user-agent is + suffixed behind it (``" "``) so callers can attribute + telemetry to their application while still tracking toolkit usage. + Otherwise, only the toolkit's user-agent is returned. + """ + custom = (custom_user_agent or "").strip() + if custom: + return f"{custom} {COSMOS_USER_AGENT}" + return COSMOS_USER_AGENT + + def new_id(memory_type: str) -> str: """Return a fresh, type-prefixed UUID-backed memory id.""" prefix_map = { diff --git a/azure/cosmos/agent_memory/aio/cosmos_memory_client.py b/azure/cosmos/agent_memory/aio/cosmos_memory_client.py index 75b8aff..ecd6557 100644 --- a/azure/cosmos/agent_memory/aio/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/aio/cosmos_memory_client.py @@ -89,6 +89,7 @@ def __init__( chat_deployment_name: str = "gpt-4o-mini", use_default_credential: bool = True, enable_turn_embeddings: Optional[bool] = None, + user_agent: Optional[str] = None, embeddings_client: Optional[Any] = None, chat_client: Optional[Any] = None, processor: Optional[AsyncMemoryProcessor] = None, @@ -115,6 +116,7 @@ def __init__( chat_deployment_name=chat_deployment_name, use_default_credential=use_default_credential, enable_turn_embeddings=enable_turn_embeddings, + user_agent=user_agent, default_credential_module="azure.identity.aio", ) self._background_tasks: set[asyncio.Task[Any]] = set() @@ -246,7 +248,9 @@ async def connect_cosmos( from azure.cosmos.aio import CosmosClient await self._drain_cosmos_client() - client = CosmosClient(self._cosmos_endpoint, credential=self._cosmos_credential) + client = CosmosClient( + self._cosmos_endpoint, credential=self._cosmos_credential, user_agent=self._cosmos_user_agent + ) db = client.get_database_client(self._cosmos_database) self._cosmos_client = client self._memories_container_client = db.get_container_client(self._cosmos_container) @@ -319,7 +323,9 @@ async def create_memory_store( from azure.cosmos.aio import CosmosClient await self._drain_cosmos_client() - client = CosmosClient(self._cosmos_endpoint, credential=self._cosmos_credential) + client = CosmosClient( + self._cosmos_endpoint, credential=self._cosmos_credential, user_agent=self._cosmos_user_agent + ) db = await client.create_database_if_not_exists(id=self._cosmos_database) partition_key = PartitionKey(path=["/user_id", "/thread_id"], kind="MultiHash") offer = _cosmos_container_offer_throughput( diff --git a/azure/cosmos/agent_memory/cosmos_memory_client.py b/azure/cosmos/agent_memory/cosmos_memory_client.py index 71a5670..6467e21 100644 --- a/azure/cosmos/agent_memory/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/cosmos_memory_client.py @@ -81,6 +81,7 @@ def __init__( chat_deployment_name: str = "gpt-4o-mini", use_default_credential: bool = True, enable_turn_embeddings: Optional[bool] = None, + user_agent: Optional[str] = None, embeddings_client: Optional[Any] = None, chat_client: Optional[Any] = None, processor: Optional[MemoryProcessor] = None, @@ -107,6 +108,7 @@ def __init__( chat_deployment_name=chat_deployment_name, use_default_credential=use_default_credential, enable_turn_embeddings=enable_turn_embeddings, + user_agent=user_agent, ) # Embeddings/chat clients may be injected (e.g. an OpenAI-compatible backend, a # caller-configured client, or a deterministic fake for offline tests). When a client @@ -219,7 +221,9 @@ def connect_cosmos( from azure.cosmos import CosmosClient self._drain_cosmos_client() - client = CosmosClient(self._cosmos_endpoint, credential=self._cosmos_credential) + client = CosmosClient( + self._cosmos_endpoint, credential=self._cosmos_credential, user_agent=self._cosmos_user_agent + ) db = client.get_database_client(self._cosmos_database) self._cosmos_client = client self._memories_container_client = db.get_container_client(self._cosmos_container) @@ -287,7 +291,9 @@ def create_memory_store( from azure.cosmos import CosmosClient, PartitionKey, ThroughputProperties self._drain_cosmos_client() - client = CosmosClient(self._cosmos_endpoint, credential=self._cosmos_credential) + client = CosmosClient( + self._cosmos_endpoint, credential=self._cosmos_credential, user_agent=self._cosmos_user_agent + ) db = client.create_database_if_not_exists(id=self._cosmos_database) partition_key = PartitionKey(path=["/user_id", "/thread_id"], kind="MultiHash") offer = _cosmos_container_offer_throughput( diff --git a/pyproject.toml b/pyproject.toml index 2ead3b0..cf7554a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,7 +12,7 @@ namespaces = true [project] name = "azure-cosmos-agent-memory" -version = "0.2.0b2" +version = "0.2.0b3" description = "Store, retrieve, and transform AI agent memories backed by Azure Cosmos DB" readme = "README.md" license = {file = "LICENSE"} diff --git a/tests/unit/test_cosmos_memory_client.py b/tests/unit/test_cosmos_memory_client.py index c50e6af..5333210 100644 --- a/tests/unit/test_cosmos_memory_client.py +++ b/tests/unit/test_cosmos_memory_client.py @@ -349,6 +349,38 @@ def test_require_cosmos_before_connect(self): mem._require_cosmos() +class TestUserAgent: + def _connect_with_mock(self, **client_kwargs): + mock_cosmos_cls = MagicMock() + mock_client = MagicMock() + mock_db = MagicMock() + mock_cosmos_cls.return_value = mock_client + mock_client.get_database_client.return_value = mock_db + + mem = _make_client(**client_kwargs) + with patch.dict( + "sys.modules", + {"azure.cosmos": MagicMock(CosmosClient=mock_cosmos_cls)}, + ): + mem.connect_cosmos( + endpoint="https://fake.documents.azure.com:443/", + credential="fake-key", + ) + return mock_cosmos_cls + + def test_default_user_agent_forwarded_to_cosmos_client(self): + mock_cosmos_cls = self._connect_with_mock() + ua = mock_cosmos_cls.call_args.kwargs["user_agent"] + assert ua == "azsdk-python-cosmos-agent-memory/" + ua.split("/", 1)[1] + assert ua.startswith("azsdk-python-cosmos-agent-memory/") + + def test_custom_user_agent_is_prefixed_before_toolkit(self): + mock_cosmos_cls = self._connect_with_mock(user_agent="MyApp/1.2.3") + ua = mock_cosmos_cls.call_args.kwargs["user_agent"] + assert ua.startswith("MyApp/1.2.3 ") + assert ua.endswith("azsdk-python-cosmos-agent-memory/" + ua.rsplit("/", 1)[1]) + + class TestValidateTopology: def test_validate_topology_succeeds_on_healthy_deploy(self): mem = _make_client() diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 0749a0a..32e6fad 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -3,6 +3,7 @@ import pytest from azure.cosmos.agent_memory._utils import ( + COSMOS_USER_AGENT, DEFAULT_TTL_BY_TYPE, _build_container_kwargs, _container_policies, @@ -11,6 +12,7 @@ _resolve_embedding_data_type, _resolve_full_text_language, _resolve_vector_index_type, + build_cosmos_user_agent, compute_content_hash, normalize_ai_foundry_endpoint, ) @@ -350,3 +352,29 @@ def test_normalize_ai_foundry_endpoint_leaves_openai_host_with_project_path(): normalize_ai_foundry_endpoint("https://my-res.openai.azure.com/api/projects/my-project") == "https://my-res.openai.azure.com/api/projects/my-project" ) + + +# --------------------------------------------------------------------------- +# User-agent helpers +# --------------------------------------------------------------------------- + + +def test_cosmos_user_agent_follows_azure_sdk_convention(): + assert COSMOS_USER_AGENT.startswith("azsdk-python-cosmos-agent-memory/") + + +def test_build_cosmos_user_agent_defaults_to_toolkit_only(): + assert build_cosmos_user_agent() == COSMOS_USER_AGENT + assert build_cosmos_user_agent(None) == COSMOS_USER_AGENT + + +@pytest.mark.parametrize("blank", ["", " ", "\t\n"]) +def test_build_cosmos_user_agent_treats_blank_as_unset(blank): + assert build_cosmos_user_agent(blank) == COSMOS_USER_AGENT + + +def test_build_cosmos_user_agent_suffixes_toolkit_behind_custom(): + result = build_cosmos_user_agent("MyApp/1.2.3") + assert result == f"MyApp/1.2.3 {COSMOS_USER_AGENT}" + assert result.startswith("MyApp/1.2.3 ") + assert result.endswith(COSMOS_USER_AGENT) diff --git a/uv.lock b/uv.lock index 610993f..8514440 100644 --- a/uv.lock +++ b/uv.lock @@ -194,7 +194,7 @@ wheels = [ [[package]] name = "azure-cosmos-agent-memory" -version = "0.2.0b2" +version = "0.2.0b3" source = { editable = "." } dependencies = [ { name = "aiohttp" }, From 1b69baa8d98c2859e530bfa507741334a032f01a Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Wed, 8 Jul 2026 15:09:29 -0700 Subject: [PATCH 2/3] Updating changelog --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 917cf41..32e7342 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,7 @@ argument on `CosmosMemoryClient` and `AsyncCosmosMemoryClient`. The toolkit's own user-agent (`azsdk-python-cosmos-agent-memory/`) is always sent to Azure Cosmos DB so usage can be tracked; when a custom value is provided it is prefixed and - the toolkit's user-agent is suffixed behind it (`" "`). + the toolkit's user-agent is suffixed behind it (`" "`). See [PR:#30](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/30) * Per-turn processing cadence can now be set in-process via the new `cadence_thresholds` constructor argument on `CosmosMemoryClient` and `AsyncCosmosMemoryClient`, instead of only through environment variables. Pass a From 83b5167871363237274a9b62b7ca29c87f7d032f Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Wed, 8 Jul 2026 18:46:10 -0700 Subject: [PATCH 3/3] Resolving comments --- azure/cosmos/agent_memory/aio/cosmos_memory_client.py | 8 ++++++-- azure/cosmos/agent_memory/cosmos_memory_client.py | 8 ++++++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/azure/cosmos/agent_memory/aio/cosmos_memory_client.py b/azure/cosmos/agent_memory/aio/cosmos_memory_client.py index ecd6557..7b24172 100644 --- a/azure/cosmos/agent_memory/aio/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/aio/cosmos_memory_client.py @@ -249,7 +249,9 @@ async def connect_cosmos( await self._drain_cosmos_client() client = CosmosClient( - self._cosmos_endpoint, credential=self._cosmos_credential, user_agent=self._cosmos_user_agent + self._cosmos_endpoint, + credential=self._cosmos_credential, + user_agent=self._cosmos_user_agent, ) db = client.get_database_client(self._cosmos_database) self._cosmos_client = client @@ -324,7 +326,9 @@ async def create_memory_store( await self._drain_cosmos_client() client = CosmosClient( - self._cosmos_endpoint, credential=self._cosmos_credential, user_agent=self._cosmos_user_agent + self._cosmos_endpoint, + credential=self._cosmos_credential, + user_agent=self._cosmos_user_agent, ) db = await client.create_database_if_not_exists(id=self._cosmos_database) partition_key = PartitionKey(path=["/user_id", "/thread_id"], kind="MultiHash") diff --git a/azure/cosmos/agent_memory/cosmos_memory_client.py b/azure/cosmos/agent_memory/cosmos_memory_client.py index 6467e21..84114cd 100644 --- a/azure/cosmos/agent_memory/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/cosmos_memory_client.py @@ -222,7 +222,9 @@ def connect_cosmos( self._drain_cosmos_client() client = CosmosClient( - self._cosmos_endpoint, credential=self._cosmos_credential, user_agent=self._cosmos_user_agent + self._cosmos_endpoint, + credential=self._cosmos_credential, + user_agent=self._cosmos_user_agent, ) db = client.get_database_client(self._cosmos_database) self._cosmos_client = client @@ -292,7 +294,9 @@ def create_memory_store( self._drain_cosmos_client() client = CosmosClient( - self._cosmos_endpoint, credential=self._cosmos_credential, user_agent=self._cosmos_user_agent + self._cosmos_endpoint, + credential=self._cosmos_credential, + user_agent=self._cosmos_user_agent, ) db = client.create_database_if_not_exists(id=self._cosmos_database) partition_key = PartitionKey(path=["/user_id", "/thread_id"], kind="MultiHash")