From 91f93ac39a27c0d3f22263de37fb165c07fd0cbf Mon Sep 17 00:00:00 2001 From: Shashwat Gupta Date: Tue, 3 Nov 2026 00:11:39 +0530 Subject: [PATCH 1/2] Python: Preserve injected Cosmos memory client configuration ownership --- python/packages/azure-cosmos-memory/README.md | 6 +- .../_context_provider.py | 60 ++++++------- .../tests/test_context_provider.py | 89 +++++++++++++++---- 3 files changed, 101 insertions(+), 54 deletions(-) diff --git a/python/packages/azure-cosmos-memory/README.md b/python/packages/azure-cosmos-memory/README.md index 87d0857ce24..96fcc3af7eb 100644 --- a/python/packages/azure-cosmos-memory/README.md +++ b/python/packages/azure-cosmos-memory/README.md @@ -420,8 +420,10 @@ For fine-grained control over memory processing: from azure.cosmos.agent_memory.aio import AsyncCosmosMemoryClient # Create a custom memory client. To disable automatic extraction, zero the cadence thresholds -# on the client you build - the provider cannot reconfigure a client you pass in, so supplying -# a memory_client together with auto_extract=False or processor_config raises ValueError. +# on the client you build. When memory_client is supplied, the provider ignores its connection, +# credential, auto_extract, processor_config and prompts_dir options and preserves the client's +# configuration. Retrieval options (top_k, min_confidence, memory_types and context_prompt) still apply. +# The caller is responsible for closing the supplied client. memory_client = AsyncCosmosMemoryClient( cosmos_endpoint=cosmos_endpoint, cosmos_database="ai_memory", 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 264fcb86a13..02baed23b5f 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 @@ -119,7 +119,9 @@ def __init__( credential: Azure credential for authentication. When provided it is used for both Cosmos DB and AI Foundry; when ``None`` the toolkit builds (and owns) a ``DefaultAzureCredential``. - memory_client: Pre-created AsyncCosmosMemoryClient. + memory_client: Pre-created AsyncCosmosMemoryClient, used as configured by the caller. + Connection settings, credential, auto_extract, processor_config and prompts_dir + are ignored when supplied. The caller remains responsible for closing the client. top_k: Number of memories to retrieve in search. min_confidence: Minimum confidence score (0.0-1.0) for retrieved memories. memory_types: Types of memories to retrieve. Default: ["fact", "procedural"]. @@ -127,18 +129,16 @@ def __init__( auto_extract: Enable automatic background memory extraction/summarization after turn writes. When ``False`` the cadence thresholds are zeroed so nothing runs automatically and callers drive processing via ``memory_client.process_now()``. - Only applied when the provider builds the client; supplying ``memory_client`` - together with ``auto_extract=False`` raises ``ValueError``. + Only applied when the provider builds the client; ignored with ``memory_client``. processor_config: Optional processor cadence configuration, forwarded to the toolkit client via ``cadence_thresholds``. Only applied when the provider builds the - client; supplying ``memory_client`` together with ``processor_config`` raises - ``ValueError`` (configure cadence on your own client instead). + client; ignored with ``memory_client`` (configure cadence on your own client instead). prompts_dir: Optional directory of Prompty templates for the memory pipeline. When set, the extraction and summarization steps read their templates (including ``extract_memories.prompty``) from this directory instead of the toolkit's bundled defaults, letting you customize what the extraction LLM produces. The - directory must contain the full template set. Applies whether the client is built - by the provider or supplied via ``memory_client``. + directory must contain the full template set. Only applied when the provider builds + the client; ignored with ``memory_client``. Raises: SettingNotFoundError: If ``cosmos_endpoint``, ``foundry_endpoint``, ``embedding_model``, @@ -163,34 +163,24 @@ def __init__( self.memory_types: list[MemoryType] = list(memory_types) if memory_types else ["fact", "procedural"] self.context_prompt = context_prompt self.auto_extract = auto_extract - self._prompts_dir = prompts_dir - - # Build the per-instance cadence override for the toolkit client. The Agent Memory Toolkit - # accepts these thresholds directly via ``cadence_thresholds=`` (v0.2.0b3+), so the provider - # configures the processor without mutating global ``os.environ``. ``auto_extract=False`` - # zeroes the extraction/summary steps so the toolkit's background auto-trigger never runs on - # turn writes; callers then drive processing explicitly via ``memory_client.process_now(...)``. - # Keys not present fall back to the toolkit's environment/defaults. - cadence_thresholds: dict[str, int] = { - str(k): int(v) for k, v in cast("Mapping[str, int]", processor_config or {}).items() - } - if not auto_extract: - cadence_thresholds["FACT_EXTRACTION_EVERY_N"] = 0 - cadence_thresholds["THREAD_SUMMARY_EVERY_N"] = 0 - cadence_thresholds["USER_SUMMARY_EVERY_N"] = 0 - - # A caller-supplied client owns its own cadence configuration; the provider cannot apply - # ``cadence_thresholds`` to an already-constructed client. Reject the combination instead of - # silently ignoring the requested configuration. - if memory_client is not None and cadence_thresholds: - raise ValueError( - "processor_config and auto_extract=False only take effect when the provider builds " - "the memory client. When supplying your own memory_client, configure cadence via " - "AsyncCosmosMemoryClient(cadence_thresholds=...) directly." - ) + self._prompts_dir = prompts_dir if memory_client is None else None - # Initialize memory client if not provided + # Caller-supplied clients retain their own connection, cadence and prompt configuration. if memory_client is None: + # Build the per-instance cadence override for the toolkit client. The Agent Memory Toolkit + # accepts these thresholds directly via ``cadence_thresholds=`` (v0.2.0b3+), so the provider + # configures the processor without mutating global ``os.environ``. ``auto_extract=False`` + # zeroes the extraction/summary steps so the toolkit's background auto-trigger never runs on + # turn writes; callers then drive processing explicitly via ``memory_client.process_now(...)``. + # Keys not present fall back to the toolkit's environment/defaults. + cadence_thresholds: dict[str, int] = { + str(k): int(v) for k, v in cast("Mapping[str, int]", processor_config or {}).items() + } + if not auto_extract: + cadence_thresholds["FACT_EXTRACTION_EVERY_N"] = 0 + cadence_thresholds["THREAD_SUMMARY_EVERY_N"] = 0 + cadence_thresholds["USER_SUMMARY_EVERY_N"] = 0 + # Resolve connection settings from explicit args, then the environment. ``load_settings`` # validates that the required endpoints are present (raising if not), replacing manual # ``os.getenv`` + ``if not ...: raise`` blocks. @@ -512,13 +502,13 @@ async def after_run( content=msg.text.strip(), ) - # Auto-extraction and processing: + # Auto-extraction and processing for provider-created clients: # When auto_extract is True (default), the turn write schedules cadence-aware background # processing (fact extraction, summaries, reconciliation) based on the configured # thresholds (FACT_EXTRACTION_EVERY_N, DEDUP_EVERY_N, etc.), so no explicit # process_now() call is needed. When auto_extract is False, those thresholds were # zeroed in __init__ so nothing runs automatically; call memory_client.process_now() - # to drive extraction manually. + # to drive extraction manually. Injected clients retain their own cadence settings. except Exception as e: logger.warning("Failed to store conversation turns: %s", e, exc_info=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 aa48689ac48..4364465bc5a 100644 --- a/python/packages/azure-cosmos-memory/tests/test_context_provider.py +++ b/python/packages/azure-cosmos-memory/tests/test_context_provider.py @@ -349,17 +349,50 @@ def test_default_cadence_thresholds_is_none(self) -> None: _, kwargs = mock_client_class.call_args assert kwargs["cadence_thresholds"] is None - def test_processor_config_with_supplied_client_raises(self, mock_memory_client: AsyncMock) -> None: - """Cadence config cannot apply to a caller-supplied client, so combining them raises.""" - with pytest.raises(ValueError, match="processor_config"): - CosmosMemoryContextProvider( - memory_client=mock_memory_client, processor_config={"FACT_EXTRACTION_EVERY_N": 10} + @pytest.mark.parametrize("auto_extract", [True, False]) + def test_supplied_client_ignores_construction_options( + self, mock_memory_client: AsyncMock, auto_extract: bool, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Injected clients bypass connection resolution and cadence configuration.""" + for name in ( + "COSMOS_ENDPOINT", + "FOUNDRY_PROJECT_ENDPOINT", + "FOUNDRY_ENDPOINT", + "EMBEDDING_MODEL", + "CHAT_MODEL", + ): + monkeypatch.delenv(name, raising=False) + with ( + patch("agent_framework_azure_cosmos_memory._context_provider.load_settings") as mock_load_settings, + patch("agent_framework_azure_cosmos_memory._context_provider.AsyncCosmosMemoryClient") as mock_client_class, + ): + provider = CosmosMemoryContextProvider( + memory_client=mock_memory_client, + processor_config={"FACT_EXTRACTION_EVERY_N": 2}, + auto_extract=auto_extract, + top_k=7, + min_confidence=0.8, + memory_types=["episodic"], + context_prompt="Selected memories:", ) - def test_auto_extract_false_with_supplied_client_raises(self, mock_memory_client: AsyncMock) -> None: - """auto_extract=False cannot apply to a caller-supplied client, so combining them raises.""" - with pytest.raises(ValueError, match="processor_config"): - CosmosMemoryContextProvider(memory_client=mock_memory_client, auto_extract=False) + assert provider.memory_client is mock_memory_client + assert provider.top_k == 7 + assert provider.min_confidence == 0.8 + assert provider.memory_types == ["episodic"] + assert provider.context_prompt == "Selected memories:" + mock_load_settings.assert_not_called() + mock_client_class.assert_not_called() + + def test_supplied_client_does_not_parse_cadence(self, mock_memory_client: AsyncMock) -> None: + """Construction-only cadence values are not evaluated for injected clients.""" + cadence = MagicMock() + cadence.items.side_effect = AssertionError("cadence should not be inspected") + + provider = CosmosMemoryContextProvider(memory_client=mock_memory_client, processor_config=cadence) + + assert provider.memory_client is mock_memory_client + cadence.items.assert_not_called() # -- before_run tests ---------------------------------------------------------- @@ -955,9 +988,12 @@ async def test_provided_client_not_closed(self, mock_memory_client: AsyncMock) - """When client is provided externally, provider should not close it.""" provider = CosmosMemoryContextProvider(memory_client=mock_memory_client) - async with provider: - pass + with patch.object(provider, "flush", new_callable=AsyncMock) as mock_flush: + async with provider: + pass + mock_flush.assert_awaited_once() + mock_memory_client.close.assert_not_awaited() # Should still enter the client mock_memory_client.__aenter__.assert_awaited_once() # But should NOT exit it (caller owns it) @@ -1052,12 +1088,19 @@ async def test_prompts_dir_redirects_pipeline_loader(self, mock_memory_client: A # _get_pipeline is synchronous on the toolkit client; return our stand-in pipeline. mock_memory_client._get_pipeline = MagicMock(return_value=mock_pipeline) - provider = CosmosMemoryContextProvider( - memory_client=mock_memory_client, - prompts_dir="/custom/prompts", - ) - async with provider: - pass + with patch( + "agent_framework_azure_cosmos_memory._context_provider.AsyncCosmosMemoryClient", + return_value=mock_memory_client, + ): + provider = CosmosMemoryContextProvider( + cosmos_endpoint="https://test.documents.azure.com:443/", + foundry_endpoint="https://test.ai.azure.com", + embedding_model="text-embedding-3-large", + chat_model="gpt-4o-mini", + prompts_dir="/custom/prompts", + ) + async with provider: + pass from azure.cosmos.agent_memory.services._pipeline_helpers import PromptyLoader @@ -1074,3 +1117,15 @@ async def test_no_prompts_dir_leaves_pipeline_untouched(self, mock_memory_client pass mock_memory_client._get_pipeline.assert_not_called() + + async def test_supplied_client_ignores_prompts_dir(self, mock_memory_client: AsyncMock) -> None: + """Custom templates must not replace the caller's existing pipeline loader.""" + mock_memory_client._get_pipeline = MagicMock() + provider = CosmosMemoryContextProvider(memory_client=mock_memory_client, prompts_dir="/custom/prompts") + + async with provider: + pass + + mock_memory_client._get_pipeline.assert_not_called() + mock_memory_client.__aexit__.assert_not_awaited() + mock_memory_client.close.assert_not_awaited() From c94f6de12343cb72cd8cca06f14136b658328086 Mon Sep 17 00:00:00 2001 From: Shashwat Gupta Date: Tue, 3 Nov 2026 00:17:52 +0530 Subject: [PATCH 2/2] Test custom emulator prompts with a provider-owned client --- .../tests/test_emulator.py | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/python/packages/azure-cosmos-memory/tests/test_emulator.py b/python/packages/azure-cosmos-memory/tests/test_emulator.py index 16800265ac2..509568a0e90 100644 --- a/python/packages/azure-cosmos-memory/tests/test_emulator.py +++ b/python/packages/azure-cosmos-memory/tests/test_emulator.py @@ -335,9 +335,22 @@ async def test_prompts_dir_overrides_extraction_prompt( custom_dir = _make_custom_prompts_dir(tmp_path / "prompts", marker) chat = _FakeChat() - client = _build_emulator_client(monkeypatch, chat) + + def _create_emulator_client(**kwargs: Any) -> AsyncCosmosMemoryClient: + return _build_emulator_client(monkeypatch, chat) + + # Keep the real emulator-backed toolkit client, but let the provider create and own it + # so its prompts_dir option applies. Foundry calls still use the deterministic fakes. + monkeypatch.setattr( + "agent_framework_azure_cosmos_memory._context_provider.AsyncCosmosMemoryClient", + _create_emulator_client, + ) provider = CosmosMemoryContextProvider( - memory_client=client, + cosmos_endpoint=_EMULATOR_ENDPOINT, + cosmos_database="test_af_mem", + foundry_endpoint="https://test.ai.azure.com", + embedding_model="test-embedding", + chat_model="test-chat", top_k=5, min_confidence=0.0, memory_types=["fact"], @@ -346,7 +359,7 @@ async def test_prompts_dir_overrides_extraction_prompt( try: await provider.__aenter__() except Exception as exc: # noqa: BLE001 - clear skip on any connectivity/setup failure - await client.close() + await provider.memory_client.close() pytest.skip(f"Cosmos DB emulator not reachable or vector search unavailable at {_EMULATOR_ENDPOINT}: {exc}") try: @@ -372,4 +385,3 @@ async def test_prompts_dir_overrides_extraction_prompt( # in the overridden template, so its presence proves prompts_dir was honored end to end. sent = "\n".join(str(msg.get("content", "")) for call in chat.calls for msg in call) assert marker in sent, "custom extract_memories.prompty was not used by the extraction pipeline" - await client.close()