From fc24f195bc69a8c8b1bc07bdee1ca19e87128cc1 Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Mon, 17 Aug 2026 13:58:56 -0700 Subject: [PATCH 1/3] Removing some unused code --- azure/cosmos/agent_memory/_utils.py | 32 ----- .../agent_memory/aio/services/pipeline.py | 114 ----------------- .../cosmos/agent_memory/services/pipeline.py | 121 ------------------ tests/integration/test_async_full_pipeline.py | 4 +- tests/integration/test_full_pipeline.py | 3 +- tests/unit/test_utils.py | 50 -------- 6 files changed, 3 insertions(+), 321 deletions(-) diff --git a/azure/cosmos/agent_memory/_utils.py b/azure/cosmos/agent_memory/_utils.py index 5e89729..50a2d82 100644 --- a/azure/cosmos/agent_memory/_utils.py +++ b/azure/cosmos/agent_memory/_utils.py @@ -341,9 +341,6 @@ def _resolve_vector_index_type(val: Optional[str]) -> str: return raw -_SIMILARITY_DESCENDING_FUNCTIONS = frozenset({"cosine", "dotproduct"}) - - def cosine_similarity(a: list[float], b: list[float]) -> float: """Cosine similarity of two equal-length vectors. @@ -386,35 +383,6 @@ def vector_centroid(vectors: list[list[float]]) -> list[float]: return [value / count for value in acc] -def vector_order_direction(distance_function: str) -> str: - """Return the ``ORDER BY VectorDistance(...)`` direction for most-similar-first. - - ``DESC`` for cosine/dotproduct (higher score = more similar), ``ASC`` for - euclidean (lower distance = more similar). - """ - return "DESC" if distance_function in _SIMILARITY_DESCENDING_FUNCTIONS else "ASC" - - -def distance_function_from_container_properties(props: Any, *, default: str = "cosine") -> str: - """Read the vector embedding's ``distanceFunction`` from container properties. - - The distance function (cosine/dotproduct/euclidean) is chosen at - ``create_memory_store`` time, written immutably into the container's vector - embedding policy, and read back here from the authoritative source - (``container.read()``) so the dedup vector-floor logic matches how the - container actually ranks. This SDK provisions exactly one vector embedding; - falls back to ``default`` (cosine) when the policy is absent or malformed - (e.g. ``__new__``-built test instances with mocked containers). - """ - policy = props.get("vectorEmbeddingPolicy") if isinstance(props, dict) else None - embeddings = policy.get("vectorEmbeddings") if isinstance(policy, dict) else None - entry = embeddings[0] if isinstance(embeddings, list) and embeddings else None - fn = entry.get("distanceFunction") if isinstance(entry, dict) else None - if isinstance(fn, str) and fn in _ALLOWED_DISTANCE_FUNCTIONS: - return fn - return default - - def _resolve_full_text_language(val: Optional[str]) -> str: """Resolve full-text language from the explicit value, defaulting to ``en-US``. diff --git a/azure/cosmos/agent_memory/aio/services/pipeline.py b/azure/cosmos/agent_memory/aio/services/pipeline.py index 6a3cfed..dbcad5a 100644 --- a/azure/cosmos/agent_memory/aio/services/pipeline.py +++ b/azure/cosmos/agent_memory/aio/services/pipeline.py @@ -29,8 +29,6 @@ from azure.cosmos.agent_memory._utils import ( DEFAULT_TTL_BY_TYPE, compute_content_hash, - distance_function_from_container_properties, - vector_order_direction, ) from azure.cosmos.agent_memory.aio.store import AsyncMemoryStore from azure.cosmos.agent_memory.exceptions import ( @@ -301,118 +299,6 @@ async def _embed_one(self, text: str) -> list[float]: async def _embed_batch(self, texts: list[str]) -> list[list[float]]: return await self._embeddings.generate_batch(texts) - async def _vector_distance_function(self) -> str: - """Return the container's configured Cosmos ``distanceFunction`` (cached). - - Read from the container's vector embedding policy (``await container.read()``) - - the authoritative, immutable source set when the container was created. - Drives the ORDER BY direction and similarity-threshold comparisons so dedup - never silently assumes cosine. Falls back to cosine when the policy can't be - read (e.g. ``__new__``-built test instances with mocked containers). - """ - fn = getattr(self, "_distance_function_cache", None) - if fn is not None: - return fn - try: - props = await self._memories_container.read() - except Exception: - # See sync pipeline: don't cache a defaulted cosine, and flag the - # failure so the destructive in-place fold path skips this run rather - # than mis-applying cosine bands to (possibly euclidean) distances. - self._distance_function_read_failed = True - logger.debug( - "vector dedup: could not read container vector policy; defaulting to cosine (not cached)", - exc_info=True, - ) - return "cosine" - fn = distance_function_from_container_properties(props) - self._distance_function_cache = fn - self._distance_function_read_failed = False - return fn - - def _warn_distance_policy_unavailable_once(self) -> None: - """One-shot WARN that in-place folding was skipped (policy unreadable).""" - if getattr(self, "_warned_distance_policy_unavailable", False): - return - self._warned_distance_policy_unavailable = True - logger.warning( - "vector dedup: container vector policy could not be read; skipping " - "near-exact auto-drop this run to avoid mis-calibrated drops. Memories are " - "written as-is and reconciled on a later run once the policy is readable." - ) - - def _warn_euclidean_autodrop_once(self, distance_function: str) -> None: - """One-shot WARN that the near-exact vector auto-drop is disabled. - - The near-exact threshold is cosine-calibrated; on euclidean - the destructive auto-drop is skipped and LLM reconcile still runs. - Logged once per pipeline instance to avoid hot-path spam. - """ - if getattr(self, "_warned_euclidean_autodrop", False): - return - self._warned_euclidean_autodrop = True - logger.warning( - "Container distanceFunction=%r: near-exact vector auto-drop is " - "cosine-calibrated and has been DISABLED for this distance function. " - "Duplicate detection falls back to borderline tagging + LLM reconcile. " - "Use cosine/dotproduct embeddings for vector-floor auto-dedup.", - distance_function, - ) - - async def _vector_candidates( - self, - *, - user_id: str, - embedding, - memory_type, - top_k, - exclude_ids, - ) -> list[dict]: - """Return active same-user vector candidates from Cosmos.""" - if not user_id or not embedding or not top_k or int(top_k) < 1: - return [] - excluded = set(exclude_ids or []) - capped_top = top_literal(int(top_k), name="_vector_candidates.top_k") - distance_function = await self._vector_distance_function() - order_direction = vector_order_direction(distance_function) - field = "embedding" - query = ( - f"SELECT TOP {capped_top} c.id, c.content, c.type, " - f"VectorDistance(c.{field}, @vec) AS score " - "FROM c WHERE c.user_id = @user_id " - "AND c.type = @memory_type " - f"AND {_ACTIVE_DOC_FILTER} " - f"AND IS_DEFINED(c.{field}) " - # Cosmos orders ORDER BY VectorDistance() most-similar-first per the - # container's distanceFunction; an explicit ASC/DESC is rejected (BadRequest). - f"ORDER BY VectorDistance(c.{field}, @vec)" - ) - rows = await self._query_items( - self._memories_container, - query=query, - parameters=[ - {"name": "@user_id", "value": user_id}, - {"name": "@memory_type", "value": memory_type}, - {"name": "@vec", "value": embedding}, - ], - ) - candidates = [ - { - "id": row.get("id"), - "content": row.get("content"), - "type": row.get("type"), - "score": float(row.get("score") or 0.0), - } - for row in rows - if row.get("id") and row.get("id") not in excluded - ] - # Most-similar-first: descending score for cosine/dotproduct, ascending for euclidean. - candidates.sort( - key=lambda item: item.get("score", 0.0), - reverse=order_direction == "DESC", - ) - return candidates - def _prompt_lineage(self, filename: str) -> dict[str, str]: """Return ``{prompt_id, prompt_version}`` for stamping a doc. diff --git a/azure/cosmos/agent_memory/services/pipeline.py b/azure/cosmos/agent_memory/services/pipeline.py index f308d4e..5a0370b 100644 --- a/azure/cosmos/agent_memory/services/pipeline.py +++ b/azure/cosmos/agent_memory/services/pipeline.py @@ -28,8 +28,6 @@ from azure.cosmos.agent_memory._utils import ( DEFAULT_TTL_BY_TYPE, compute_content_hash, - distance_function_from_container_properties, - vector_order_direction, ) from azure.cosmos.agent_memory.exceptions import ( ValidationError, @@ -283,125 +281,6 @@ def _build_transcript( include_timestamp=include_timestamp, ) - def _vector_distance_function(self) -> str: - """Return the container's configured Cosmos ``distanceFunction`` (cached). - - Read from the container's vector embedding policy (``container.read()``) - - the authoritative, immutable source set when the container was created. - Drives the ORDER BY direction and similarity-threshold comparisons so dedup - never silently assumes cosine. Falls back to cosine when the policy can't be - read (e.g. ``__new__``-built test instances with mocked containers). - """ - fn = getattr(self, "_distance_function_cache", None) - if fn is not None: - return fn - try: - props = self._memories_container.read() - except Exception: - # Transient read failure (429/503/connection) is indistinguishable from - # "no policy" once we drop to None - so DON'T cache here. Returning an - # uncached cosine default lets the next call self-heal; caching it would - # pin cosine for the instance's life and silently mis-handle a euclidean - # container (cosine bands applied to euclidean distances -> data loss). - # Flag the failure so the *destructive* in-place fold path can skip - # entirely (a defaulted cosine on a euclidean container would fold and - # overwrite unrelated memories). - self._distance_function_read_failed = True - logger.debug( - "vector dedup: could not read container vector policy; defaulting to cosine (not cached)", - exc_info=True, - ) - return "cosine" - fn = distance_function_from_container_properties(props) - self._distance_function_cache = fn - self._distance_function_read_failed = False - return fn - - def _warn_euclidean_autodrop_once(self, distance_function: str) -> None: - """One-shot WARN that the near-exact vector auto-drop is disabled. - - The near-exact threshold is cosine-calibrated; on euclidean - the destructive auto-drop is skipped and LLM reconcile still runs. - Logged once per pipeline instance to avoid hot-path spam. - """ - if getattr(self, "_warned_euclidean_autodrop", False): - return - self._warned_euclidean_autodrop = True - logger.warning( - "Container distanceFunction=%r: near-exact vector auto-drop is " - "cosine-calibrated and has been DISABLED for this distance function. " - "Duplicate detection falls back to borderline tagging + LLM reconcile. " - "Use cosine/dotproduct embeddings for vector-floor auto-dedup.", - distance_function, - ) - - def _warn_distance_policy_unavailable_once(self) -> None: - """One-shot WARN that in-place folding was skipped (policy unreadable).""" - if getattr(self, "_warned_distance_policy_unavailable", False): - return - self._warned_distance_policy_unavailable = True - logger.warning( - "vector dedup: container vector policy could not be read; skipping " - "near-exact auto-drop this run to avoid mis-calibrated drops. Memories are " - "written as-is and reconciled on a later run once the policy is readable." - ) - - def _vector_candidates( - self, - *, - user_id: str, - embedding: list[float], - memory_type: str, - top_k: int, - exclude_ids: set[str], - ) -> list[dict[str, Any]]: - """Return nearest active same-type memories using Cosmos VectorDistance.""" - if not user_id or not embedding or top_k < 1: - return [] - capped_top_k = top_literal(top_k, name="_vector_candidates.top_k") - distance_function = self._vector_distance_function() - order_direction = vector_order_direction(distance_function) - field = "embedding" - query = ( - f"SELECT TOP {capped_top_k} c.id, c.content, c.type, " - f"VectorDistance(c.{field}, @vec) AS score " - "FROM c WHERE c.user_id = @user_id " - "AND c.type = @memory_type " - f"AND {_ACTIVE_DOC_FILTER} " - f"AND IS_DEFINED(c.{field}) " - # Cosmos orders ORDER BY VectorDistance() most-similar-first per the - # container's distanceFunction; an explicit ASC/DESC is rejected (BadRequest). - f"ORDER BY VectorDistance(c.{field}, @vec)" - ) - rows = list( - self._memories_container.query_items( - query=query, - parameters=[ - {"name": "@user_id", "value": user_id}, - {"name": "@memory_type", "value": memory_type}, - {"name": "@vec", "value": embedding}, - ], - enable_cross_partition_query=True, - ) - ) - excluded = set(exclude_ids or set()) - candidates = [ - { - "id": row.get("id"), - "content": row.get("content"), - "type": row.get("type"), - "score": float(row.get("score") or 0.0), - } - for row in rows - if row.get("id") and row.get("id") not in excluded - ] - # Most-similar-first: descending score for cosine/dotproduct, ascending for euclidean. - candidates.sort( - key=lambda row: row.get("score", 0.0), - reverse=order_direction == "DESC", - ) - return candidates - def _query_active_memories( self, user_id: str, diff --git a/tests/integration/test_async_full_pipeline.py b/tests/integration/test_async_full_pipeline.py index a1a3af4..306d10b 100644 --- a/tests/integration/test_async_full_pipeline.py +++ b/tests/integration/test_async_full_pipeline.py @@ -156,8 +156,8 @@ async def _async_wait_vector_searchable( timeout: float = 20.0, ) -> None: """Poll vector search until the user's seeded fact is retrievable (DiskANN - index caught up), so the subsequent ``_vector_candidates`` lookup is - deterministic rather than racing the async index.""" + index caught up), so the subsequent retrieval is deterministic rather than + racing the async index.""" deadline = time.time() + timeout while time.time() < deadline: try: diff --git a/tests/integration/test_full_pipeline.py b/tests/integration/test_full_pipeline.py index 7091c4f..19ff1d4 100644 --- a/tests/integration/test_full_pipeline.py +++ b/tests/integration/test_full_pipeline.py @@ -171,8 +171,7 @@ def _wait_vector_searchable( ``upsert_memory`` stores the embedding synchronously, but Cosmos's DiskANN vector index catches up asynchronously (~1-2s). Gating on a real vector search makes - the subsequent ``_vector_candidates`` lookup deterministic instead of racing - the index.""" + the subsequent retrieval deterministic instead of racing the index.""" deadline = time.time() + timeout while time.time() < deadline: try: diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 7d3af33..6c1c27b 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -15,10 +15,8 @@ _resolve_vector_index_type, build_cosmos_user_agent, compute_content_hash, - distance_function_from_container_properties, extract_keywords, normalize_ai_foundry_endpoint, - vector_order_direction, ) from azure.cosmos.agent_memory.exceptions import ConfigurationError, ValidationError @@ -226,54 +224,6 @@ def test_resolve_distance_function_invalid_raises(): _resolve_distance_function("manhattan") -def test_vector_order_direction_per_function(): - # cosine/dotproduct: higher VectorDistance == more similar -> DESC for nearest-first. - assert vector_order_direction("cosine") == "DESC" - assert vector_order_direction("dotproduct") == "DESC" - # euclidean: lower distance == more similar -> ASC for nearest-first. - assert vector_order_direction("euclidean") == "ASC" - - -def test_distance_function_from_container_properties_reads_policy(): - props = { - "id": "memories", - "vectorEmbeddingPolicy": { - "vectorEmbeddings": [ - {"path": "/embedding", "dataType": "float32", "distanceFunction": "euclidean", "dimensions": 1536} - ] - }, - } - assert distance_function_from_container_properties(props) == "euclidean" - - -def test_distance_function_from_container_properties_reads_single_embedding(): - # This SDK provisions a single vector embedding; the resolver reads its - # distanceFunction directly (the path value is irrelevant here). - props = { - "vectorEmbeddingPolicy": { - "vectorEmbeddings": [ - {"path": "/embedding", "distanceFunction": "dotproduct"}, - ] - } - } - assert distance_function_from_container_properties(props) == "dotproduct" - - -@pytest.mark.parametrize( - "props", - [ - None, - {}, - {"vectorEmbeddingPolicy": {}}, - {"vectorEmbeddingPolicy": {"vectorEmbeddings": []}}, - {"vectorEmbeddingPolicy": {"vectorEmbeddings": [{"path": "/embedding", "distanceFunction": "manhattan"}]}}, - "not-a-dict", - ], -) -def test_distance_function_from_container_properties_falls_back_to_cosine(props): - assert distance_function_from_container_properties(props) == "cosine" - - def test_extract_keywords_basic_and_stopwords(): # Stopwords removed, lowercased, de-duplicated, first-seen order preserved. assert extract_keywords("The user LOVES hiking and hiking trails") == [ From dbf620265c5d30b1da74659c2ddede8f75b367be Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Mon, 17 Aug 2026 15:06:26 -0700 Subject: [PATCH 2/3] Fixing issues 39 and 40 --- .github/workflows/_smoke.yml | 42 ++++++ .github/workflows/ci.yml | 7 +- .../agent_memory/aio/services/pipeline.py | 3 +- azure/cosmos/agent_memory/exceptions.py | 15 +++ .../cosmos/agent_memory/services/pipeline.py | 3 +- function_app/local.settings.json.template | 4 + function_app/orchestrators/user_summary.py | 59 +++++++-- function_app/shared/config.py | 24 ++++ tests/unit/function_app/test_orchestrators.py | 121 ++++++++++++++++++ 9 files changed, 264 insertions(+), 14 deletions(-) create mode 100644 .github/workflows/_smoke.yml diff --git a/.github/workflows/_smoke.yml b/.github/workflows/_smoke.yml new file mode 100644 index 0000000..3d19ef4 --- /dev/null +++ b/.github/workflows/_smoke.yml @@ -0,0 +1,42 @@ +name: smoke + +on: + workflow_call: + +permissions: + contents: read + +jobs: + function-app-import: + name: "function app import #${{ matrix.python-version }}" + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + # The Function app deploys on the Flex Consumption Python 3.11 runtime. + python-version: ["3.11"] + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: pip + cache-dependency-path: pyproject.toml + + # Install the toolkit from this same commit so the Function app is + # validated against the source it ships with, rather than a possibly-older + # published package. The toolkit pinned in requirements.txt is then already + # satisfied and is not re-fetched. + - name: Install toolkit from source + run: pip install . + + - name: Install Function app requirements + working-directory: function_app + run: pip install -r requirements.txt + + # Fails if the Function source uses a toolkit API the in-repo toolkit does + # not provide - the drift that left the host unable to load the app. + - name: Import function_app + working-directory: function_app + run: PYTHONPATH=. python -c "import function_app; print('import function_app OK')" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c826cbd..38e9756 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -25,12 +25,17 @@ jobs: permissions: contents: read + smoke: + uses: ./.github/workflows/_smoke.yml + permissions: + contents: read + ci-success: # Single required-status-check target. Branch protection should require # only this job; it summarizes lint + test matrix results so a failure # in any matrix cell blocks merge. name: "CI Success" - needs: [lint, test] + needs: [lint, test, smoke] if: always() runs-on: ubuntu-latest env: diff --git a/azure/cosmos/agent_memory/aio/services/pipeline.py b/azure/cosmos/agent_memory/aio/services/pipeline.py index dbcad5a..ead339a 100644 --- a/azure/cosmos/agent_memory/aio/services/pipeline.py +++ b/azure/cosmos/agent_memory/aio/services/pipeline.py @@ -32,6 +32,7 @@ ) from azure.cosmos.agent_memory.aio.store import AsyncMemoryStore from azure.cosmos.agent_memory.exceptions import ( + NoSourceMemoriesError, ValidationError, ) from azure.cosmos.agent_memory.logging import get_logger @@ -1520,7 +1521,7 @@ async def generate_user_summary_durable( user_doc.pop("embedding", None) return user_doc if not existing_summary and not items: - raise ValidationError(f"No memories found for user_id={user_id!r}") + raise NoSourceMemoriesError(f"No memories found for user_id={user_id!r}") items.sort(key=lambda m: m.get("created_at", ""), reverse=True) if recent_k is not None: diff --git a/azure/cosmos/agent_memory/exceptions.py b/azure/cosmos/agent_memory/exceptions.py index 36d2259..ac5439c 100644 --- a/azure/cosmos/agent_memory/exceptions.py +++ b/azure/cosmos/agent_memory/exceptions.py @@ -39,6 +39,21 @@ class ValidationError(AgentMemoryError): error_code = "validation" +class NoSourceMemoriesError(ValidationError): + """Raised by the durable summary path when a user has no source memories yet. + + A brand-new user can cross the user-summary threshold before fact extraction + has persisted anything, because the change feed starts the extract and + user-summary orchestrations independently. Subclasses :class:`ValidationError` + so existing ``except ValidationError`` handlers keep catching it, while + letting the Durable user-summary orchestrator tell this expected "extraction + has not landed yet" race apart from a real validation failure and wait/retry + instead of failing after its short activity-retry window. + """ + + error_code = "no_source_memories" + + class CosmosNotConnectedError(AgentMemoryError): """Raised when a Cosmos DB operation is attempted without an active connection.""" diff --git a/azure/cosmos/agent_memory/services/pipeline.py b/azure/cosmos/agent_memory/services/pipeline.py index 5a0370b..c5908e5 100644 --- a/azure/cosmos/agent_memory/services/pipeline.py +++ b/azure/cosmos/agent_memory/services/pipeline.py @@ -30,6 +30,7 @@ compute_content_hash, ) from azure.cosmos.agent_memory.exceptions import ( + NoSourceMemoriesError, ValidationError, ) from azure.cosmos.agent_memory.logging import get_logger @@ -1586,7 +1587,7 @@ def generate_user_summary_durable( user_doc.pop("embedding", None) return user_doc if not existing_summary and not items: - raise ValidationError(f"No memories found for user_id={user_id!r}") + raise NoSourceMemoriesError(f"No memories found for user_id={user_id!r}") items.sort(key=lambda m: m.get("created_at", ""), reverse=True) if recent_k is not None: diff --git a/function_app/local.settings.json.template b/function_app/local.settings.json.template index 2d65bde..fa55e9f 100644 --- a/function_app/local.settings.json.template +++ b/function_app/local.settings.json.template @@ -31,6 +31,10 @@ "// --- Batch knob ---": "", "MAX_BATCH_SIZE": "20", + "// --- UserSummaryOrchestrator waits this long (seconds) for extraction to persist the first memories of a brand-new user before skipping the cadence; polls every interval. ---": "", + "USER_SUMMARY_WAIT_SECONDS": "120", + "USER_SUMMARY_WAIT_INTERVAL_SECONDS": "10", + "// --- Turn vector search: embed raw turns on write so search_turns() works. Turns container is always vector-indexed; default false. ---": "", "ENABLE_TURN_EMBEDDINGS": "false" } diff --git a/function_app/orchestrators/user_summary.py b/function_app/orchestrators/user_summary.py index e4024b5..cd6d335 100644 --- a/function_app/orchestrators/user_summary.py +++ b/function_app/orchestrators/user_summary.py @@ -7,17 +7,24 @@ from __future__ import annotations import logging +from datetime import timedelta import azure.durable_functions as df from shared import config from shared.pipeline_factory import get_pipeline +from azure.cosmos.agent_memory.exceptions import NoSourceMemoriesError + from ._retry import default_retry_options logger = logging.getLogger(__name__) bp = df.Blueprint() +# Sentinel returned by ``us_Extract`` when the user has no persisted memories +# yet, so the orchestrator waits for extraction to land instead of failing. +_NO_MEMORIES_YET_STATUS = "no_memories_yet" + @bp.orchestration_trigger(context_name="context") def UserSummaryOrchestrator(context: df.DurableOrchestrationContext): @@ -25,12 +32,33 @@ def UserSummaryOrchestrator(context: df.DurableOrchestrationContext): user_id = payload["user_id"] thread_ids = payload.get("thread_ids") or None retry = default_retry_options() + extract_payload = { + "user_id": user_id, + "limit": config.get_max_batch_size(), + "thread_ids": thread_ids, + } - user_summary = yield context.call_activity_with_retry( - "us_Extract", - retry, - {"user_id": user_id, "limit": config.get_max_batch_size(), "thread_ids": thread_ids}, - ) + # A brand-new user can cross the user-summary threshold before fact + # extraction (started independently by the change feed) has persisted + # anything. Rather than fail after the short activity-retry window, poll on a + # replay-safe Durable timer until memories land or a bounded budget is + # exhausted, then skip this cadence (a later threshold retries). + wait_budget = config.get_user_summary_wait_seconds() + poll_interval = config.get_user_summary_wait_interval_seconds() + deadline = context.current_utc_datetime + timedelta(seconds=wait_budget) + + user_summary = yield context.call_activity_with_retry("us_Extract", retry, extract_payload) + while isinstance(user_summary, dict) and user_summary.get("status") == _NO_MEMORIES_YET_STATUS: + if context.current_utc_datetime >= deadline: + logger.warning( + "UserSummary no memories persisted within %ss for user=%s; skipping this " + "cadence (a later user-summary threshold will retry)", + wait_budget, + user_id, + ) + return {"persisted": False, "user_summary_id": None, "skipped": _NO_MEMORIES_YET_STATUS} + yield context.create_timer(context.current_utc_datetime + timedelta(seconds=poll_interval)) + user_summary = yield context.call_activity_with_retry("us_Extract", retry, extract_payload) persisted = yield context.call_activity_with_retry( "us_PersistUserSummary", @@ -46,13 +74,22 @@ def UserSummaryOrchestrator(context: df.DurableOrchestrationContext): @bp.activity_trigger(input_name="payload") def us_Extract(payload: dict) -> dict: - """Generate a cross-thread user summary body only.""" + """Generate a cross-thread user summary body only. + + Returns a ``{"status": "no_memories_yet"}`` sentinel instead of raising when + the user has no persisted memories yet, so the orchestrator can wait for + extraction to land rather than exhausting its activity retries. + """ user_id = payload["user_id"] - summary = get_pipeline().generate_user_summary_durable( - user_id=user_id, - recent_k=payload.get("limit"), - thread_ids=payload.get("thread_ids") or None, - ) + try: + summary = get_pipeline().generate_user_summary_durable( + user_id=user_id, + recent_k=payload.get("limit"), + thread_ids=payload.get("thread_ids") or None, + ) + except NoSourceMemoriesError: + logger.info("UserSummary no source memories yet user=%s; will wait and retry", user_id) + return {"status": _NO_MEMORIES_YET_STATUS} logger.info("UserSummary extracted user=%s", user_id) return summary diff --git a/function_app/shared/config.py b/function_app/shared/config.py index 2ace7a6..2733d59 100644 --- a/function_app/shared/config.py +++ b/function_app/shared/config.py @@ -81,6 +81,14 @@ DEFAULT_MAX_BATCH_SIZE = 20 +# UserSummaryOrchestrator can cross the user-summary threshold before fact +# extraction has persisted anything for a brand-new user, because the change +# feed starts the extract and user-summary orchestrations independently. These +# bound a replay-safe Durable wait so the summary is still produced once +# extraction lands, instead of failing after the short activity-retry window. +DEFAULT_USER_SUMMARY_WAIT_SECONDS = 120 +DEFAULT_USER_SUMMARY_WAIT_INTERVAL_SECONDS = 10 + def _parse_threshold(name: str, default: int) -> int: """Parse an integer threshold env var. @@ -172,6 +180,22 @@ def get_max_batch_size() -> int: return _parse_int("MAX_BATCH_SIZE", DEFAULT_MAX_BATCH_SIZE) +def get_user_summary_wait_seconds() -> int: + """Total seconds ``UserSummaryOrchestrator`` waits for extraction to persist + the first memories before giving up for this cadence. ``0`` disables waiting + (a single attempt).""" + return _parse_threshold("USER_SUMMARY_WAIT_SECONDS", DEFAULT_USER_SUMMARY_WAIT_SECONDS) + + +def get_user_summary_wait_interval_seconds() -> int: + """Seconds between ``UserSummaryOrchestrator`` readiness polls, floored at 1 + so a misconfigured ``0`` cannot busy-loop.""" + return max( + 1, + _parse_threshold("USER_SUMMARY_WAIT_INTERVAL_SECONDS", DEFAULT_USER_SUMMARY_WAIT_INTERVAL_SECONDS), + ) + + def get_thread_summary_every_n() -> int: """Threshold for triggering ``ThreadSummaryOrchestrator``. ``0`` disables.""" return _parse_threshold( diff --git a/tests/unit/function_app/test_orchestrators.py b/tests/unit/function_app/test_orchestrators.py index 582ba18..61b915a 100644 --- a/tests/unit/function_app/test_orchestrators.py +++ b/tests/unit/function_app/test_orchestrators.py @@ -9,6 +9,7 @@ from __future__ import annotations +from datetime import datetime, timedelta, timezone from unittest.mock import AsyncMock, MagicMock, call, patch import pytest @@ -16,6 +17,8 @@ from orchestrators import thread_summary as ts_mod from orchestrators import user_summary as us_mod +from azure.cosmos.agent_memory.exceptions import NoSourceMemoriesError + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -77,6 +80,41 @@ def _drive(gen, activity_results): return stop.value, yields +class _TimerContext: + """DurableOrchestrationContext double for the user-summary wait/poll loop. + + Unlike the ``MagicMock``-based :func:`_make_context`, this exposes a real + ``current_utc_datetime`` that advances by ``step_seconds`` on every read + (so the wait budget can elapse deterministically) and records + ``create_timer`` calls. Activity calls yield the same ``("__call__", ...)`` + sentinel that :func:`_drive` feeds pre-canned results into. + """ + + def __init__(self, payload, *, start=None, step_seconds=0): + self._payload = payload + self._now = start or datetime(2026, 1, 1, tzinfo=timezone.utc) + self._step = timedelta(seconds=step_seconds) + self.activity_calls: list[tuple] = [] + self.timer_fire_times: list = [] + + def get_input(self): + return self._payload + + @property + def current_utc_datetime(self): + now = self._now + self._now = now + self._step + return now + + def call_activity_with_retry(self, name, retry, activity_payload): + self.activity_calls.append((name, activity_payload)) + return ("__call__", name, activity_payload) + + def create_timer(self, fire_at): + self.timer_fire_times.append(fire_at) + return ("__timer__", fire_at) + + # --------------------------------------------------------------------------- # Shared env fixture: ensure MAX_BATCH_SIZE has a deterministic value across # tests (other tests may set it via @patch.dict). @@ -500,3 +538,86 @@ def test_missing_user_id_raises(self): gen = self._orchestrator()(ctx) with pytest.raises(KeyError): next(gen) + + +class TestUserSummaryWaitForExtraction: + """A brand-new user can cross the user-summary threshold before extraction + has persisted anything (the change feed starts the two orchestrations + independently). The orchestrator must wait on a replay-safe Durable timer + and still produce the summary, instead of failing after the short + activity-retry window.""" + + def _orchestrator(self): + return _user_function(us_mod.UserSummaryOrchestrator) + + @patch.object(us_mod, "default_retry_options", return_value=MagicMock()) + def test_waits_then_persists_when_memories_arrive_late(self, _retry, monkeypatch): + monkeypatch.setenv("USER_SUMMARY_WAIT_SECONDS", "120") + monkeypatch.setenv("USER_SUMMARY_WAIT_INTERVAL_SECONDS", "10") + # Static clock: the budget never elapses, so the loop ends only when + # memories finally land. + ctx = _TimerContext({"user_id": "u1"}, step_seconds=0) + gen = self._orchestrator()(ctx) + result, _ = _drive( + gen, + [ + {"status": "no_memories_yet"}, # us_Extract #1 - extraction not done + None, # create_timer #1 + {"status": "no_memories_yet"}, # us_Extract #2 - still not done + None, # create_timer #2 + {"id": "user-sum-1"}, # us_Extract #3 - memories landed + {"id": "user-sum-1"}, # us_PersistUserSummary + ], + ) + assert result == {"persisted": True, "user_summary_id": "user-sum-1"} + assert len(ctx.timer_fire_times) == 2 # waited twice before succeeding + assert [name for name, _ in ctx.activity_calls] == [ + "us_Extract", + "us_Extract", + "us_Extract", + "us_PersistUserSummary", + ] + + @patch.object(us_mod, "default_retry_options", return_value=MagicMock()) + def test_gives_up_after_budget_without_persisting(self, _retry, monkeypatch): + monkeypatch.setenv("USER_SUMMARY_WAIT_SECONDS", "15") + monkeypatch.setenv("USER_SUMMARY_WAIT_INTERVAL_SECONDS", "10") + # Clock advances 10s per read, so the 15s budget elapses after one wait. + ctx = _TimerContext({"user_id": "u1"}, step_seconds=10) + gen = self._orchestrator()(ctx) + result, _ = _drive( + gen, + [ + {"status": "no_memories_yet"}, # us_Extract #1 + None, # create_timer #1 + {"status": "no_memories_yet"}, # us_Extract #2 - budget now exceeded + ], + ) + assert result == {"persisted": False, "user_summary_id": None, "skipped": "no_memories_yet"} + assert "us_PersistUserSummary" not in [name for name, _ in ctx.activity_calls] + + @patch.object(us_mod, "default_retry_options", return_value=MagicMock()) + def test_persists_immediately_when_memories_already_present(self, _retry, monkeypatch): + monkeypatch.setenv("USER_SUMMARY_WAIT_SECONDS", "120") + ctx = _TimerContext({"user_id": "u1"}, step_seconds=0) + gen = self._orchestrator()(ctx) + result, _ = _drive(gen, [{"id": "us"}, {"id": "us"}]) + assert result == {"persisted": True, "user_summary_id": "us"} + assert ctx.timer_fire_times == [] # no waiting when memories already exist + + +class TestUserSummaryExtractActivity: + def test_returns_sentinel_when_no_source_memories(self): + pipeline = MagicMock() + pipeline.generate_user_summary_durable.side_effect = NoSourceMemoriesError("No memories found for user_id='u1'") + with patch.object(us_mod, "get_pipeline", return_value=pipeline): + result = us_mod.us_Extract({"user_id": "u1", "limit": 20, "thread_ids": None}) + assert result == {"status": "no_memories_yet"} + + def test_passes_through_summary_doc_when_memories_present(self): + pipeline = MagicMock() + pipeline.generate_user_summary_durable.return_value = {"id": "user_summary_u1", "type": "user_summary"} + with patch.object(us_mod, "get_pipeline", return_value=pipeline): + result = us_mod.us_Extract({"user_id": "u1", "limit": 20, "thread_ids": ["t1"]}) + assert result == {"id": "user_summary_u1", "type": "user_summary"} + pipeline.generate_user_summary_durable.assert_called_once_with(user_id="u1", recent_k=20, thread_ids=["t1"]) From fd2b7862d340f0320cb27c0b2ed792783c49fa5f Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Mon, 17 Aug 2026 15:26:41 -0700 Subject: [PATCH 3/3] Fixing issues 39 and 40 --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index bf04185..2d353a1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -72,6 +72,7 @@ Changelog = "https://github.com/AzureCosmosDB/AgentMemoryToolkit/blob/main/CHANG [project.optional-dependencies] dev = [ + "httpx>=0.27", "pytest>=8.0", "pytest-asyncio>=0.23", "pytest-cov>=5.0",