diff --git a/CHANGELOG.md b/CHANGELOG.md index 2f9f74c..7bc3f45 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,18 @@ ## Release History +## [Unreleased] + +#### Features Added +* 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. + ## [0.2.0b2] (2026-07-01) #### Features Added diff --git a/azure/cosmos/agent_memory/aio/cosmos_memory_client.py b/azure/cosmos/agent_memory/aio/cosmos_memory_client.py index 7caa832..75b8aff 100644 --- a/azure/cosmos/agent_memory/aio/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/aio/cosmos_memory_client.py @@ -4,7 +4,7 @@ import asyncio from datetime import datetime -from typing import TYPE_CHECKING, Any, Iterable, Optional +from typing import TYPE_CHECKING, Any, Iterable, Mapping, Optional from azure.cosmos.agent_memory._base import _BaseMemoryClient from azure.cosmos.agent_memory._base.base_client import is_transient_tail_step_error @@ -29,7 +29,10 @@ from azure.cosmos.agent_memory.aio.store import AsyncMemoryStore from azure.cosmos.agent_memory.exceptions import CosmosNotConnectedError, CosmosOperationError, ValidationError from azure.cosmos.agent_memory.logging import get_logger -from azure.cosmos.agent_memory.services._pipeline_helpers import _normalize_metadata_keys +from azure.cosmos.agent_memory.services._pipeline_helpers import ( + _normalize_cadence_thresholds, + _normalize_metadata_keys, +) from azure.cosmos.agent_memory.thresholds import DEFAULT_TTL_BY_TYPE if TYPE_CHECKING: # pragma: no cover - typing-only import @@ -90,6 +93,7 @@ def __init__( chat_client: Optional[Any] = None, processor: Optional[AsyncMemoryProcessor] = None, transcript_metadata_keys: Optional[Iterable[str]] = None, + cadence_thresholds: Optional[Mapping[str, int]] = None, ) -> None: self._init_base_config( cosmos_endpoint=cosmos_endpoint, @@ -146,6 +150,14 @@ def __init__( self._processor: Optional[AsyncMemoryProcessor] = processor self._processor_explicit = processor is not None self._transcript_metadata_keys: Optional[tuple[str, ...]] = _normalize_metadata_keys(transcript_metadata_keys) + # Optional per-turn cadence override, 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``). When provided, the auto-trigger uses these values instead of + # reading ``os.environ``; any key not present falls back to the environment/defaults. + # ``None`` preserves the env-only behavior. Normalized to a defensive ``dict[str, int]`` + # so later mutation of the caller's mapping cannot change client behavior, and invalid + # values fail fast at construction rather than deep inside the trigger path. + self._cadence_thresholds: Optional[dict[str, int]] = _normalize_cadence_thresholds(cadence_thresholds) logger.info("AsyncCosmosMemoryClient initialized") async def __aenter__(self) -> "AsyncCosmosMemoryClient": @@ -539,7 +551,12 @@ def _get_counter_container(self) -> Any: async def _maybe_auto_trigger(self, turn_counts: dict[tuple[str, str], int]) -> None: if not turn_counts: return - await maybe_trigger_steps(self._get_processor(), self._get_counter_container(), turn_counts) + await maybe_trigger_steps( + self._get_processor(), + self._get_counter_container(), + turn_counts, + thresholds=self._cadence_thresholds, + ) def _container_for_type(self, memory_type: str) -> Any: """Return the Cosmos container client that owns ``memory_type``.""" diff --git a/azure/cosmos/agent_memory/cosmos_memory_client.py b/azure/cosmos/agent_memory/cosmos_memory_client.py index 4bd03a6..71a5670 100644 --- a/azure/cosmos/agent_memory/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/cosmos_memory_client.py @@ -3,7 +3,7 @@ from __future__ import annotations from datetime import datetime -from typing import TYPE_CHECKING, Any, Iterable, Optional +from typing import TYPE_CHECKING, Any, Iterable, Mapping, Optional from azure.cosmos.agent_memory.logging import get_logger @@ -27,7 +27,7 @@ from .embeddings import EmbeddingsClient from .exceptions import CosmosOperationError, ValidationError from .processors import InProcessProcessor, MemoryProcessor -from .services._pipeline_helpers import _normalize_metadata_keys +from .services._pipeline_helpers import _normalize_cadence_thresholds, _normalize_metadata_keys from .services.pipeline import PipelineService from .store import MemoryStore from .thresholds import DEFAULT_TTL_BY_TYPE @@ -85,6 +85,7 @@ def __init__( chat_client: Optional[Any] = None, processor: Optional[MemoryProcessor] = None, transcript_metadata_keys: Optional[Iterable[str]] = None, + cadence_thresholds: Optional[Mapping[str, int]] = None, ) -> None: self._init_base_config( cosmos_endpoint=cosmos_endpoint, @@ -138,6 +139,14 @@ def __init__( self._processor: Optional[MemoryProcessor] = processor self._processor_explicit = processor is not None self._transcript_metadata_keys: Optional[tuple[str, ...]] = _normalize_metadata_keys(transcript_metadata_keys) + # Optional per-turn cadence override, 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``). When provided, the auto-trigger uses these values instead of + # reading ``os.environ``; any key not present falls back to the environment/defaults. + # ``None`` preserves the env-only behavior. Normalized to a defensive ``dict[str, int]`` + # so later mutation of the caller's mapping cannot change client behavior, and invalid + # values fail fast at construction rather than deep inside the trigger path. + self._cadence_thresholds: Optional[dict[str, int]] = _normalize_cadence_thresholds(cadence_thresholds) if self._cosmos_endpoint: self.create_memory_store() logger.info("CosmosMemoryClient initialized") @@ -494,7 +503,12 @@ def _get_counter_container(self) -> Any: def _maybe_auto_trigger(self, turn_counts: dict[tuple[str, str], int]) -> None: if not turn_counts: return - maybe_trigger_steps(self._get_processor(), self._get_counter_container(), turn_counts) + maybe_trigger_steps( + self._get_processor(), + self._get_counter_container(), + turn_counts, + thresholds=self._cadence_thresholds, + ) def _container_for_type(self, memory_type: str) -> Any: """Return the Cosmos container client that owns ``memory_type``.""" diff --git a/azure/cosmos/agent_memory/services/_pipeline_helpers.py b/azure/cosmos/agent_memory/services/_pipeline_helpers.py index 9748168..14ddf8f 100644 --- a/azure/cosmos/agent_memory/services/_pipeline_helpers.py +++ b/azure/cosmos/agent_memory/services/_pipeline_helpers.py @@ -14,7 +14,7 @@ import re from collections import defaultdict from pathlib import Path -from typing import Any, Iterable, Optional +from typing import Any, Iterable, Mapping, Optional from azure.cosmos.agent_memory.exceptions import LLMError @@ -202,6 +202,35 @@ def _normalize_metadata_keys( return keys or None +def _normalize_cadence_thresholds( + value: Optional[Mapping[str, int]], +) -> Optional[dict[str, int]]: + """Validate + defensively copy a ``cadence_thresholds`` argument. + + Coerces each value to ``int`` and rejects negatives (``0`` disables the + corresponding step). Returns a new ``dict`` so later mutation of the + caller's mapping cannot change client behavior. Returns ``None`` for + missing or empty input. + """ + if value is None: + return None + if not isinstance(value, Mapping): + raise TypeError( + "cadence_thresholds must be a mapping of env-var name to int " + f"(e.g. {{'FACT_EXTRACTION_EVERY_N': 5}}), not {type(value).__name__}." + ) + normalized: dict[str, int] = {} + for key, raw in value.items(): + try: + coerced = int(raw) + except (TypeError, ValueError) as exc: + raise ValueError(f"cadence_thresholds[{key!r}] must be an int, got {raw!r}.") from exc + if coerced < 0: + raise ValueError(f"cadence_thresholds[{key!r}] must be >= 0 (0 disables the step), got {coerced}.") + normalized[str(key)] = coerced + return normalized or None + + def _format_metadata_segment( metadata: Any, metadata_keys: Optional[tuple[str, ...]], diff --git a/tests/unit/aio/test_auto_trigger.py b/tests/unit/aio/test_auto_trigger.py index edcf382..f244563 100644 --- a/tests/unit/aio/test_auto_trigger.py +++ b/tests/unit/aio/test_auto_trigger.py @@ -8,7 +8,7 @@ from __future__ import annotations import asyncio -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -145,3 +145,67 @@ async def capture(turn_counts): await client.push_to_cosmos() await asyncio.gather(*list(client._background_tasks), return_exceptions=True) assert captured == [{("u1", "t1"): 1}] + + +class TestCadenceThresholdsForwarding: + """A ``cadence_thresholds`` mapping on the client is forwarded to the auto-trigger. + + This lets callers set per-turn cadence in-process instead of mutating ``os.environ``. + """ + + @pytest.mark.asyncio + async def test_cadence_thresholds_forwarded(self): + thresholds = {"FACT_EXTRACTION_EVERY_N": 3, "DEDUP_EVERY_N": 2} + client = AsyncCosmosMemoryClient(use_default_credential=False, cadence_thresholds=thresholds) + client._get_processor = MagicMock(return_value=MagicMock()) + client._get_counter_container = MagicMock(return_value=MagicMock()) + + with patch( + "azure.cosmos.agent_memory.aio.cosmos_memory_client.maybe_trigger_steps", + new=AsyncMock(), + ) as mock_trigger: + await client._maybe_auto_trigger({("u1", "t1"): 1}) + + mock_trigger.assert_awaited_once() + assert mock_trigger.await_args.kwargs["thresholds"] == thresholds + + @pytest.mark.asyncio + async def test_defaults_to_none_when_unset(self): + client = AsyncCosmosMemoryClient(use_default_credential=False) + client._get_processor = MagicMock(return_value=MagicMock()) + client._get_counter_container = MagicMock(return_value=MagicMock()) + + with patch( + "azure.cosmos.agent_memory.aio.cosmos_memory_client.maybe_trigger_steps", + new=AsyncMock(), + ) as mock_trigger: + await client._maybe_auto_trigger({("u1", "t1"): 1}) + + # None preserves the env-only behavior (the auto-trigger treats None as defaults). + assert mock_trigger.await_args.kwargs["thresholds"] is None + + +class TestCadenceThresholdsNormalization: + """The async client normalizes ``cadence_thresholds`` at construction time.""" + + def test_defensive_copy_isolates_later_mutation(self): + thresholds = {"FACT_EXTRACTION_EVERY_N": 3} + client = AsyncCosmosMemoryClient(use_default_credential=False, cadence_thresholds=thresholds) + thresholds["FACT_EXTRACTION_EVERY_N"] = 99 + assert client._cadence_thresholds == {"FACT_EXTRACTION_EVERY_N": 3} + + def test_string_values_are_coerced_to_int(self): + client = AsyncCosmosMemoryClient(use_default_credential=False, cadence_thresholds={"DEDUP_EVERY_N": "5"}) + assert client._cadence_thresholds == {"DEDUP_EVERY_N": 5} + + def test_negative_value_rejected(self): + with pytest.raises(ValueError): + AsyncCosmosMemoryClient(use_default_credential=False, cadence_thresholds={"DEDUP_EVERY_N": -1}) + + def test_non_int_value_rejected(self): + with pytest.raises(ValueError): + AsyncCosmosMemoryClient(use_default_credential=False, cadence_thresholds={"DEDUP_EVERY_N": "x"}) + + def test_non_mapping_rejected(self): + with pytest.raises(TypeError): + AsyncCosmosMemoryClient(use_default_credential=False, cadence_thresholds=[("DEDUP_EVERY_N", 5)]) diff --git a/tests/unit/test_cosmos_memory_client.py b/tests/unit/test_cosmos_memory_client.py index f84277d..c50e6af 100644 --- a/tests/unit/test_cosmos_memory_client.py +++ b/tests/unit/test_cosmos_memory_client.py @@ -997,3 +997,59 @@ def test_list_tags_delegates_to_store(): kwargs = container.query_items.call_args.kwargs assert "SELECT VALUE c.tags" in kwargs["query"] assert kwargs["parameters"] == [{"name": "@user_id", "value": "u1"}] + + +class TestSyncCadenceThresholdsForwarding: + """A ``cadence_thresholds`` mapping on the sync client is forwarded to the auto-trigger. + + This lets callers set per-turn cadence in-process instead of mutating ``os.environ``. + """ + + def test_cadence_thresholds_forwarded(self): + thresholds = {"FACT_EXTRACTION_EVERY_N": 3, "DEDUP_EVERY_N": 2} + mem = CosmosMemoryClient(use_default_credential=False, cadence_thresholds=thresholds) + mem._get_processor = MagicMock(return_value=MagicMock()) + mem._get_counter_container = MagicMock(return_value=MagicMock()) + + with patch("azure.cosmos.agent_memory.cosmos_memory_client.maybe_trigger_steps") as mock_trigger: + mem._maybe_auto_trigger({("u1", "t1"): 1}) + + mock_trigger.assert_called_once() + assert mock_trigger.call_args.kwargs["thresholds"] == thresholds + + def test_defaults_to_none_when_unset(self): + mem = CosmosMemoryClient(use_default_credential=False) + mem._get_processor = MagicMock(return_value=MagicMock()) + mem._get_counter_container = MagicMock(return_value=MagicMock()) + + with patch("azure.cosmos.agent_memory.cosmos_memory_client.maybe_trigger_steps") as mock_trigger: + mem._maybe_auto_trigger({("u1", "t1"): 1}) + + # None preserves the env-only behavior (the auto-trigger treats None as defaults). + assert mock_trigger.call_args.kwargs["thresholds"] is None + + +class TestSyncCadenceThresholdsNormalization: + """The sync client normalizes ``cadence_thresholds`` at construction time.""" + + def test_defensive_copy_isolates_later_mutation(self): + thresholds = {"FACT_EXTRACTION_EVERY_N": 3} + mem = CosmosMemoryClient(use_default_credential=False, cadence_thresholds=thresholds) + thresholds["FACT_EXTRACTION_EVERY_N"] = 99 + assert mem._cadence_thresholds == {"FACT_EXTRACTION_EVERY_N": 3} + + def test_string_values_are_coerced_to_int(self): + mem = CosmosMemoryClient(use_default_credential=False, cadence_thresholds={"DEDUP_EVERY_N": "5"}) + assert mem._cadence_thresholds == {"DEDUP_EVERY_N": 5} + + def test_negative_value_rejected(self): + with pytest.raises(ValueError): + CosmosMemoryClient(use_default_credential=False, cadence_thresholds={"DEDUP_EVERY_N": -1}) + + def test_non_int_value_rejected(self): + with pytest.raises(ValueError): + CosmosMemoryClient(use_default_credential=False, cadence_thresholds={"DEDUP_EVERY_N": "x"}) + + def test_non_mapping_rejected(self): + with pytest.raises(TypeError): + CosmosMemoryClient(use_default_credential=False, cadence_thresholds=[("DEDUP_EVERY_N", 5)])