Adding user agent

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/<version>`) 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 (`"<custom> <toolkit>"`). 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
   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)
 

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

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 (``"<custom> <toolkit>"``) 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 = {

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,11 @@ 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 +325,11 @@ 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(

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,11 @@ 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 +293,11 @@ 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(

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"}

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

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)