From 2347428db4fb4e59257d90cd43ac931067dff80d Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Wed, 1 Jul 2026 09:29:26 -0700 Subject: [PATCH 1/8] Rebasing --- .env.template | 17 - CHANGELOG.md | 24 + Docs/concepts.md | 40 +- Docs/design_patterns.md | 1 - Docs/public_api.md | 8 +- Docs/troubleshooting.md | 10 +- Samples/Advanced/advanced_search_patterns.py | 13 +- Samples/Notebooks/Demo_async.ipynb | 2 +- azure/cosmos/agent_memory/_counters.py | 78 +- azure/cosmos/agent_memory/_utils.py | 206 ++++- azure/cosmos/agent_memory/aio/auto_trigger.py | 61 +- .../agent_memory/aio/cosmos_memory_client.py | 23 +- .../agent_memory/aio/processors/base.py | 2 + .../agent_memory/aio/processors/durable.py | 3 +- .../agent_memory/aio/processors/inprocess.py | 28 +- .../agent_memory/aio/services/pipeline.py | 750 ++++++++++++----- .../agent_memory/aio/store/memory_store.py | 27 +- azure/cosmos/agent_memory/auto_trigger.py | 62 +- .../agent_memory/cosmos_memory_client.py | 29 +- azure/cosmos/agent_memory/processors/base.py | 2 + .../cosmos/agent_memory/processors/durable.py | 3 +- .../agent_memory/processors/inprocess.py | 29 +- azure/cosmos/agent_memory/prompts/_schemas.py | 35 +- .../cosmos/agent_memory/prompts/dedup.prompty | 2 +- .../prompts/dedup_episodic.prompty | 160 ++++ .../prompts/extract_memories.prompty | 293 +------ .../services/_pipeline_helpers.py | 61 +- .../cosmos/agent_memory/services/pipeline.py | 764 ++++++++++++++---- .../agent_memory/store/_search_helpers.py | 21 +- .../cosmos/agent_memory/store/memory_store.py | 27 +- azure/cosmos/agent_memory/thresholds.py | 76 ++ function_app/local.settings.json.template | 1 + .../orchestrators/extract_memories.py | 76 +- function_app/shared/config.py | 11 - function_app/shared/counters.py | 39 + function_app/triggers/change_feed.py | 29 + tests/integration/test_async_full_pipeline.py | 274 +++++++ tests/integration/test_full_pipeline.py | 129 ++- .../integration/test_processor_integration.py | 9 +- .../test_processor_integration_async.py | 17 +- tests/unit/aio/processors/test_inprocess.py | 22 +- .../processors/test_protocol_satisfaction.py | 1 + .../aio/services/test_dedup_vector_async.py | 446 ++++++++++ tests/unit/aio/test_auto_trigger.py | 259 +++++- tests/unit/aio/test_cosmos_memory_client.py | 106 ++- tests/unit/aio/test_process_now.py | 10 +- tests/unit/aio/test_reconcile_telemetry.py | 8 + tests/unit/function_app/test_change_feed.py | 83 +- tests/unit/function_app/test_orchestrators.py | 198 ++++- tests/unit/processors/test_inprocess.py | 23 +- .../processors/test_protocol_satisfaction.py | 1 + .../services/test_chaos_extract_persist.py | 20 + tests/unit/services/test_dedup_vector.py | 383 +++++++++ tests/unit/services/test_extract_dry.py | 568 +++++-------- tests/unit/services/test_pipeline_service.py | 33 +- tests/unit/store/test_memory_store.py | 57 ++ tests/unit/test_auto_trigger.py | 214 ++++- tests/unit/test_cosmos_memory_client.py | 113 ++- tests/unit/test_pipeline_confidence.py | 218 ++--- tests/unit/test_procedural_synthesis.py | 12 + tests/unit/test_process_now.py | 10 +- tests/unit/test_reconcile.py | 404 +-------- tests/unit/test_thresholds.py | 147 ++++ tests/unit/test_utils.py | 146 +++- 64 files changed, 5029 insertions(+), 1895 deletions(-) create mode 100644 azure/cosmos/agent_memory/prompts/dedup_episodic.prompty create mode 100644 tests/integration/test_async_full_pipeline.py create mode 100644 tests/unit/aio/services/test_dedup_vector_async.py create mode 100644 tests/unit/services/test_dedup_vector.py diff --git a/.env.template b/.env.template index 68bef3c..a264e6e 100644 --- a/.env.template +++ b/.env.template @@ -18,15 +18,6 @@ COSMOS_DB_SUMMARIES_CONTAINER="memories_summaries" COSMOS_DB_TURNS_CONTAINER="memories_turns" COSMOS_DB_COUNTERS_CONTAINER=counter COSMOS_DB_LEASE_CONTAINER=leases -# Throughput mode for all required Cosmos DB containers created by the toolkit -# (memories, counter, and lease). -# - serverless: default. The toolkit does not send container RU/s settings. -# Use this only with a Cosmos DB account configured for serverless. -# - autoscale: the toolkit provisions all required containers with autoscale -# throughput using COSMOS_DB_AUTOSCALE_MAX_RU as the max RU/s cap. -# Default max RU/s is 1000. -COSMOS_DB_THROUGHPUT_MODE=serverless -COSMOS_DB_AUTOSCALE_MAX_RU=1000 # ---- Processing thresholds (set to 0 to disable) ---- THREAD_SUMMARY_EVERY_N=10 @@ -58,14 +49,6 @@ AI_FOUNDRY_ENDPOINT=https://.openai.azure.com/ AI_FOUNDRY_API_KEY= AI_FOUNDRY_EMBEDDING_DEPLOYMENT_NAME=text-embedding-3-large AI_FOUNDRY_EMBEDDING_DIMENSIONS=1536 -AI_FOUNDRY_EMBEDDING_DATA_TYPE=float32 -AI_FOUNDRY_EMBEDDING_DISTANCE_FUNCTION=cosine -# Optional. Vector index type for the memories container: quantizedFlat -# (default), diskANN, or flat. quantizedFlat works on any Cosmos DB account -# (including the classic emulator); diskANN requires the DiskANN capability on -# the Cosmos DB account, so opt into it explicitly when available. -AI_FOUNDRY_EMBEDDING_VECTOR_INDEX_TYPE=quantizedFlat -COSMOS_DB_FULL_TEXT_LANGUAGE=en-US # Embed raw conversation turns on write so they can be vector-searched via # search_turns(). The turns container is always provisioned with a diff --git a/CHANGELOG.md b/CHANGELOG.md index efc4558..b822456 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,29 @@ ## Release History +## [0.3.0b1] (Unreleased) + +#### Features Added +* Contradiction-aware vector dedup and reconciliation. Extraction now runs a + vector-similarity dedup ladder: near-exact duplicates are auto-dropped and + borderline matches are tagged for review before they are written. A periodic + LLM reconcile pass (driven by `dedup.prompty` for facts and + `dedup_episodic.prompty` for episodic memories) then merges duplicate groups + and resolves contradictions, soft-deleting the losers with a supersede reason. + Reconciliation is distance-function aware (the destructive near-exact + auto-drop is disabled for `euclidean` containers). In-process reconcile now + covers both facts and episodic memories, matching the Durable backend, and its + cadence is derived from the persisted message counter. See [PR:#26](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/26) +* Extraction watermark. The window of recent turns sent to extraction is sized + from a persisted per-thread watermark (`last_extract_count`) that only advances + after a successful extract, so under normal operation no turns are skipped when + extraction lags or transiently fails. See [PR:#26](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/26) + +#### Other Changes +* `search_cosmos` and `search_turns` now always fuse vector similarity with + BM25 (full-text) ranking, falling back to vector-only for all-stopword + queries. The `hybrid_search` flag has been removed — hybrid ranking is the + default and requires no opt-in. See [PR:#26](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/26) + ## [0.2.0b1] (2026-06-30) #### Features Added diff --git a/Docs/concepts.md b/Docs/concepts.md index e3b646f..b72807c 100644 --- a/Docs/concepts.md +++ b/Docs/concepts.md @@ -118,26 +118,52 @@ Prompts for summarization and fact extraction live in `azure_functions/prompts/` ## Memory Reconciliation -The `reconcile_memories(user_id, n=50)` pipeline step reads up to N most-recent active facts for a user and asks the LLM to identify two orthogonal outcomes in one pass: +Reconciliation runs in **two complementary tiers**: a cheap, LLM-free **vector-floor dedup ladder** applied to freshly-extracted memories before they persist, and a periodic **LLM reconcile** that runs in a **dual mode** (cheap candidate clusters most sweeps, a full-pool backstop occasionally). -- **Duplicates** — two or more facts that restate the same claim in different words. Resolution: collapse into one merged fact; the originals are soft-deleted with `supersede_reason="duplicate"` and `superseded_by` set to the merged fact's id. -- **Contradictions** — two facts that assert opposing claims about the same subject. Resolution: keep the winner (more recent first, higher confidence as tiebreaker), soft-delete the loser with `supersede_reason="contradict"` and `superseded_by` set to the winner. +### Vector-floor dedup ladder (write path, LLM-free) -### Why one pass +Between extraction and persist, `dedup_extracted_memories` compares each new fact/episodic memory against the user's existing active memories of the same type using Cosmos `VectorDistance` (pure vector, no hybrid). Each new memory takes one rung of a similarity ladder: -Detecting contradictions semantically requires the LLM to see the candidate pool as a whole — paraphrased ("user prefers aisle seats") and contradictory ("user is vegetarian" vs "user loves steak") facts often have very different embedding vectors and would never co-occur in any cosine cluster. Putting all N candidates into one prompt lets the LLM do the semantic reasoning across both axes simultaneously. The pipeline returns `{"kept": int, "merged": int, "contradicted": int}`. +| band | condition (cosine) | action | +|------|--------------------|--------| +| exact | `content_hash` hit | skip (Stage 0, free) | +| near-exact | `s ≥ DEDUP_SIM_HIGH` (0.97) | **auto-skip** the new memory (no LLM); logged for audit | +| borderline | `DEDUP_SIM_LOW ≤ s < DEDUP_SIM_HIGH` (0.80–0.97) | persist, tag `sys:dup-candidate` + stash `dup_of`/`dup_score` for the LLM reconcile | +| novel | `s < DEDUP_SIM_LOW` | persist clean | + +The thresholds are calibrated for **cosine/dotproduct** on normalized embeddings. On a container whose `distanceFunction` is **euclidean**, the destructive near-exact auto-skip is **disabled** (one-shot warning) and those memories fall through to borderline tagging so the LLM adjudicates — euclidean distances aren't a bounded [0,1] similarity and would mis-fire the cosine-tuned drop. + +### Dual-mode LLM reconcile + +`reconcile_memories(user_id, n=50, *, memory_type="fact", full_rebuild=False)` identifies two orthogonal outcomes: + +- **Duplicates** — facts restating the same claim. Resolution: collapse into one merged fact; originals soft-deleted with `supersede_reason="duplicate"` and `superseded_by` set to the merged fact. +- **Contradictions** — facts asserting opposing claims about the same subject. Resolution: keep the winner (more recent first, higher confidence as tiebreaker), soft-delete the loser with `supersede_reason="contradict"`. + +It runs in one of two modes: + +- **Candidate mode** (default auto sweeps) — builds connected-component clusters from the `sys:dup-candidate` seeds + their vector neighbors (edge threshold `DEDUP_CLUSTER_SIM`, 0.60) and sends **only those clusters** to the LLM. Cheap, but keyed on near-duplicate similarity. Tagged seeds that never join a cluster have their stale tag cleared so they aren't re-scanned forever. +- **Full-pool backstop** — every `DEDUP_FULL_RECLUSTER_EVERY_N`-th sweep (default 12), and on any explicit `reconcile(full_rebuild=True)`, the **entire** active pool goes into one LLM pass. This is the only path that catches **dissimilar contradictions** — paraphrased ("prefers aisle seats") and contradictory ("vegetarian" vs "loves steak") facts have very different embedding vectors and would never co-occur in a cosine cluster, so candidate mode alone can't link them. + +Both modes return `{"kept": int, "merged": int, "contradicted": int}`. In-process and durable backends reconcile **both** facts and episodic memories so episodic duplicates don't accrue forever. ### Loser preservation -Soft-deleted facts stay in the container with their `supersede_reason`, `superseded_at`, and `superseded_by` fields populated. Default reads (`get_memories`, `search_cosmos`) filter them out via `superseded_by IS NULL`. To inspect the audit trail (e.g. "show everything that ever applied to this user"), opt out of the filter at the query level. +Soft-deleted facts stay in the container with their `supersede_reason`, `superseded_at`, and `superseded_by` fields populated. Default reads (`get_memories`, `search_cosmos`) filter them out via `superseded_by IS NULL`. To inspect the audit trail, opt out of the filter at the query level. ### Write-time exact dedup Each fact written by `extract_memories` carries a `content_hash` (SHA-256 of normalized content, truncated to 32 hex chars; lowercase, whitespace-collapsed). Before upserting a freshly-extracted fact, the pipeline checks the hash against existing active facts and short-circuits if a match exists, incrementing the `exact_dedup_skipped` metric. This catches identical re-extractions cheaply without an LLM call. +### Extraction watermark (`recent_k`) + +The auto-trigger paths size `recent_k` (how many recent turns extraction reads) from a per-thread **watermark** (`last_extract_count` on the counter doc): `recent_k = current_count − last_extract_count` (with `last_extract_count` treated as `0` before the first successful extract). The newest-`recent_k` turns are exactly the turns added since the last successful extract, and the watermark advances **only after a successful extract** — so under normal operation no turns are skipped when extraction lags or transiently fails: a failed run leaves the watermark put and the full backlog is retried next sweep. The window is deliberately **not** capped by `DEDUP_POOL_SIZE` (that knob governs the reconcile prompt, not the extraction window) — capping would extract only the newest N and silently strand the oldest backlog turns. + +> **Caveat (rare):** the SDK's inline counter increment is best-effort — under sustained optimistic-concurrency contention it can drop an increment rather than block the user's write path (see `increment_counter_sync`). A dropped increment leaves `current_count` lagging the true turn count, which can in turn under-cover a later extraction window. This is the one case where the "no turns skipped" property does not hold; the Function App backend avoids it by raising to force change-feed redelivery. + ### Tunable -`DEDUP_EVERY_N` (default 5) controls how often `reconcile_memories` runs in the auto-trigger path. Set to `0` to disable. The candidate cap `n` (default 50) is tunable per call; larger values give the LLM a wider view at higher token cost. +`DEDUP_EVERY_N` (default 5) controls how often reconcile runs in the auto-trigger path. Set to `0` to disable. The candidate cap `n` (default `DEDUP_POOL_SIZE`, 50) is tunable per call; larger values give the LLM a wider view at higher token cost. `DEDUP_FULL_RECLUSTER_EVERY_N` (default 12) sets how often the full-pool backstop fires. > **Indexing note.** The reconcile pool query orders by `created_at` (matching the prompt's "more recent first" tiebreaker). Cosmos's default indexing policy includes every property, so this works out of the box. If you customize the indexing policy to reduce write RU, ensure `/created_at/?` remains indexed or the query will fail with a 400 (`Order-by over a non-indexed path`). diff --git a/Docs/design_patterns.md b/Docs/design_patterns.md index ce8e72b..0e520c3 100644 --- a/Docs/design_patterns.md +++ b/Docs/design_patterns.md @@ -170,7 +170,6 @@ facts = await mem.search_cosmos( results = await mem.search_cosmos( search_terms="PostgreSQL to Cosmos DB", user_id="user-1", - hybrid_search=True, top_k=5, ) ``` diff --git a/Docs/public_api.md b/Docs/public_api.md index dd17452..bec37e4 100644 --- a/Docs/public_api.md +++ b/Docs/public_api.md @@ -37,8 +37,8 @@ ### Retrieval -- `search_cosmos(search_terms, memory_id=None, user_id=None, role=None, memory_types=None, thread_id=None, hybrid_search=False, top_k=5, tags_all=None, tags_any=None, exclude_tags=None, include_superseded=False, min_salience=None, min_confidence=None, created_after=None, created_before=None) -> list[dict]` — vector or hybrid search derived memories (facts/episodic/procedural). -- `search_turns(search_terms, user_id, thread_id=None, role=None, hybrid_search=False, top_k=5, tags_all=None, tags_any=None, exclude_tags=None, created_after=None, created_before=None) -> list[dict]` — vector or hybrid search the raw conversation log instead of facts/episodic/procedural (requires turn embeddings; see `enable_turn_embeddings`). `user_id` is required so the search is scoped to one partition instead of scanning every user's turns. +- `search_cosmos(search_terms, memory_id=None, user_id=None, role=None, memory_types=None, thread_id=None, top_k=5, tags_all=None, tags_any=None, exclude_tags=None, include_superseded=False, min_salience=None, min_confidence=None, created_after=None, created_before=None) -> list[dict]` — hybrid vector/full-text search memories, falling back to vector-only for all-stopword queries. +- `search_turns(search_terms, user_id, thread_id=None, role=None, top_k=5, tags_all=None, tags_any=None, exclude_tags=None, created_after=None, created_before=None) -> list[dict]` — hybrid vector/full-text search the raw conversation log instead of facts/episodic/procedural (requires turn embeddings; see `enable_turn_embeddings`). `user_id` is required so the search is scoped to one partition instead of scanning every user's turns. - `get_procedural_prompt(user_id) -> Optional[str]` — read the active procedural prompt. - `get_procedural_history(user_id, limit=10) -> list[dict]` — read procedural prompt history. - `get_procedural_memories(user_id, priority=None, category=None, min_salience=None, include_superseded=False) -> list[dict]` — retrieve procedural memory documents. @@ -91,8 +91,8 @@ Local-buffer methods remain synchronous in-memory operations; Cosmos, retrieval, ### Retrieval -- `async search_cosmos(search_terms, memory_id=None, user_id=None, role=None, memory_types=None, thread_id=None, hybrid_search=False, top_k=5, tags_all=None, tags_any=None, exclude_tags=None, include_superseded=False, min_salience=None, min_confidence=None, created_after=None, created_before=None) -> list[dict]` — vector or hybrid search derived memories (facts/episodic/procedural). -- `async search_turns(search_terms, user_id, thread_id=None, role=None, hybrid_search=False, top_k=5, tags_all=None, tags_any=None, exclude_tags=None, created_after=None, created_before=None) -> list[dict]` — vector or hybrid search the raw conversation log instead of facts/episodic/procedural (requires turn embeddings; see `enable_turn_embeddings`). `user_id` is required so the search is scoped to one partition instead of scanning every user's turns. +- `async search_cosmos(search_terms, memory_id=None, user_id=None, role=None, memory_types=None, thread_id=None, top_k=5, tags_all=None, tags_any=None, exclude_tags=None, include_superseded=False, min_salience=None, min_confidence=None, created_after=None, created_before=None) -> list[dict]` — hybrid vector/full-text search memories, falling back to vector-only for all-stopword queries. +- `async search_turns(search_terms, user_id, thread_id=None, role=None, top_k=5, tags_all=None, tags_any=None, exclude_tags=None, created_after=None, created_before=None) -> list[dict]` — hybrid vector/full-text search the raw conversation log instead of facts/episodic/procedural (requires turn embeddings; see `enable_turn_embeddings`). `user_id` is required so the search is scoped to one partition instead of scanning every user's turns. - `async get_procedural_prompt(user_id) -> Optional[str]` — read the active procedural prompt. - `async get_procedural_history(user_id, limit=10) -> list[dict]` — read procedural prompt history. - `async get_procedural_memories(user_id, priority=None, category=None, min_salience=None, include_superseded=False) -> list[dict]` — retrieve procedural memory documents. diff --git a/Docs/troubleshooting.md b/Docs/troubleshooting.md index 27f500f..a64d1cd 100644 --- a/Docs/troubleshooting.md +++ b/Docs/troubleshooting.md @@ -50,15 +50,11 @@ COSMOS_DB_DATABASE=ai_memory COSMOS_DB_MEMORIES_CONTAINER=memories COSMOS_DB_COUNTERS_CONTAINER=counter COSMOS_DB_LEASE_CONTAINER=leases -COSMOS_DB_THROUGHPUT_MODE=serverless -COSMOS_DB_AUTOSCALE_MAX_RU=1000 AI_FOUNDRY_ENDPOINT=https://.openai.azure.com/ AI_FOUNDRY_API_KEY= AI_FOUNDRY_EMBEDDING_DEPLOYMENT_NAME=text-embedding-3-large AI_FOUNDRY_EMBEDDING_DIMENSIONS=1536 -AI_FOUNDRY_EMBEDDING_DATA_TYPE=float32 -AI_FOUNDRY_EMBEDDING_DISTANCE_FUNCTION=cosine AI_FOUNDRY_CHAT_DEPLOYMENT_NAME= ``` @@ -77,8 +73,6 @@ The notebooks and samples pass these values into the client like this: | `AI_FOUNDRY_EMBEDDING_DEPLOYMENT_NAME` | `embedding_deployment_name` | | `AI_FOUNDRY_CHAT_DEPLOYMENT_NAME` | `chat_deployment_name` | -`AI_FOUNDRY_EMBEDDING_DIMENSIONS`, `AI_FOUNDRY_EMBEDDING_DATA_TYPE`, and `AI_FOUNDRY_EMBEDDING_DISTANCE_FUNCTION` are read by the toolkit when creating the Cosmos DB vector policy. The Function App also reads `COSMOS_DB__accountEndpoint` for its identity-based Cosmos DB trigger binding; set it to the same value as `COSMOS_DB_ENDPOINT`. - Run `az login` before using `DefaultAzureCredential`. Required roles: @@ -104,7 +98,7 @@ The memories container is created with: If vector or full-text search fails after changing dimensions or indexing settings, create a fresh container with the desired configuration. Cosmos container vector policies are creation-time infrastructure choices. -Use `COSMOS_DB_THROUGHPUT_MODE=serverless` for the default setup. Use `autoscale` with `COSMOS_DB_AUTOSCALE_MAX_RU` when you need provisioned autoscale throughput. +Pass `cosmos_throughput_mode="serverless"` (the default) when creating the client. Use `cosmos_throughput_mode="autoscale"` with `cosmos_autoscale_max_ru` when you need provisioned autoscale throughput. --- @@ -117,7 +111,7 @@ Embedding failures usually mean one of these is wrong: - `AI_FOUNDRY_EMBEDDING_DIMENSIONS` - Azure OpenAI / AI Services RBAC -For hybrid search, `search_terms` is required when `hybrid_search=True`. +Search always uses hybrid vector/full-text ranking when keyword extraction finds terms; all-stopword queries fall back to vector-only ranking. If search returns documents but scores look poor, check that records have an `embedding` field and that the query uses similar language to the stored memory content. diff --git a/Samples/Advanced/advanced_search_patterns.py b/Samples/Advanced/advanced_search_patterns.py index 2d9a518..dff042d 100644 --- a/Samples/Advanced/advanced_search_patterns.py +++ b/Samples/Advanced/advanced_search_patterns.py @@ -89,10 +89,11 @@ def seed_memories(mem: CosmosMemoryClient, user_id: str, thread_id: str) -> None # --------------------------------------------------------------------------- def vector_search(mem: CosmosMemoryClient, user_id: str) -> None: - """Pattern 1 — Pure vector (semantic similarity) search.""" - print_header("1. Vector Search (semantic similarity)") + """Pattern 1 — Semantic-style query (natural language, low keyword overlap).""" + print_header("1. Semantic Search (natural-language query)") print(" Query: 'outdoor activities'") - print(" Finds semantically related memories even without exact keyword matches.\n") + print(" Hybrid ranking leans on embedding similarity when there are few exact") + print(" keyword matches, so semantically related memories still surface.\n") results = mem.search_cosmos( search_terms="outdoor activities", @@ -103,15 +104,15 @@ def vector_search(mem: CosmosMemoryClient, user_id: str) -> None: def hybrid_search(mem: CosmosMemoryClient, user_id: str) -> None: - """Pattern 2 — Hybrid search (vector + full-text).""" + """Pattern 2 — Hybrid search (vector + full-text) is the default.""" print_header("2. Hybrid Search (vector + full-text)") print(" Query: 'hiking trails Pacific Northwest'") - print(" Combines embedding similarity with BM25 keyword matching.\n") + print(" Every search_cosmos call fuses embedding similarity with BM25 keyword") + print(" matching automatically — no flag required.\n") results = mem.search_cosmos( search_terms="hiking trails Pacific Northwest", user_id=user_id, - hybrid_search=True, top_k=5, ) print_results(results) diff --git a/Samples/Notebooks/Demo_async.ipynb b/Samples/Notebooks/Demo_async.ipynb index bf413f8..3aca2d1 100644 --- a/Samples/Notebooks/Demo_async.ipynb +++ b/Samples/Notebooks/Demo_async.ipynb @@ -872,7 +872,7 @@ "results_search_async = await memory.search_cosmos(\n", " search_terms=\"What did the user ask about the weather?\",\n", " user_id=USER_ID,\n", - " top_k=3, hybrid_search= True\n", + " top_k=3\n", ")" ] }, diff --git a/azure/cosmos/agent_memory/_counters.py b/azure/cosmos/agent_memory/_counters.py index 2223d06..65d888c 100644 --- a/azure/cosmos/agent_memory/_counters.py +++ b/azure/cosmos/agent_memory/_counters.py @@ -170,7 +170,7 @@ def increment_counter_sync( if exc.status_code == 412 and attempt < MAX_RETRIES - 1: continue logger.warning( - "Counter increment failed counter_id=%s status=%s — auto-trigger skipped", + "Counter increment failed counter_id=%s status=%s — auto-trigger skipped (increment dropped)", counter_id, exc.status_code, ) @@ -277,6 +277,10 @@ def _build_counter_doc( doc["last_batch_old_count"] = existing.get("last_batch_old_count", old_count) else: doc["last_batch_lsn"] = None + # Preserve the extraction watermark (count value at the last successful + # extract) so recent_k can cover every turn since, not just this batch. + if existing is not None and "last_extract_count" in existing: + doc["last_extract_count"] = existing.get("last_extract_count") # Carry over auto-trigger failure breadcrumbs so they aren't blown away # by a successful write. ``stamp_failure_*`` helpers refresh them on # failure; operators can monitor ``last_failure_at`` directly. @@ -346,6 +350,74 @@ async def stamp_failure_async( logger.debug("stamp_failure_async failed counter_id=%s: %s", counter_id, exc) +def read_extract_watermark_sync( + container: Any, + counter_id: str, + user_id: str, + thread_id: str, +) -> Optional[int]: + """Return the count value at the last successful extract, or ``None``. + + The watermark lets recent_k cover every turn since the previous extract + succeeded, instead of just the current batch — so turns are never skipped + when extraction lags or transiently fails. Best-effort: returns ``None`` + on any read error so callers fall back to a batch-based recent_k. + """ + try: + doc = container.read_item(item=counter_id, partition_key=[user_id, thread_id]) + value = doc.get("last_extract_count") + return int(value) if value is not None else None + except Exception as exc: # pragma: no cover - best-effort + logger.debug("read_extract_watermark_sync failed counter_id=%s: %s", counter_id, exc) + return None + + +def advance_extract_watermark_sync( + container: Any, + counter_id: str, + user_id: str, + thread_id: str, + count: int, +) -> None: + """Stamp ``last_extract_count=count`` after a successful extract (sync).""" + patch_ops = [{"op": "add", "path": "/last_extract_count", "value": int(count)}] + try: + container.patch_item(item=counter_id, partition_key=[user_id, thread_id], patch_operations=patch_ops) + except Exception as exc: # pragma: no cover - best-effort + logger.debug("advance_extract_watermark_sync failed counter_id=%s: %s", counter_id, exc) + + +async def read_extract_watermark_async( + container: Any, + counter_id: str, + user_id: str, + thread_id: str, +) -> Optional[int]: + """Async version of :func:`read_extract_watermark_sync`.""" + try: + doc = await container.read_item(item=counter_id, partition_key=[user_id, thread_id]) + value = doc.get("last_extract_count") + return int(value) if value is not None else None + except Exception as exc: # pragma: no cover - best-effort + logger.debug("read_extract_watermark_async failed counter_id=%s: %s", counter_id, exc) + return None + + +async def advance_extract_watermark_async( + container: Any, + counter_id: str, + user_id: str, + thread_id: str, + count: int, +) -> None: + """Stamp ``last_extract_count=count`` after a successful extract (async).""" + patch_ops = [{"op": "add", "path": "/last_extract_count", "value": int(count)}] + try: + await container.patch_item(item=counter_id, partition_key=[user_id, thread_id], patch_operations=patch_ops) + except Exception as exc: # pragma: no cover - best-effort + logger.debug("advance_extract_watermark_async failed counter_id=%s: %s", counter_id, exc) + + __all__ = [ "USER_COUNTER_THREAD_ID", "thread_counter_id", @@ -355,4 +427,8 @@ async def stamp_failure_async( "increment_counter_async", "stamp_failure_sync", "stamp_failure_async", + "read_extract_watermark_sync", + "advance_extract_watermark_sync", + "read_extract_watermark_async", + "advance_extract_watermark_async", ] diff --git a/azure/cosmos/agent_memory/_utils.py b/azure/cosmos/agent_memory/_utils.py index 849ef48..ac32f9e 100644 --- a/azure/cosmos/agent_memory/_utils.py +++ b/azure/cosmos/agent_memory/_utils.py @@ -224,11 +224,12 @@ def normalize_ai_foundry_endpoint(endpoint: Optional[str]) -> Optional[str]: def _resolve_embedding_data_type(val: Optional[str]) -> str: - """Resolve embedding data type from explicit value or ``AI_FOUNDRY_EMBEDDING_DATA_TYPE`` env var. + """Resolve embedding data type from the explicit value, defaulting to ``float32``. - Defaults to ``float32``. Raises :class:`ConfigurationError` for unknown values. + Provided by the caller at memory-client creation. Raises :class:`ConfigurationError` + for unknown values. """ - raw = (val if val is not None else os.environ.get("AI_FOUNDRY_EMBEDDING_DATA_TYPE") or "float32").strip() + raw = (val if val is not None else "float32").strip() if raw not in _ALLOWED_EMBEDDING_DATA_TYPES: raise ConfigurationError( message=( @@ -241,11 +242,12 @@ def _resolve_embedding_data_type(val: Optional[str]) -> str: def _resolve_distance_function(val: Optional[str]) -> str: - """Resolve distance function from explicit value or ``AI_FOUNDRY_EMBEDDING_DISTANCE_FUNCTION`` env var. + """Resolve distance function from the explicit value, defaulting to ``cosine``. - Defaults to ``cosine``. Raises :class:`ConfigurationError` for unknown values. + Provided by the caller at memory-client creation. Raises :class:`ConfigurationError` + for unknown values. """ - raw = (val if val is not None else os.environ.get("AI_FOUNDRY_EMBEDDING_DISTANCE_FUNCTION") or "cosine").strip() + raw = (val if val is not None else "cosine").strip() if raw not in _ALLOWED_DISTANCE_FUNCTIONS: raise ConfigurationError( message=( @@ -258,17 +260,16 @@ def _resolve_distance_function(val: Optional[str]) -> str: def _resolve_vector_index_type(val: Optional[str]) -> str: - """Resolve vector index type from explicit value or ``AI_FOUNDRY_EMBEDDING_VECTOR_INDEX_TYPE`` env var. + """Resolve the vector index type from the explicit value, defaulting to ``quantizedFlat``. - Defaults to ``quantizedFlat``. Raises :class:`ConfigurationError` for unknown values. + Provided by the caller at memory-client creation. Raises :class:`ConfigurationError` + for unknown values. ``quantizedFlat`` works on any Cosmos DB account (including the classic emulator). ``diskANN`` requires the Cosmos DB account to have the DiskANN vector index capability enabled; opt into it explicitly when available. """ - raw = ( - val if val is not None else os.environ.get("AI_FOUNDRY_EMBEDDING_VECTOR_INDEX_TYPE") or "quantizedFlat" - ).strip() + raw = (val if val is not None else "quantizedFlat").strip() if raw not in _ALLOWED_VECTOR_INDEX_TYPES: raise ConfigurationError( message=( @@ -280,21 +281,79 @@ def _resolve_vector_index_type(val: Optional[str]) -> str: return raw +_SIMILARITY_DESCENDING_FUNCTIONS = frozenset({"cosine", "dotproduct"}) + + +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 vector_similarity_at_least(score: float, threshold: float, distance_function: str) -> bool: + """Return ``True`` when ``score`` meets/exceeds ``threshold`` similarity. + + For cosine/dotproduct (higher = more similar) this is ``score >= threshold``; + for euclidean (lower = more similar) it inverts to ``score <= threshold``. The + dedup thresholds (``DEDUP_SIM_*``) are calibrated for cosine/dotproduct on + normalized embeddings; euclidean gets the correct *direction* but its + thresholds would need separate calibration. + """ + if distance_function in _SIMILARITY_DESCENDING_FUNCTIONS: + return score >= threshold + return score <= threshold + + +def vector_autodrop_supported(distance_function: str) -> bool: + """Whether the cosine-calibrated near-exact auto-drop is safe to apply. + + The destructive ``DEDUP_SIM_HIGH`` auto-skip drops a new memory without an + LLM check, relying on thresholds (~0.97) calibrated for cosine/dotproduct + on normalized embeddings. Euclidean returns an *unbounded distance* (not a + [0,1] similarity), so those thresholds mis-fire — auto-drop is disabled for + euclidean and the borderline tagging path (LLM-adjudicated) is used instead. + """ + return distance_function != "euclidean" + + +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 explicit value or ``COSMOS_DB_FULL_TEXT_LANGUAGE`` env var. + """Resolve full-text language from the explicit value, defaulting to ``en-US``. - Defaults to ``en-US``. Empty values fall back to the default. + Provided by the caller at memory-client creation. Empty values fall back to the default. """ - raw = (val if val is not None else os.environ.get("COSMOS_DB_FULL_TEXT_LANGUAGE") or "en-US").strip() + raw = (val if val is not None else "en-US").strip() return raw or "en-US" def _resolve_cosmos_throughput_mode(val: Optional[str]) -> str: - """Resolve throughput mode from explicit value or env var. + """Resolve throughput mode from the explicit value, defaulting to ``serverless``. Allowed values are ``serverless`` and ``autoscale``. """ - raw = (val if val is not None else os.environ.get("COSMOS_DB_THROUGHPUT_MODE") or "serverless").strip().lower() + raw = (val if val is not None else "serverless").strip().lower() if raw not in {"serverless", "autoscale"}: raise ConfigurationError( @@ -307,28 +366,15 @@ def _resolve_cosmos_throughput_mode(val: Optional[str]) -> str: def _resolve_cosmos_autoscale_max_ru(val: Optional[int]) -> int: - """Resolve autoscale max RU from explicit value or env var.""" - if val is not None: - if not isinstance(val, int) or isinstance(val, bool) or val <= 0: - raise ConfigurationError( - message=f"Invalid configuration for cosmos_autoscale_max_ru: expected a positive integer, got '{val}'", - parameter="cosmos_autoscale_max_ru", - ) - return val - raw = (os.environ.get("COSMOS_DB_AUTOSCALE_MAX_RU") or "1000").strip() - try: - parsed = int(raw) - except ValueError as exc: + """Resolve autoscale max RU from the explicit value, defaulting to 1000.""" + if val is None: + return 1000 + if not isinstance(val, int) or isinstance(val, bool) or val <= 0: raise ConfigurationError( - message=(f"Invalid configuration for cosmos_autoscale_max_ru: expected an integer, got '{raw}'"), - parameter="cosmos_autoscale_max_ru", - ) from exc - if parsed <= 0: - raise ConfigurationError( - message=(f"Invalid configuration for cosmos_autoscale_max_ru: expected a positive integer, got '{raw}'"), + message=f"Invalid configuration for cosmos_autoscale_max_ru: expected a positive integer, got '{val}'", parameter="cosmos_autoscale_max_ru", ) - return parsed + return val def _resolve_cosmos_provisioning_autoscale_max_ru( @@ -490,10 +536,86 @@ def _container_policies( return vector_embedding_policy, indexing_policy, full_text_policy -def _validate_hybrid_search( - hybrid_search: bool, - search_terms: Optional[str], -) -> None: - """Raise :class:`ValidationError` if hybrid search is requested without search terms.""" - if hybrid_search and not search_terms: - raise ValidationError("search_terms is required when hybrid_search is True") +_FULLTEXT_STOPWORDS: frozenset[str] = frozenset( + """ + 0 1 2 3 4 5 6 7 8 9 a a's able about above according accordingly across actually + after afterwards again against ain't all allow allows almost alone along already + also although always am among amongst an and another any anybody anyhow anyone + anything anyway anyways anywhere apart appear appreciate appropriate are aren't + around as aside ask asking associated at available away awfully b be became + because become becomes becoming been before beforehand behind being believe below + beside besides best better between beyond both brief but by c c'mon c's came can + can't cannot cant cause causes certain certainly changes clearly co com come comes + concerning consequently consider considering contain containing contains + corresponding could couldn't course currently d definitely described despite did + didn't different do does doesn't doing don don't done down downwards during e each + edu eg eight either else elsewhere enough entirely especially et etc even ever + every everybody everyone everything everywhere ex exactly example except f far few + fifth first five followed following follows for former formerly forth four from + further furthermore g get gets getting given gives go goes going gone got gotten + greetings h had hadn't happens hardly has hasn't have haven't having he he's hello + help hence her here here's hereafter hereby herein hereupon hers herself hi him + himself his hither hopefully how howbeit however i i'd i'll i'm i've ie if ignored + immediate in inasmuch inc indeed indicate indicated indicates inner insofar instead + into inward is isn't it it'd it'll it's its itself j just k keep keeps kept know + known knows l last lately later latter latterly least less lest let let's like + liked likely little ll look looking looks ltd m mainly make many may maybe me mean + meanwhile merely might more moreover most mostly mr mrs ms much must my myself n + name namely nd near nearly necessary need needs neither never nevertheless new next + nine no nobody non none noone nor normally not nothing novel now nowhere o obviously + of off often oh ok okay old on once one ones only onto or other others otherwise + ought our ours ourselves out outside over overall own p particular particularly per + perhaps placed please plus possible presumably probably provides q que quite qv r + rather rd re really reasonably regarding regardless regards relatively respectively + right s said same saw say saying says second secondly see seeing seem seemed seeming + seems seen self selves sensible sent serious seriously seven several shall she + should shouldn't since six so some somebody somehow someone something sometime + sometimes somewhat somewhere soon sorry specified specify specifying still sub such + sup sure t t's take taken tell tends th than thank thanks thanx that that's thats + the their theirs them themselves then thence there there's thereafter thereby + therefore therein theres thereupon these they they'd they'll they're they've think + third this thorough thoroughly those though three through throughout thru thus to + together too took toward towards tried tries truly try trying twice two u un under + unfortunately unless unlikely until unto up upon us use used useful uses using + usually v value various ve very via viz vs w want wants was wasn't way we we'd we'll + we're we've welcome well went were weren't what what's whatever when whence whenever + where where's whereafter whereas whereby wherein whereupon wherever whether which + while whither who who's whoever whole whom whose why will willing wish with within + without won't wonder would wouldn't x y yes yet you you'd you'll you're you've your + yours yourself yourselves z zero + """.split() +) + +_KEYWORD_TOKEN_RE = re.compile(r"[a-z0-9]+") + +# Azure Cosmos DB ``FullTextScore`` accepts at most 30 search terms; a query with +# 31+ terms is rejected with ``BadRequest: One of the input values is invalid``. +# Keyword extraction is therefore capped here so the hybrid search SQL can never +# exceed the limit. The full query text is still embedded uncapped for the vector +# half of the hybrid rank, so trimming the BM25 keyword tail does not lose semantics. +MAX_FULLTEXT_TERMS = 30 + + +def extract_keywords(text: Optional[str]) -> list[str]: + """Extract de-duplicated, stopword-filtered keyword terms for full-text search. + + Lowercases, tokenizes on alphanumeric runs (apostrophes/punctuation split into + fragments that are themselves stopwords), removes stopwords, and de-duplicates + while preserving first-seen order. The result is capped at ``MAX_FULLTEXT_TERMS`` + (30) — the hard limit on terms Azure Cosmos DB ``FullTextScore`` accepts — so the + hybrid search query is always valid. Returns ``[]`` when the text is empty or all + stopwords, which the search layer treats as a signal to fall back to pure vector + ranking. + """ + if not text: + return [] + seen: set[str] = set() + keywords: list[str] = [] + for token in _KEYWORD_TOKEN_RE.findall(text.lower()): + if token in _FULLTEXT_STOPWORDS or token in seen: + continue + seen.add(token) + keywords.append(token) + if len(keywords) >= MAX_FULLTEXT_TERMS: + break + return keywords diff --git a/azure/cosmos/agent_memory/aio/auto_trigger.py b/azure/cosmos/agent_memory/aio/auto_trigger.py index b0d4c47..78f8716 100644 --- a/azure/cosmos/agent_memory/aio/auto_trigger.py +++ b/azure/cosmos/agent_memory/aio/auto_trigger.py @@ -64,6 +64,10 @@ async def maybe_trigger_steps( return n_dedup_turns = n_facts * n_dedup if n_facts > 0 and n_dedup > 0 else 0 + # Persisted-counter full-pool backstop cadence (durable-safe, mirrors the + # change-feed): every DEDUP_FULL_RECLUSTER_EVERY_N-th reconcile. + n_full_recluster = _threshold_int(thresholds, "get_dedup_full_recluster_every_n", "DEDUP_FULL_RECLUSTER_EVERY_N") + n_full_turns = n_dedup_turns * n_full_recluster if (n_dedup_turns > 0 and n_full_recluster > 0) else 0 user_batch_counts = await _trigger_thread_steps( processor, counter_container, @@ -71,6 +75,7 @@ async def maybe_trigger_steps( n_facts=n_facts, n_summary=n_summary, n_dedup_turns=n_dedup_turns, + n_full_turns=n_full_turns, thresholds=thresholds, ) await _trigger_user_steps(processor, counter_container, user_batch_counts, n_user=n_user) @@ -84,6 +89,7 @@ async def _trigger_thread_steps( n_facts: int, n_summary: int, n_dedup_turns: int, + n_full_turns: int, thresholds: Any = None, ) -> dict[str, int]: user_batch_counts: dict[str, int] = {} @@ -110,14 +116,38 @@ async def _trigger_thread_steps( counter_id, user_id, thread_id, + new_count=new_count, fire_extract=n_facts > 0 and _counters.crosses_threshold(old_count, new_count, n_facts), fire_summary=n_summary > 0 and _counters.crosses_threshold(old_count, new_count, n_summary), fire_dedup=n_dedup_turns > 0 and _counters.crosses_threshold(old_count, new_count, n_dedup_turns), + fire_full_rebuild=n_full_turns > 0 and _counters.crosses_threshold(old_count, new_count, n_full_turns), thresholds=thresholds, ) return user_batch_counts +async def _watermark_recent_k( + counter_container: Any, + counter_id: str, + user_id: str, + thread_id: str, + *, + new_count: int, +) -> int: + """Async: recent_k covering every turn since the last successful extract. + + Not capped — ``new_count - watermark`` is exactly the unextracted backlog and + the newest-``recent_k`` slice covers precisely those turns, so the watermark + can advance to ``new_count`` with no stranded turns. **Bootstrap:** with no + watermark yet the base is ``0`` (``recent_k = new_count``), so turns added + during earlier failed extracts aren't stranded when the watermark first + advances to ``new_count``. + """ + watermark = await _counters.read_extract_watermark_async(counter_container, counter_id, user_id, thread_id) + base = watermark if watermark is not None else 0 + return max(new_count - base, 1) + + async def _fire_thread_steps( processor: AsyncInProcessProcessor, counter_container: Any, @@ -125,9 +155,11 @@ async def _fire_thread_steps( user_id: str, thread_id: str, *, + new_count: int, fire_extract: bool, fire_summary: bool, fire_dedup: bool, + fire_full_rebuild: bool = False, thresholds: Any = None, ) -> None: fire_procedural = fire_dedup and bool( @@ -138,14 +170,33 @@ async def _fire_thread_steps( default=True, ) ) + if fire_extract: + recent_k = await _watermark_recent_k( + counter_container, + counter_id, + user_id, + thread_id, + new_count=new_count, + ) + try: + await _call_async_compatible( + processor.process_extract_memories, user_id=user_id, thread_id=thread_id, recent_k=recent_k + ) + await _counters.advance_extract_watermark_async( + counter_container, counter_id, user_id, thread_id, new_count + ) + except Exception as exc: + logger.warning("Async auto-trigger process_extract_memories failed for %s/%s: %s", user_id, thread_id, exc) + await _counters.stamp_failure_async( + counter_container, counter_id, user_id, thread_id, f"process_extract_memories: {exc!r}" + ) calls = ( ( - fire_extract, - "process_extract_memories", - processor.process_extract_memories, - {"user_id": user_id, "thread_id": thread_id}, + fire_dedup, + "process_reconcile", + processor.process_reconcile, + {"user_id": user_id, "full_rebuild": fire_full_rebuild}, ), - (fire_dedup, "process_reconcile", processor.process_reconcile, {"user_id": user_id}), ( fire_procedural, "synthesize_procedural", diff --git a/azure/cosmos/agent_memory/aio/cosmos_memory_client.py b/azure/cosmos/agent_memory/aio/cosmos_memory_client.py index 4c5b4a7..6bfd626 100644 --- a/azure/cosmos/agent_memory/aio/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/aio/cosmos_memory_client.py @@ -663,7 +663,6 @@ async def search_cosmos( role: Optional[str] = None, memory_types: Optional[list[str]] = None, thread_id: Optional[str] = None, - hybrid_search: bool = False, top_k: int = 5, tags_all: Optional[list[str]] = None, tags_any: Optional[list[str]] = None, @@ -681,7 +680,6 @@ async def search_cosmos( role=role, memory_types=memory_types, thread_id=thread_id, - hybrid_search=hybrid_search, top_k=top_k, tags_all=tags_all, tags_any=tags_any, @@ -699,7 +697,6 @@ async def search_turns( user_id: str, thread_id: Optional[str] = None, role: Optional[str] = None, - hybrid_search: bool = False, top_k: int = 5, tags_all: Optional[list[str]] = None, tags_any: Optional[list[str]] = None, @@ -721,7 +718,6 @@ async def search_turns( user_id=user_id, thread_id=thread_id, role=role, - hybrid_search=hybrid_search, top_k=top_k, tags_all=tags_all, tags_any=tags_any, @@ -836,12 +832,23 @@ async def search_episodic_memories( min_salience: Optional[float] = None, include_superseded: bool = False, ) -> list[dict[str, Any]]: - return await self._get_store().search_episodic(user_id, search_terms, top_k, min_salience, include_superseded) + return await self._get_store().search_episodic( + user_id, + search_terms, + top_k, + min_salience, + include_superseded, + ) async def build_procedural_context(self, user_id: str) -> str: return await self._get_pipeline().build_procedural_context(user_id) - async def build_episodic_context(self, user_id: str, query: str, top_k: int = 3) -> str: + async def build_episodic_context( + self, + user_id: str, + query: str, + top_k: int = 3, + ) -> str: return await self._get_store().build_episodic_context(user_id, query, top_k) async def extract_memories(self, user_id: str, thread_id: str, recent_k: Optional[int] = None) -> dict[str, int]: @@ -879,7 +886,9 @@ async def generate_user_summary( async def reconcile(self, user_id: str, n: Optional[int] = None) -> dict[str, int]: from azure.cosmos.agent_memory.thresholds import get_dedup_pool_size - return await self._get_pipeline().reconcile_memories(user_id, n if n is not None else get_dedup_pool_size()) + return await self._get_pipeline().reconcile_memories( + user_id, n if n is not None else get_dedup_pool_size(), full_rebuild=True + ) async def process_now(self, *, user_id: str, thread_id: str) -> "ProcessThreadResult": """Force the processor to run the full pipeline RIGHT NOW for one thread. diff --git a/azure/cosmos/agent_memory/aio/processors/base.py b/azure/cosmos/agent_memory/aio/processors/base.py index cdab0c0..421a5e5 100644 --- a/azure/cosmos/agent_memory/aio/processors/base.py +++ b/azure/cosmos/agent_memory/aio/processors/base.py @@ -33,6 +33,7 @@ async def process_extract_memories( *, user_id: str, thread_id: str, + recent_k: Optional[int] = None, ) -> dict[str, int]: ... async def process_thread_summary( @@ -53,6 +54,7 @@ async def process_reconcile( self, *, user_id: str, + full_rebuild: bool = False, ) -> int: ... async def generate_user_summary( diff --git a/azure/cosmos/agent_memory/aio/processors/durable.py b/azure/cosmos/agent_memory/aio/processors/durable.py index 677855d..59fef03 100644 --- a/azure/cosmos/agent_memory/aio/processors/durable.py +++ b/azure/cosmos/agent_memory/aio/processors/durable.py @@ -40,6 +40,7 @@ async def process_extract_memories( *, user_id: str, thread_id: str, + recent_k: Optional[int] = None, ) -> dict[str, int]: logger.debug( "AsyncDurableFunctionProcessor.process_extract_memories no-op user_id=%s thread_id=%s", @@ -73,7 +74,7 @@ async def process_user_summary( ) return UserSummaryResult(summary=None) - async def process_reconcile(self, *, user_id: str) -> int: + async def process_reconcile(self, *, user_id: str, full_rebuild: bool = False) -> int: logger.debug( "AsyncDurableFunctionProcessor.process_reconcile no-op user_id=%s", user_id, diff --git a/azure/cosmos/agent_memory/aio/processors/inprocess.py b/azure/cosmos/agent_memory/aio/processors/inprocess.py index 9e10071..62cac70 100644 --- a/azure/cosmos/agent_memory/aio/processors/inprocess.py +++ b/azure/cosmos/agent_memory/aio/processors/inprocess.py @@ -70,9 +70,7 @@ async def process_thread( thread_summary = await self._pipeline.generate_thread_summary(user_id, thread_id) extracted = await self._pipeline.extract_memories(user_id, thread_id) - reconciled = await self._pipeline.reconcile_memories(user_id, get_dedup_pool_size()) - - deduped_count = self._extract_reconcile_count(reconciled) + deduped_count = await self._reconcile_fact_and_episodic(user_id, get_dedup_pool_size()) extracted_counts: dict[str, int] = ( {k: v for k, v in extracted.items() if isinstance(v, int)} if isinstance(extracted, dict) else {} @@ -91,8 +89,9 @@ async def process_extract_memories( *, user_id: str, thread_id: str, + recent_k: Optional[int] = None, ) -> dict[str, int]: - extracted = await self._pipeline.extract_memories(user_id, thread_id) + extracted = await self._pipeline.extract_memories(user_id, thread_id, recent_k=recent_k) return {k: v for k, v in extracted.items() if isinstance(v, int)} if isinstance(extracted, dict) else {} async def process_thread_summary( @@ -113,11 +112,26 @@ async def process_user_summary( summary = await self._pipeline.generate_user_summary(user_id, thread_ids) return UserSummaryResult(summary=summary if isinstance(summary, dict) else None) - async def process_reconcile(self, *, user_id: str) -> int: + async def process_reconcile(self, *, user_id: str, full_rebuild: bool = False) -> int: from ...thresholds import get_dedup_pool_size - reconciled = await self._pipeline.reconcile_memories(user_id, get_dedup_pool_size()) - return self._extract_reconcile_count(reconciled) + return await self._reconcile_fact_and_episodic(user_id, get_dedup_pool_size(), full_rebuild=full_rebuild) + + async def _reconcile_fact_and_episodic(self, user_id: str, n: int, *, full_rebuild: bool = False) -> int: + """Reconcile facts and episodic memories; sum merged+contradicted counts. + + SDK in-process processing reconciles both types (matching the Durable + backend) so episodic dups don't accrue forever. ``full_rebuild`` (set by + the auto-trigger on its persisted-counter full-recluster cadence) forces + the full-pool LLM pass that catches dissimilar-embedding contradictions. + """ + total = 0 + for memory_type in ("fact", "episodic"): + reconciled = await self._pipeline.reconcile_memories( + user_id, n=n, memory_type=memory_type, full_rebuild=full_rebuild + ) + total += self._extract_reconcile_count(reconciled) + return total @staticmethod def _extract_reconcile_count(reconciled: Any) -> int: diff --git a/azure/cosmos/agent_memory/aio/services/pipeline.py b/azure/cosmos/agent_memory/aio/services/pipeline.py index 0db0475..4e5de1e 100644 --- a/azure/cosmos/agent_memory/aio/services/pipeline.py +++ b/azure/cosmos/agent_memory/aio/services/pipeline.py @@ -10,7 +10,6 @@ from __future__ import annotations -import asyncio import hashlib import inspect import json @@ -20,13 +19,19 @@ from typing import Any, Iterable, Literal, Optional from azure.cosmos.exceptions import ( - CosmosHttpResponseError, CosmosResourceExistsError, CosmosResourceNotFoundError, ) from azure.cosmos.agent_memory._container_routing import ContainerKey -from azure.cosmos.agent_memory._utils import DEFAULT_TTL_BY_TYPE, compute_content_hash +from azure.cosmos.agent_memory._utils import ( + DEFAULT_TTL_BY_TYPE, + compute_content_hash, + distance_function_from_container_properties, + vector_autodrop_supported, + vector_order_direction, + vector_similarity_at_least, +) from azure.cosmos.agent_memory.aio.store import AsyncMemoryStore from azure.cosmos.agent_memory.exceptions import ( LLMError, @@ -58,9 +63,6 @@ coerce_valence, parse_llm_json, ) -from azure.cosmos.agent_memory.services._pipeline_helpers import ( - format_existing_episodics as _format_existing_episodics, -) from azure.cosmos.agent_memory.services._pipeline_helpers import ( is_real_number as _is_real_number, ) @@ -68,6 +70,16 @@ max_or_none as _max_or_none, ) from azure.cosmos.agent_memory.store._search_helpers import top_literal +from azure.cosmos.agent_memory.thresholds import ( + get_dedup_candidate_topk, + get_dedup_cluster_sim, + get_dedup_context_topk, + get_dedup_context_vector_enabled, + get_dedup_reconcile_mode, + get_dedup_sim_high, + get_dedup_sim_low, + get_dedup_vector_enabled, +) logger = get_logger("azure.cosmos.agent_memory.pipeline.aio") @@ -252,6 +264,106 @@ 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: + # Transient read failure is indistinguishable from "no policy" once we + # drop to None — so DON'T cache. An uncached cosine default self-heals on + # the next call; caching it would pin cosine and silently mis-handle a + # euclidean container (cosine bands on euclidean distances → data loss). + 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 + 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 ``DEDUP_SIM_HIGH`` thresholds are cosine-calibrated; on euclidean + the destructive auto-drop is skipped (borderline tagging + LLM reconcile + still run). 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. @@ -354,58 +466,6 @@ def _stable_source_timestamp(items: list[dict[str, Any]]) -> str: return max(timestamps) return datetime.now(timezone.utc).isoformat() - async def _mark_extracted_superseded( - self, - *, - user_id: str, - thread_id: str, - supersedes_id: str, - superseder_id: str, - reason: Literal["update", "contradict"], - ) -> bool: - try: - old_mem = await self._read_item( - self._memories_container, - item=supersedes_id, - partition_key=[user_id, thread_id], - ) - if old_mem.get("superseded_by"): - logger.debug( - "extract_memories: skipping UPDATE — target %s already superseded by %s", - supersedes_id, - old_mem.get("superseded_by"), - ) - return False - return await self._mark_superseded(old_mem, superseder_id, reason=reason) - except CosmosResourceNotFoundError: - logger.debug( - "extract_memories: %s not found at (user_id, thread_id) — retrying cross-partition", - supersedes_id, - ) - except Exception as exc: - # Includes 429s, 503s, transient connection errors — surface at WARNING - # so they're not masked by the silent cross-partition fallback below. - logger.warning( - "extract_memories: read_item failed for %s (%s); retrying cross-partition", - supersedes_id, - type(exc).__name__, - ) - try: - q = f"SELECT * FROM c WHERE c.id = @id AND c.user_id = @uid AND {_ACTIVE_DOC_FILTER}" - results = await self._query_items( - self._memories_container, - query=q, - parameters=[ - {"name": "@id", "value": supersedes_id}, - {"name": "@uid", "value": user_id}, - ], - ) - if results and not results[0].get("superseded_by"): - return await self._mark_superseded(results[0], superseder_id, reason=reason) - except CosmosHttpResponseError as exc: - logger.warning("Failed to mark superseded memory %s: %s", supersedes_id, exc) - return False - async def _mark_superseded( self, old_doc: dict[str, Any], @@ -464,30 +524,35 @@ async def extract_memories_dry( logger.warning("extract_memories_dry no memories found user_id=%s thread_id=%s", user_id, thread_id) return {"facts": [], "episodic": [], "updates": [], "processed_turn_docs": []} - existing_facts, existing_episodics = await asyncio.gather( - self._load_existing_memories(user_id, ["fact"]), - self._load_existing_memories(user_id, ["episodic"]), - ) + transcript = self._build_transcript(items) + existing_for_hash = await self._load_existing_memories(user_id, ["fact"]) existing_fact_hashes: set[str] = { - m["content_hash"] for m in existing_facts if m.get("type") == "fact" and m.get("content_hash") + m["content_hash"] for m in existing_for_hash if m.get("type") == "fact" and m.get("content_hash") } - if existing_facts: + if get_dedup_context_vector_enabled(): + user_turns_text = "\n".join( + str(it.get("content", "")) for it in items if it.get("role") == "user" + ).strip() + context_query = user_turns_text or transcript + existing = await self._store.search( + search_terms=context_query, + user_id=user_id, + memory_types=["fact"], + top_k=get_dedup_context_topk(), + ) + else: + existing = existing_for_hash + if existing: existing_text = "\n".join( - f"- [ID: {mem['id']}] {mem.get('content', '')} (type=fact, salience={mem.get('salience', 'N/A')})" - for mem in existing_facts + f"- [ID: {mem['id']}] {mem.get('content', '')} " + f"(type={mem.get('type', 'fact')}, salience={mem.get('salience', 'N/A')})" + for mem in existing ) else: existing_text = "(none)" - existing_episodics_text = _format_existing_episodics(existing_episodics) - - transcript = self._build_transcript(items) response_text = await self._run_prompty( "extract_memories.prompty", - inputs={ - "existing_facts": existing_text, - "existing_episodics": existing_episodics_text, - "transcript": transcript, - }, + inputs={"existing_facts": existing_text, "transcript": transcript}, ) parsed = self._parse_llm_json(response_text) facts = parsed.get("facts", []) @@ -502,14 +567,13 @@ async def extract_memories_dry( dropped_episodic_count = 0 for fact in facts: - action = fact.get("action", "ADD").upper() text = fact.get("text") if not text: logger.warning("extract_memories: dropping malformed fact (missing 'text'): %r", fact) continue new_content_hash = compute_content_hash(text) - if action == "ADD" and new_content_hash in existing_fact_hashes: + if new_content_hash in existing_fact_hashes: logger.debug( "extract_memories: skipping exact-dup fact hash=%s user_id=%s thread_id=%s", new_content_hash, @@ -546,22 +610,6 @@ async def extract_memories_dry( "updated_at": doc_timestamp, } - if action in {"UPDATE", "CONTRADICT"} and fact.get("supersedes_id"): - reason: Literal["update", "contradict"] = "contradict" if action == "CONTRADICT" else "update" - if det_id == fact["supersedes_id"]: - logger.debug("extract_memories: skipping UPDATE — det_id == supersedes_id (%s)", det_id) - continue - doc["supersedes_ids"] = [fact["supersedes_id"]] - updates.append( - { - "op": "supersede", - "supersedes_id": fact["supersedes_id"], - "superseder_id": det_id, - "thread_id": thread_id, - "reason": reason, - } - ) - fact_docs.append(self._validate_extracted_doc(doc)) existing_fact_hashes.add(new_content_hash) @@ -581,28 +629,16 @@ async def extract_memories_dry( dropped_episodic_count += 1 continue - text_raw = ep.get("text") - text = text_raw.strip() if isinstance(text_raw, str) else None - if not text: - logger.error( - "extract_memories: dropping episodic with empty/missing text field " - "(LLM extraction did not populate the required `text` field — likely a " - "weaker extraction model that needs upgrading or a prompt-compliance issue). " - "scope_type=%s scope_value=%s user_id=%s thread_id=%s reason=missing_text", - scope_type, - scope_value, - user_id, - thread_id, - ) - dropped_episodic_count += 1 - continue - situation = ep.get("situation") action_taken = ep.get("action_taken") outcome = ep.get("outcome") + if situation and action_taken and outcome: + text = f"{situation} → {action_taken} → {outcome}" + else: + text = f"For the user's {scope_value} {scope_type}, intent recorded." content_hash = compute_content_hash(text) - seed = _ID_SEED_SEP.join((user_id, scope_type, scope_value)) + seed = _ID_SEED_SEP.join((user_id, thread_id, content_hash)) det_id = f"ep_{hashlib.sha256(seed.encode()).hexdigest()[:32]}" topic_tags = build_topic_tags(ep.get("tags", [])) confidence = ep.get("confidence") @@ -619,7 +655,7 @@ async def extract_memories_dry( doc = { "id": det_id, "user_id": user_id, - "thread_id": "__episodic__", + "thread_id": thread_id, "role": "system", "type": "episodic", "content": text, @@ -630,7 +666,6 @@ async def extract_memories_dry( "metadata": { "scope_type": scope_type, "scope_value": scope_value, - "originating_thread_id": thread_id, "situation": situation, "action_taken": action_taken, "outcome": outcome, @@ -694,7 +729,7 @@ async def extract_memories_dry( check_extracted_fact_grounding( fact_docs, items, - existing_facts, + existing, user_id=user_id, thread_id=thread_id, logger=logger, @@ -716,6 +751,100 @@ async def extract_memories_dry( ) return result + async def dedup_extracted_memories(self, user_id: str, extracted: dict) -> dict: + """Apply gated vector-floor deduplication to extracted facts/episodes.""" + if not get_dedup_vector_enabled(): + return extracted + if not user_id: + raise ValidationError("user_id is required") + if not isinstance(extracted, dict): + raise ValidationError("extracted must be a dict") + + high = get_dedup_sim_high() + low = get_dedup_sim_low() + top_k = get_dedup_candidate_topk() + distance_function = await self._vector_distance_function() + autodrop_ok = vector_autodrop_supported(distance_function) + if not autodrop_ok: + self._warn_euclidean_autodrop_once(distance_function) + result = { + "facts": [dict(doc) for doc in extracted.get("facts", [])], + "episodic": [dict(doc) for doc in extracted.get("episodic", [])], + "updates": [dict(op) for op in extracted.get("updates", [])], + } + docs = [doc for doc in result["facts"] + result["episodic"] if doc.get("content")] + missing_embeddings = [doc for doc in docs if not doc.get("embedding")] + if missing_embeddings: + embeddings = await self._embed_batch([str(doc["content"]) for doc in missing_embeddings]) + for doc, embedding in zip(missing_embeddings, embeddings): + doc["embedding"] = embedding + + vector_dedup_skipped = 0 + dup_candidates_tagged = 0 + kept_ids: set[str] = set() + dropped_ids: set[str] = set() + filtered_by_key: dict[str, list[dict[str, Any]]] = {"facts": [], "episodic": []} + for key in ("facts", "episodic"): + for doc in result[key]: + if not doc.get("content"): + filtered_by_key[key].append(doc) + continue + doc_id = str(doc.get("id") or "") + memory_type = str(doc.get("type") or "") + if not doc_id or memory_type not in {"fact", "episodic"}: + # Parity with sync: under-specified docs (no id / unknown type) + # skip dedup and pass through verbatim. + filtered_by_key[key].append(doc) + continue + exclude_ids = kept_ids | dropped_ids | {doc_id, *(doc.get("supersedes_ids") or [])} + candidates = await self._vector_candidates( + user_id=user_id, + embedding=doc.get("embedding"), + memory_type=memory_type, + top_k=top_k, + exclude_ids=exclude_ids, + ) + best: dict[str, Any] | None = candidates[0] if candidates else None + score = float(best.get("score") or 0.0) if best else 0.0 + if best and autodrop_ok and vector_similarity_at_least(score, high, distance_function): + vector_dedup_skipped += 1 + dropped_ids.add(doc_id) + logger.info( + "dedup_extracted_memories: vector skip user_id=%s dropped=%r " + "surviving_id=%s surviving=%r score=%.4f", + user_id, + doc.get("content"), + best.get("id"), + best.get("content"), + score, + ) + continue + if best and vector_similarity_at_least(score, low, distance_function): + tags = list(doc.get("tags") or []) + if "sys:dup-candidate" not in tags: + tags.append("sys:dup-candidate") + doc["tags"] = tags + metadata = dict(doc.get("metadata") or {}) + metadata["dup_of"] = best.get("id") + metadata["dup_score"] = score + doc["metadata"] = metadata + dup_candidates_tagged += 1 + + kept_ids.add(doc_id) + filtered_by_key[key].append(doc) + + if vector_dedup_skipped or dup_candidates_tagged: + result["updates"].append( + { + "op": "stats", + "vector_dedup_skipped": vector_dedup_skipped, + "dup_candidates_tagged": dup_candidates_tagged, + } + ) + result["facts"] = filtered_by_key["facts"] + result["episodic"] = filtered_by_key["episodic"] + return result + async def persist_extracted_memories( self, user_id: str, @@ -763,27 +892,14 @@ async def persist_extracted_memories( if op.get("op") == "stats": result["exact_dedup_skipped"] += int(op.get("exact_dedup_skipped") or 0) result["dropped_episodic_count"] += int(op.get("dropped_episodic_count") or 0) - continue - if op.get("op") != "supersede": - continue - reason = op.get("reason") - op_thread_id = op.get("thread_id") - supersedes_id = op.get("supersedes_id") - superseder_id = op.get("superseder_id") - if reason not in {"update", "contradict"} or not op_thread_id or not supersedes_id or not superseder_id: - continue - marked = await self._mark_extracted_superseded( - user_id=user_id, - thread_id=op_thread_id, - supersedes_id=supersedes_id, - superseder_id=superseder_id, - reason=reason, - ) - if marked: - if reason == "contradict": - result["contradicted_count"] += 1 - else: - result["updated_count"] += 1 + if "vector_dedup_skipped" in op: + result["vector_dedup_skipped"] = result.get("vector_dedup_skipped", 0) + int( + op.get("vector_dedup_skipped") or 0 + ) + if "dup_candidates_tagged" in op: + result["dup_candidates_tagged"] = result.get("dup_candidates_tagged", 0) + int( + op.get("dup_candidates_tagged") or 0 + ) logger.info("persist_extracted_memories completed user_id=%s counts=%s", user_id, result) @@ -834,6 +950,8 @@ async def extract_memories( ) -> dict[str, int]: """Extract facts and episodic memories from a thread and persist them.""" extracted = await self.extract_memories_dry(user_id, thread_id, recent_k, turns=turns) + if get_dedup_vector_enabled(): + extracted = await self.dedup_extracted_memories(user_id, extracted) return await self.persist_extracted_memories(user_id, extracted) async def synthesize_procedural( @@ -1031,7 +1149,7 @@ def _render_bullets(values: list[str]) -> str: } validated = construct_internal(ProceduralRecord, new_doc).to_doc() try: - await self._create_item(self._memories_container, body=validated) + await self._create_item(self._memories_container, body=dict(validated)) written_doc = validated break except CosmosResourceExistsError: @@ -1392,7 +1510,228 @@ def _emit_reconcile_outcome( }, ) - async def reconcile_memories(self, user_id: str, n: int = 50) -> dict[str, int]: + async def _active_memories_for_reconcile(self, user_id: str, memory_type: str, n: int) -> list[dict[str, Any]]: + capped_n = top_literal(n, name="reconcile_memories.n") + query = ( + f"SELECT TOP {capped_n} * FROM c " + "WHERE c.user_id = @user_id " + "AND c.type = @memory_type " + f"AND {_ACTIVE_DOC_FILTER} " + "ORDER BY c.created_at DESC" + ) + return await self._query_items( + self._memories_container, + query=query, + parameters=[ + {"name": "@user_id", "value": user_id}, + {"name": "@memory_type", "value": memory_type}, + ], + ) + + async def _load_memories_by_ids( + self, + user_id: str, + memory_type: str, + ids: Iterable[str], + ) -> list[dict[str, Any]]: + id_list = [mid for mid in dict.fromkeys(ids) if mid] + if not id_list: + return [] + placeholders = ", ".join(f"@id{i}" for i in range(len(id_list))) + query = ( + "SELECT * FROM c WHERE c.user_id = @user_id " + "AND c.type = @memory_type " + f"AND c.id IN ({placeholders}) " + f"AND {_ACTIVE_DOC_FILTER}" + ) + parameters = [ + {"name": "@user_id", "value": user_id}, + {"name": "@memory_type", "value": memory_type}, + ] + parameters.extend({"name": f"@id{i}", "value": mid} for i, mid in enumerate(id_list)) + return await self._query_items(self._memories_container, query=query, parameters=parameters) + + async def _build_candidate_clusters( + self, + user_id: str, + memory_type: str, + n: int, + ) -> tuple[list[list[dict[str, Any]]], int, list[dict[str, Any]]]: + """Cluster dup-candidate seeds (+ vector neighbors) into connected components. + + Returns ``(clusters, node_count, seeds)`` where each cluster has >= 2 members, + ``node_count`` is the total distinct memories pulled into the graph (used to + report ``reconcile_llm_calls_saved``), and ``seeds`` is the tagged seed scan + (so the caller can clear stale tags on orphan seeds that never clustered). + The seed scan is bounded to ``n`` so a single cluster can never exceed the + reconcile prompt's pool cap. + """ + cluster_sim = get_dedup_cluster_sim() + top_k = get_dedup_candidate_topk() + distance_function = await self._vector_distance_function() + query = ( + f"SELECT TOP {top_literal(n, name='reconcile_memories.candidate_n')} * FROM c " + "WHERE c.user_id = @user_id " + "AND c.type = @memory_type " + "AND ARRAY_CONTAINS(c.tags, @tag) " + f"AND {_ACTIVE_DOC_FILTER} " + "ORDER BY c.created_at DESC" + ) + seeds = await self._query_items( + self._memories_container, + query=query, + parameters=[ + {"name": "@user_id", "value": user_id}, + {"name": "@memory_type", "value": memory_type}, + {"name": "@tag", "value": "sys:dup-candidate"}, + ], + ) + nodes_by_id: dict[str, dict[str, Any]] = {doc["id"]: doc for doc in seeds if doc.get("id")} + edges: set[tuple[str, str]] = set() + for seed in seeds: + sid = seed.get("id") + if not sid: + continue + dup_of = (seed.get("metadata") or {}).get("dup_of") if isinstance(seed.get("metadata"), dict) else None + if dup_of: + for doc in await self._load_memories_by_ids(user_id, memory_type, [dup_of]): + nodes_by_id[doc["id"]] = doc + edges.add(tuple(sorted((sid, doc["id"])))) + for cand in await self._vector_candidates( + user_id=user_id, + embedding=seed.get("embedding"), + memory_type=memory_type, + top_k=top_k, + exclude_ids={sid}, + ): + if vector_similarity_at_least(float(cand.get("score") or 0.0), cluster_sim, distance_function): + for doc in await self._load_memories_by_ids(user_id, memory_type, [cand.get("id")]): + nodes_by_id[doc["id"]] = doc + edges.add(tuple(sorted((sid, doc["id"])))) + node_ids = set(nodes_by_id) + for doc in list(nodes_by_id.values()): + did = doc.get("id") + if not did: + continue + for cand in await self._vector_candidates( + user_id=user_id, + embedding=doc.get("embedding"), + memory_type=memory_type, + top_k=top_k, + exclude_ids={did}, + ): + cid = cand.get("id") + if cid in node_ids and vector_similarity_at_least( + float(cand.get("score") or 0.0), cluster_sim, distance_function + ): + edges.add(tuple(sorted((did, cid)))) + + adjacency: dict[str, set[str]] = {node_id: set() for node_id in nodes_by_id} + for left, right in edges: + if left != right and left in adjacency and right in adjacency: + adjacency[left].add(right) + adjacency[right].add(left) + clusters: list[list[dict[str, Any]]] = [] + seen: set[str] = set() + for node_id in adjacency: + if node_id in seen: + continue + stack = [node_id] + component: list[str] = [] + seen.add(node_id) + while stack: + current = stack.pop() + component.append(current) + for nxt in adjacency[current]: + if nxt not in seen: + seen.add(nxt) + stack.append(nxt) + if len(component) >= 2: + # Cap cluster size at the reconcile pool limit: lowering the cluster + # threshold can chain many facts into one giant transitive component + # that would blow the prompt cap; keep the most-recent ``n``. + if len(component) > n: + component = component[:n] + clusters.append([nodes_by_id[cid] for cid in component]) + return clusters, len(nodes_by_id), seeds + + async def _reconcile_candidate_mode( + self, user_id: str, *, n: int, memory_type: str, started_at: float + ) -> dict[str, int]: + # Candidate clustering only. The periodic full-pool backstop that catches + # dissimilar-embedding contradictions ("vegetarian" vs "loves steak") is + # driven by the caller via ``full_rebuild`` on a PERSISTED-counter cadence + # (in-process auto-trigger + durable change-feed), not an in-memory sweep + # counter — the latter reset per worker/process and never fired reliably on + # the Function-App backend. + clusters, node_count, seeds = await self._build_candidate_clusters(user_id, memory_type, n) + aggregate = {"kept": 0, "merged": 0, "contradicted": 0} + clustered_ids: set[str] = set() + for cluster in clusters: + # Mark members as clustered BEFORE the LLM call so a failed cluster's + # seeds are not treated as orphans below (which would clear their tags + # and prevent a retry). A truncated/malformed LLM response on one + # cluster must not abort the sweep or starve the remaining clusters. + clustered_ids.update(doc["id"] for doc in cluster if doc.get("id")) + try: + counts, consumed = await self._reconcile_pool(user_id, memory_type, cluster) + except Exception as exc: + logger.warning( + "reconcile_memories: cluster reconcile failed user_id=%s memory_type=%s; " + "skipping cluster, tags retained for next sweep: %s", + user_id, + memory_type, + exc, + ) + continue + for key in aggregate: + aggregate[key] += int(counts.get(key, 0)) + # Clear dup-candidate tags only on survivors. Re-upserting a doc that + # was just superseded (duplicate source or contradiction loser) would + # resurrect it, since the in-memory cluster copy lacks superseded_by. + survivors = [doc for doc in cluster if doc.get("id") and doc["id"] not in consumed] + await self._clear_dup_candidate_tags(survivors) + # Orphan seeds: tagged dup-candidates that never joined a cluster have no + # near-duplicate, so clear the stale tag — otherwise every future sweep + # re-scans them as seeds and they accumulate forever. + orphan_seeds = [seed for seed in seeds if seed.get("id") and seed["id"] not in clustered_ids] + await self._clear_dup_candidate_tags(orphan_seeds) + aggregate["reconcile_clusters_sent"] = len(clusters) + aggregate["reconcile_llm_calls_saved"] = max(0, node_count - len(clusters)) + logger.info( + "reconcile_memories candidate completed user_id=%s memory_type=%s result=%s", + user_id, + memory_type, + aggregate, + ) + self._emit_reconcile_outcome( + started_at=started_at, + user_id=user_id, + candidates=node_count, + result=aggregate, + ) + return aggregate + + async def _clear_dup_candidate_tags(self, docs: Iterable[dict[str, Any]]) -> None: + for doc in docs: + tags = [tag for tag in (doc.get("tags") or []) if tag != "sys:dup-candidate"] + if tags == (doc.get("tags") or []): + continue + updated = dict(doc) + updated["tags"] = tags + metadata = dict(updated.get("metadata") or {}) + metadata.pop("dup_of", None) + metadata.pop("dup_score", None) + updated["metadata"] = metadata + updated["updated_at"] = datetime.now(timezone.utc).isoformat() + try: + await self._upsert_memory(updated) + except Exception: + logger.exception("reconcile_memories: failed to clear dup-candidate tag id=%s", doc.get("id")) + + async def reconcile_memories( + self, user_id: str, n: int = 50, *, memory_type: str = "fact", full_rebuild: bool = False + ) -> dict[str, int]: """Reconcile a user's active facts in a single LLM pass. Loads the most recent ``n`` active (non-superseded) facts for @@ -1417,44 +1756,64 @@ async def reconcile_memories(self, user_id: str, n: int = 50) -> dict[str, int]: raise ValidationError(f"n must be a positive integer, got {n!r}") if n > 500: raise ValidationError(f"n must be <= 500 to bound prompt size and LLM cost, got {n}") + if memory_type not in {"fact", "episodic", "procedural"}: + raise ValidationError(f"memory_type must be one of fact, episodic, procedural, got {memory_type!r}") + if memory_type == "procedural": + result = { + "kept": 0, + "merged": 0, + "contradicted": 0, + "reconcile_clusters_sent": 0, + "reconcile_llm_calls_saved": 0, + } + logger.info("reconcile_memories procedural no-op user_id=%s result=%s", user_id, result) + return result started_at = time.monotonic() - logger.info("reconcile_memories started user_id=%s n=%d", user_id, n) + logger.info("reconcile_memories started user_id=%s n=%d memory_type=%s", user_id, n, memory_type) + + # Explicit user-triggered reconcile (full_rebuild) always takes the + # full-pool single-LLM-pass path: it sees every active fact together, so it + # catches contradictions that aren't vector-similar (e.g. "vegetarian" vs + # "loves steak") — which candidate clustering, keyed on near-duplicate + # similarity, would never group. Automatic sweeps use cheap candidate mode. + if get_dedup_reconcile_mode() == "candidate" and not full_rebuild: + return await self._reconcile_candidate_mode( + user_id, n=n, memory_type=memory_type, started_at=started_at + ) - # ---- 1. Load up to N most recent active facts ---- - # ORDER BY c.created_at DESC keeps the TOP cap deterministic across - # physical partitions and matches the dedup prompt's tiebreaker - # ("more recent created_at first"). Cosmos's _ts is the last-write - # timestamp, which would diverge from created_at after any UPDATE. - query = ( - f"SELECT TOP {n} * FROM c " - "WHERE c.user_id = @user_id " - "AND c.type = 'fact' " - f"AND {_ACTIVE_DOC_FILTER} " - "ORDER BY c.created_at DESC" - ) - parameters: list[dict[str, Any]] = [ - {"name": "@user_id", "value": user_id}, - ] - facts = await self._query_items( - self._memories_container, - query=query, - parameters=parameters, + facts = await self._active_memories_for_reconcile(user_id, memory_type, n) + result, consumed = await self._reconcile_pool(user_id, memory_type, facts) + # Clear dup-candidate tags on survivors so an explicit reconcile(full_rebuild=True) + # doesn't leave stale sys:dup-candidate/dup_of metadata on user-visible memories. + survivors = [doc for doc in facts if doc.get("id") and doc["id"] not in consumed] + await self._clear_dup_candidate_tags(survivors) + self._emit_reconcile_outcome( + started_at=started_at, + user_id=user_id, + candidates=len(facts), + result=result, ) + return result + async def _reconcile_pool( + self, user_id: str, memory_type: str, facts: list[dict[str, Any]] + ) -> tuple[dict[str, int], set[str]]: + """Reconcile an explicit pool of same-type memories in one LLM pass. + + Returns ``({"kept", "merged", "contradicted"}, consumed_ids)`` where + ``consumed_ids`` are the source/loser ids that were actually superseded, + so callers can skip them when clearing dup-candidate tags (re-upserting a + superseded source would resurrect it). Does not emit telemetry — the + caller owns the ``reconcile.outcome`` line. + """ if len(facts) <= 1: logger.info( - "reconcile_memories: %d facts, nothing to reconcile", + "reconcile_memories: %d %s memories, nothing to reconcile", len(facts), + memory_type, ) - early_result = {"kept": len(facts), "merged": 0, "contradicted": 0} - self._emit_reconcile_outcome( - started_at=started_at, - user_id=user_id, - candidates=len(facts), - result=early_result, - ) - return early_result + return {"kept": len(facts), "merged": 0, "contradicted": 0}, set() # ---- 2. Format the facts pool for the prompt ---- # ``json.dumps`` escapes embedded quotes and pipes inside content so @@ -1481,14 +1840,16 @@ async def reconcile_memories(self, user_id: str, n: int = 50) -> dict[str, int]: facts_text = "\n".join(lines) # ---- 3. Single LLM call over the entire pool ---- - response_text = await self._run_prompty( - "dedup.prompty", - inputs={"facts_text": facts_text}, - ) + # Polarity keyed on an explicit predicate so a future third type (e.g. + # procedural, were it ever routed here) can't silently diverge from sync. + is_episodic = memory_type == "episodic" + prompt_name = "dedup_episodic.prompty" if is_episodic else "dedup.prompty" + prompt_inputs = {"episodics_text": facts_text} if is_episodic else {"facts_text": facts_text} + response_text = await self._run_prompty(prompt_name, inputs=prompt_inputs) parsed = self._parse_llm_json(response_text) duplicate_groups = parsed.get("duplicate_groups", []) or [] - contradicted_pairs = parsed.get("contradicted_pairs", []) or [] + contradicted_pairs = [] if is_episodic else (parsed.get("contradicted_pairs", []) or []) # ``kept_ids`` from the LLM is used below as a cross-check for # accounting drift (hallucinated IDs, double-counting). The actual # kept count is computed from facts minus consumed losers. @@ -1564,11 +1925,13 @@ async def reconcile_memories(self, user_id: str, n: int = 50) -> dict[str, int]: seen_tags: set[str] = set() for src in source_docs: for t in src.get("tags", []) or []: + if t == "sys:dup-candidate": + continue if t not in seen_tags: seen_tags.add(t) merged_tags.append(t) if not merged_tags: - merged_tags = ["sys:fact"] + merged_tags = [f"sys:{memory_type}"] # Union source_memory_ids across all source docs (provenance chain). merged_source_memory_ids: list[str] = [] @@ -1625,16 +1988,39 @@ async def reconcile_memories(self, user_id: str, n: int = 50) -> dict[str, int]: # see the same id rather than chaining through a new UUID. merged_content_hash = compute_content_hash(merged_content) merged_id_seed = _ID_SEED_SEP.join((user_id, "merged", merged_content_hash)) - merged_id = "fact_" + hashlib.sha256(merged_id_seed.encode()).hexdigest()[:32] - + merged_prefix = "ep_" if memory_type == "episodic" else "fact_" + merged_id = merged_prefix + hashlib.sha256(merged_id_seed.encode()).hexdigest()[:32] try: + if memory_type == "episodic": + base_metadata = dict(source_docs[0].get("metadata") or {}) + base_metadata.update( + { + "lesson": merged_content, + "merged_via": "reconcile", + "merged_from_count": len(valid_source_ids), + } + ) + base_metadata.setdefault("scope_type", base_metadata.get("scope_type") or "general") + base_metadata.setdefault("scope_value", base_metadata.get("scope_value") or "general") + base_metadata.setdefault("outcome_valence", base_metadata.get("outcome_valence") or "neutral") + record_cls = EpisodicRecord + prompt_lineage = self._prompt_lineage("dedup_episodic.prompty") + metadata = base_metadata + else: + record_cls = FactRecord + prompt_lineage = self._prompt_lineage("dedup.prompty") + metadata = { + "category": "preference", + "merged_via": "reconcile", + "merged_from_count": len(valid_source_ids), + } merged_record = construct_internal( - FactRecord, + record_cls, { "id": merged_id, "user_id": user_id, "role": "system", - "type": "fact", + "type": memory_type, "content": merged_content, "thread_id": recent_thread_id or f"__reconciled__:{user_id}", "confidence": confidence_val if confidence_val is not None else 0.5, @@ -1643,12 +2029,8 @@ async def reconcile_memories(self, user_id: str, n: int = 50) -> dict[str, int]: "source_memory_ids": merged_source_memory_ids, "tags": merged_tags, "content_hash": merged_content_hash, - "metadata": { - "category": "preference", - "merged_via": "reconcile", - "merged_from_count": len(valid_source_ids), - }, - **self._prompt_lineage("dedup.prompty"), + "metadata": metadata, + **prompt_lineage, "created_at": datetime.now(timezone.utc).isoformat(), "updated_at": datetime.now(timezone.utc).isoformat(), }, @@ -1820,13 +2202,7 @@ async def reconcile_memories(self, user_id: str, n: int = 50) -> dict[str, int]: ) result = {"kept": kept, "merged": merged, "contradicted": contradicted} logger.info("reconcile_memories completed: %s", result) - self._emit_reconcile_outcome( - started_at=started_at, - user_id=user_id, - candidates=len(facts), - result=result, - ) - return result + return result, consumed_ids async def build_procedural_context(self, user_id: str) -> str: """Return the active synthesized procedural prompt for system injection.""" diff --git a/azure/cosmos/agent_memory/aio/store/memory_store.py b/azure/cosmos/agent_memory/aio/store/memory_store.py index 4f725f9..b7aec2e 100644 --- a/azure/cosmos/agent_memory/aio/store/memory_store.py +++ b/azure/cosmos/agent_memory/aio/store/memory_store.py @@ -17,8 +17,8 @@ from azure.cosmos.agent_memory._utils import ( _build_memory_query_builder, _coerce_datetime_iso, - _validate_hybrid_search, compute_content_hash, + extract_keywords, new_id, ) from azure.cosmos.agent_memory.exceptions import ( @@ -804,7 +804,6 @@ async def search( role: Optional[str] = None, memory_types: Optional[list[str]] = None, thread_id: Optional[str] = None, - hybrid_search: bool = False, top_k: int = 5, tags_all: Optional[list[str]] = None, tags_any: Optional[list[str]] = None, @@ -823,9 +822,9 @@ async def search( :meth:`search_turns` to vector-search the raw conversation log instead. """ terms = require_search_terms(search_terms, query) - _validate_hybrid_search(hybrid_search, terms) top = top_literal(top_k, name="top_k") query_vector = await self._embed(terms) + keywords = extract_keywords(terms) qb = _build_memory_query_builder( memory_id=memory_id, @@ -845,11 +844,16 @@ async def search( ) add_salience_filter(qb, min_salience) - sql = build_search_sql(qb=qb, top=top, hybrid_search=hybrid_search, include_superseded=include_superseded) + sql = build_search_sql( + qb=qb, + top=top, + keyword_count=len(keywords), + include_superseded=include_superseded, + ) parameters = qb.get_parameters() parameters.append({"name": "@embedding", "value": query_vector}) - if hybrid_search: - parameters.append({"name": "@key_terms", "value": terms}) + for i, kw in enumerate(keywords): + parameters.append({"name": f"@kw{i}", "value": kw}) partition_key, _ = query_scope(user_id, thread_id) if thread_id is not None and (not memory_types or set(memory_types) & USER_SCOPED_MEMORIES_TYPES): @@ -868,7 +872,6 @@ async def search_turns( user_id: Optional[str] = None, thread_id: Optional[str] = None, role: Optional[str] = None, - hybrid_search: bool = False, top_k: int = 5, tags_all: Optional[list[str]] = None, tags_any: Optional[list[str]] = None, @@ -878,7 +881,7 @@ async def search_turns( *, query: Optional[str] = None, ) -> list[dict[str, Any]]: - """Search raw conversation turns using vector similarity with optional hybrid ranking. + """Search raw conversation turns using vector similarity with hybrid ranking. Only vector-searchable when turn embeddings were enabled at write time (see ``enable_turn_embeddings``). ``user_id`` is required and always @@ -890,9 +893,9 @@ async def search_turns( if not user_id: raise ValidationError("user_id is required for search_turns") terms = require_search_terms(search_terms, query) - _validate_hybrid_search(hybrid_search, terms) top = top_literal(top_k, name="top_k") query_vector = await self._embed(terms) + keywords = extract_keywords(terms) qb = _QueryBuilder() qb.add_filter("c.user_id", "@user_id", user_id) @@ -907,11 +910,11 @@ async def search_turns( before_param="@created_before", ) - sql = build_search_sql(qb=qb, top=top, hybrid_search=hybrid_search, include_superseded=False) + sql = build_search_sql(qb=qb, top=top, keyword_count=len(keywords), include_superseded=False) parameters = qb.get_parameters() parameters.append({"name": "@embedding", "value": query_vector}) - if hybrid_search: - parameters.append({"name": "@key_terms", "value": terms}) + for i, kw in enumerate(keywords): + parameters.append({"name": f"@kw{i}", "value": kw}) partition_key, _ = query_scope(user_id, thread_id) logger.debug("AsyncMemoryStore.search_turns query: %s", sql) diff --git a/azure/cosmos/agent_memory/auto_trigger.py b/azure/cosmos/agent_memory/auto_trigger.py index 10ef595..a494181 100644 --- a/azure/cosmos/agent_memory/auto_trigger.py +++ b/azure/cosmos/agent_memory/auto_trigger.py @@ -108,6 +108,12 @@ def maybe_trigger_steps( return n_dedup_turns = n_facts * n_dedup if n_facts > 0 and n_dedup > 0 else 0 + # Full-pool reconcile backstop cadence, derived from the PERSISTED counter so + # it fires reliably regardless of process/worker lifetime (the old in-memory + # per-instance sweep counter reset and under-fired). Every + # DEDUP_FULL_RECLUSTER_EVERY_N-th reconcile = every n_dedup_turns * N turns. + n_full_recluster = _threshold_int(thresholds, "get_dedup_full_recluster_every_n", "DEDUP_FULL_RECLUSTER_EVERY_N") + n_full_turns = n_dedup_turns * n_full_recluster if (n_dedup_turns > 0 and n_full_recluster > 0) else 0 user_batch_counts = _trigger_thread_steps( processor, counter_container, @@ -115,6 +121,7 @@ def maybe_trigger_steps( n_facts=n_facts, n_summary=n_summary, n_dedup_turns=n_dedup_turns, + n_full_turns=n_full_turns, thresholds=thresholds, ) _trigger_user_steps(processor, counter_container, user_batch_counts, n_user=n_user) @@ -128,6 +135,7 @@ def _trigger_thread_steps( n_facts: int, n_summary: int, n_dedup_turns: int, + n_full_turns: int, thresholds: Any = None, ) -> dict[str, int]: user_batch_counts: dict[str, int] = {} @@ -154,14 +162,43 @@ def _trigger_thread_steps( counter_id, user_id, thread_id, + new_count=new_count, fire_extract=n_facts > 0 and _counters.crosses_threshold(old_count, new_count, n_facts), fire_summary=n_summary > 0 and _counters.crosses_threshold(old_count, new_count, n_summary), fire_dedup=n_dedup_turns > 0 and _counters.crosses_threshold(old_count, new_count, n_dedup_turns), + fire_full_rebuild=n_full_turns > 0 and _counters.crosses_threshold(old_count, new_count, n_full_turns), thresholds=thresholds, ) return user_batch_counts +def _watermark_recent_k( + counter_container: Any, + counter_id: str, + user_id: str, + thread_id: str, + *, + new_count: int, +) -> int: + """recent_k covering every turn since the last successful extract. + + ``new_count - watermark`` is exactly the count of turns added since the last + successful extract, and the newest-``recent_k`` slice in the pipeline picks + precisely those turns — so the whole backlog is covered and the watermark can + safely advance to ``new_count``. **Not capped:** capping would extract only the + newest N and strand the oldest ``backlog - N`` turns. + + **Bootstrap:** before a thread's first successful extract there is no watermark, + so we treat it as ``0`` — ``recent_k = new_count`` covers every turn the thread + has so far. Using only the current batch size here would strand turns added + during earlier *failed* extracts, because the watermark still advances to + ``new_count`` on the first success. + """ + watermark = _counters.read_extract_watermark_sync(counter_container, counter_id, user_id, thread_id) + base = watermark if watermark is not None else 0 + return max(new_count - base, 1) + + def _fire_thread_steps( processor: InProcessProcessor, counter_container: Any, @@ -169,9 +206,11 @@ def _fire_thread_steps( user_id: str, thread_id: str, *, + new_count: int, fire_extract: bool, fire_summary: bool, fire_dedup: bool, + fire_full_rebuild: bool = False, thresholds: Any = None, ) -> None: fire_procedural = fire_dedup and bool( @@ -182,13 +221,28 @@ def _fire_thread_steps( default=True, ) ) + if fire_extract: + recent_k = _watermark_recent_k( + counter_container, + counter_id, + user_id, + thread_id, + new_count=new_count, + ) + try: + processor.process_extract_memories(user_id=user_id, thread_id=thread_id, recent_k=recent_k) + _counters.advance_extract_watermark_sync(counter_container, counter_id, user_id, thread_id, new_count) + except Exception as exc: + logger.warning("Auto-trigger process_extract_memories failed for %s/%s: %s", user_id, thread_id, exc) + _counters.stamp_failure_sync( + counter_container, counter_id, user_id, thread_id, f"process_extract_memories: {exc!r}" + ) for enabled, label, call in ( ( - fire_extract, - "process_extract_memories", - lambda: processor.process_extract_memories(user_id=user_id, thread_id=thread_id), + fire_dedup, + "process_reconcile", + lambda: processor.process_reconcile(user_id=user_id, full_rebuild=fire_full_rebuild), ), - (fire_dedup, "process_reconcile", lambda: processor.process_reconcile(user_id=user_id)), ( fire_procedural, "synthesize_procedural", diff --git a/azure/cosmos/agent_memory/cosmos_memory_client.py b/azure/cosmos/agent_memory/cosmos_memory_client.py index f4924c3..4bad143 100644 --- a/azure/cosmos/agent_memory/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/cosmos_memory_client.py @@ -621,7 +621,6 @@ def search_cosmos( role: Optional[str] = None, memory_types: Optional[list[str]] = None, thread_id: Optional[str] = None, - hybrid_search: bool = False, top_k: int = 5, tags_all: Optional[list[str]] = None, tags_any: Optional[list[str]] = None, @@ -644,7 +643,6 @@ def search_cosmos( role=role, memory_types=memory_types, thread_id=thread_id, - hybrid_search=hybrid_search, top_k=top_k, tags_all=tags_all, tags_any=tags_any, @@ -662,7 +660,6 @@ def search_turns( user_id: str, thread_id: Optional[str] = None, role: Optional[str] = None, - hybrid_search: bool = False, top_k: int = 5, tags_all: Optional[list[str]] = None, tags_any: Optional[list[str]] = None, @@ -684,7 +681,6 @@ def search_turns( user_id=user_id, thread_id=thread_id, role=role, - hybrid_search=hybrid_search, top_k=top_k, tags_all=tags_all, tags_any=tags_any, @@ -801,15 +797,30 @@ def search_episodic_memories( include_superseded: bool = False, ) -> list[dict[str, Any]]: """Semantic search across episodic memories for a user.""" - return self._get_store().search_episodic(user_id, search_terms, top_k, min_salience, include_superseded) + return self._get_store().search_episodic( + user_id=user_id, + search_terms=search_terms, + top_k=top_k, + min_salience=min_salience, + include_superseded=include_superseded, + ) def build_procedural_context(self, user_id: str) -> str: """Build formatted procedural context for prompt injection.""" return self._get_pipeline().build_procedural_context(user_id) - def build_episodic_context(self, user_id: str, query: str, top_k: int = 3) -> str: + def build_episodic_context( + self, + user_id: str, + query: str, + top_k: int = 3, + ) -> str: """Build formatted context of relevant past experiences.""" - return self._get_store().build_episodic_context(user_id, query, top_k) + return self._get_store().build_episodic_context( + user_id=user_id, + query=query, + top_k=top_k, + ) def extract_memories( self, @@ -855,7 +866,9 @@ def reconcile(self, user_id: str, n: Optional[int] = None) -> dict[str, int]: """Reconcile a user's facts via the contradiction-aware dedup pass.""" from .thresholds import get_dedup_pool_size - return self._get_pipeline().reconcile_memories(user_id, n if n is not None else get_dedup_pool_size()) + return self._get_pipeline().reconcile_memories( + user_id, n if n is not None else get_dedup_pool_size(), full_rebuild=True + ) def process_now(self, *, user_id: str, thread_id: str) -> "ProcessThreadResult": """Force the processor to run the full pipeline RIGHT NOW for one thread. diff --git a/azure/cosmos/agent_memory/processors/base.py b/azure/cosmos/agent_memory/processors/base.py index cd257b7..2fbf9b2 100644 --- a/azure/cosmos/agent_memory/processors/base.py +++ b/azure/cosmos/agent_memory/processors/base.py @@ -79,6 +79,7 @@ def process_extract_memories( *, user_id: str, thread_id: str, + recent_k: Optional[int] = None, ) -> dict[str, int]: ... def process_thread_summary( @@ -99,6 +100,7 @@ def process_reconcile( self, *, user_id: str, + full_rebuild: bool = False, ) -> int: ... def generate_user_summary( diff --git a/azure/cosmos/agent_memory/processors/durable.py b/azure/cosmos/agent_memory/processors/durable.py index d63b24f..9834dcf 100644 --- a/azure/cosmos/agent_memory/processors/durable.py +++ b/azure/cosmos/agent_memory/processors/durable.py @@ -46,6 +46,7 @@ def process_extract_memories( *, user_id: str, thread_id: str, + recent_k: Optional[int] = None, ) -> dict[str, int]: logger.debug( "DurableFunctionProcessor.process_extract_memories no-op user_id=%s thread_id=%s", @@ -79,7 +80,7 @@ def process_user_summary( ) return UserSummaryResult(summary=None) - def process_reconcile(self, *, user_id: str) -> int: + def process_reconcile(self, *, user_id: str, full_rebuild: bool = False) -> int: logger.debug( "DurableFunctionProcessor.process_reconcile no-op user_id=%s", user_id, diff --git a/azure/cosmos/agent_memory/processors/inprocess.py b/azure/cosmos/agent_memory/processors/inprocess.py index 9684ce0..aac1760 100644 --- a/azure/cosmos/agent_memory/processors/inprocess.py +++ b/azure/cosmos/agent_memory/processors/inprocess.py @@ -81,9 +81,7 @@ def process_thread( thread_summary = self._pipeline.generate_thread_summary(user_id, thread_id) extracted = self._pipeline.extract_memories(user_id, thread_id) - reconciled = self._pipeline.reconcile_memories(user_id, get_dedup_pool_size()) - - deduped_count = self._extract_reconcile_count(reconciled) + deduped_count = self._reconcile_fact_and_episodic(user_id, get_dedup_pool_size()) extracted_counts: dict[str, int] = ( {k: v for k, v in extracted.items() if isinstance(v, int)} if isinstance(extracted, dict) else {} @@ -102,8 +100,9 @@ def process_extract_memories( *, user_id: str, thread_id: str, + recent_k: Optional[int] = None, ) -> dict[str, int]: - extracted = self._pipeline.extract_memories(user_id, thread_id) + extracted = self._pipeline.extract_memories(user_id, thread_id, recent_k=recent_k) return {k: v for k, v in extracted.items() if isinstance(v, int)} if isinstance(extracted, dict) else {} def process_thread_summary( @@ -126,16 +125,32 @@ def process_user_summary( summary = self._pipeline.generate_user_summary(user_id, thread_ids) return UserSummaryResult(summary=summary if isinstance(summary, dict) else None) - def process_reconcile(self, *, user_id: str) -> int: + def process_reconcile(self, *, user_id: str, full_rebuild: bool = False) -> int: """Run reconciliation standalone. Returns count of facts merged + contradicted. Pool size is read from ``DEDUP_POOL_SIZE`` (env-tunable, default 50, capped at 500) so the auto-trigger and the standalone path agree. + ``full_rebuild`` (set by the auto-trigger on its persisted-counter + full-recluster cadence) forces the full-pool LLM pass that catches + dissimilar-embedding contradictions. """ from ..thresholds import get_dedup_pool_size - reconciled = self._pipeline.reconcile_memories(user_id, n=get_dedup_pool_size()) - return self._extract_reconcile_count(reconciled) + return self._reconcile_fact_and_episodic(user_id, get_dedup_pool_size(), full_rebuild=full_rebuild) + + def _reconcile_fact_and_episodic(self, user_id: str, n: int, *, full_rebuild: bool = False) -> int: + """Reconcile facts and episodic memories; sum merged+contradicted counts. + + SDK in-process processing reconciles both types (matching the Durable + backend) so episodic dups don't accrue forever. + """ + total = 0 + for memory_type in ("fact", "episodic"): + reconciled = self._pipeline.reconcile_memories( + user_id, n=n, memory_type=memory_type, full_rebuild=full_rebuild + ) + total += self._extract_reconcile_count(reconciled) + return total @staticmethod def _extract_reconcile_count(reconciled: Any) -> int: diff --git a/azure/cosmos/agent_memory/prompts/_schemas.py b/azure/cosmos/agent_memory/prompts/_schemas.py index 1b34973..6c3306c 100644 --- a/azure/cosmos/agent_memory/prompts/_schemas.py +++ b/azure/cosmos/agent_memory/prompts/_schemas.py @@ -65,6 +65,34 @@ } +# --------------------------------------------------------------------------- +# dedup_episodic.prompty — reconcile a pool of active episodic memories +# (MERGE-ONLY: same-event duplicates collapse; no contradiction/deletion) +# --------------------------------------------------------------------------- +DEDUP_EPISODIC_SCHEMA: dict[str, Any] = { + "type": "object", + "properties": { + "duplicate_groups": { + "type": "array", + "items": { + "type": "object", + "properties": { + "merged_content": {"type": "string"}, + "source_ids": {"type": "array", "items": {"type": "string"}}, + "confidence": {"type": ["number", "null"]}, + "salience": {"type": ["number", "null"]}, + }, + "required": ["merged_content", "source_ids", "confidence", "salience"], + "additionalProperties": False, + }, + }, + "kept_ids": {"type": "array", "items": {"type": "string"}}, + }, + "required": ["duplicate_groups", "kept_ids"], + "additionalProperties": False, +} + + # --------------------------------------------------------------------------- # extract_memories.prompty — extract facts + episodic + unclassified # --------------------------------------------------------------------------- @@ -91,8 +119,6 @@ "salience": {"type": "number"}, "temporal_context": {"type": ["string", "null"]}, "tags": {"type": "array", "items": {"type": "string"}}, - "action": {"type": "string", "enum": ["ADD", "UPDATE", "CONTRADICT"]}, - "supersedes_id": {"type": ["string", "null"]}, }, "required": [ "text", @@ -104,8 +130,6 @@ "salience", "temporal_context", "tags", - "action", - "supersedes_id", ], "additionalProperties": False, } @@ -115,7 +139,6 @@ "properties": { "scope_type": {"type": "string"}, "scope_value": {"type": "string"}, - "text": {"type": "string"}, "situation": {"type": ["string", "null"]}, "action_taken": {"type": ["string", "null"]}, "outcome": {"type": ["string", "null"]}, @@ -133,7 +156,6 @@ "required": [ "scope_type", "scope_value", - "text", "situation", "action_taken", "outcome", @@ -282,6 +304,7 @@ # --------------------------------------------------------------------------- PROMPTY_SCHEMAS: dict[str, tuple[str, dict[str, Any]]] = { "dedup.prompty": ("DedupOutput", DEDUP_SCHEMA), + "dedup_episodic.prompty": ("DedupEpisodicOutput", DEDUP_EPISODIC_SCHEMA), "extract_memories.prompty": ("ExtractMemoriesOutput", EXTRACT_MEMORIES_SCHEMA), "summarize.prompty": ("SummarizeOutput", SUMMARIZE_SCHEMA), "summarize_update.prompty": ("SummarizeUpdateOutput", SUMMARIZE_UPDATE_SCHEMA), diff --git a/azure/cosmos/agent_memory/prompts/dedup.prompty b/azure/cosmos/agent_memory/prompts/dedup.prompty index c83a05a..d45bb60 100644 --- a/azure/cosmos/agent_memory/prompts/dedup.prompty +++ b/azure/cosmos/agent_memory/prompts/dedup.prompty @@ -6,7 +6,7 @@ model: apiType: chat options: seed: 43 - maxOutputTokens: 2000 + maxOutputTokens: 16384 additionalProperties: response_format: type: json_object diff --git a/azure/cosmos/agent_memory/prompts/dedup_episodic.prompty b/azure/cosmos/agent_memory/prompts/dedup_episodic.prompty new file mode 100644 index 0000000..6008dd2 --- /dev/null +++ b/azure/cosmos/agent_memory/prompts/dedup_episodic.prompty @@ -0,0 +1,160 @@ +--- +name: dedup_episodic +version: v1 +description: Reconcile a pool of active episodic memories — collapse only true same-event duplicates. +model: + apiType: chat + options: + seed: 43 + maxOutputTokens: 16384 + additionalProperties: + response_format: + type: json_object +inputs: + episodics_text: + type: string +--- + +system: +You are a precision episodic-memory reconciliation system. You receive a pool of active episodic memories (each with an ID, content, confidence, salience, and creation timestamp) and must collapse only true same-event duplicates while leaving distinct experiences untouched. + +## Your Goal +Produce a clean merge-only reconciliation that: +1. Collapses repeated extractions of the same past experience into a single merged episode. +2. Preserves distinct occurrences, even when they involve the same action, tool, project, or outcome. +3. Leaves everything else untouched. + +## What is an EPISODIC MEMORY + +An episodic memory is a **past experience**: a specific event or bounded experience in the shape `situation → action_taken → outcome`, with a scope and an `outcome_valence`. It records what happened, not an atomic claim about what is generally true. + +Examples: +- "During the Q3 launch, Redis-backed sessions timed out under load → the team switched to DynamoDB → sessions stabilized" → episodic. +- "On Monday's load test, increasing worker count to 20 still failed with timeouts" → episodic. + +## Two Orthogonal Outcomes (mutually exclusive, exhaustive) + +Every input episode must end up in exactly one of these two places: +- `duplicate_groups[*].source_ids` — the episode is a re-extraction of the same event as one or more other episodes. +- `kept_ids` — the episode is not a true same-event duplicate. + +No input episode may be omitted from the output. Do NOT emit a `contradicted_pairs` key or any deletion bucket. + +## What is a DUPLICATE + +Two or more episodic memories are duplicates only when they describe the **same event re-extracted**: the same situation, the same action taken, and the same outcome. + +The duplicate bar is HIGH. Same topic, same tool, same action pattern, or same outcome is not enough. + +Resolution: emit one entry in `duplicate_groups` with: +- `merged_content` — a clean, self-contained same-event restatement. Do NOT concatenate the originals; synthesize. +- `source_ids` — every original ID that participates in the duplicate group. It must contain at least 2 IDs. +- `confidence` — the **max** confidence among the source episodes. +- `salience` — the **max** salience among the source episodes. + +## What is KEPT + +Episodes that are not true same-event duplicates go in `kept_ids`. + +Distinct occurrences MUST NOT merge. If the user tried the same action on Monday and again on Tuesday with the same outcome, those are separate experiences. Frequency and repetition are evidence and must be preserved. + +## No Contradiction Handling + +Past events are append-only ground truth. There is no contradiction handling and no deletion of episodes in this prompt. + +- Do NOT pick winners or losers. +- Do NOT soft-delete an episode because another episode appears to conflict with it. +- Do NOT emit `contradicted_pairs` or any contradiction/deletion bucket. +- Conflicting lessons are resolved elsewhere, not here. + +## Decision Guidelines + +1. **Conservative bias.** If you cannot confidently prove two episodes are the same event, put them both in `kept_ids`. Over-merging distinct experiences is worse than retaining redundancy. + +2. **Same event test.** Merge only when the memories align on the concrete situation/scope, action taken, and outcome. If the date, run, incident, session, customer, environment, or other occurrence marker differs, keep them separate. + +3. **Don't drop information.** A `merged_content` must preserve every material detail from its sources. If you cannot preserve everything in one clean same-event restatement, the episodes probably are not duplicates. + +4. **Don't fabricate.** Never introduce dates, outcomes, causes, tools, or entities that are not present in at least one source. + +## Input Format + +You will receive a numbered list of episodic memories. Each line has the form: +``` +N. ID: | Content: "" | Confidence: 0.85 | Salience: 0.7 | Created: +``` + +## Missing-Field Handling + +Some episodes may show `Confidence: N/A`, `Salience: N/A`, or `Created: N/A`. + +- Treat `N/A` confidence/salience as **unknown** — set those fields to `null` + in the `duplicate_groups` output (do NOT omit the keys); the pipeline will + fall back to `max(source.confidence)` / `max(source.salience)`. +- Created timestamps are for traceability only. Do not use recency to delete or prefer episodes. + +## Output Format + +You must output ONLY valid JSON matching this exact schema. No preamble, no explanation, no markdown fences — just the JSON object. + +```json +{ + "duplicate_groups": [ + {"merged_content": "", "source_ids": ["", ""], "confidence": 0.0, "salience": 0.0} + ], + "kept_ids": ["", ""] +} +``` + + +The `confidence` and `salience` values **must** be real numbers between 0 and 1 in the actual output — the `0.0` above is a structural placeholder, not a literal value to echo. Compute them as the maximum across source episodes in the group. If you cannot determine confidence or salience, set the field to `null` (the pipeline will fall back to `max(source.*)` from the source records). **Do not omit the keys** — the response schema requires both fields to appear on every group, with either a number or `null` as the value. + +If a bucket is empty, emit it as an empty array (`[]`) rather than omitting the key. Never include `contradicted_pairs`. + +## Worked Examples + +### Example 1: True same-event duplicate + +**Input pool:** +``` +1. ID: E1 | Content: "During the Q3 checkout launch, Redis sessions timed out under load; the team switched session storage to DynamoDB and checkout stabilized." | Confidence: 0.9 | Salience: 0.8 | Created: 2024-01-06T00:00:00Z +2. ID: E2 | Content: "In the Q3 checkout launch incident, Redis-backed sessions failed under load, so the team moved sessions to DynamoDB, which stabilized checkout." | Confidence: 0.85 | Salience: 0.9 | Created: 2024-01-07T00:00:00Z +3. ID: E3 | Content: "On Tuesday's checkout load test, raising Redis connection limits still produced session timeouts." | Confidence: 0.8 | Salience: 0.6 | Created: 2024-01-08T00:00:00Z +``` + +**Expected output:** +```json +{ + "duplicate_groups": [ + {"merged_content": "During the Q3 checkout launch, Redis-backed sessions timed out under load; the team moved session storage to DynamoDB, and checkout stabilized.", "source_ids": ["E1", "E2"], "confidence": 0.9, "salience": 0.9} + ], + "kept_ids": ["E3"] +} +``` + +E1 and E2 describe the same situation, action, and outcome. E3 is related but a different occurrence, so it stays kept. + +### Example 2: Same action on different days is not a duplicate + +**Input pool:** +``` +1. ID: E4 | Content: "On Monday's ingestion test, the team increased workers to 20, but the pipeline still timed out." | Confidence: 0.9 | Salience: 0.7 | Created: 2024-02-01T00:00:00Z +2. ID: E5 | Content: "On Tuesday's ingestion test, the team increased workers to 20, but the pipeline still timed out." | Confidence: 0.9 | Salience: 0.7 | Created: 2024-02-02T00:00:00Z +``` + +**Expected output:** +```json +{ + "duplicate_groups": [], + "kept_ids": ["E4", "E5"] +} +``` + +These episodes share the same action and outcome, but they occurred on different days. Preserve both because repetition is evidence. + +user: +Reconcile the following pool of active episodic memories: + +{{episodics_text}} + +Return JSON exactly matching the schema above. diff --git a/azure/cosmos/agent_memory/prompts/extract_memories.prompty b/azure/cosmos/agent_memory/prompts/extract_memories.prompty index d266f1d..723b458 100644 --- a/azure/cosmos/agent_memory/prompts/extract_memories.prompty +++ b/azure/cosmos/agent_memory/prompts/extract_memories.prompty @@ -6,7 +6,7 @@ model: apiType: chat options: seed: 42 - maxOutputTokens: 2500 + maxOutputTokens: 16384 additionalProperties: response_format: type: json_object @@ -14,9 +14,6 @@ inputs: existing_facts: type: string default: '(none)' - existing_episodics: - type: string - default: '(none)' transcript: type: string --- @@ -34,11 +31,11 @@ Every memory must be explicitly grounded in the conversation — never inferred, ## Speaker Discrimination — Where Memories May Come From -The transcript below is line-tagged: `[user]:` lines are the human's own words, `[assistant]:` lines are the agent's response. These two sources are NOT interchangeable. +The transcript below is line-tagged: `[user]:` lines are the human's own words, `[agent]:` lines are the agent's response. These two sources are NOT interchangeable. -- **Facts about the user may ONLY come from `[user]:` lines.** The assistant may restate, paraphrase, confirm, or make assumptions on the user's behalf ("Got it, you don't eat meat", "I assume you want a luxury hotel") — those are the agent's response, NOT the user's assertion. Never treat assistant text as a new source of user facts. If a fact appears only in `[assistant]:` content and is not asserted by the user, do not extract it. -- **Episodic memories may use both speakers' content** — the user's stated intent or scope is the anchor (and must be present in `[user]:`), but the assistant's content may help fill in the `action_taken` or `outcome` of a `situation → action_taken → outcome` arc when the agent carried out the action on the user's behalf. -- The assistant's general world-knowledge answers (e.g. "Python 3.13 was released in October 2024") are never user facts and never user episodic memories. +- **Facts about the user may ONLY come from `[user]:` lines.** The agent may restate, paraphrase, confirm, or make assumptions on the user's behalf ("Got it, you don't eat meat", "I assume you want a luxury hotel") — those are the agent's response, NOT the user's assertion. Never treat agent text as a new source of user facts. If a fact appears only in `[agent]:` content and is not asserted by the user, do not extract it. +- **Episodic memories may use both speakers' content** — the user's stated intent or scope is the anchor (and must be present in `[user]:`), but the agent's content may help fill in the `action_taken` or `outcome` of a `situation → action_taken → outcome` arc when the agent carried out the action on the user's behalf. +- The agent's general world-knowledge answers (e.g. "Python 3.13 was released in October 2024") are never user facts and never user episodic memories. ## Confidence Scoring @@ -73,94 +70,12 @@ Extract concrete, factual statements that fall into these categories: - Each fact must be self-contained and intelligible without context — no pronouns like "it" or "they" without antecedents - Write in third person ("The user...", "The project...") - Keep each fact concise — under 40 words -- Consolidate closely related items **within the same category** into a single fact (e.g., multiple search results on the same topic) -- **Never merge across categories.** A single user turn that combines, say, a `preference` ("I don't eat meat") and a `requirement` ("I need wheelchair-accessible restaurants") MUST produce two separate facts. Compound user statements regularly cross category boundaries; extract every category that applies. Silently folding one into the other drops information. -- Only split *within* a category when the claims are about genuinely different topics +- Consolidate closely related items into a single fact (e.g., multiple search results on the same topic) +- Only split into separate facts when the claims are about genuinely different topics - Each fact will be stored as its own document with its own vector embedding — self-contained phrasing is essential for retrieval -### Fact Reconciliation -Before adding a fact, check it against the existing memories provided below. Use the `action` field: -- **ADD** — This is a genuinely new fact not present in existing memories -- **UPDATE** — This fact refines, narrows, or supplies more detail to an existing memory while remaining *compatible* with it (e.g. "user lives in Seattle" → "user lives in Seattle, WA"). Set `supersedes_id` to the ID of the memory being refined -- **CONTRADICT** — This new fact is the opposite, a negation, or otherwise mutually exclusive with an existing memory (e.g. existing: "user is vegetarian"; new: "user now eats meat"). The old fact is no longer true — emit `CONTRADICT`, **not** UPDATE. Set `supersedes_id` to the ID of the memory being contradicted -- **NONE** — This fact is already captured in existing memories with no meaningful change. Do not include NONE entries in your output — simply omit them - -**UPDATE vs CONTRADICT — the decisive test:** If the old fact and the new fact could both be true at the same moment about the same subject, use UPDATE. If they cannot both be true (one negates, reverses, or excludes the other), use CONTRADICT. - -### ADD / CONTRADICT / UPDATE — hard constraints - -These rules are absolute. Violating them silently corrupts the user's memory store. - -1. **Every emitted fact (ADD, UPDATE, or CONTRADICT) must paraphrase a claim directly stated by the user in a `[user]:` line in this extraction's transcript.** Do not invent supporting, clarifying, or "explicit-negation" facts. Do not synthesize facts by combining, restating, or "consolidating" entries from the existing-facts list. If the [user]: lines in this transcript do not assert claim X, do NOT emit claim X — regardless of what the existing-facts list contains. Merging existing facts is the job of a separate reconciliation pass, not yours. -2. **One user statement → one fact, even when it implies a contradiction.** If a single user statement both adds new information AND opposes an existing fact, emit EXACTLY ONE fact: `text` paraphrases what the user actually said, `action="CONTRADICT"`, `supersedes_id` points at the opposed fact. The `supersedes_id` field by itself encodes the semantic opposition — do NOT also emit a second "explicit-negation" fact that restates the opposite of the prior claim. See the worked examples below. -3. **Fact text must describe the world, never describe a memory operation.** Strings like `"X is contradicted by Y"`, `"The user's prior preference is no longer accurate"`, or `"Previous fact superseded"` are meta-commentary about prior reconciliations and are **never** valid fact content. If you find yourself writing one, drop the item entirely. -4. **`supersedes_id` must point at a fact that is semantically about the same subject and property as the new fact.** A new dietary fact may only contradict an existing dietary fact; a new accessibility requirement may only update an existing accessibility requirement. Cross-subject supersedes (e.g. contradicting a "wheelchair access" fact with a new "loves seafood" fact) are always wrong — emit `ADD` instead. -5. **When in doubt, emit `ADD`.** A spurious `ADD` is recoverable (exact-content-hash deduping catches it; reconciliation can later merge or supersede it). A spurious `CONTRADICT` or `UPDATE` corrupts the audit trail and silently marks a still-valid fact as superseded. - -#### Worked example A — explicit CONTRADICT (user states the negation directly) - -- Existing memory: `[ID: fact_abc123] User is vegetarian.` -- New turn: *"I started eating meat again last month."* -- Correct output (one fact, text paraphrases what the user said): - ```json - { - "text": "User eats meat.", - "category": "preference", - "action": "CONTRADICT", - "supersedes_id": "fact_abc123", - "confidence": 0.95, - "salience": 0.8 - } - ``` -- **Wrong**: emitting `"action": "UPDATE"` here would be a silent bug — the pipeline treats UPDATE as a compatible refinement and CONTRADICT as an opposing claim, and downstream telemetry / belief-revision logic depends on the distinction. - -#### Worked example B — implicit CONTRADICT (user states a claim that semantically opposes an existing fact) - -- Existing memory: `[ID: fact_xyz789] The user does not eat meat.` -- New turn: *"Actually, I love steak and seafood."* -- Correct output (ONE fact — `text` is what the user said, `supersedes_id` carries the opposition): - ```json - { - "text": "The user loves steak and seafood.", - "category": "preference", - "action": "CONTRADICT", - "supersedes_id": "fact_xyz789", - "confidence": 0.95, - "salience": 0.8 - } - ``` -- **Wrong** — emitting two facts: - ```json - // BAD: phantom explicit-negation fact alongside the literal user claim - [ - {"text": "The user loves steak and seafood.", "action": "ADD", ...}, - {"text": "The user eats meat.", "action": "CONTRADICT", "supersedes_id": "fact_xyz789", ...} - ] - ``` - The phantom `"The user eats meat"` fact was never said by the user — it is an invented restatement to make the contradiction "explicit". This pollutes the fact store with claims the user did not make. The CONTRADICT relation on the literal fact is sufficient. - -#### Worked example C — do NOT synthesize ADDs from existing facts - -- Existing memories: - - `[ID: fact_111] The user eats meat.` - - `[ID: fact_222] The user loves steak and seafood.` -- New turn: *"Normally, I prefer moderate hotels."* -- Correct output (ONE fact — only the hotel preference; the existing-facts list is reference-only): - ```json - { - "text": "The user normally prefers moderate hotels.", - "category": "preference", - "action": "ADD", - "confidence": 0.9, - "salience": 0.7 - } - ``` -- **Wrong** — emitting a synthesized "consolidation" fact: - ```json - // BAD: this fact is a paraphrase-merge of fact_111 + fact_222; the user never said this in this turn - {"text": "The user loves steak and seafood, indicating they eat meat.", "action": "ADD", ...} - ``` - Merging existing facts is the job of a separate reconciliation pass. Your job here is to extract claims from the new [user] turn only. +### Avoiding Duplicates +Before adding a fact, check it against the existing memories provided below. **Only emit facts that are genuinely new.** If a fact is already captured in existing memories with no meaningful change, simply omit it. Do not try to update, merge, or flag contradictions with existing facts — that is handled later by a separate reconciliation step. Your only job here is to extract new facts and episodic memories from the transcript. --- @@ -175,12 +90,8 @@ Extract memories that are tied to a **specific situation, scope, or context** th ### Required Fields - **scope_type** — short, free-form noun describing the kind of context (e.g. `trip`, `project`, `event`, `session`, `release`, `campaign`). Pick whatever vocabulary fits the user's domain. Do not invent a value if one is not implied — if no scope is present, the memory probably belongs in facts. - **scope_value** — the specific instance of that scope (e.g. `Paris 2025`, `Acme revamp`, `Q3 launch`). -- **text** — short, self-contained one-liner (under 25 words) describing what this memory is *about*. Required, non-empty. This is the field that gets embedded and full-text-indexed — vague text like "intent recorded" silently kills retrieval. Write it like a subject line: - - *Planned / in-flight* intents — capture the goal **plus the key constraints** the user has stated. Example: `"Planning a Tokyo trip with vegetarian and wheelchair-accessible-restaurant constraints."` - - *Past events* — summarize situation→action→outcome in one sentence. Example: `"Resolved Q3 K8s OOM outage by raising pod memory limits to 1GB."` - - *Ongoing context* — describe the current state. Example: `"User is heads-down on the Acme launch this week and wants short answers."` -All three of `scope_type`, `scope_value`, and `text` must be non-empty. (The `text` field plays the same role for episodic that it does for facts — self-contained phrasing intended for embedding.) +Both must be non-empty. ### Optional Fields (include only when applicable) - **situation** — the context or problem faced (present for past/in-flight events) @@ -191,7 +102,7 @@ All three of `scope_type`, `scope_value`, and `text` must be non-empty. (The `te - **lesson** — a transferable takeaway - **domain** — topic area -For planned/in-flight or ongoing-context memories, leave `situation`, `action_taken`, `outcome`, `outcome_valence`, `reasoning`, and `lesson` as `null`. The scope fields plus `text` carry the meaning. +For planned/in-flight or ongoing-context memories, leave `situation`, `action_taken`, `outcome`, `outcome_valence`, `reasoning`, and `lesson` as `null`. The scope fields alone carry the meaning. --- @@ -214,34 +125,18 @@ When in doubt: --- ## Existing Memories +The following facts already exist for this user. They are provided so you can avoid re-extracting facts that are already known. If a fact from the new transcript is already captured here, omit it. -### Existing facts (REFERENCE ONLY — never source ADDs from this list) +### Existing Memories Are Context, Not a Source -The list below shows facts already stored for this user. Use it ONLY for these three purposes: +The existing memories block is provided ONLY for grounding and for deciding whether a fact from the new transcript is already known and should be omitted. It is NOT part of the conversation transcript. -1. **Deduplication** — if a fact you would otherwise ADD is already captured (semantically equivalent to an existing entry), omit it (action=NONE; don't include NONE entries in your output). -2. **UPDATE** — if a NEW `[user]:` line in this transcript refines or adds detail to an existing fact in a compatible way, emit one fact with `action=UPDATE` and `supersedes_id` pointing at the existing fact. -3. **CONTRADICT** — if a NEW `[user]:` line in this transcript opposes or negates an existing fact, emit one fact with `action=CONTRADICT` and `supersedes_id` pointing at the existing fact. The new fact's `text` paraphrases what the user actually said (not an invented explicit-negation). - -**Do NOT** treat this list as source material for ADD. Never combine, merge, restate, or "consolidate" entries from this list into a new ADD fact. If the new `[user]:` lines do not directly assert claim X, do NOT emit claim X — regardless of what the existing list contains. Cross-fact merging is the job of a separate reconciliation pass, not this one. +- New memories must come exclusively from the new transcript below. +- Never extract a fact or episodic memory solely because it appears in existing memories. +- If information appears only in existing memories and not in the new transcript, do not emit it. {{existing_facts}} -### Existing episodics - -The following episodic memories already exist for this user, grouped by scope. **Each scope (the pair of `scope_type` + `scope_value`) is the unique identity of an episodic memory** — there is one episodic per scope. Storage is upsert-by-scope: whatever you emit for an existing scope **replaces** the prior record. There is no ADD/UPDATE/CONTRADICT vocabulary for episodics — emit the full current state in `text` and the pipeline does the right thing. - -Decision rule when the transcript mentions a scope: - -- **Already captured, nothing new** — if the same scope is already covered and the transcript adds no new information (e.g. the user just re-states "for the Tokyo trip, I want luxury hotels" and that intent is already in the existing record), **omit it from the output**. Re-emitting it accomplishes nothing. -- **Refinement / extension** — the transcript adds new details to a scope already present (e.g. existing: "Planning Tokyo trip with luxury hotels"; new turn: "let's also book a 5-night Shinjuku stay"). Emit the episodic with **the merged, richer `text`** that includes both the prior intent and the new detail. The new text replaces the old record by scope. -- **Reversal** — the transcript negates or replaces the prior intent for the same scope (e.g. existing: "Planning Tokyo trip with luxury hotels"; new turn: "switching to budget hostels"). Emit the episodic with the **new** `text` describing the updated intent. The reversal replaces the old record by scope. -- **New scope** — the transcript introduces a scope that is not in the list below. Emit a fresh episodic with that scope. - -If two genuinely-distinct events would share the same `(scope_type, scope_value)` (e.g. "lost wallet in Tokyo" and "booked Tokyo hotel" both as `(trip, Tokyo)`), differentiate them via `scope_value` (`Tokyo lost-wallet incident` vs `Tokyo trip`) — not by emitting two records under the same scope. - -{{existing_episodics}} - --- ## Salience Scoring Rubric @@ -283,7 +178,7 @@ A fact is a standing claim that holds outside any specific context. The test: if 1. Could someone act on or reference this fact without reading the original thread? 2. Is this fact stated explicitly, not inferred? 3. Is each fact truly atomic — one claim per entry? -4. Is every emitted fact grounded in a `[user]:` line in this transcript? (Existing-facts entries are reference-only; never source new ADDs from them.) +4. Can any facts be merged because they describe variants of the same thing? **For Episodic:** 1. Did this event actually happen, or is it hypothetical? @@ -312,9 +207,7 @@ A fact is a standing claim that holds outside any specific context. The test: if "confidence": 1.0, "salience": 0.9, "temporal_context": null, - "tags": ["topic:identity"], - "action": "ADD", - "supersedes_id": null + "tags": ["topic:identity"] }, { "text": "Alex is a data engineer at Acme Corp.", @@ -325,9 +218,7 @@ A fact is a standing claim that holds outside any specific context. The test: if "confidence": 1.0, "salience": 0.8, "temporal_context": null, - "tags": ["topic:identity", "topic:career"], - "action": "ADD", - "supersedes_id": null + "tags": ["topic:identity", "topic:career"] }, { "text": "Alex's new ETL pipeline project has a deadline at end of Q2.", @@ -338,9 +229,7 @@ A fact is a standing claim that holds outside any specific context. The test: if "confidence": 0.95, "salience": 0.9, "temporal_context": "end of Q2", - "tags": ["topic:project", "topic:ETL"], - "action": "ADD", - "supersedes_id": null + "tags": ["topic:project", "topic:ETL"] } ], "episodic": [] @@ -365,16 +254,13 @@ A fact is a standing claim that holds outside any specific context. The test: if "confidence": 0.95, "salience": 0.7, "temporal_context": "last month", - "tags": ["topic:kubernetes", "topic:infrastructure"], - "action": "ADD", - "supersedes_id": null + "tags": ["topic:kubernetes", "topic:infrastructure"] } ], "episodic": [ { "scope_type": "incident", "scope_value": "Q3 K8s OOM outage", - "text": "Resolved Q3 K8s OOM outage by raising pod memory limits to 1GB and adding per-namespace resource quotas.", "situation": "Kubernetes pods in production were repeatedly OOM-killed, causing an outage.", "action_taken": "Bumped pod memory limits from 512MB to 1GB and added resource quotas per namespace.", "outcome": "The OOM-killing stopped and production stabilized.", @@ -408,9 +294,7 @@ A fact is a standing claim that holds outside any specific context. The test: if "confidence": 0.95, "salience": 0.7, "temporal_context": null, - "tags": ["topic:database", "topic:BigQuery"], - "action": "ADD", - "supersedes_id": null + "tags": ["topic:database", "topic:BigQuery"] }, { "text": "The marketing team's budget is $50,000 for the current quarter.", @@ -421,9 +305,7 @@ A fact is a standing claim that holds outside any specific context. The test: if "confidence": 0.95, "salience": 0.9, "temporal_context": "current quarter", - "tags": ["topic:budget", "topic:marketing"], - "action": "ADD", - "supersedes_id": null + "tags": ["topic:budget", "topic:marketing"] } ], "episodic": [] @@ -451,16 +333,13 @@ The first statement is a standing preference and belongs in `facts`. The second "confidence": 0.95, "salience": 0.7, "temporal_context": null, - "tags": ["topic:travel", "topic:hotels"], - "action": "ADD", - "supersedes_id": null + "tags": ["topic:travel", "topic:hotels"] } ], "episodic": [ { "scope_type": "trip", "scope_value": "Paris", - "text": "Planning a Paris trip with a luxury-accommodations preference.", "situation": null, "action_taken": null, "outcome": null, @@ -476,126 +355,11 @@ The first statement is a standing preference and belongs in `facts`. The second } ``` -### Example 6: Episodic — same scope, same intent already captured (OMIT) - -**Existing episodics:** -``` -- trip = Tokyo (1 episodic) - - [ID: ep_a1b2c3] (salience 0.8) Planning a Tokyo trip with a luxury hotel preference. -``` - -**Conversation:** -> User: "For this Tokyo trip, I want luxury hotels." -> User: "Normally, I prefer moderate hotels." - -The first user turn re-states an intent that is already captured for `(trip, Tokyo)`. Omit it — re-emitting the same intent for the same scope accomplishes nothing (the pipeline upserts by scope). The second user turn is a standing preference and belongs in `facts`. - -**Output:** -```json -{ - "facts": [ - { - "text": "The user normally prefers moderate hotels.", - "category": "preference", - "subject": "user", - "predicate": "hotel_preference", - "object": "moderate hotels", - "confidence": 0.95, - "salience": 0.7, - "temporal_context": null, - "tags": ["topic:travel", "topic:hotels"], - "action": "ADD", - "supersedes_id": null - } - ], - "episodic": [] -} -``` - -### Example 7: Episodic — refinement (emit merged text for the same scope) - -**Existing episodics:** -``` -- trip = Tokyo (1 episodic) - - [ID: ep_a1b2c3] (salience 0.8) Planning a Tokyo trip with a luxury hotel preference. -``` - -**Conversation:** -> User: "For the Tokyo trip, let's also book a 5-night stay in Shinjuku." - -The transcript adds a new detail (5-night Shinjuku stay) to the existing Tokyo scope. Emit one episodic for `(trip, Tokyo)` with the **merged richer text** that carries both the prior luxury-hotel intent and the new accommodation detail. The pipeline upserts by scope, so this replaces the prior record. - -**Output:** -```json -{ - "facts": [], - "episodic": [ - { - "scope_type": "trip", - "scope_value": "Tokyo", - "text": "Planning a Tokyo trip with a luxury hotel preference and a 5-night stay in Shinjuku.", - "situation": null, - "action_taken": null, - "outcome": null, - "outcome_valence": null, - "reasoning": null, - "lesson": null, - "domain": "travel", - "confidence": 0.95, - "salience": 0.85, - "tags": ["topic:travel", "topic:hotels", "topic:itinerary"] - } - ] -} -``` - -### Example 5: Compound user statement crossing fact categories - -**Conversation:** -> User: "I don't eat meat and I need wheelchair-accessible restaurants." - -A single user turn that combines a `preference` (diet) and a `requirement` (accessibility). These are different categories — they MUST be emitted as separate facts. Collapsing both into one "restaurant preferences" fact silently loses the accessibility constraint, which is the more operationally critical of the two. - -**Output:** -```json -{ - "facts": [ - { - "text": "The user does not eat meat.", - "category": "preference", - "subject": "user", - "predicate": "dietary_restriction", - "object": "no meat", - "confidence": 1.0, - "salience": 0.9, - "temporal_context": null, - "tags": ["topic:diet", "topic:food"], - "action": "ADD", - "supersedes_id": null - }, - { - "text": "The user requires wheelchair-accessible restaurants.", - "category": "requirement", - "subject": "user", - "predicate": "accessibility_requirement", - "object": "wheelchair-accessible restaurants", - "confidence": 1.0, - "salience": 0.95, - "temporal_context": null, - "tags": ["topic:accessibility", "topic:restaurants"], - "action": "ADD", - "supersedes_id": null - } - ], - "episodic": [] -} -``` - --- ## Output Format -You must output ONLY valid JSON matching the schema below. No preamble, no explanation, no closing remarks — just the JSON object. Each array can be empty if no memories of that type are found. Omit entries with action=NONE entirely. +You must output ONLY valid JSON matching the schema below. No preamble, no explanation, no closing remarks — just the JSON object. Each array can be empty if no memories of that type are found. Omit any fact that is already captured in the existing memories. ```json { @@ -609,16 +373,13 @@ You must output ONLY valid JSON matching the schema below. No preamble, no expla "confidence": 0.95, "salience": 0.8, "temporal_context": "optional time ref or null", - "tags": ["topic:x"], - "action": "ADD|UPDATE|CONTRADICT", - "supersedes_id": "id or null" + "tags": ["topic:x"] } ], "episodic": [ { "scope_type": "trip|project|event|session|release|campaign|... (free-form, required, non-empty)", "scope_value": "specific instance, e.g. Paris 2025 (required, non-empty)", - "text": "self-contained one-liner describing the memory — required, non-empty", "situation": "context/problem, or null", "action_taken": "what was tried, or null", "outcome": "what happened, or null", diff --git a/azure/cosmos/agent_memory/services/_pipeline_helpers.py b/azure/cosmos/agent_memory/services/_pipeline_helpers.py index 9748168..e94ce91 100644 --- a/azure/cosmos/agent_memory/services/_pipeline_helpers.py +++ b/azure/cosmos/agent_memory/services/_pipeline_helpers.py @@ -280,47 +280,6 @@ def build_transcript( return "\n".join(parts) -def format_existing_episodics(memories: list[dict[str, Any]]) -> str: - """Render existing episodic memories for the extract_memories prompt. - - Groups by ``(scope_type, scope_value)`` so the LLM can see, per-scope, - which intent is already captured. Episodics use **scope-as-identity**: - the deterministic id is seeded from ``(user_id, scope_type, scope_value)``, - so any re-emission for the same scope (paraphrased intent, added detail, - or a reversal) collides and overwrites the prior record via upsert. The - LLM does NOT make ``ADD``/``UPDATE``/``CONTRADICT`` decisions on - episodics — that vocabulary is not in the episodic schema. - - What this rendering gives the model is per-scope context so it can: - - 1. Emit a single coherent ``text`` that reflects the *current* intent - for the scope (the upsert will overwrite the prior one). - 2. Avoid re-emitting an episodic when the new turn carries no - additional signal beyond what the existing one already records. - - Distinct events under the same umbrella (e.g. hotel booking vs lost - wallet, both under a Tokyo trip) belong under distinct ``scope_value`` - strings so they don't collide on the deterministic id. - """ - if not memories: - return "(none)" - grouped: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list) - for mem in memories: - meta = mem.get("metadata") or {} - scope_type = (meta.get("scope_type") or "(none)").strip() or "(none)" - scope_value = (meta.get("scope_value") or "(none)").strip() or "(none)" - grouped[(scope_type, scope_value)].append(mem) - lines: list[str] = [] - for (scope_type, scope_value), bucket in grouped.items(): - lines.append(f"- {scope_type} = {scope_value} ({len(bucket)} episodic{'s' if len(bucket) != 1 else ''})") - for mem in bucket: - mem_id = mem.get("id", "(no-id)") - salience = mem.get("salience", "N/A") - content = (mem.get("content") or "").strip() or "(empty content)" - lines.append(f" - [ID: {mem_id}] (salience {salience}) {content}") - return "\n".join(lines) - - # Stopwords stripped from grounding checks. Keep this list short and focused # on tokens that carry no factual content; any word a memory might legitimately # differ on (e.g. "not", "no") must NOT be added here. @@ -502,13 +461,31 @@ def parse_llm_json(text: str | None) -> dict[str, Any]: cleaned = cleaned.lstrip("`").lstrip() if cleaned.endswith("```"): cleaned = cleaned[:-3] + cleaned = cleaned.strip() try: - return json.loads(cleaned.strip()) + return json.loads(cleaned) except json.JSONDecodeError as exc: preview = (text or "")[:200].replace("\n", " ") + if _looks_truncated(cleaned, exc): + raise LLMError( + "LLM JSON output appears TRUNCATED (decode error at the very end of a " + f"{len(cleaned)}-char body — the model almost certainly hit its output-token " + "cap mid-object). Increase 'maxOutputTokens' in the calling prompty, or reduce " + "the amount of input per call (e.g. lower the fact-extraction batch size / " + f"recent_k, or split oversized turns). Decode error: {exc}. preview={preview!r}" + ) from exc raise LLMError(f"LLM returned invalid JSON (preview={preview!r}): {exc}") from exc +def _looks_truncated(cleaned: str, exc: json.JSONDecodeError) -> bool: + """Heuristic: did the JSON fail because the model ran out of output tokens?""" + if not cleaned: + return False + unbalanced = cleaned.count("{") > cleaned.count("}") or cleaned.count("[") > cleaned.count("]") + unterminated_string = "Unterminated string" in str(exc) + return unbalanced or unterminated_string + + def default_prompts_dir() -> str: """Default ``prompts/`` directory location: under ``azure/cosmos/agent_memory/``.""" pkg_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) diff --git a/azure/cosmos/agent_memory/services/pipeline.py b/azure/cosmos/agent_memory/services/pipeline.py index 2211669..b216816 100644 --- a/azure/cosmos/agent_memory/services/pipeline.py +++ b/azure/cosmos/agent_memory/services/pipeline.py @@ -17,13 +17,20 @@ from typing import Any, Iterable, Literal, Optional from azure.cosmos.exceptions import ( - CosmosHttpResponseError, CosmosResourceExistsError, CosmosResourceNotFoundError, ) +from azure.cosmos.agent_memory import thresholds as threshold_config from azure.cosmos.agent_memory._container_routing import ContainerKey -from azure.cosmos.agent_memory._utils import DEFAULT_TTL_BY_TYPE, compute_content_hash +from azure.cosmos.agent_memory._utils import ( + DEFAULT_TTL_BY_TYPE, + compute_content_hash, + distance_function_from_container_properties, + vector_autodrop_supported, + vector_order_direction, + vector_similarity_at_least, +) from azure.cosmos.agent_memory.exceptions import ( LLMError, MemoryConflictError, @@ -55,9 +62,6 @@ coerce_valence, parse_llm_json, ) -from azure.cosmos.agent_memory.services._pipeline_helpers import ( - format_existing_episodics as _format_existing_episodics, -) from azure.cosmos.agent_memory.services._pipeline_helpers import ( is_real_number as _is_real_number, ) @@ -276,6 +280,183 @@ def _load_existing_memories( ) return items + 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). + 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 + 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 ``DEDUP_SIM_HIGH`` thresholds are cosine-calibrated; on euclidean + the destructive auto-drop is skipped (borderline tagging + LLM reconcile + still run). 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 _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, + memory_type: str, + *, + limit: int | None = None, + tagged_only: bool = False, + ) -> list[dict[str, Any]]: + top_clause = f"TOP {top_literal(limit, name='_query_active_memories.limit')} " if limit else "" + tag_clause = "AND ARRAY_CONTAINS(c.tags, @tag) " if tagged_only else "" + query = ( + f"SELECT {top_clause}* FROM c WHERE c.user_id = @user_id " + "AND c.type = @memory_type " + f"AND {_ACTIVE_DOC_FILTER} " + f"{tag_clause}" + "ORDER BY c.created_at DESC" + ) + parameters = [ + {"name": "@user_id", "value": user_id}, + {"name": "@memory_type", "value": memory_type}, + ] + if tagged_only: + parameters.append({"name": "@tag", "value": "sys:dup-candidate"}) + return list( + self._memories_container.query_items( + query=query, + parameters=parameters, + enable_cross_partition_query=True, + ) + ) + + def _load_memories_by_ids( + self, user_id: str, memory_type: str, ids: Iterable[str] + ) -> list[dict[str, Any]]: + ids = [mid for mid in dict.fromkeys(ids) if mid] + if not ids: + return [] + placeholders = ", ".join(f"@id{i}" for i in range(len(ids))) + parameters = [ + {"name": "@user_id", "value": user_id}, + {"name": "@memory_type", "value": memory_type}, + ] + parameters.extend({"name": f"@id{i}", "value": mid} for i, mid in enumerate(ids)) + query = ( + "SELECT * FROM c WHERE c.user_id = @user_id " + "AND c.type = @memory_type " + f"AND c.id IN ({placeholders}) " + f"AND {_ACTIVE_DOC_FILTER}" + ) + return list( + self._memories_container.query_items( + query=query, + parameters=parameters, + enable_cross_partition_query=True, + ) + ) + + def _clear_dup_candidate_tags(self, docs: Iterable[dict[str, Any]]) -> None: + for doc in docs: + tags = [tag for tag in (doc.get("tags") or []) if tag != "sys:dup-candidate"] + if tags == (doc.get("tags") or []): + continue + updated = dict(doc) + updated["tags"] = tags + metadata = dict(updated.get("metadata") or {}) + metadata.pop("dup_of", None) + metadata.pop("dup_score", None) + updated["metadata"] = metadata + updated["updated_at"] = datetime.now(timezone.utc).isoformat() + try: + self._upsert_memory(updated) + except Exception: + logger.exception("reconcile_memories: failed to clear dup-candidate tag id=%s", doc.get("id")) + def _upsert_memory(self, doc: dict[str, Any]) -> dict[str, Any]: """Upsert a fact, episodic, or procedural document to the memories container.""" response = self._memories_container.upsert_item(body=doc) @@ -314,56 +495,6 @@ def _stable_source_timestamp(items: list[dict[str, Any]]) -> str: return max(timestamps) return datetime.now(timezone.utc).isoformat() - def _mark_extracted_superseded( - self, - *, - user_id: str, - thread_id: str, - supersedes_id: str, - superseder_id: str, - reason: Literal["update", "contradict"], - ) -> bool: - try: - old_mem = self._memories_container.read_item(item=supersedes_id, partition_key=[user_id, thread_id]) - if old_mem.get("superseded_by"): - logger.debug( - "extract_memories: skipping UPDATE — target %s already superseded by %s", - supersedes_id, - old_mem.get("superseded_by"), - ) - return False - return self._mark_superseded(old_mem, superseder_id, reason=reason) - except CosmosResourceNotFoundError: - logger.debug( - "extract_memories: %s not found at (user_id, thread_id) — retrying cross-partition", - supersedes_id, - ) - except Exception as exc: - # Includes 429s, 503s, transient connection errors — surface at WARNING - # so they're not masked by the silent cross-partition fallback below. - logger.warning( - "extract_memories: read_item failed for %s (%s); retrying cross-partition", - supersedes_id, - type(exc).__name__, - ) - try: - q = f"SELECT * FROM c WHERE c.id = @id AND c.user_id = @uid AND {_ACTIVE_DOC_FILTER}" - results = list( - self._memories_container.query_items( - query=q, - parameters=[ - {"name": "@id", "value": supersedes_id}, - {"name": "@uid", "value": user_id}, - ], - enable_cross_partition_query=True, - ) - ) - if results and not results[0].get("superseded_by"): - return self._mark_superseded(results[0], superseder_id, reason=reason) - except CosmosHttpResponseError as exc: - logger.warning("Failed to mark superseded memory %s: %s", supersedes_id, exc) - return False - def _mark_superseded( self, old_doc: dict[str, Any], @@ -462,28 +593,35 @@ def extract_memories_dry( logger.warning("extract_memories_dry no memories found user_id=%s thread_id=%s", user_id, thread_id) return {"facts": [], "episodic": [], "updates": [], "processed_turn_docs": []} - existing_facts = self._load_existing_memories(user_id, ["fact"]) - existing_episodics = self._load_existing_memories(user_id, ["episodic"]) + existing_for_hashes = self._load_existing_memories(user_id, ["fact"]) existing_fact_hashes: set[str] = { - m["content_hash"] for m in existing_facts if m.get("type") == "fact" and m.get("content_hash") + m["content_hash"] for m in existing_for_hashes if m.get("type") == "fact" and m.get("content_hash") } - if existing_facts: + transcript = self._build_transcript(items) + existing = existing_for_hashes + if threshold_config.get_dedup_context_vector_enabled(): + user_turns_text = "\n".join( + str(it.get("content", "")) for it in items if it.get("role") == "user" + ).strip() + context_query = user_turns_text or transcript + existing = self._store.search( + search_terms=context_query, + user_id=user_id, + memory_types=["fact"], + top_k=threshold_config.get_dedup_context_topk(), + ) + if existing: existing_text = "\n".join( - f"- [ID: {mem['id']}] {mem.get('content', '')} (type=fact, salience={mem.get('salience', 'N/A')})" - for mem in existing_facts + f"- [ID: {mem['id']}] {mem.get('content', '')} " + f"(type={mem.get('type', 'fact')}, salience={mem.get('salience', 'N/A')})" + for mem in existing ) else: existing_text = "(none)" - existing_episodics_text = _format_existing_episodics(existing_episodics) - transcript = self._build_transcript(items) response_text = self._run_prompty( "extract_memories.prompty", - inputs={ - "existing_facts": existing_text, - "existing_episodics": existing_episodics_text, - "transcript": transcript, - }, + inputs={"existing_facts": existing_text, "transcript": transcript}, ) parsed = self._parse_llm_json(response_text) facts = parsed.get("facts", []) @@ -498,14 +636,13 @@ def extract_memories_dry( dropped_episodic_count = 0 for fact in facts: - action = fact.get("action", "ADD").upper() text = fact.get("text") if not text: logger.warning("extract_memories: dropping malformed fact (missing 'text'): %r", fact) continue new_content_hash = compute_content_hash(text) - if action == "ADD" and new_content_hash in existing_fact_hashes: + if new_content_hash in existing_fact_hashes: logger.debug( "extract_memories: skipping exact-dup fact hash=%s user_id=%s thread_id=%s", new_content_hash, @@ -542,22 +679,6 @@ def extract_memories_dry( "updated_at": doc_timestamp, } - if action in {"UPDATE", "CONTRADICT"} and fact.get("supersedes_id"): - reason: Literal["update", "contradict"] = "contradict" if action == "CONTRADICT" else "update" - if det_id == fact["supersedes_id"]: - logger.debug("extract_memories: skipping UPDATE — det_id == supersedes_id (%s)", det_id) - continue - doc["supersedes_ids"] = [fact["supersedes_id"]] - updates.append( - { - "op": "supersede", - "supersedes_id": fact["supersedes_id"], - "superseder_id": det_id, - "thread_id": thread_id, - "reason": reason, - } - ) - fact_docs.append(self._validate_extracted_doc(doc)) existing_fact_hashes.add(new_content_hash) @@ -577,28 +698,16 @@ def extract_memories_dry( dropped_episodic_count += 1 continue - text_raw = ep.get("text") - text = text_raw.strip() if isinstance(text_raw, str) else None - if not text: - logger.error( - "extract_memories: dropping episodic with empty/missing text field " - "(LLM extraction did not populate the required `text` field — likely a " - "weaker extraction model that needs upgrading or a prompt-compliance issue). " - "scope_type=%s scope_value=%s user_id=%s thread_id=%s reason=missing_text", - scope_type, - scope_value, - user_id, - thread_id, - ) - dropped_episodic_count += 1 - continue - situation = ep.get("situation") action_taken = ep.get("action_taken") outcome = ep.get("outcome") + if situation and action_taken and outcome: + text = f"{situation} → {action_taken} → {outcome}" + else: + text = f"For the user's {scope_value} {scope_type}, intent recorded." content_hash = compute_content_hash(text) - seed = _ID_SEED_SEP.join((user_id, scope_type, scope_value)) + seed = _ID_SEED_SEP.join((user_id, thread_id, content_hash)) det_id = f"ep_{hashlib.sha256(seed.encode()).hexdigest()[:32]}" topic_tags = build_topic_tags(ep.get("tags", [])) confidence = ep.get("confidence") @@ -615,7 +724,7 @@ def extract_memories_dry( doc = { "id": det_id, "user_id": user_id, - "thread_id": "__episodic__", + "thread_id": thread_id, "role": "system", "type": "episodic", "content": text, @@ -626,7 +735,6 @@ def extract_memories_dry( "metadata": { "scope_type": scope_type, "scope_value": scope_value, - "originating_thread_id": thread_id, "situation": situation, "action_taken": action_taken, "outcome": outcome, @@ -690,7 +798,7 @@ def extract_memories_dry( check_extracted_fact_grounding( fact_docs, items, - existing_facts, + existing, user_id=user_id, thread_id=thread_id, logger=logger, @@ -712,6 +820,104 @@ def extract_memories_dry( ) return result + def dedup_extracted_memories( + self, + user_id: str, + extracted: dict[str, list[dict[str, Any]]], + ) -> dict[str, list[dict[str, Any]]]: + """Apply the gated Stage-3 vector dedup ladder to extracted docs.""" + if not threshold_config.get_dedup_vector_enabled(): + return extracted + if not user_id: + raise ValidationError("user_id is required") + if not isinstance(extracted, dict): + raise ValidationError("extracted must be a dict") + + high = threshold_config.get_dedup_sim_high() + low = threshold_config.get_dedup_sim_low() + top_k = threshold_config.get_dedup_candidate_topk() + distance_function = self._vector_distance_function() + autodrop_ok = vector_autodrop_supported(distance_function) + if not autodrop_ok: + self._warn_euclidean_autodrop_once(distance_function) + vector_dedup_skipped = 0 + dup_candidates_tagged = 0 + + result: dict[str, list[dict[str, Any]]] = { + "facts": [dict(doc) for doc in extracted.get("facts", [])], + "episodic": [dict(doc) for doc in extracted.get("episodic", [])], + "updates": [dict(op) for op in extracted.get("updates", [])], + } + docs = [doc for bucket in ("facts", "episodic") for doc in result[bucket] if doc.get("content")] + if not docs: + return result + + missing_embeddings = [doc for doc in docs if not doc.get("embedding")] + if missing_embeddings: + embeddings = self._embed_batch([str(doc["content"]) for doc in missing_embeddings]) + for doc, embedding in zip(missing_embeddings, embeddings): + doc["embedding"] = embedding + + kept_ids: set[str] = set() + dropped_ids: set[str] = set() + for doc in docs: + doc_id = str(doc.get("id") or "") + memory_type = str(doc.get("type") or "") + if not doc_id or memory_type not in {"fact", "episodic"}: + continue + + best: dict[str, Any] | None = None + candidates = self._vector_candidates( + user_id=user_id, + embedding=doc.get("embedding") or [], + memory_type=memory_type, + top_k=top_k, + exclude_ids=kept_ids | dropped_ids | {doc_id} | set(doc.get("supersedes_ids") or []), + ) + if candidates: + best = candidates[0] + + score = float(best.get("score") or 0.0) if best else 0.0 + if best and autodrop_ok and vector_similarity_at_least(score, high, distance_function): + logger.info( + "vector dedup skipped new memory id=%s type=%s score=%.4f surviving_id=%s " + "new_content=%r surviving_content=%r", + doc_id, + memory_type, + score, + best.get("id"), + doc.get("content"), + best.get("content"), + ) + vector_dedup_skipped += 1 + dropped_ids.add(doc_id) + continue + + if best and vector_similarity_at_least(score, low, distance_function): + tags = list(doc.get("tags") or []) + if "sys:dup-candidate" not in tags: + tags.append("sys:dup-candidate") + doc["tags"] = tags + metadata = dict(doc.get("metadata") or {}) + metadata["dup_of"] = best.get("id") + metadata["dup_score"] = score + doc["metadata"] = metadata + dup_candidates_tagged += 1 + + kept_ids.add(doc_id) + + for bucket in ("facts", "episodic"): + result[bucket] = [doc for doc in result[bucket] if doc.get("id") not in dropped_ids] + if vector_dedup_skipped or dup_candidates_tagged: + result["updates"].append( + { + "op": "stats", + "vector_dedup_skipped": vector_dedup_skipped, + "dup_candidates_tagged": dup_candidates_tagged, + } + ) + return result + def persist_extracted_memories( self, user_id: str, @@ -764,27 +970,9 @@ def persist_extracted_memories( if op.get("op") == "stats": result["exact_dedup_skipped"] += int(op.get("exact_dedup_skipped") or 0) result["dropped_episodic_count"] += int(op.get("dropped_episodic_count") or 0) - continue - if op.get("op") != "supersede": - continue - reason = op.get("reason") - op_thread_id = op.get("thread_id") - supersedes_id = op.get("supersedes_id") - superseder_id = op.get("superseder_id") - if reason not in {"update", "contradict"} or not op_thread_id or not supersedes_id or not superseder_id: - continue - marked = self._mark_extracted_superseded( - user_id=user_id, - thread_id=op_thread_id, - supersedes_id=supersedes_id, - superseder_id=superseder_id, - reason=reason, - ) - if marked: - if reason == "contradict": - result["contradicted_count"] += 1 - else: - result["updated_count"] += 1 + for key in ("vector_dedup_skipped", "dup_candidates_tagged"): + if key in op: + result[key] = result.get(key, 0) + int(op.get(key) or 0) logger.info("persist_extracted_memories completed user_id=%s counts=%s", user_id, result) @@ -839,6 +1027,8 @@ def extract_memories( ) -> dict[str, int]: """Extract facts and episodic memories from a thread and persist them.""" extracted = self.extract_memories_dry(user_id, thread_id, recent_k, turns=turns) + if threshold_config.get_dedup_vector_enabled(): + extracted = self.dedup_extracted_memories(user_id, extracted) return self.persist_extracted_memories(user_id, extracted) def synthesize_procedural( @@ -1403,7 +1593,171 @@ def _emit_reconcile_outcome( }, ) - def reconcile_memories(self, user_id: str, n: int = 50) -> dict[str, int]: + def _build_candidate_clusters( + self, user_id: str, memory_type: str, n: int + ) -> tuple[list[list[dict[str, Any]]], int, list[dict[str, Any]]]: + """Cluster dup-candidate seeds (+ vector neighbors) into connected components. + + Returns ``(clusters, node_count, seeds)`` where each cluster has >= 2 members, + ``node_count`` is the total distinct memories pulled into the graph (used to + report ``reconcile_llm_calls_saved``), and ``seeds`` is the tagged seed scan + (so the caller can clear stale tags on orphan seeds that never clustered). + The seed scan is bounded to ``n`` so a single cluster can never exceed the + reconcile prompt's pool cap. + """ + cluster_sim = threshold_config.get_dedup_cluster_sim() + top_k = threshold_config.get_dedup_candidate_topk() + distance_function = self._vector_distance_function() + seeds = self._query_active_memories(user_id, memory_type, limit=n, tagged_only=True) + nodes_by_id: dict[str, dict[str, Any]] = {doc["id"]: doc for doc in seeds if doc.get("id")} + edge_pairs: set[tuple[str, str]] = set() + + dup_of_ids = [ + (doc.get("metadata") or {}).get("dup_of") + for doc in seeds + if isinstance(doc.get("metadata"), dict) and (doc.get("metadata") or {}).get("dup_of") + ] + for doc in self._load_memories_by_ids(user_id, memory_type, dup_of_ids): + if doc.get("id"): + nodes_by_id[doc["id"]] = doc + + for seed in seeds: + seed_id = seed.get("id") + if not seed_id: + continue + dup_of = (seed.get("metadata") or {}).get("dup_of") if isinstance(seed.get("metadata"), dict) else None + if dup_of: + edge_pairs.add(tuple(sorted((seed_id, dup_of)))) + if not seed.get("embedding"): + continue + candidates = self._vector_candidates( + user_id=user_id, + embedding=seed.get("embedding") or [], + memory_type=memory_type, + top_k=top_k, + exclude_ids={seed_id}, + ) + candidate_ids = [ + row["id"] + for row in candidates + if row.get("id") + and vector_similarity_at_least(row.get("score", 0.0), cluster_sim, distance_function) + ] + for cid in candidate_ids: + edge_pairs.add(tuple(sorted((seed_id, cid)))) + for doc in self._load_memories_by_ids(user_id, memory_type, candidate_ids): + if doc.get("id"): + nodes_by_id[doc["id"]] = doc + + node_ids = list(nodes_by_id) + node_id_set = set(node_ids) + for node_id in node_ids: + node = nodes_by_id[node_id] + if not node.get("embedding"): + continue + for row in self._vector_candidates( + user_id=user_id, + embedding=node.get("embedding") or [], + memory_type=memory_type, + top_k=top_k, + exclude_ids={node_id}, + ): + neighbor_id = row.get("id") + if neighbor_id in node_id_set and vector_similarity_at_least( + float(row.get("score") or 0.0), cluster_sim, distance_function + ): + edge_pairs.add(tuple(sorted((node_id, neighbor_id)))) + + adjacency: dict[str, set[str]] = {node_id: set() for node_id in node_ids} + for left_id, right_id in edge_pairs: + if left_id != right_id and left_id in adjacency and right_id in adjacency: + adjacency[left_id].add(right_id) + adjacency[right_id].add(left_id) + + clusters: list[list[dict[str, Any]]] = [] + seen: set[str] = set() + for node_id in node_ids: + if node_id in seen: + continue + stack = [node_id] + component: list[str] = [] + seen.add(node_id) + while stack: + current = stack.pop() + component.append(current) + for neighbor in adjacency.get(current, set()): + if neighbor not in seen: + seen.add(neighbor) + stack.append(neighbor) + if len(component) >= 2: + # Cap cluster size at the reconcile pool limit: lowering the cluster + # threshold can chain many facts into one giant transitive component + # that would blow the prompt cap; keep the most-recent ``n``. + if len(component) > n: + component = component[:n] + clusters.append([nodes_by_id[cid] for cid in component]) + return clusters, len(nodes_by_id), seeds + + def _reconcile_candidate_mode( + self, user_id: str, *, n: int, memory_type: str, started_at: float + ) -> dict[str, int]: + # Candidate clustering only. The periodic full-pool backstop that catches + # dissimilar-embedding contradictions ("vegetarian" vs "loves steak") is + # driven by the caller via ``full_rebuild`` on a PERSISTED-counter cadence + # (in-process auto-trigger + durable change-feed), not an in-memory sweep + # counter — the latter reset per worker/process and never fired reliably on + # the Function-App backend. + clusters, node_count, seeds = self._build_candidate_clusters(user_id, memory_type, n) + aggregate = {"kept": 0, "merged": 0, "contradicted": 0} + clustered_ids: set[str] = set() + for cluster in clusters: + # Mark members as clustered BEFORE the LLM call so a failed cluster's + # seeds are not treated as orphans below (which would clear their tags + # and prevent a retry). A truncated/malformed LLM response on one + # cluster must not abort the sweep or starve the remaining clusters. + clustered_ids.update(doc["id"] for doc in cluster if doc.get("id")) + try: + counts, consumed = self._reconcile_pool(user_id, memory_type, cluster) + except Exception as exc: + logger.warning( + "reconcile_memories: cluster reconcile failed user_id=%s memory_type=%s; " + "skipping cluster, tags retained for next sweep: %s", + user_id, + memory_type, + exc, + ) + continue + for key in aggregate: + aggregate[key] += int(counts.get(key, 0)) + # Clear dup-candidate tags only on survivors. Re-upserting a doc that + # was just superseded (duplicate source or contradiction loser) would + # resurrect it, since the in-memory cluster copy lacks superseded_by. + survivors = [doc for doc in cluster if doc.get("id") and doc["id"] not in consumed] + self._clear_dup_candidate_tags(survivors) + # Orphan seeds: tagged dup-candidates that never joined a cluster have no + # near-duplicate, so clear the stale tag — otherwise every future sweep + # re-scans them as seeds and they accumulate forever. + orphan_seeds = [seed for seed in seeds if seed.get("id") and seed["id"] not in clustered_ids] + self._clear_dup_candidate_tags(orphan_seeds) + aggregate["reconcile_clusters_sent"] = len(clusters) + aggregate["reconcile_llm_calls_saved"] = max(0, node_count - len(clusters)) + logger.info( + "reconcile_memories candidate completed user_id=%s memory_type=%s result=%s", + user_id, + memory_type, + aggregate, + ) + self._emit_reconcile_outcome( + started_at=started_at, + user_id=user_id, + candidates=node_count, + result=aggregate, + ) + return aggregate + + def reconcile_memories( + self, user_id: str, n: int = 50, *, memory_type: str = "fact", full_rebuild: bool = False + ) -> dict[str, int]: """Reconcile a user's active facts in a single LLM pass. Loads the most recent ``n`` active (non-superseded) facts for @@ -1418,6 +1772,12 @@ def reconcile_memories(self, user_id: str, n: int = 50) -> dict[str, int]: the winner. Dangling references are resolved transparently when a contradicted id was just absorbed into a duplicate group. + ``full_rebuild`` (set by the explicit public ``reconcile()`` entrypoint) + makes candidate mode cluster the whole active pool instead of only + Stage-3-tagged seeds, so a user-triggered reconcile dedups every active + fact. Automatic background sweeps leave it ``False`` for the cheap + tagged-only path. + Returns ``{"kept": int, "merged": int, "contradicted": int}`` where ``merged`` and ``contradicted`` count the *losers* that were soft-deleted (duplicates and contradictions respectively). @@ -1428,46 +1788,97 @@ def reconcile_memories(self, user_id: str, n: int = 50) -> dict[str, int]: raise ValidationError(f"n must be a positive integer, got {n!r}") if n > 500: raise ValidationError(f"n must be <= 500 to bound prompt size and LLM cost, got {n}") + if memory_type not in {"fact", "episodic", "procedural"}: + raise ValidationError(f"memory_type must be one of fact, episodic, procedural; got {memory_type!r}") + if memory_type == "procedural": + result = { + "kept": 0, + "merged": 0, + "contradicted": 0, + "reconcile_clusters_sent": 0, + "reconcile_llm_calls_saved": 0, + } + logger.info("reconcile_memories procedural no-op user_id=%s result=%s", user_id, result) + return result started_at = time.monotonic() - logger.info("reconcile_memories started user_id=%s n=%d", user_id, n) + logger.info("reconcile_memories started user_id=%s n=%d memory_type=%s", user_id, n, memory_type) + + # Explicit user-triggered reconcile (full_rebuild) always takes the + # full-pool single-LLM-pass path: it sees every active fact together, so it + # catches contradictions that aren't vector-similar (e.g. "vegetarian" vs + # "loves steak") — which candidate clustering, keyed on near-duplicate + # similarity, would never group. Automatic sweeps use cheap candidate mode. + if threshold_config.get_dedup_reconcile_mode() == "candidate" and not full_rebuild: + return self._reconcile_candidate_mode( + user_id, n=n, memory_type=memory_type, started_at=started_at + ) - # ---- 1. Load up to N most recent active facts ---- + facts = self._active_memories_for_reconcile(user_id, memory_type, n) + result, consumed = self._reconcile_pool(user_id, memory_type, facts) + # Clear dup-candidate tags on survivors so an explicit reconcile(full_rebuild=True) + # doesn't leave stale sys:dup-candidate/dup_of metadata on user-visible memories. + survivors = [doc for doc in facts if doc.get("id") and doc["id"] not in consumed] + self._clear_dup_candidate_tags(survivors) + self._emit_reconcile_outcome( + started_at=started_at, + user_id=user_id, + candidates=len(facts), + result=result, + ) + return result + + def _active_memories_for_reconcile( + self, user_id: str, memory_type: str, n: int + ) -> list[dict[str, Any]]: + # ---- Load up to N most recent active memories ---- # ORDER BY c.created_at DESC keeps the TOP cap deterministic across # physical partitions and matches the dedup prompt's tiebreaker # ("more recent created_at first"). Cosmos's _ts is the last-write # timestamp, which would diverge from created_at after any UPDATE. query = ( - f"SELECT TOP {n} * FROM c " + f"SELECT TOP {top_literal(n, name='reconcile_memories.n')} * FROM c " "WHERE c.user_id = @user_id " - "AND c.type = 'fact' " + "AND c.type = @memory_type " f"AND {_ACTIVE_DOC_FILTER} " "ORDER BY c.created_at DESC" ) - parameters: list[dict[str, Any]] = [ - {"name": "@user_id", "value": user_id}, - ] - facts = list( + return list( self._memories_container.query_items( query=query, - parameters=parameters, + parameters=[ + {"name": "@user_id", "value": user_id}, + {"name": "@memory_type", "value": memory_type}, + ], enable_cross_partition_query=True, ) ) + def _reconcile_pool( + self, user_id: str, memory_type: str, facts: list[dict[str, Any]] + ) -> tuple[dict[str, int], set[str]]: + """Reconcile an explicit pool of same-type memories in one LLM pass. + + Returns ``({"kept", "merged", "contradicted"}, consumed_ids)`` where + ``consumed_ids`` are the source/loser ids that were actually superseded, + so callers can skip them when clearing dup-candidate tags (re-upserting a + superseded source would resurrect it). Does not emit telemetry — the + caller owns the ``reconcile.outcome`` line. + """ + # Polarity keyed on an explicit predicate so a future third type (e.g. + # procedural, were it ever routed here) can't silently diverge from async. + is_episodic = memory_type == "episodic" + prompt_filename = "dedup_episodic.prompty" if is_episodic else "dedup.prompty" + prompt_input_key = "episodics_text" if is_episodic else "facts_text" + allow_contradictions = not is_episodic + if len(facts) <= 1: logger.info( - "reconcile_memories: %d facts, nothing to reconcile", + "reconcile_memories: %d %s memories, nothing to reconcile", len(facts), + memory_type, ) - early_result = {"kept": len(facts), "merged": 0, "contradicted": 0} - self._emit_reconcile_outcome( - started_at=started_at, - user_id=user_id, - candidates=len(facts), - result=early_result, - ) - return early_result + return {"kept": len(facts), "merged": 0, "contradicted": 0}, set() # ---- 2. Format the facts pool for the prompt ---- # ``json.dumps`` escapes embedded quotes and pipes inside content so @@ -1495,13 +1906,13 @@ def reconcile_memories(self, user_id: str, n: int = 50) -> dict[str, int]: # ---- 3. Single LLM call over the entire pool ---- response_text = self._run_prompty( - "dedup.prompty", - inputs={"facts_text": facts_text}, + prompt_filename, + inputs={prompt_input_key: facts_text}, ) parsed = self._parse_llm_json(response_text) duplicate_groups = parsed.get("duplicate_groups", []) or [] - contradicted_pairs = parsed.get("contradicted_pairs", []) or [] + contradicted_pairs = (parsed.get("contradicted_pairs", []) or []) if allow_contradictions else [] # ``kept_ids`` from the LLM is used below as a cross-check for # accounting drift (hallucinated IDs, double-counting). The actual # kept count is computed from facts minus consumed losers. @@ -1573,15 +1984,19 @@ def reconcile_memories(self, user_id: str, n: int = 50) -> dict[str, int]: source_docs.sort(key=lambda d: d.get("_ts", 0), reverse=True) # Union tags across all source docs (preserve order, dedupe). + # Strip sys:dup-candidate so the merged doc isn't reborn as a + # permanent reconcile seed. merged_tags: list[str] = [] seen_tags: set[str] = set() for src in source_docs: for t in src.get("tags", []) or []: + if t == "sys:dup-candidate": + continue if t not in seen_tags: seen_tags.add(t) merged_tags.append(t) if not merged_tags: - merged_tags = ["sys:fact"] + merged_tags = [f"sys:{memory_type}"] # Union source_memory_ids across all source docs (provenance chain). merged_source_memory_ids: list[str] = [] @@ -1638,16 +2053,15 @@ def reconcile_memories(self, user_id: str, n: int = 50) -> dict[str, int]: # see the same id rather than chaining through a new UUID. merged_content_hash = compute_content_hash(merged_content) merged_id_seed = _ID_SEED_SEP.join((user_id, "merged", merged_content_hash)) - merged_id = "fact_" + hashlib.sha256(merged_id_seed.encode()).hexdigest()[:32] + id_prefix = "fact_" if memory_type == "fact" else "ep_" + merged_id = id_prefix + hashlib.sha256(merged_id_seed.encode()).hexdigest()[:32] try: - merged_record = construct_internal( - FactRecord, - { + merged_payload: dict[str, Any] = { "id": merged_id, "user_id": user_id, "role": "system", - "type": "fact", + "type": memory_type, "content": merged_content, "thread_id": recent_thread_id or f"__reconciled__:{user_id}", "confidence": confidence_val if confidence_val is not None else 0.5, @@ -1656,16 +2070,32 @@ def reconcile_memories(self, user_id: str, n: int = 50) -> dict[str, int]: "source_memory_ids": merged_source_memory_ids, "tags": merged_tags, "content_hash": merged_content_hash, - "metadata": { - "category": "preference", - "merged_via": "reconcile", - "merged_from_count": len(valid_source_ids), - }, - **self._prompt_lineage("dedup.prompty"), + **self._prompt_lineage(prompt_filename), "created_at": datetime.now(timezone.utc).isoformat(), "updated_at": datetime.now(timezone.utc).isoformat(), - }, - ) + } + if memory_type == "fact": + merged_payload["metadata"] = { + "category": "preference", + "merged_via": "reconcile", + "merged_from_count": len(valid_source_ids), + } + record_cls = FactRecord + else: + source_meta = dict(source_docs[0].get("metadata") or {}) + source_meta.update( + { + "lesson": merged_content, + "merged_via": "reconcile", + "merged_from_count": len(valid_source_ids), + } + ) + source_meta.setdefault("scope_type", source_meta.get("scope_type") or "general") + source_meta.setdefault("scope_value", source_meta.get("scope_value") or "general") + source_meta.setdefault("outcome_valence", source_meta.get("outcome_valence") or "neutral") + merged_payload["metadata"] = source_meta + record_cls = EpisodicRecord + merged_record = construct_internal(record_cls, merged_payload) except Exception: logger.exception( "reconcile_memories: failed to build merged record for group %r", @@ -1833,13 +2263,7 @@ def reconcile_memories(self, user_id: str, n: int = 50) -> dict[str, int]: ) result = {"kept": kept, "merged": merged, "contradicted": contradicted} logger.info("reconcile_memories completed: %s", result) - self._emit_reconcile_outcome( - started_at=started_at, - user_id=user_id, - candidates=len(facts), - result=result, - ) - return result + return result, consumed_ids def build_procedural_context(self, user_id: str) -> str: """Return the active synthesized procedural prompt for system injection.""" diff --git a/azure/cosmos/agent_memory/store/_search_helpers.py b/azure/cosmos/agent_memory/store/_search_helpers.py index b0f6c5e..032965c 100644 --- a/azure/cosmos/agent_memory/store/_search_helpers.py +++ b/azure/cosmos/agent_memory/store/_search_helpers.py @@ -94,16 +94,27 @@ def build_search_sql( *, qb: _QueryBuilder, top: int, - hybrid_search: bool, + keyword_count: int, include_superseded: bool, ) -> str: + """Build the search SQL. + + When ``keyword_count > 0`` the query is hybrid: ``RANK RRF`` fuses the vector + distance with ``FullTextScore`` over ``keyword_count`` individual keyword + parameters (``@kw0``..``@kw{n-1}``). When there are no keywords (e.g. an + all-stopword query) it falls back to pure vector ranking. ``similarity_score`` + is always the vector distance and is *not* the RRF ranking basis under hybrid. + """ if not include_superseded: qb.add_is_null_or_undefined("c.superseded_by") - order_by = "ORDER BY VectorDistance(c.embedding, @embedding)" - if hybrid_search: - order_by = "ORDER BY RANK RRF(VectorDistance(c.embedding, @embedding), FullTextScore(c.content, @key_terms))" + vector_distance = "VectorDistance(c.embedding, @embedding)" + if keyword_count > 0: + keyword_params = ", ".join(f"@kw{i}" for i in range(keyword_count)) + order_by = f"ORDER BY RANK RRF({vector_distance}, FullTextScore(c.content, {keyword_params}))" + else: + order_by = f"ORDER BY {vector_distance}" return ( f"SELECT TOP {top} {MEMORY_PROJECTION}, " - "VectorDistance(c.embedding, @embedding) AS similarity_score " + f"{vector_distance} AS similarity_score " f"FROM c{qb.build_where()} {order_by}" ) diff --git a/azure/cosmos/agent_memory/store/memory_store.py b/azure/cosmos/agent_memory/store/memory_store.py index 5755311..dd5d345 100644 --- a/azure/cosmos/agent_memory/store/memory_store.py +++ b/azure/cosmos/agent_memory/store/memory_store.py @@ -15,8 +15,8 @@ from azure.cosmos.agent_memory._utils import ( _build_memory_query_builder, _coerce_datetime_iso, - _validate_hybrid_search, compute_content_hash, + extract_keywords, new_id, ) from azure.cosmos.agent_memory.exceptions import ( @@ -840,7 +840,6 @@ def search( role: Optional[str] = None, memory_types: Optional[list[str]] = None, thread_id: Optional[str] = None, - hybrid_search: bool = False, top_k: int = 5, tags_all: Optional[list[str]] = None, tags_any: Optional[list[str]] = None, @@ -859,9 +858,9 @@ def search( :meth:`search_turns` to vector-search the raw conversation log instead. """ terms = require_search_terms(search_terms, query) - _validate_hybrid_search(hybrid_search, terms) top = top_literal(top_k, name="top_k") query_vector = self._embed(terms) + keywords = extract_keywords(terms) qb = _build_memory_query_builder( memory_id=memory_id, @@ -881,11 +880,16 @@ def search( ) add_salience_filter(qb, min_salience) - sql = build_search_sql(qb=qb, top=top, hybrid_search=hybrid_search, include_superseded=include_superseded) + sql = build_search_sql( + qb=qb, + top=top, + keyword_count=len(keywords), + include_superseded=include_superseded, + ) parameters = qb.get_parameters() parameters.append({"name": "@embedding", "value": query_vector}) - if hybrid_search: - parameters.append({"name": "@key_terms", "value": terms}) + for i, kw in enumerate(keywords): + parameters.append({"name": f"@kw{i}", "value": kw}) partition_key, cross_partition = query_scope(user_id, thread_id) if thread_id is not None and (not memory_types or set(memory_types) & USER_SCOPED_MEMORIES_TYPES): @@ -905,7 +909,6 @@ def search_turns( user_id: Optional[str] = None, thread_id: Optional[str] = None, role: Optional[str] = None, - hybrid_search: bool = False, top_k: int = 5, tags_all: Optional[list[str]] = None, tags_any: Optional[list[str]] = None, @@ -915,7 +918,7 @@ def search_turns( *, query: Optional[str] = None, ) -> list[dict[str, Any]]: - """Search raw conversation turns using vector similarity with optional hybrid ranking. + """Search raw conversation turns using vector similarity with hybrid ranking. Only vector-searchable when turn embeddings were enabled at write time (see ``enable_turn_embeddings``). ``user_id`` is required and always @@ -927,9 +930,9 @@ def search_turns( if not user_id: raise ValidationError("user_id is required for search_turns") terms = require_search_terms(search_terms, query) - _validate_hybrid_search(hybrid_search, terms) top = top_literal(top_k, name="top_k") query_vector = self._embed(terms) + keywords = extract_keywords(terms) qb = _QueryBuilder() qb.add_filter("c.user_id", "@user_id", user_id) @@ -944,11 +947,11 @@ def search_turns( before_param="@created_before", ) - sql = build_search_sql(qb=qb, top=top, hybrid_search=hybrid_search, include_superseded=False) + sql = build_search_sql(qb=qb, top=top, keyword_count=len(keywords), include_superseded=False) parameters = qb.get_parameters() parameters.append({"name": "@embedding", "value": query_vector}) - if hybrid_search: - parameters.append({"name": "@key_terms", "value": terms}) + for i, kw in enumerate(keywords): + parameters.append({"name": f"@kw{i}", "value": kw}) partition_key, cross_partition = query_scope(user_id, thread_id) logger.debug("MemoryStore.search_turns query: %s", sql) diff --git a/azure/cosmos/agent_memory/thresholds.py b/azure/cosmos/agent_memory/thresholds.py index 7fffc60..ca4c508 100644 --- a/azure/cosmos/agent_memory/thresholds.py +++ b/azure/cosmos/agent_memory/thresholds.py @@ -28,6 +28,23 @@ # of 500 (enforced by the pipeline) bounds prompt size and LLM cost. DEFAULT_DEDUP_POOL_SIZE = 50 +# --------------------------------------------------------------------------- +# INTERNAL dedup/search tuning — NOT customer-configurable. +# These ship as fixed feature constants (no env vars, not in any settings +# template). They are maintainer-tunable here in code only; if a knob ever +# needs to become operator-facing we add the env plumbing back deliberately. +# The dedup + hybrid-search features ship ON via these values. +# --------------------------------------------------------------------------- +DEDUP_CONTEXT_VECTOR_ENABLED = True # Stage-1 relevance-ranked extraction context +DEDUP_CONTEXT_TOPK = 10 +DEDUP_VECTOR_ENABLED = True # Stage-3 vector near-dup ladder +DEDUP_SIM_HIGH = 0.97 # >= -> auto-skip near-exact +DEDUP_SIM_LOW = 0.80 # < -> novel; between -> tag candidate +DEDUP_CANDIDATE_TOPK = 10 +DEDUP_RECONCILE_MODE = "candidate" # clustered candidate reconcile (vs legacy full_pool) +DEDUP_CLUSTER_SIM = 0.60 # Stage-5 clustering edge threshold +DEDUP_FULL_RECLUSTER_EVERY_N = 12 # full re-cluster safety net cadence + DEFAULT_TTL_BY_TYPE: dict[str, int] = { "turn": 2_592_000, "episodic": 7_776_000, @@ -138,6 +155,56 @@ def get_dedup_pool_size() -> int: return raw +# --------------------------------------------------------------------------- +# Internal dedup/search feature accessors — return fixed constants (no env). +# Kept as thin functions so call sites stay stable and the values can be +# changed in one place; NOT customer-configurable. +# --------------------------------------------------------------------------- +def get_dedup_context_vector_enabled() -> bool: + """Whether Stage-1 extraction context uses vector retrieval (internal; on).""" + return DEDUP_CONTEXT_VECTOR_ENABLED + + +def get_dedup_context_topk() -> int: + """Top-K memories to retrieve for Stage-1 extraction context (internal).""" + return DEDUP_CONTEXT_TOPK + + +def get_dedup_vector_enabled() -> bool: + """Whether Stage-3 vector deduplication is enabled (internal; on).""" + return DEDUP_VECTOR_ENABLED + + +def get_dedup_sim_high() -> float: + """Similarity at/above which near-exact memories are auto-skipped (internal).""" + return DEDUP_SIM_HIGH + + +def get_dedup_sim_low() -> float: + """Similarity below which memories are treated as novel (internal).""" + return DEDUP_SIM_LOW + + +def get_dedup_candidate_topk() -> int: + """Top-K existing memories pulled per new memory in Stage-3 (internal).""" + return DEDUP_CANDIDATE_TOPK + + +def get_dedup_reconcile_mode() -> str: + """Reconcile mode: ``candidate`` clustering (internal; the shipped feature).""" + return DEDUP_RECONCILE_MODE + + +def get_dedup_cluster_sim() -> float: + """Similarity threshold for Stage-5 clustering edges (internal).""" + return DEDUP_CLUSTER_SIM + + +def get_dedup_full_recluster_every_n() -> int: + """Full re-cluster safety-net cadence, every Nth reconcile sweep (internal).""" + return DEDUP_FULL_RECLUSTER_EVERY_N + + def get_procedural_synthesis_auto() -> bool: """Whether procedural synthesis auto-fires after extract. @@ -214,6 +281,15 @@ def get_processor_owner() -> Optional[str]: "get_user_summary_every_n", "get_dedup_every_n", "get_dedup_pool_size", + "get_dedup_context_vector_enabled", + "get_dedup_context_topk", + "get_dedup_vector_enabled", + "get_dedup_sim_high", + "get_dedup_sim_low", + "get_dedup_candidate_topk", + "get_dedup_reconcile_mode", + "get_dedup_cluster_sim", + "get_dedup_full_recluster_every_n", "get_procedural_synthesis_auto", "get_enable_turn_embeddings", "get_processor_owner", diff --git a/function_app/local.settings.json.template b/function_app/local.settings.json.template index e2f2fa6..2d65bde 100644 --- a/function_app/local.settings.json.template +++ b/function_app/local.settings.json.template @@ -22,6 +22,7 @@ "// --- Threshold knobs (set to 0 to disable that orchestrator) ---": "", "THREAD_SUMMARY_EVERY_N": "10", "FACT_EXTRACTION_EVERY_N": "1", + "DEDUP_EVERY_N": "5", "USER_SUMMARY_EVERY_N": "20", "// --- Processor ownership: 'durable' (FA owns) or 'inprocess' (SDK owns). FA skips processing if not set to 'durable' to avoid double-fire when an SDK install runs against the same Cosmos. ---": "", diff --git a/function_app/orchestrators/extract_memories.py b/function_app/orchestrators/extract_memories.py index a911f20..ec859bc 100644 --- a/function_app/orchestrators/extract_memories.py +++ b/function_app/orchestrators/extract_memories.py @@ -1,7 +1,8 @@ """Memory-extraction orchestrator + activities. -Chain: ``Extract`` → ``Persist`` followed by an optional ``ReconcileMemories`` -activity, then a best-effort ``SynthesizeProceduralOrchestrator`` sub-call. +Chain: ``Extract`` → ``Dedup`` → ``Persist`` followed by an optional +``ReconcileMemories`` activity, then a best-effort +``SynthesizeProceduralOrchestrator`` sub-call. Reconciliation is gated by the change-feed trigger (which tracks the per-user/thread turn counter) and signaled to the orchestrator via the ``reconcile`` flag on its input payload. Procedural synthesis fires only @@ -33,26 +34,44 @@ def ExtractMemoriesOrchestrator(context: df.DurableOrchestrationContext): user_id = payload["user_id"] thread_id = payload["thread_id"] should_reconcile = bool(payload.get("reconcile", False)) + full_rebuild = bool(payload.get("full_rebuild", False)) + recent_k = payload.get("recent_k") retry = default_retry_options() + extract_payload = {"user_id": user_id, "thread_id": thread_id} + if recent_k is not None: + extract_payload["recent_k"] = recent_k extracted = yield context.call_activity_with_retry( "em_Extract", retry, - {"user_id": user_id, "thread_id": thread_id, "limit": config.get_max_batch_size()}, + extract_payload, + ) + deduped = yield context.call_activity_with_retry( + "em_Dedup", + retry, + {"user_id": user_id, "extracted": extracted}, ) persisted = yield context.call_activity_with_retry( "em_Persist", retry, - {"user_id": user_id, "extracted": extracted}, + {"user_id": user_id, "extracted": deduped}, ) + count = payload.get("count") + if count is not None: + yield context.call_activity_with_retry( + "em_AdvanceExtractWatermark", + retry, + {"user_id": user_id, "thread_id": thread_id, "count": count}, + ) + reconciled = None procedural = None if should_reconcile: reconciled = yield context.call_activity_with_retry( "em_ReconcileMemories", retry, - {"user_id": user_id}, + {"user_id": user_id, "full_rebuild": full_rebuild}, ) if config.get_procedural_synthesis_auto(): count = payload.get("count") @@ -86,11 +105,13 @@ def em_Extract(payload: dict) -> dict: """Load recent turns and run LLM extraction without embeddings or writes.""" user_id = payload["user_id"] thread_id = payload["thread_id"] - limit = payload.get("limit") + recent_k = payload.get("recent_k") + if recent_k is None: + recent_k = config.get_max_batch_size() extracted = get_pipeline().extract_memories_dry( user_id=user_id, thread_id=thread_id, - recent_k=limit, + recent_k=recent_k, ) logger.info( "ExtractMemories extracted user=%s thread=%s facts=%d episodic=%d updates=%d", @@ -103,6 +124,15 @@ def em_Extract(payload: dict) -> dict: return extracted +@bp.activity_trigger(input_name="payload") +def em_Dedup(payload: dict) -> dict: + """vector-floor dedup ladder (gated; passthrough when disabled).""" + return get_pipeline().dedup_extracted_memories( + user_id=payload["user_id"], + extracted=payload["extracted"], + ) or payload["extracted"] + + @bp.activity_trigger(input_name="payload") def em_Persist(payload: dict) -> dict: """Persist extracted docs with embeddings and deterministic create semantics.""" @@ -115,12 +145,40 @@ def em_Persist(payload: dict) -> dict: return counts or {} +@bp.activity_trigger(input_name="payload") +async def em_AdvanceExtractWatermark(payload: dict) -> bool: + """Advance the extraction watermark after a successful extract→persist. + + Stamps ``last_extract_count`` on the thread counter so the next batch's + recent_k spans only turns added since this run, never skipping any. + Runs only on success (after persist) so failed extracts re-process. + """ + from shared.cosmos_clients import get_counter_container_async + from shared.counters import advance_extract_watermark, thread_counter_id + + user_id = payload["user_id"] + thread_id = payload["thread_id"] + count = payload["count"] + container = await get_counter_container_async() + await advance_extract_watermark(container, thread_counter_id(user_id, thread_id), user_id, thread_id, count) + return True + + @bp.activity_trigger(input_name="payload") def em_ReconcileMemories(payload: dict) -> dict: # GA keeps reconcile single-activity: its LLM dedup decisions and supersession - # operations are larger/more coupled than the extract→persist split handled here. + # operations are larger/more coupled than the extract→dedup→persist split handled here. user_id = payload["user_id"] + # full_rebuild forces the full-pool single-LLM-pass path (catches dissimilar + # contradictions). The change-feed sets it on a persisted-counter cadence so it + # fires reliably on FA, where the in-memory candidate-mode sweep counter can't. + full_rebuild = bool(payload.get("full_rebuild", False)) pipeline = get_pipeline() from azure.cosmos.agent_memory.thresholds import get_dedup_pool_size - return pipeline.reconcile_memories(user_id=user_id, n=get_dedup_pool_size()) or {} + n = get_dedup_pool_size() + facts = pipeline.reconcile_memories(user_id=user_id, n=n, memory_type="fact", full_rebuild=full_rebuild) or {} + episodic = ( + pipeline.reconcile_memories(user_id=user_id, n=n, memory_type="episodic", full_rebuild=full_rebuild) or {} + ) + return {"fact": facts, "episodic": episodic} diff --git a/function_app/shared/config.py b/function_app/shared/config.py index f15ed35..da4f4fa 100644 --- a/function_app/shared/config.py +++ b/function_app/shared/config.py @@ -116,17 +116,6 @@ def _parse_int(name: str, default: int) -> int: return default -def _parse_float(name: str, default: float) -> float: - raw = os.environ.get(name) - if raw is None or raw == "": - return default - try: - return float(raw) - except (ValueError, TypeError): - logger.warning("Invalid value for %s=%r, using default %f", name, raw, default) - return default - - def _parse_bool(name: str, default: bool) -> bool: raw = os.environ.get(name) if raw is None or raw == "": diff --git a/function_app/shared/counters.py b/function_app/shared/counters.py index 1ab09cc..1adb0e5 100644 --- a/function_app/shared/counters.py +++ b/function_app/shared/counters.py @@ -177,6 +177,8 @@ async def increment_counter_by( new_doc["last_failure_at"] = existing_doc.get("last_failure_at") if "last_failure_reason" in existing_doc: new_doc["last_failure_reason"] = existing_doc.get("last_failure_reason") + if "last_extract_count" in existing_doc: + new_doc["last_extract_count"] = existing_doc.get("last_extract_count") # Stamp the writing backend (advisory only — not enforced server-side). if owner is not None: new_doc["last_owner"] = owner @@ -259,3 +261,40 @@ def crosses_threshold(old_count: int, new_count: int, n: int) -> bool: if n <= 0: raise ValueError("n must be > 0") return old_count // n != new_count // n + + +async def read_extract_watermark( + container: Any, + counter_id: str, + user_id: str, + thread_id: str, +) -> int | None: + """Return the count value at the last successful extract, or ``None``. + + Lets recent_k cover every turn since the previous extract succeeded so + turns are never skipped when extraction lags or transiently fails. + Best-effort: returns ``None`` on any read error so callers fall back to a + batch-based recent_k. + """ + try: + doc = await container.read_item(item=counter_id, partition_key=[user_id, thread_id]) + value = doc.get("last_extract_count") + return int(value) if value is not None else None + except Exception as exc: # pragma: no cover - best-effort + logger.debug("read_extract_watermark failed counter_id=%s: %s", counter_id, exc) + return None + + +async def advance_extract_watermark( + container: Any, + counter_id: str, + user_id: str, + thread_id: str, + count: int, +) -> None: + """Stamp ``last_extract_count=count`` after a successful extract.""" + patch_ops = [{"op": "add", "path": "/last_extract_count", "value": int(count)}] + try: + await container.patch_item(item=counter_id, partition_key=[user_id, thread_id], patch_operations=patch_ops) + except Exception as exc: # pragma: no cover - best-effort + logger.debug("advance_extract_watermark failed counter_id=%s: %s", counter_id, exc) diff --git a/function_app/triggers/change_feed.py b/function_app/triggers/change_feed.py index c7b2c8b..4630a46 100644 --- a/function_app/triggers/change_feed.py +++ b/function_app/triggers/change_feed.py @@ -26,6 +26,7 @@ from shared.counters import ( crosses_threshold, increment_counter_by, + read_extract_watermark, thread_counter_id, user_counter_id, ) @@ -128,6 +129,16 @@ async def process_changefeed_batch( # Reconcile fires every (n_facts * n_dedup) turns, matching the SDK # auto-trigger contract. Disabled when either knob is 0. n_dedup_turns = n_facts * n_dedup if (n_facts > 0 and n_dedup > 0) else 0 + # Full-pool reconcile backstop cadence. The candidate-mode "every Nth sweep" + # gate is an in-memory per-worker counter, which on the FA backend (per-worker + # singleton pipeline) resets on recycle/scale and never reaches N for a user — + # so the full-pool pass that catches dissimilar-embedding contradictions almost + # never fires. Derive it from the PERSISTED counter instead: every + # DEDUP_FULL_RECLUSTER_EVERY_N-th reconcile = every n_dedup_turns * N turns. + from azure.cosmos.agent_memory.thresholds import get_dedup_full_recluster_every_n + + n_full_recluster = get_dedup_full_recluster_every_n() + n_full_turns = n_dedup_turns * n_full_recluster if (n_dedup_turns > 0 and n_full_recluster > 0) else 0 if n_thread == 0 and n_facts == 0 and n_user == 0: return # all orchestrators disabled @@ -206,6 +217,22 @@ async def process_changefeed_batch( if n_facts > 0 and crosses_threshold(old_count, new_count, n_facts): instance_id = f"extract:{user_id}:{thread_id}:{new_count}" should_reconcile = bool(n_dedup_turns > 0 and crosses_threshold(old_count, new_count, n_dedup_turns)) + # Persisted-counter backstop: every n_full_turns turns, force a + # full-pool reconcile so dissimilar-embedding contradictions are + # caught reliably on FA (not gated by the in-memory sweep counter). + should_full_reconcile = bool( + n_full_turns > 0 and crosses_threshold(old_count, new_count, n_full_turns) + ) + watermark = await read_extract_watermark(counter_container, cid, user_id, thread_id) + # Not capped: new_count - watermark is exactly the unextracted backlog + # and the orchestrator advances the watermark to new_count, so capping + # would strand the oldest turns (DEDUP_POOL_SIZE is the reconcile knob, + # not the extraction window). Bootstrap: with no watermark yet, base=0 + # so recent_k = new_count covers every turn so far — using only this + # batch (new_count - old_count) would strand turns added during earlier + # failed extracts once the watermark first advances to new_count. + base = watermark if watermark is not None else 0 + recent_k = max(new_count - base, 1) await _safe_start( starter, "ExtractMemoriesOrchestrator", @@ -215,6 +242,8 @@ async def process_changefeed_batch( "thread_id": thread_id, "count": new_count, "reconcile": should_reconcile, + "full_rebuild": should_full_reconcile, + "recent_k": recent_k, }, orchestration_errors, ) diff --git a/tests/integration/test_async_full_pipeline.py b/tests/integration/test_async_full_pipeline.py new file mode 100644 index 0000000..88f66c4 --- /dev/null +++ b/tests/integration/test_async_full_pipeline.py @@ -0,0 +1,274 @@ +"""Async live integration smoke for ``AsyncCosmosMemoryClient``. + +The async pipeline mirrors every sync dedup/reconcile code path line-for-line +(watermark, euclidean guard, ``exclude_ids`` parity, reconcile cadence, the +vector-floor dedup ladder). The sync suite covers all of that against a live +backend, but the async client had **no** live coverage — its processor test is +fully mocked. This module exercises the real async end-to-end flow (write turns +→ extract → reconcile → search) so the async mirror can't silently diverge from +sync without a test failing. + +The Azure Function host is **not** required: the same ``AsyncPipelineService`` +the change-feed trigger drives is exposed directly on the client. + +Enable by setting:: + + AGENT_MEMORY_RUN_INTEGRATION=true + +Auth: ``COSMOS_DB_KEY`` is used when present; otherwise ``DefaultAzureCredential``. +""" + +from __future__ import annotations + +import asyncio +import time +import uuid + +import pytest + +from azure.cosmos.agent_memory.aio.cosmos_memory_client import AsyncCosmosMemoryClient +from tests.conftest import INTEGRATION_ENABLED + +pytestmark = [ + pytest.mark.integration, + pytest.mark.skipif( + not INTEGRATION_ENABLED, + reason="Set AGENT_MEMORY_RUN_INTEGRATION=true", + ), +] + + +# --------------------------------------------------------------------------- +# Fixtures / helpers +# --------------------------------------------------------------------------- + + +@pytest.fixture +async def async_agent_memory( + cosmos_endpoint, + cosmos_key, + cosmos_database, + cosmos_container, + ai_foundry_endpoint, + ai_foundry_api_key, + embedding_deployment_name, + embedding_dimensions, + chat_deployment_name, +): + """A live AsyncCosmosMemoryClient with its containers provisioned/connected.""" + if not cosmos_endpoint or not ai_foundry_endpoint: + pytest.skip("COSMOS_DB_ENDPOINT / AI_FOUNDRY_ENDPOINT not set") + + client = AsyncCosmosMemoryClient( + cosmos_endpoint=cosmos_endpoint, + cosmos_key=cosmos_key or None, + cosmos_database=cosmos_database, + cosmos_container=cosmos_container, + ai_foundry_endpoint=ai_foundry_endpoint, + ai_foundry_api_key=ai_foundry_api_key or None, + embedding_deployment_name=embedding_deployment_name, + embedding_dimensions=embedding_dimensions, + chat_deployment_name=chat_deployment_name, + ) + # The async client cannot auto-connect in __init__; do it explicitly. + await client.create_memory_store() + try: + yield client + finally: + await client.close() + + +async def _async_add_turns( + mem: AsyncCosmosMemoryClient, + user_id: str, + thread_id: str, + turns: list[tuple[str, str]], +) -> None: + for role, content in turns: + await mem.add_cosmos( + user_id=user_id, + role=role, + content=content, + memory_type="turn", + thread_id=thread_id, + ) + + +async def _async_cleanup(mem: AsyncCosmosMemoryClient, user_id: str) -> None: + """Best-effort delete of every document for *user_id* across all containers.""" + sql = "SELECT c.id, c.thread_id FROM c WHERE c.user_id = @uid" + params = [{"name": "@uid", "value": user_id}] + for container in ( + mem._turns_container_client, + mem._memories_container_client, + mem._summaries_container_client, + ): + if container is None: + continue + try: + docs = [doc async for doc in container.query_items(query=sql, parameters=params)] + except Exception: + continue + for doc in docs: + try: + await container.delete_item( + item=doc["id"], + partition_key=[user_id, doc.get("thread_id", "")], + ) + except Exception: + pass + + +async def _async_seed_fact_with_embedding( + mem: AsyncCosmosMemoryClient, + user_id: str, + thread_id: str, + content: str, + *, + retries: int = 4, +) -> None: + """Seed a fact and confirm it stored *with* an embedding (async mirror of the + sync helper). Retries through transient embedding-service blips so the + extract-time vector floor always has a neighbour; skips honestly if the + embedding service is genuinely unavailable.""" + check = ( + "SELECT c.id FROM c WHERE c.user_id = @uid AND c.content = @content " + "AND IS_DEFINED(c.embedding)" + ) + params = [{"name": "@uid", "value": user_id}, {"name": "@content", "value": content}] + for _ in range(retries): + await mem.add_cosmos( + user_id=user_id, + role="user", + content=content, + memory_type="fact", + thread_id=thread_id, + salience=0.7, + ) + embedded = [ + doc async for doc in mem._memories_container_client.query_items(query=check, parameters=params) + ] + if embedded: + return + await asyncio.sleep(1) + pytest.skip(f"embedding service unavailable — could not seed an embedded fact for {content!r}") + + +async def _async_wait_vector_searchable( + mem: AsyncCosmosMemoryClient, + user_id: str, + search_terms: str, + *, + 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.""" + deadline = time.time() + timeout + while time.time() < deadline: + try: + if await mem.search_cosmos(search_terms=search_terms, user_id=user_id, top_k=5): + return + except Exception: + pass + await asyncio.sleep(1) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestAsyncEndToEnd: + async def test_extract_reconcile_search( + self, + async_agent_memory, + unique_user_id, + unique_thread_id, + ): + """Full async round-trip: turns → extract → reconcile → search.""" + mem = async_agent_memory + try: + await _async_add_turns( + mem, + unique_user_id, + unique_thread_id, + [ + ("user", "I just adopted a golden retriever puppy named Cosmo."), + ("agent", "Congrats on Cosmo! How old is the puppy?"), + ("user", "Cosmo is 10 weeks old and already loves swimming in the lake."), + ], + ) + await asyncio.sleep(1) + + extract_stats = await mem.extract_memories( + user_id=unique_user_id, + thread_id=unique_thread_id, + ) + assert isinstance(extract_stats, dict) + + facts = await mem.get_memories(user_id=unique_user_id, memory_types=["fact"]) + assert len(facts) >= 1, "Async extraction should produce at least one fact" + + reconcile_stats = await mem.reconcile(user_id=unique_user_id) + assert isinstance(reconcile_stats, dict) + assert {"kept", "merged", "contradicted"} <= set(reconcile_stats), ( + f"reconcile should return kept/merged/contradicted, got {reconcile_stats}" + ) + + results = await mem.search_cosmos( + search_terms="golden retriever puppy", + user_id=unique_user_id, + top_k=5, + ) + assert len(results) >= 1, "Async search should return at least one result" + finally: + await _async_cleanup(mem, unique_user_id) + + async def test_dedup_extracted_memories_flags_near_duplicate_of_stored_fact( + self, + async_agent_memory, + unique_user_id, + unique_thread_id, + ): + """Async extract-time vector floor drops/tags a near-duplicate fact. + + Parity check with the sync ``TestExtractTimeVectorDedup`` — guards the + async ``dedup_extracted_memories`` mirror (``_vector_candidates`` + + similarity bands) against a live backend. Driven with a controlled + near-duplicate (no LLM variance) so the assertion is deterministic. + """ + mem = async_agent_memory + try: + await _async_seed_fact_with_embedding( + mem, unique_user_id, unique_thread_id, "The user has a cat named Whiskers." + ) + await _async_wait_vector_searchable(mem, unique_user_id, "cat named Whiskers") + + extracted = { + "facts": [ + { + "id": f"fact_{uuid.uuid4().hex}", + "type": "fact", + "user_id": unique_user_id, + "thread_id": unique_thread_id, + "content": "The user's cat is called Whiskers.", + "tags": [], + } + ], + "episodic": [], + "updates": [], + } + result = await mem._get_pipeline().dedup_extracted_memories(unique_user_id, extracted) + + stats = next((op for op in result.get("updates", []) if op.get("op") == "stats"), {}) + suppressed = int(stats.get("vector_dedup_skipped", 0)) + int(stats.get("dup_candidates_tagged", 0)) + surviving = result.get("facts", []) + was_dropped = len(surviving) == 0 + was_tagged = any("sys:dup-candidate" in (f.get("tags") or []) for f in surviving) + assert suppressed >= 1 and (was_dropped or was_tagged), ( + "Async vector floor should drop or tag the near-duplicate of the stored " + f"'cat named Whiskers' fact; surviving={surviving} stats={stats}" + ) + finally: + await _async_cleanup(mem, unique_user_id) diff --git a/tests/integration/test_full_pipeline.py b/tests/integration/test_full_pipeline.py index 6e80bd9..dcb217e 100644 --- a/tests/integration/test_full_pipeline.py +++ b/tests/integration/test_full_pipeline.py @@ -122,6 +122,70 @@ def _delete(container, doc: dict) -> None: _delete(container, doc) +def _seed_fact_with_embedding( + mem: CosmosMemoryClient, + user_id: str, + thread_id: str, + content: str, + *, + retries: int = 4, +) -> None: + """Seed a fact and confirm it was stored *with* an embedding. + + ``add_cosmos`` generates the embedding synchronously; a transient + embedding-service blip logs "proceeding without embedding" and stores the doc + without a vector, which would leave the extract-time vector floor with no + neighbour to match. Retry until an embedded copy exists (indexing is fast — + the doc is vector-searchable within ~2s), and skip honestly if the embedding + service is genuinely unavailable rather than reporting a false failure.""" + check = ( + "SELECT c.id FROM c WHERE c.user_id = @uid AND c.content = @content " + "AND IS_DEFINED(c.embedding)" + ) + params = [{"name": "@uid", "value": user_id}, {"name": "@content", "value": content}] + for _ in range(retries): + mem.add_cosmos( + user_id=user_id, + role="user", + content=content, + memory_type="fact", + thread_id=thread_id, + salience=0.7, + ) + embedded = list( + mem._memories_container_client.query_items( + query=check, parameters=params, enable_cross_partition_query=True + ) + ) + if embedded: + return + time.sleep(1) + pytest.skip(f"embedding service unavailable — could not seed an embedded fact for {content!r}") + + +def _wait_vector_searchable( + mem: CosmosMemoryClient, + user_id: str, + search_terms: str, + *, + timeout: float = 20.0, +) -> None: + """Poll vector search until the user's seeded fact is retrievable. + + ``add_cosmos`` 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.""" + deadline = time.time() + timeout + while time.time() < deadline: + try: + if mem.search_cosmos(search_terms=search_terms, user_id=user_id, top_k=5): + return + except Exception: + pass + time.sleep(1) + + # --------------------------------------------------------------------------- # Tests # --------------------------------------------------------------------------- @@ -243,7 +307,7 @@ def test_multi_thread_user_summary(self, agent_memory, unique_user_id): class TestSearchAfterExtraction: - def test_vector_and_hybrid_search(self, agent_memory, unique_user_id, unique_thread_id): + def test_search_after_extraction(self, agent_memory, unique_user_id, unique_thread_id): try: _add_turns( agent_memory, @@ -268,13 +332,12 @@ def test_vector_and_hybrid_search(self, agent_memory, unique_user_id, unique_thr ) assert len(vec) >= 1, "Vector search should return at least 1 result" - hyb = agent_memory.search_cosmos( + hybrid = agent_memory.search_cosmos( search_terms="Buddy the dog park", user_id=unique_user_id, - hybrid_search=True, top_k=5, ) - assert len(hyb) >= 1, "Hybrid search should return at least 1 result" + assert len(hybrid) >= 1, "Hybrid search should return at least 1 result" finally: _cleanup(agent_memory, unique_user_id) @@ -504,3 +567,61 @@ def test_reconcile_writes_supersede_metadata(self, agent_memory, unique_user_id, assert any(m["id"] == survivor_id for m in live), "supersede_by must point at a live record" finally: _cleanup(agent_memory, unique_user_id) + + +class TestExtractTimeVectorDedup: + """Extract-time vector floor (``dedup_extracted_memories``), distinct from the + ``reconcile`` path. A freshly-extracted fact that near-duplicates an + *already-stored* fact is either auto-dropped (``vector_dedup_skipped``, + sim >= DEDUP_SIM_HIGH) or tagged ``sys:dup-candidate`` + (``dup_candidates_tagged``, DEDUP_SIM_LOW <= sim < DEDUP_SIM_HIGH). + + The ladder is driven directly with a controlled extracted fact rather than + through the LLM: extraction phrasing varies run-to-run and often lands the + fact below the 0.80 floor (or produces unrelated facts), which is a property + of the model, not the dedup code. Feeding a fixed near-duplicate keeps the + assertion deterministic while still exercising the real embedding call, the + live Cosmos ``VectorDistance`` query, and the similarity bands.""" + + def test_dedup_extracted_memories_flags_near_duplicate_of_stored_fact( + self, agent_memory, unique_user_id, unique_thread_id + ): + try: + # Seed a stored fact (embedded + vector-indexed) to dedup against. + # Concrete, minimally-reworded facts embed ~0.93-0.98 cosine — well + # inside the DEDUP_SIM_LOW (0.80) / DEDUP_SIM_HIGH (0.97) bands. + _seed_fact_with_embedding( + agent_memory, unique_user_id, unique_thread_id, "The user has a cat named Whiskers." + ) + _wait_vector_searchable(agent_memory, unique_user_id, "cat named Whiskers") + + # A controlled "extracted" near-duplicate (not byte-identical to the + # seed, so this is the vector floor rather than an exact-hash match). + extracted = { + "facts": [ + { + "id": f"fact_{uuid.uuid4().hex}", + "type": "fact", + "user_id": unique_user_id, + "thread_id": unique_thread_id, + "content": "The user's cat is called Whiskers.", + "tags": [], + } + ], + "episodic": [], + "updates": [], + } + result = agent_memory._get_pipeline().dedup_extracted_memories(unique_user_id, extracted) + + stats = next((op for op in result.get("updates", []) if op.get("op") == "stats"), {}) + suppressed = int(stats.get("vector_dedup_skipped", 0)) + int(stats.get("dup_candidates_tagged", 0)) + surviving = result.get("facts", []) + was_dropped = len(surviving) == 0 + was_tagged = any("sys:dup-candidate" in (f.get("tags") or []) for f in surviving) + assert suppressed >= 1 and (was_dropped or was_tagged), ( + "Vector floor should drop or tag the near-duplicate of the stored " + f"'cat named Whiskers' fact; surviving={surviving} stats={stats}" + ) + finally: + _cleanup(agent_memory, unique_user_id) + diff --git a/tests/integration/test_processor_integration.py b/tests/integration/test_processor_integration.py index 1995626..b5a0be2 100644 --- a/tests/integration/test_processor_integration.py +++ b/tests/integration/test_processor_integration.py @@ -9,6 +9,7 @@ from __future__ import annotations from unittest.mock import MagicMock +from unittest.mock import call as mock_call from azure.cosmos.agent_memory.cosmos_memory_client import CosmosMemoryClient from azure.cosmos.agent_memory.processors import ( @@ -70,7 +71,13 @@ def test_process_now_invokes_pipeline_with_correct_args(self): client.get_thread.assert_called_once_with(thread_id="thread-paris", user_id="u-paris") pipeline.generate_thread_summary.assert_called_once_with("u-paris", "thread-paris") pipeline.extract_memories.assert_called_once_with("u-paris", "thread-paris") - pipeline.reconcile_memories.assert_called_once_with("u-paris", 50) + assert pipeline.reconcile_memories.call_count == 2 + pipeline.reconcile_memories.assert_has_calls( + [ + mock_call("u-paris", n=50, memory_type="fact", full_rebuild=False), + mock_call("u-paris", n=50, memory_type="episodic", full_rebuild=False), + ] + ) # --------------------------------------------------------------------------- diff --git a/tests/integration/test_processor_integration_async.py b/tests/integration/test_processor_integration_async.py index f437715..5769c3e 100644 --- a/tests/integration/test_processor_integration_async.py +++ b/tests/integration/test_processor_integration_async.py @@ -3,6 +3,7 @@ from __future__ import annotations from unittest.mock import AsyncMock, MagicMock +from unittest.mock import call as mock_call import pytest @@ -29,9 +30,9 @@ class TestAsyncInProcessProcessNowEndToEnd: @pytest.mark.asyncio async def test_process_now_invokes_pipeline_with_correct_args(self): pipeline = MagicMock() - # ``AsyncInProcessProcessor`` awaits these three pipeline methods, - # so the mocks must be ``AsyncMock`` — a plain ``MagicMock`` returns - # the dict synchronously, which ``await`` cannot consume. + # ``AsyncInProcessProcessor`` awaits these pipeline methods, so the + # mocks must be ``AsyncMock`` — a plain ``MagicMock`` returns the dict + # synchronously, which ``await`` cannot consume. pipeline.generate_thread_summary = AsyncMock( return_value={ "id": "summary-1", @@ -47,6 +48,8 @@ async def test_process_now_invokes_pipeline_with_correct_args(self): } ) pipeline.reconcile_memories = AsyncMock(return_value={"kept": 0, "merged": 0, "contradicted": 0}) + pipeline.synthesize_procedural = AsyncMock(return_value={"id": "proc-1", "type": "procedural"}) + pipeline.generate_user_summary = AsyncMock(return_value={"id": "us-1", "type": "user_summary"}) processor = AsyncInProcessProcessor(pipeline=pipeline) client = _build_client(processor=processor) @@ -73,7 +76,13 @@ async def test_process_now_invokes_pipeline_with_correct_args(self): client.get_thread.assert_awaited_once_with(thread_id="thread-paris", user_id="u-paris") pipeline.generate_thread_summary.assert_awaited_once_with("u-paris", "thread-paris") pipeline.extract_memories.assert_awaited_once_with("u-paris", "thread-paris") - pipeline.reconcile_memories.assert_awaited_once_with("u-paris", 50) + assert pipeline.reconcile_memories.await_count == 2 + pipeline.reconcile_memories.assert_has_awaits( + [ + mock_call("u-paris", n=50, memory_type="fact", full_rebuild=False), + mock_call("u-paris", n=50, memory_type="episodic", full_rebuild=False), + ] + ) # --------------------------------------------------------------------------- diff --git a/tests/unit/aio/processors/test_inprocess.py b/tests/unit/aio/processors/test_inprocess.py index 2d0852b..ec22819 100644 --- a/tests/unit/aio/processors/test_inprocess.py +++ b/tests/unit/aio/processors/test_inprocess.py @@ -24,10 +24,11 @@ async def test_process_thread_calls_summarize_extract_reconcile_in_order(): "generate_thread_summary", "extract_memories", "reconcile_memories", + "reconcile_memories", ] assert isinstance(result, ProcessThreadResult) assert result.thread_summary == {"id": "summary", "type": "thread_summary"} - assert result.reconciled_count == 2 + assert result.reconciled_count == 4 @pytest.mark.asyncio @@ -62,8 +63,19 @@ async def test_process_reconcile_invokes_pipeline_with_env_pool_size(monkeypatch proc = AsyncInProcessProcessor(pipeline=pipeline) count = await proc.process_reconcile(user_id="u") - pipeline.reconcile_memories.assert_called_once_with("u", 37) - assert count == 5 # merged + contradicted + # Reconciles fact + episodic; both forward the env pool size. + assert pipeline.reconcile_memories.await_count == 2 + assert pipeline.reconcile_memories.await_args_list[0].kwargs == { + "n": 37, + "memory_type": "fact", + "full_rebuild": False, + } + assert pipeline.reconcile_memories.await_args_list[1].kwargs == { + "n": 37, + "memory_type": "episodic", + "full_rebuild": False, + } + assert count == 10 # (merged+contradicted) x2 types @pytest.mark.asyncio @@ -75,9 +87,9 @@ async def test_process_extract_memories_invokes_pipeline_and_filters_to_ints(): } proc = AsyncInProcessProcessor(pipeline=pipeline) - result = await proc.process_extract_memories(user_id="u", thread_id="t") + result = await proc.process_extract_memories(user_id="u", thread_id="t", recent_k=3) - pipeline.extract_memories.assert_called_once_with("u", "t") + pipeline.extract_memories.assert_called_once_with("u", "t", recent_k=3) assert result == {"fact_count": 3} diff --git a/tests/unit/aio/processors/test_protocol_satisfaction.py b/tests/unit/aio/processors/test_protocol_satisfaction.py index 04cda06..2a23515 100644 --- a/tests/unit/aio/processors/test_protocol_satisfaction.py +++ b/tests/unit/aio/processors/test_protocol_satisfaction.py @@ -32,6 +32,7 @@ async def process_extract_memories( *, user_id: str, thread_id: str, + recent_k: Optional[int] = None, ) -> dict[str, int]: return {} diff --git a/tests/unit/aio/services/test_dedup_vector_async.py b/tests/unit/aio/services/test_dedup_vector_async.py new file mode 100644 index 0000000..be62553 --- /dev/null +++ b/tests/unit/aio/services/test_dedup_vector_async.py @@ -0,0 +1,446 @@ +from __future__ import annotations + +import asyncio +import json +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from azure.cosmos.agent_memory.aio.services.pipeline import AsyncPipelineService + + +def _service() -> AsyncPipelineService: + p = AsyncPipelineService.__new__(AsyncPipelineService) + p._memories_container = MagicMock() + p._embed_batch = AsyncMock() + p._embed_one = AsyncMock(return_value=[0.1, 0.2]) + p._upsert_memory = AsyncMock(side_effect=lambda doc: doc) + p._mark_superseded = AsyncMock(return_value=True) + return p + + +def _fact(fid: str, content: str, embedding=None, tags=None, metadata=None) -> dict: + return { + "id": fid, + "user_id": "u1", + "thread_id": "t1", + "type": "fact", + "role": "system", + "content": content, + "content_hash": "0" * 32, + "confidence": 0.8, + "salience": 0.7, + "tags": list(tags or ["sys:fact"]), + "metadata": dict(metadata or {"category": "preference"}), + "created_at": "2025-01-01T00:00:00+00:00", + "embedding": embedding or [1.0, 0.0], + } + + +def _episode(eid: str, content: str) -> dict: + return { + "id": eid, + "user_id": "u1", + "thread_id": "t1", + "type": "episodic", + "role": "system", + "content": content, + "content_hash": "1" * 32, + "confidence": 0.8, + "salience": 0.7, + "tags": ["sys:episodic", "sys:dup-candidate"], + "metadata": { + "scope_type": "project", + "scope_value": "CI", + "lesson": content, + "outcome_valence": "positive", + }, + "created_at": "2025-01-01T00:00:00+00:00", + "embedding": [1.0, 0.0], + } + + +@pytest.mark.asyncio +async def test_vector_distance_function_reads_container_policy(): + # The distance function comes from the container's vector embedding policy + # (read once, cached), NOT an env var. + p = _service() + p._memories_container.read = AsyncMock( + return_value={ + "vectorEmbeddingPolicy": {"vectorEmbeddings": [{"path": "/embedding", "distanceFunction": "dotproduct"}]} + } + ) + assert await p._vector_distance_function() == "dotproduct" + assert await p._vector_distance_function() == "dotproduct" + assert p._memories_container.read.await_count == 1 + + +@pytest.mark.asyncio +async def test_vector_candidates_orders_nearest_first_by_distance_function(): + # Regression: async _vector_candidates must order most-similar-first per the + # container's distanceFunction. For cosine/dotproduct higher score = more + # similar (DESC); for euclidean lower distance = more similar (ASC). A missing + # DESC silently fetched the LEAST-similar rows when the pool exceeded top_k. + p = _service() + captured: dict[str, str] = {} + + async def fake_query_items(_container, *, query, parameters): + captured["query"] = query + return [ + {"id": "near", "content": "a", "type": "fact", "score": 0.95}, + {"id": "far", "content": "b", "type": "fact", "score": 0.10}, + ] + + p._query_items = AsyncMock(side_effect=fake_query_items) + + p._distance_function_cache = "cosine" + out = await p._vector_candidates( + user_id="u1", embedding=[1.0, 0.0], memory_type="fact", top_k=2, exclude_ids=set() + ) + # Cosmos rejects an explicit ASC/DESC on ORDER BY VectorDistance(); it orders + # most-similar-first server-side. Direction-awareness lives in the Python sort. + assert "ORDER BY VectorDistance(c.embedding, @vec)" in captured["query"] + assert "VectorDistance(c.embedding, @vec) DESC" not in captured["query"] + assert "VectorDistance(c.embedding, @vec) ASC" not in captured["query"] + assert [c["id"] for c in out] == ["near", "far"] + + p._distance_function_cache = "euclidean" + out = await p._vector_candidates( + user_id="u1", embedding=[1.0, 0.0], memory_type="fact", top_k=2, exclude_ids=set() + ) + assert "VectorDistance(c.embedding, @vec) ASC" not in captured["query"] + # euclidean: lower distance = more similar, so 0.10 ("far" label) sorts first. + # euclidean: lower distance = more similar, so 0.10 ("far" label) sorts first. + assert [c["id"] for c in out] == ["far", "near"] + + +@pytest.mark.asyncio +async def test_candidate_mode_clears_tags_only_for_survivors(): + # Latent-bug regression (async mirror): consumed sources must not be + # re-upserted by tag-clearing, which would resurrect them without superseded_by. + p = _service() + f1 = _fact("f1", "a", tags=["sys:fact", "sys:dup-candidate"]) + f2 = _fact("f2", "b", tags=["sys:fact", "sys:dup-candidate"]) + f3 = _fact("f3", "c", tags=["sys:fact", "sys:dup-candidate"]) + p._build_candidate_clusters = AsyncMock(return_value=([[f1, f2, f3]], 3, [f1, f2, f3])) + p._reconcile_pool = AsyncMock(return_value=({"kept": 1, "merged": 2, "contradicted": 0}, {"f1", "f2"})) + cleared: list[str] = [] + + async def clear(docs): + cleared.extend(d["id"] for d in docs) + + p._clear_dup_candidate_tags = AsyncMock(side_effect=clear) + p._emit_reconcile_outcome = MagicMock() + + result = await p._reconcile_candidate_mode("u1", n=50, memory_type="fact", started_at=0.0) + + assert cleared == ["f3"] + assert result == { + "kept": 1, + "merged": 2, + "contradicted": 0, + "reconcile_clusters_sent": 1, + "reconcile_llm_calls_saved": 2, + } + + +@pytest.mark.asyncio +async def test_candidate_mode_clears_tags_on_orphan_seeds(): + # Async mirror: orphan dup-candidate seeds (no cluster) get their stale tag cleared. + p = _service() + orphan = _fact("orphan", "lonely", tags=["sys:fact", "sys:dup-candidate"]) + c1 = _fact("c1", "a", tags=["sys:fact", "sys:dup-candidate"]) + c2 = _fact("c2", "b", tags=["sys:fact", "sys:dup-candidate"]) + p._build_candidate_clusters = AsyncMock(return_value=([[c1, c2]], 3, [orphan, c1, c2])) + p._reconcile_pool = AsyncMock(return_value=({"kept": 2, "merged": 0, "contradicted": 0}, set())) + cleared: list[str] = [] + + async def clear(docs): + cleared.extend(d["id"] for d in docs) + + p._clear_dup_candidate_tags = AsyncMock(side_effect=clear) + p._emit_reconcile_outcome = MagicMock() + + await p._reconcile_candidate_mode("u1", n=50, memory_type="fact", started_at=0.0) + + assert set(cleared) == {"c1", "c2", "orphan"} + + +@pytest.mark.asyncio +async def test_sweep_survives_one_cluster_failure(): + # Async mirror: one failing cluster must not abort the sweep; failed cluster's + # tags are retained, remaining cluster + orphan are still cleared. + p = _service() + c1 = [ + _fact("a1", "x", tags=["sys:fact", "sys:dup-candidate"]), + _fact("a2", "y", tags=["sys:fact", "sys:dup-candidate"]), + ] + c2 = [ + _fact("b1", "p", tags=["sys:fact", "sys:dup-candidate"]), + _fact("b2", "q", tags=["sys:fact", "sys:dup-candidate"]), + ] + orphan = _fact("o1", "lonely", tags=["sys:fact", "sys:dup-candidate"]) + p._build_candidate_clusters = AsyncMock(return_value=([c1, c2], 5, [*c1, *c2, orphan])) + p._reconcile_pool = AsyncMock( + side_effect=[RuntimeError("truncated LLM response"), ({"kept": 2, "merged": 0, "contradicted": 0}, set())] + ) + cleared: list[str] = [] + + async def clear(docs): + cleared.extend(d["id"] for d in docs) + + p._clear_dup_candidate_tags = AsyncMock(side_effect=clear) + p._emit_reconcile_outcome = MagicMock() + + result = await p._reconcile_candidate_mode("u1", n=50, memory_type="fact", started_at=0.0) + + assert "a1" not in cleared and "a2" not in cleared + assert {"b1", "b2", "o1"} <= set(cleared) + assert result["reconcile_clusters_sent"] == 2 + + +@pytest.mark.asyncio +async def test_full_rebuild_clears_survivor_tags(monkeypatch): + # Async mirror: full_rebuild full-pool path clears survivor dup-candidate tags. + monkeypatch.setattr( + "azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_reconcile_mode", lambda: "candidate" + ) + p = _service() + pool = [ + _fact("f1", "a", tags=["sys:fact", "sys:dup-candidate"]), + _fact("f2", "b", tags=["sys:fact", "sys:dup-candidate"]), + ] + p._active_memories_for_reconcile = AsyncMock(return_value=pool) + p._reconcile_pool = AsyncMock(return_value=({"kept": 1, "merged": 1, "contradicted": 0}, {"f1"})) + cleared: list[str] = [] + + async def clear(docs): + cleared.extend(d["id"] for d in docs) + + p._clear_dup_candidate_tags = AsyncMock(side_effect=clear) + p._emit_reconcile_outcome = MagicMock() + + await p.reconcile_memories("u1", n=50, memory_type="fact", full_rebuild=True) + + assert cleared == ["f2"] + + +@pytest.mark.asyncio +async def test_dedup_extracted_memories_flag_off_is_noop(monkeypatch): + monkeypatch.setattr("azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_vector_enabled", lambda: False) + p = _service() + extracted = {"facts": [_fact("f1", "User likes tea")], "episodic": [], "updates": []} + + out = await p.dedup_extracted_memories("u1", extracted) + + assert out is extracted + p._embed_batch.assert_not_called() + + +@pytest.mark.asyncio +async def test_euclidean_disables_near_exact_autodrop(): + # Async mirror: euclidean disables the cosine-calibrated near-exact auto-drop; + # the near-identical new memory is kept + tagged for LLM reconcile. + p = _service() + p._vector_distance_function = AsyncMock(return_value="euclidean") + fact = _fact("f-new", "near identical") + fact.pop("embedding", None) + p._embed_batch.return_value = [[1.0, 0.0]] + p._vector_candidates = AsyncMock( + return_value=[{"id": "existing", "content": "same", "type": "fact", "score": 0.05}] + ) + + out = await p.dedup_extracted_memories("u1", {"facts": [fact], "episodic": [], "updates": []}) + + assert [doc["id"] for doc in out["facts"]] == ["f-new"] + assert out["facts"][0]["tags"][-1] == "sys:dup-candidate" + assert out["updates"][-1]["vector_dedup_skipped"] == 0 + assert out["updates"][-1]["dup_candidates_tagged"] == 1 + + +@pytest.mark.asyncio +async def test_dedup_extracted_memories_passes_user_id_per_concurrent_call(): + p = _service() + seen: list[tuple[str, str]] = [] + + async def vector_candidates(**kwargs): + await asyncio.sleep(0) + exclude_ids = kwargs["exclude_ids"] + doc_id = next(iter(exclude_ids)) + seen.append((doc_id, kwargs["user_id"])) + return [] + + p._vector_candidates = AsyncMock(side_effect=vector_candidates) + user_a_doc = _fact("user-a-doc", "User A likes tea") + user_b_doc = _fact("user-b-doc", "User B likes coffee") + + await asyncio.gather( + p.dedup_extracted_memories("user-a", {"facts": [user_a_doc], "episodic": [], "updates": []}), + p.dedup_extracted_memories("user-b", {"facts": [user_b_doc], "episodic": [], "updates": []}), + ) + + assert dict(seen) == {"user-a-doc": "user-a", "user-b-doc": "user-b"} + + +@pytest.mark.asyncio +async def test_dedup_extracted_memories_vector_ladder_and_intra_batch(): + p = _service() + docs = [ + _fact("drop-existing", "User likes tea"), + _fact("tag-existing", "User likes coffee"), + _fact("clean", "User likes water"), + _fact("batch-keeper", "User likes green tea"), + _fact("drop-batch", "User likes green tea too"), + ] + for doc in docs: + doc.pop("embedding", None) + p._embed_batch.return_value = [ + [1.0, 0.0], + [0.0, 1.0], + [1.0, -1.0], + [0.7, 0.7], + [0.7, 0.7], + ] + p._vector_candidates = AsyncMock( + side_effect=[ + [{"id": "old-high", "content": "User likes tea.", "type": "fact", "score": 0.99}], + [{"id": "old-mid", "content": "User enjoys coffee.", "type": "fact", "score": 0.90}], + [{"id": "old-low", "content": "Unrelated", "type": "fact", "score": 0.20}], + [], + [], + ] + ) + + out = await p.dedup_extracted_memories("u1", {"facts": docs, "episodic": [], "updates": []}) + + # Intra-batch new-vs-new dedup was removed: drop-batch (no existing match) + # is now kept; same-batch near-dups are deferred to reconcile. + ids = [doc["id"] for doc in out["facts"]] + assert ids == ["tag-existing", "clean", "batch-keeper", "drop-batch"] + assert all("embedding" in doc for doc in out["facts"]) + tagged = out["facts"][0] + assert "sys:dup-candidate" in tagged["tags"] + assert tagged["metadata"]["dup_of"] == "old-mid" + assert tagged["metadata"]["dup_score"] == 0.90 + assert "sys:dup-candidate" not in out["facts"][1]["tags"] + stats = out["updates"][-1] + assert stats["vector_dedup_skipped"] == 1 + assert stats["dup_candidates_tagged"] == 1 + + +@pytest.mark.asyncio +async def test_dedup_skips_underspecified_doc_verbatim(): + # Parity with sync: a doc with no/unknown type is passed through untouched + # and never runs vector dedup (async previously defaulted type to the bucket). + p = _service() + p._vector_candidates = AsyncMock( + return_value=[{"id": "x", "content": "c", "type": "fact", "score": 0.99}] + ) + doc = _fact("f1", "content") + doc.pop("type") + doc.pop("embedding", None) + p._embed_batch.return_value = [[1.0, 0.0]] + + out = await p.dedup_extracted_memories("u1", {"facts": [doc], "episodic": [], "updates": []}) + + assert [d["id"] for d in out["facts"]] == ["f1"] + assert "sys:dup-candidate" not in out["facts"][0].get("tags", []) + p._vector_candidates.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_reconcile_memory_type_routes_episodic_merge_only(monkeypatch): + monkeypatch.setattr("azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_reconcile_mode", lambda: "full_pool") + p = _service() + episodes = [_episode("ep_1", "CI failed then retries fixed it"), _episode("ep_2", "CI failed; retries fixed it")] + p._active_memories_for_reconcile = AsyncMock(return_value=episodes) + p._run_prompty = AsyncMock( + return_value=json.dumps( + { + "duplicate_groups": [ + { + "merged_content": "CI failed, then retries fixed it", + "source_ids": ["ep_1", "ep_2"], + "confidence": 0.9, + "salience": 0.8, + } + ], + "kept_ids": [], + "contradicted_pairs": [{"winner_id": "ep_1", "loser_id": "ep_2"}], + } + ) + ) + + result = await p.reconcile_memories("u1", memory_type="episodic") + + assert result == {"kept": 0, "merged": 2, "contradicted": 0} + assert p._run_prompty.await_args.kwargs["inputs"].keys() == {"episodics_text"} + assert p._run_prompty.await_args.args[0] == "dedup_episodic.prompty" + + +@pytest.mark.asyncio +async def test_candidate_reconcile_builds_connected_component(): + p = _service() + docs = { + "f1": _fact("f1", "User likes aisle seats", tags=["sys:fact", "sys:dup-candidate"]), + "f2": _fact("f2", "User prefers aisle seats"), + "f3": _fact("f3", "User enjoys aisle seats"), + } + + async def query_items(_container, **kwargs): + params = {p["name"]: p["value"] for p in kwargs.get("parameters", [])} + if params.get("@tag") == "sys:dup-candidate": + return [docs["f1"]] + ids = [value for name, value in params.items() if name.startswith("@id")] + return [docs[mid] for mid in ids if mid in docs] + + p._query_items = AsyncMock(side_effect=query_items) + p._vector_candidates = AsyncMock( + side_effect=[ + [ + {"id": "f2", "content": docs["f2"]["content"], "type": "fact", "score": 0.90}, + {"id": "f3", "content": docs["f3"]["content"], "type": "fact", "score": 0.88}, + ], + [ + {"id": "f2", "content": docs["f2"]["content"], "type": "fact", "score": 0.90}, + {"id": "f3", "content": docs["f3"]["content"], "type": "fact", "score": 0.88}, + ], + [ + {"id": "f1", "content": docs["f1"]["content"], "type": "fact", "score": 0.90}, + {"id": "f3", "content": docs["f3"]["content"], "type": "fact", "score": 0.89}, + ], + [ + {"id": "f1", "content": docs["f1"]["content"], "type": "fact", "score": 0.88}, + {"id": "f2", "content": docs["f2"]["content"], "type": "fact", "score": 0.89}, + ], + ] + ) + p._run_prompty = AsyncMock( + return_value=json.dumps( + { + "duplicate_groups": [ + { + "merged_content": "User prefers aisle seats", + "source_ids": ["f1", "f2", "f3"], + "confidence": 0.9, + "salience": 0.8, + } + ], + "contradicted_pairs": [], + "kept_ids": [], + } + ) + ) + + result = await p.reconcile_memories("u1") + + assert result == { + "kept": 0, + "merged": 3, + "contradicted": 0, + "reconcile_clusters_sent": 1, + "reconcile_llm_calls_saved": 2, + } + p._run_prompty.assert_awaited_once() + facts_text = p._run_prompty.await_args.kwargs["inputs"]["facts_text"] + assert all(fid in facts_text for fid in ["f1", "f2", "f3"]) diff --git a/tests/unit/aio/test_auto_trigger.py b/tests/unit/aio/test_auto_trigger.py index edcf382..adb9f88 100644 --- a/tests/unit/aio/test_auto_trigger.py +++ b/tests/unit/aio/test_auto_trigger.py @@ -8,14 +8,51 @@ from __future__ import annotations import asyncio -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest +from azure.cosmos.exceptions import CosmosResourceNotFoundError +from azure.cosmos.agent_memory import _counters from azure.cosmos.agent_memory.aio.cosmos_memory_client import AsyncCosmosMemoryClient from azure.cosmos.agent_memory.aio.processors import AsyncInProcessProcessor +class _AsyncFakeCounterContainer: + """Async in-memory counter container exercising the REAL increment / + watermark-read / watermark-advance helpers end-to-end (no constant mocks).""" + + def __init__(self) -> None: + self.store: dict[str, dict] = {} + self._etag = 0 + + async def read_item(self, *, item, partition_key): + if item not in self.store: + raise CosmosResourceNotFoundError(message="404") + return dict(self.store[item]) + + async def create_item(self, *, body): + self._etag += 1 + body = dict(body) + body["_etag"] = f"e{self._etag}" + self.store[body["id"]] = body + return dict(body) + + async def upsert_item(self, *, body, **_kwargs): + self._etag += 1 + body = dict(body) + body["_etag"] = f"e{self._etag}" + self.store[body["id"]] = body + return dict(body) + + async def patch_item(self, *, item, partition_key, patch_operations): + doc = self.store.setdefault(item, {"id": item}) + for op in patch_operations: + doc[op["path"].lstrip("/")] = op["value"] + return dict(doc) + + + class TestAsyncAutoTriggerNonBlocking: @pytest.mark.asyncio async def test_push_to_cosmos_does_not_await_auto_trigger(self, monkeypatch): @@ -25,7 +62,7 @@ async def test_push_to_cosmos_does_not_await_auto_trigger(self, monkeypatch): processor = AsyncInProcessProcessor(pipeline=MagicMock()) - async def slow_trigger(user_id, thread_id): + async def slow_trigger(user_id, thread_id, recent_k=None): # If push_to_cosmos awaited the trigger inline, the test would # block here for half a second before returning. await asyncio.sleep(0.5) @@ -63,6 +100,224 @@ async def fake_upsert(body): await asyncio.gather(*list(client._background_tasks), return_exceptions=True) +class TestAsyncExtractRecentK: + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("n_facts", "batch_count", "counter_result", "expected_recent_k"), + [ + (1, 1, (0, 1), 1), + (1, 3, (0, 3), 3), + (5, 1, (4, 5), 5), + ], + ) + async def test_extract_recent_k_uses_max_threshold_and_batch_count( + self, + monkeypatch, + n_facts, + batch_count, + counter_result, + expected_recent_k, + ): + monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", str(n_facts)) + monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") + monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") + + processor = AsyncInProcessProcessor(pipeline=MagicMock()) + processor.process_extract_memories = MagicMock(return_value={}) + + client = AsyncCosmosMemoryClient(use_default_credential=False, processor=processor) + + async def fake_upsert(body): + return body + + client._memories_container_client = MagicMock() + client._memories_container_client.upsert_item = MagicMock(side_effect=fake_upsert) + client._turns_container_client = client._memories_container_client + client._summaries_container_client = client._memories_container_client + client._counter_container_client = MagicMock() + + with patch( + "azure.cosmos.agent_memory._counters.increment_counter_async", + return_value=counter_result, + ): + for i in range(batch_count): + client.add_local(user_id="u1", role="user", thread_id="t1", content=f"hi {i}") + await client.push_to_cosmos() + await asyncio.gather(*list(client._background_tasks), return_exceptions=True) + + processor.process_extract_memories.assert_called_once_with( + user_id="u1", + thread_id="t1", + recent_k=expected_recent_k, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("counter_result", "watermark", "expected_recent_k"), + [ + ((5, 10), 5, 5), # backlog = new - watermark + ((0, 1), 1, 1), # new == watermark -> floored to 1 + ((98, 100), 0, 100), # large backlog is NOT capped + ((20, 30), None, 30), # BOOTSTRAP: no watermark -> base=0 -> recent_k = new_count + ], + ) + async def test_extract_recent_k_uses_watermark( + self, monkeypatch, counter_result, watermark, expected_recent_k + ): + monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") + monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") + monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") + + processor = AsyncInProcessProcessor(pipeline=MagicMock()) + processor.process_extract_memories = MagicMock(return_value={}) + client = AsyncCosmosMemoryClient(use_default_credential=False, processor=processor) + client._memories_container_client = MagicMock() + client._memories_container_client.upsert_item = AsyncMock(side_effect=lambda body: body) + client._turns_container_client = client._memories_container_client + client._summaries_container_client = client._memories_container_client + client._counter_container_client = MagicMock() + + with patch( + "azure.cosmos.agent_memory._counters.increment_counter_async", + return_value=counter_result, + ), patch( + "azure.cosmos.agent_memory._counters.read_extract_watermark_async", + new=AsyncMock(return_value=watermark), + ), patch( + "azure.cosmos.agent_memory._counters.advance_extract_watermark_async", + new=AsyncMock(), + ) as advance: + client.add_local(user_id="u1", role="user", thread_id="t1", content="hi") + await client.push_to_cosmos() + await asyncio.gather(*list(client._background_tasks), return_exceptions=True) + + processor.process_extract_memories.assert_called_once_with( + user_id="u1", thread_id="t1", recent_k=expected_recent_k + ) + advance.assert_awaited_once() + + @pytest.mark.asyncio + async def test_watermark_not_advanced_when_extract_fails(self, monkeypatch): + monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") + monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") + monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") + + processor = AsyncInProcessProcessor(pipeline=MagicMock()) + processor.process_extract_memories = MagicMock(side_effect=RuntimeError("llm down")) + client = AsyncCosmosMemoryClient(use_default_credential=False, processor=processor) + client._memories_container_client = MagicMock() + client._memories_container_client.upsert_item = AsyncMock(side_effect=lambda body: body) + client._turns_container_client = client._memories_container_client + client._summaries_container_client = client._memories_container_client + client._counter_container_client = MagicMock() + + with patch( + "azure.cosmos.agent_memory._counters.increment_counter_async", + return_value=(0, 1), + ), patch( + "azure.cosmos.agent_memory._counters.read_extract_watermark_async", + new=AsyncMock(return_value=None), + ), patch( + "azure.cosmos.agent_memory._counters.advance_extract_watermark_async", + new=AsyncMock(), + ) as advance, patch( + "azure.cosmos.agent_memory._counters.stamp_failure_async", + new=AsyncMock(), + ) as stamp: + client.add_local(user_id="u1", role="user", thread_id="t1", content="hi") + await client.push_to_cosmos() + await asyncio.gather(*list(client._background_tasks), return_exceptions=True) + + advance.assert_not_awaited() + stamp.assert_awaited_once() + + @pytest.mark.asyncio + async def test_watermark_round_trip_fail_then_succeed_no_strand(self, monkeypatch): + """Async stateful round-trip against a REAL in-memory counter: first + extract fails, second succeeds and must cover EVERY turn so far (20), + not just its own batch (10) — the bootstrap strand regression.""" + monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") + monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") + monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") + + counter = _AsyncFakeCounterContainer() + recorded: list[int] = [] + + def extract(*, user_id, thread_id, recent_k): + recorded.append(recent_k) + if len(recorded) == 1: + raise RuntimeError("transient LLM outage") + return {} + + processor = AsyncInProcessProcessor(pipeline=MagicMock()) + processor.process_extract_memories = MagicMock(side_effect=extract) + client = AsyncCosmosMemoryClient(use_default_credential=False, processor=processor) + client._memories_container_client = MagicMock() + client._memories_container_client.upsert_item = AsyncMock(side_effect=lambda body: body) + client._turns_container_client = client._memories_container_client + client._summaries_container_client = client._memories_container_client + client._counter_container_client = counter + + for i in range(10): + client.add_local(user_id="u1", role="user", thread_id="t1", content=f"a{i}") + await client.push_to_cosmos() + await asyncio.gather(*list(client._background_tasks), return_exceptions=True) + + for i in range(10): + client.add_local(user_id="u1", role="user", thread_id="t1", content=f"b{i}") + await client.push_to_cosmos() + await asyncio.gather(*list(client._background_tasks), return_exceptions=True) + + assert recorded == [10, 20] + cid = _counters.thread_counter_id("u1", "t1") + assert await _counters.read_extract_watermark_async(counter, cid, "u1", "t1") == 20 + + @pytest.mark.asyncio + async def test_reconcile_full_rebuild_on_persisted_counter_cadence(self, monkeypatch): + """Async symmetry: in-process auto-trigger requests full_rebuild on the + persisted-counter cadence (every 2 turns here), like the durable backend.""" + monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") + monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") + monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") + monkeypatch.setenv("DEDUP_EVERY_N", "1") + monkeypatch.setattr( + "azure.cosmos.agent_memory.thresholds.get_dedup_full_recluster_every_n", lambda: 2 + ) + + rebuilds: list[bool] = [] + processor = AsyncInProcessProcessor(pipeline=MagicMock()) + processor.process_extract_memories = MagicMock(return_value={}) + processor.synthesize_procedural = MagicMock(return_value=None) + processor.process_reconcile = MagicMock( + side_effect=lambda *, user_id, full_rebuild=False: rebuilds.append(full_rebuild) + ) + client = AsyncCosmosMemoryClient(use_default_credential=False, processor=processor) + client._memories_container_client = MagicMock() + client._memories_container_client.upsert_item = AsyncMock(side_effect=lambda body: body) + client._turns_container_client = client._memories_container_client + client._summaries_container_client = client._memories_container_client + client._counter_container_client = MagicMock() + + with patch( + "azure.cosmos.agent_memory._counters.increment_counter_async", + new=AsyncMock(side_effect=[(0, 1), (1, 2)]), + ), patch( + "azure.cosmos.agent_memory._counters.read_extract_watermark_async", + new=AsyncMock(return_value=None), + ), patch( + "azure.cosmos.agent_memory._counters.advance_extract_watermark_async", + new=AsyncMock(), + ): + client.add_local(user_id="u1", role="user", thread_id="t1", content="a") + await client.push_to_cosmos() + await asyncio.gather(*list(client._background_tasks), return_exceptions=True) + client.add_local(user_id="u1", role="user", thread_id="t1", content="b") + await client.push_to_cosmos() + await asyncio.gather(*list(client._background_tasks), return_exceptions=True) + + assert rebuilds == [False, True] + + class TestPushToCosmosUnflushedDelta: """``push_to_cosmos`` must use the unflushed-add delta, not a recount of ``local_memory``, so callers that retain the buffer don't re-fire diff --git a/tests/unit/aio/test_cosmos_memory_client.py b/tests/unit/aio/test_cosmos_memory_client.py index 0010495..8eb9f2c 100644 --- a/tests/unit/aio/test_cosmos_memory_client.py +++ b/tests/unit/aio/test_cosmos_memory_client.py @@ -9,7 +9,9 @@ import pytest +from azure.cosmos.agent_memory._container_routing import ContainerKey from azure.cosmos.agent_memory.aio.cosmos_memory_client import AsyncCosmosMemoryClient +from azure.cosmos.agent_memory.aio.store import AsyncMemoryStore from azure.cosmos.agent_memory.exceptions import ( ConfigurationError, CosmosNotConnectedError, @@ -369,7 +371,8 @@ async def test_create_memory_store_turns_container_uses_30_day_ttl(self): assert mem._turns_container_client is mock_turns_container async def test_create_memory_store_defaults_to_serverless(self): - mem = _make_client(cosmos_throughput_mode="serverless") + # serverless mode ignores autoscale config entirely, even an invalid value. + mem = _make_client(cosmos_throughput_mode="serverless", cosmos_autoscale_max_ru="not-an-int") mock_cosmos_cls = MagicMock() mock_client = MagicMock() mock_db = AsyncMock() @@ -391,29 +394,28 @@ async def test_create_memory_store_defaults_to_serverless(self): ] ) - with patch.dict("os.environ", {"COSMOS_DB_AUTOSCALE_MAX_RU": "not-an-int"}, clear=False): - with patch.dict( - "sys.modules", - { - "azure.cosmos.aio": MagicMock(CosmosClient=mock_cosmos_cls), - "azure.cosmos": MagicMock( - PartitionKey=MagicMock(), - ThroughputProperties=MagicMock(), - ), - }, - ): - await mem.create_memory_store( - endpoint="https://fake.documents.azure.com:443/", - credential="fake-key", - throughput_mode="serverless", - ) + with patch.dict( + "sys.modules", + { + "azure.cosmos.aio": MagicMock(CosmosClient=mock_cosmos_cls), + "azure.cosmos": MagicMock( + PartitionKey=MagicMock(), + ThroughputProperties=MagicMock(), + ), + }, + ): + await mem.create_memory_store( + endpoint="https://fake.documents.azure.com:443/", + credential="fake-key", + throughput_mode="serverless", + ) for call in mock_db.create_container_if_not_exists.await_args_list: assert "offer_throughput" not in call.kwargs - def test_constructor_ignores_invalid_autoscale_env_in_serverless_mode(self): - with patch.dict("os.environ", {"COSMOS_DB_AUTOSCALE_MAX_RU": "not-an-int"}, clear=False): - mem = _make_client(cosmos_throughput_mode="serverless") + def test_constructor_ignores_autoscale_in_serverless_mode(self): + # Even an invalid autoscale value is ignored in serverless mode. + mem = _make_client(cosmos_throughput_mode="serverless", cosmos_autoscale_max_ru="not-an-int") assert mem._cosmos_autoscale_max_ru is None @@ -738,6 +740,37 @@ async def test_search_cosmos(self): query = container.query_items.call_args.kwargs["query"] assert "VectorDistance" in query + async def test_search_uses_keyword_params_for_hybrid_sql(self): + mem, container = _connected_client() + container.query_items = MagicMock(return_value=AsyncIterator([_make_doc()])) + + mem._embeddings_client = AsyncMock() + mem._embeddings_client.generate = AsyncMock(return_value=[0.1]) + + await mem.search_cosmos(search_terms="weather Seattle", top_k=3) + + query = container.query_items.call_args.kwargs["query"] + parameters = container.query_items.call_args.kwargs["parameters"] + assert "ORDER BY RANK RRF" in query + assert "FullTextScore(c.content, @kw0, @kw1)" in query + assert {"name": "@kw0", "value": "weather"} in parameters + assert {"name": "@kw1", "value": "seattle"} in parameters + + async def test_search_all_stopwords_uses_vector_sql(self): + mem, container = _connected_client() + container.query_items = MagicMock(return_value=AsyncIterator([_make_doc()])) + + mem._embeddings_client = AsyncMock() + mem._embeddings_client.generate = AsyncMock(return_value=[0.1]) + + await mem.search_cosmos(search_terms="what is the", top_k=3) + + query = container.query_items.call_args.kwargs["query"] + parameters = container.query_items.call_args.kwargs["parameters"] + assert "ORDER BY VectorDistance" in query + assert "RRF" not in query + assert not any(parameter["name"].startswith("@kw") for parameter in parameters) + async def test_search_hybrid(self): mem, container = _connected_client() docs = [_make_doc()] @@ -748,7 +781,6 @@ async def test_search_hybrid(self): results = await mem.search_cosmos( search_terms="weather Seattle", - hybrid_search=True, top_k=5, ) @@ -788,6 +820,38 @@ async def test_search_whitespace_only_terms(self): with pytest.raises(ValidationError, match="search_terms must be a non-empty string"): await mem.search_cosmos(search_terms=" ") + async def test_search_episodic_forwards_search_options(self): + containers = {key: MagicMock() for key in ContainerKey} + store = AsyncMemoryStore(containers=containers) + store.search = AsyncMock(return_value=[]) + + await store.search_episodic( + user_id="u1", + search_terms="weather", + top_k=2, + min_salience=0.4, + include_superseded=True, + ) + + store.search.assert_awaited_once_with( + search_terms="weather", + user_id="u1", + memory_types=["episodic"], + top_k=2, + min_salience=0.4, + include_superseded=True, + ) + + async def test_build_episodic_context_forwards_search_options(self): + containers = {key: MagicMock() for key in ContainerKey} + store = AsyncMemoryStore(containers=containers) + store.search_episodic = AsyncMock(return_value=[]) + + context = await store.build_episodic_context("u1", "weather", top_k=2) + + assert context == "" + store.search_episodic.assert_awaited_once_with("u1", "weather", top_k=2) + # =================================================================== # Processing delegation (async) diff --git a/tests/unit/aio/test_process_now.py b/tests/unit/aio/test_process_now.py index 25ae108..d0aa22e 100644 --- a/tests/unit/aio/test_process_now.py +++ b/tests/unit/aio/test_process_now.py @@ -2,7 +2,7 @@ from __future__ import annotations -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, call, patch import pytest @@ -50,7 +50,13 @@ async def test_process_now_with_inprocess_invokes_full_pipeline(): assert isinstance(client._processor, AsyncInProcessProcessor) pipeline.generate_thread_summary.assert_awaited_once_with("u", "t") pipeline.extract_memories.assert_awaited_once_with("u", "t") - pipeline.reconcile_memories.assert_awaited_once_with("u", 50) + assert pipeline.reconcile_memories.await_count == 2 + pipeline.reconcile_memories.assert_has_awaits( + [ + call("u", n=50, memory_type="fact", full_rebuild=False), + call("u", n=50, memory_type="episodic", full_rebuild=False), + ] + ) pipeline.synthesize_procedural.assert_awaited_once_with("u", force=False) pipeline.generate_user_summary.assert_awaited_once_with("u", None) assert result.procedural == {"id": "proc1", "type": "procedural"} diff --git a/tests/unit/aio/test_reconcile_telemetry.py b/tests/unit/aio/test_reconcile_telemetry.py index 858e328..cea2326 100644 --- a/tests/unit/aio/test_reconcile_telemetry.py +++ b/tests/unit/aio/test_reconcile_telemetry.py @@ -17,6 +17,14 @@ ASYNC_LOGGER_NAME = "azure.cosmos.agent_memory.pipeline.aio" +@pytest.fixture(autouse=True) +def _pin_async_legacy_reconcile(monkeypatch): + monkeypatch.setattr( + "azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_reconcile_mode", + lambda: "full_pool", + ) + + def _make_async_pipeline() -> AsyncPipelineService: p = AsyncPipelineService.__new__(AsyncPipelineService) p._embeddings = MagicMock() diff --git a/tests/unit/function_app/test_change_feed.py b/tests/unit/function_app/test_change_feed.py index f5b7dd8..a70bcb3 100644 --- a/tests/unit/function_app/test_change_feed.py +++ b/tests/unit/function_app/test_change_feed.py @@ -73,9 +73,16 @@ async def upsert_item(*, body, **_kwargs): state[body["id"]] = dict(body) return body + async def patch_item(*, item, partition_key, patch_operations): + doc = state.setdefault(item, {"id": item}) + for op in patch_operations: + doc[op["path"].lstrip("/")] = op["value"] + return dict(doc) + container.read_item = AsyncMock(side_effect=read_item) container.create_item = AsyncMock(side_effect=create_item) container.upsert_item = AsyncMock(side_effect=upsert_item) + container.patch_item = AsyncMock(side_effect=patch_item) container._state = state # exposed for assertions return container @@ -618,10 +625,19 @@ def test_reconcile_flag_set_only_when_n_facts_times_n_dedup_threshold_crosses(): asyncio.run(process_changefeed_batch([_turn() for _ in range(4)], starter, counter_container=container)) extract_calls = [c for c in starter.start_new.await_args_list if c.args[0] == "ExtractMemoriesOrchestrator"] assert len(extract_calls) == 1 - assert _extract_payload(extract_calls[0]).get("reconcile") is False + first_payload = _extract_payload(extract_calls[0]) + assert first_payload.get("reconcile") is False + # Bootstrap: no watermark yet -> recent_k spans all 4 turns so far. + assert first_payload.get("recent_k") == 4 + + # Simulate the extract orchestrator's em_AdvanceExtractWatermark activity + # advancing the watermark to the processed count (4) after batch 1. + from shared.counters import advance_extract_watermark, thread_counter_id + + asyncio.run(advance_extract_watermark(container, thread_counter_id("u1", "t1"), "u1", "t1", 4)) # Next batch: counter 4 -> 5. Reconcile threshold crossed, so the same - # extract dispatch carries reconcile=True. + # extract dispatch carries reconcile=True, and recent_k = 5 - watermark(4) = 1. starter.start_new.reset_mock() asyncio.run(process_changefeed_batch([_turn()], starter, counter_container=container)) extract_calls = [c for c in starter.start_new.await_args_list if c.args[0] == "ExtractMemoriesOrchestrator"] @@ -629,6 +645,69 @@ def test_reconcile_flag_set_only_when_n_facts_times_n_dedup_threshold_crosses(): payload = _extract_payload(extract_calls[0]) assert payload.get("reconcile") is True assert payload.get("user_id") == "u1" + assert payload.get("recent_k") == 1 + + +@patch.dict( + os.environ, + { + "THREAD_SUMMARY_EVERY_N": "0", + "FACT_EXTRACTION_EVERY_N": "1", + "USER_SUMMARY_EVERY_N": "0", + "DEDUP_EVERY_N": "1", + }, + clear=False, +) +def test_full_rebuild_flag_set_on_persisted_counter_cadence(): + """The full-pool backstop is driven by the PERSISTED counter (durable-safe), + not the in-memory per-worker sweep counter: full_rebuild=True every + (n_facts * n_dedup * DEDUP_FULL_RECLUSTER_EVERY_N) turns. Here that's + 1 * 1 * 2 = every 2 turns.""" + with patch( + "azure.cosmos.agent_memory.thresholds.get_dedup_full_recluster_every_n", + return_value=2, + ): + starter = _make_starter() + container = _make_counter_container_starting_at() + + # Turn 1: counter 0->1. Reconcile crosses (n=1), full does NOT (n=2). + asyncio.run(process_changefeed_batch([_turn()], starter, counter_container=container)) + p1 = _extract_payload( + next(c for c in starter.start_new.await_args_list if c.args[0] == "ExtractMemoriesOrchestrator") + ) + assert p1.get("reconcile") is True + assert p1.get("full_rebuild") is False + + # Turn 2: counter 1->2. Full backstop threshold (2) crossed. + starter.start_new.reset_mock() + asyncio.run(process_changefeed_batch([_turn()], starter, counter_container=container)) + p2 = _extract_payload( + next(c for c in starter.start_new.await_args_list if c.args[0] == "ExtractMemoriesOrchestrator") + ) + assert p2.get("reconcile") is True + assert p2.get("full_rebuild") is True + + +@patch.dict( + os.environ, + { + "THREAD_SUMMARY_EVERY_N": "0", + "FACT_EXTRACTION_EVERY_N": "1", + "USER_SUMMARY_EVERY_N": "0", + "DEDUP_EVERY_N": "0", + }, + clear=False, +) +def test_extract_payload_recent_k_uses_max_threshold_and_batch_delta(): + starter = _make_starter() + container = _make_counter_container_starting_at() + + asyncio.run(process_changefeed_batch([_turn() for _ in range(3)], starter, counter_container=container)) + + extract_calls = [c for c in starter.start_new.await_args_list if c.args[0] == "ExtractMemoriesOrchestrator"] + assert len(extract_calls) == 1 + payload = _extract_payload(extract_calls[0]) + assert payload["recent_k"] == 3 @patch.dict( diff --git a/tests/unit/function_app/test_orchestrators.py b/tests/unit/function_app/test_orchestrators.py index 4f8f60f..9fc629f 100644 --- a/tests/unit/function_app/test_orchestrators.py +++ b/tests/unit/function_app/test_orchestrators.py @@ -9,7 +9,7 @@ from __future__ import annotations -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, call, patch import pytest from orchestrators import extract_memories as em_mod @@ -193,6 +193,7 @@ def test_extract_only_when_reconcile_flag_absent(self, _retry): gen, [ {"facts": [{"id": "f1"}], "episodic": [], "updates": []}, + {"facts": [{"id": "f1", "deduped": True}], "episodic": [], "updates": []}, { "fact_count": 2, "episodic_count": 0, @@ -201,7 +202,15 @@ def test_extract_only_when_reconcile_flag_absent(self, _retry): ], ) - assert [c[0] for c in ctx._yielded_calls] == ["em_Extract", "em_Persist"] + assert [c[0] for c in ctx._yielded_calls] == ["em_Extract", "em_Dedup", "em_Persist"] + assert ctx._yielded_calls[1][2] == { + "user_id": "u1", + "extracted": {"facts": [{"id": "f1"}], "episodic": [], "updates": []}, + } + assert ctx._yielded_calls[2][2] == { + "user_id": "u1", + "extracted": {"facts": [{"id": "f1", "deduped": True}], "episodic": [], "updates": []}, + } assert result["persisted"] is True assert result["extracted"]["fact_count"] == 2 assert result["reconciled"] is None @@ -213,21 +222,28 @@ def test_chains_reconcile_when_flag_true(self, _retry): result, _ = _drive( gen, [ + {"facts": [{"id": "f1"}], "episodic": [], "updates": []}, {"facts": [{"id": "f1"}], "episodic": [], "updates": []}, {"fact_count": 2, "episodic_count": 0, "updated_count": 0}, - {"kept": 0, "merged": 1, "contradicted": 0}, + { + "fact": {"kept": 0, "merged": 1, "contradicted": 0}, + "episodic": {"kept": 1, "merged": 0, "contradicted": 0}, + }, {"status": "synthesized", "version": 3}, ], ) names = [c[0] for c in ctx._yielded_calls] - assert names == ["em_Extract", "em_Persist", "em_ReconcileMemories"] - assert ctx._yielded_calls[2][2] == {"user_id": "u1"} + assert names == ["em_Extract", "em_Dedup", "em_Persist", "em_ReconcileMemories"] + assert ctx._yielded_calls[3][2] == {"user_id": "u1", "full_rebuild": False} assert [s[0] for s in ctx._yielded_sub_orchestrators] == [ "SynthesizeProceduralOrchestrator", ] assert ctx._yielded_sub_orchestrators[0][2] == {"user_id": "u1", "force": False} - assert result["reconciled"] == {"kept": 0, "merged": 1, "contradicted": 0} + assert result["reconciled"] == { + "fact": {"kept": 0, "merged": 1, "contradicted": 0}, + "episodic": {"kept": 1, "merged": 0, "contradicted": 0}, + } assert result["procedural"] == {"status": "synthesized", "version": 3} @patch.object(em_mod, "default_retry_options", return_value=MagicMock()) @@ -240,17 +256,24 @@ def boom_after_sub(name, retry, sub_payload, *args, **kwargs): return ("__sub_boom__", name, sub_payload) ctx.call_sub_orchestrator_with_retry.side_effect = boom_after_sub - # We'll send activity results normally for the first 3 yields, then throw - # an exception into the 4th yield (the sub-orchestrator call). + # We'll send activity results normally, then throw an exception into the + # sub-orchestrator yield. gen = self._orchestrator()(ctx) # Yield 1: em_Extract gen.send(None) - # Yield 2: em_Persist + # Yield 2: em_Dedup + gen.send({"facts": [{"id": "f1"}], "episodic": [], "updates": []}) + # Yield 3: em_Persist gen.send({"facts": [{"id": "f1"}], "episodic": [], "updates": []}) - # Yield 3: em_ReconcileMemories + # Yield 4: em_ReconcileMemories gen.send({"fact_count": 2, "episodic_count": 0, "updated_count": 0}) - # Yield 4: SynthesizeProceduralOrchestrator — throw an exception - gen.send({"kept": 0, "merged": 1, "contradicted": 0}) + # Yield 5: SynthesizeProceduralOrchestrator — throw an exception + gen.send( + { + "fact": {"kept": 0, "merged": 1, "contradicted": 0}, + "episodic": {"kept": 1, "merged": 0, "contradicted": 0}, + } + ) try: gen.throw(RuntimeError("procedural blew up")) except StopIteration as stop: @@ -259,7 +282,10 @@ def boom_after_sub(name, retry, sub_payload, *args, **kwargs): pytest.fail("orchestrator did not return after procedural exception") assert result["persisted"] is True - assert result["reconciled"] == {"kept": 0, "merged": 1, "contradicted": 0} + assert result["reconciled"] == { + "fact": {"kept": 0, "merged": 1, "contradicted": 0}, + "episodic": {"kept": 1, "merged": 0, "contradicted": 0}, + } assert result["procedural"] is None @patch.object(em_mod, "default_retry_options", return_value=MagicMock()) @@ -269,23 +295,44 @@ def test_procedural_not_called_when_reconcile_skipped(self, _retry): result, _ = _drive( gen, [ + {"facts": [], "episodic": [], "updates": []}, {"facts": [], "episodic": [], "updates": []}, {"fact_count": 0, "episodic_count": 0, "updated_count": 0}, ], ) - assert [c[0] for c in ctx._yielded_calls] == ["em_Extract", "em_Persist"] + assert [c[0] for c in ctx._yielded_calls] == ["em_Extract", "em_Dedup", "em_Persist"] assert ctx._yielded_sub_orchestrators == [] assert result["procedural"] is None @patch.object(em_mod, "default_retry_options", return_value=MagicMock()) - def test_extract_payload_carries_user_thread_and_limit(self, _retry): + def test_extract_payload_carries_user_thread_without_recent_k_when_absent(self, _retry): ctx = _make_context({"user_id": "u", "thread_id": "t"}) gen = self._orchestrator()(ctx) - _drive(gen, [{"facts": []}, {"fact_count": 0}]) + _drive(gen, [{"facts": []}, {"facts": []}, {"fact_count": 0}]) + + extract_payload = ctx._yielded_calls[0][2] + assert extract_payload == {"user_id": "u", "thread_id": "t"} + + @patch.object(em_mod, "default_retry_options", return_value=MagicMock()) + def test_extract_payload_carries_recent_k_when_provided(self, _retry): + ctx = _make_context({"user_id": "u", "thread_id": "t", "recent_k": 7}) + gen = self._orchestrator()(ctx) + _drive(gen, [{"facts": []}, {"facts": []}, {"fact_count": 0}]) extract_payload = ctx._yielded_calls[0][2] - assert extract_payload == {"user_id": "u", "thread_id": "t", "limit": 20} + assert extract_payload == {"user_id": "u", "thread_id": "t", "recent_k": 7} + + @patch.object(em_mod, "default_retry_options", return_value=MagicMock()) + def test_dedup_output_flows_to_persist(self, _retry): + extracted = {"facts": [{"id": "f1"}], "episodic": [], "updates": []} + deduped = {"facts": [{"id": "f1", "embedding": [0.1]}], "episodic": [], "updates": []} + ctx = _make_context({"user_id": "u", "thread_id": "t"}) + gen = self._orchestrator()(ctx) + _drive(gen, [extracted, deduped, {"fact_count": 1}]) + + assert ctx._yielded_calls[1][2] == {"user_id": "u", "extracted": extracted} + assert ctx._yielded_calls[2][2] == {"user_id": "u", "extracted": deduped} @patch.object(em_mod, "default_retry_options", return_value=MagicMock()) def test_activity_failure_propagates(self, _retry): @@ -295,6 +342,24 @@ def test_activity_failure_propagates(self, _retry): with pytest.raises(ValueError, match="kaboom"): gen.throw(ValueError("kaboom")) + @patch.object(em_mod, "default_retry_options", return_value=MagicMock()) + def test_advance_watermark_after_persist_when_count_present(self, _retry): + ctx = _make_context({"user_id": "u1", "thread_id": "t1", "count": 42}) + gen = self._orchestrator()(ctx) + _drive(gen, [{"facts": []}, {"facts": []}, {"fact_count": 0}, True]) + + names = [c[0] for c in ctx._yielded_calls] + assert names == ["em_Extract", "em_Dedup", "em_Persist", "em_AdvanceExtractWatermark"] + assert ctx._yielded_calls[3][2] == {"user_id": "u1", "thread_id": "t1", "count": 42} + + @patch.object(em_mod, "default_retry_options", return_value=MagicMock()) + def test_no_watermark_advance_when_count_absent(self, _retry): + ctx = _make_context({"user_id": "u1", "thread_id": "t1"}) + gen = self._orchestrator()(ctx) + _drive(gen, [{"facts": []}, {"facts": []}, {"fact_count": 0}]) + names = [c[0] for c in ctx._yielded_calls] + assert "em_AdvanceExtractWatermark" not in names + def test_missing_thread_id_raises(self): with patch.object(em_mod, "default_retry_options", return_value=MagicMock()): ctx = _make_context({"user_id": "u"}) @@ -304,10 +369,107 @@ def test_missing_thread_id_raises(self): # --------------------------------------------------------------------------- -# UserSummaryOrchestrator +# Extract memory activities # --------------------------------------------------------------------------- +class TestExtractMemoryActivities: + def test_em_extract_uses_payload_recent_k(self): + pipeline = MagicMock() + pipeline.extract_memories_dry.return_value = {"facts": [], "episodic": [], "updates": []} + + with patch.object(em_mod, "get_pipeline", return_value=pipeline): + result = em_mod.em_Extract({"user_id": "u1", "thread_id": "t1", "recent_k": 3}) + + pipeline.extract_memories_dry.assert_called_once_with(user_id="u1", thread_id="t1", recent_k=3) + assert result == {"facts": [], "episodic": [], "updates": []} + + def test_em_extract_falls_back_to_max_batch_size_when_recent_k_absent(self): + pipeline = MagicMock() + pipeline.extract_memories_dry.return_value = {"facts": [], "episodic": [], "updates": []} + + with patch.object(em_mod, "get_pipeline", return_value=pipeline): + em_mod.em_Extract({"user_id": "u1", "thread_id": "t1"}) + + pipeline.extract_memories_dry.assert_called_once_with(user_id="u1", thread_id="t1", recent_k=20) + + def test_em_dedup_delegates_to_pipeline_and_returns_deduped_dict(self): + extracted = {"facts": [{"id": "f1"}], "episodic": [], "updates": []} + deduped = {"facts": [{"id": "f1", "embedding": [0.1]}], "episodic": [], "updates": []} + pipeline = MagicMock() + pipeline.dedup_extracted_memories.return_value = deduped + + with patch.object(em_mod, "get_pipeline", return_value=pipeline): + result = em_mod.em_Dedup({"user_id": "u1", "extracted": extracted}) + + pipeline.dedup_extracted_memories.assert_called_once_with(user_id="u1", extracted=extracted) + assert result == deduped + + def test_em_dedup_falls_back_to_input_when_pipeline_returns_none(self): + extracted = {"facts": [], "episodic": [], "updates": []} + pipeline = MagicMock() + pipeline.dedup_extracted_memories.return_value = None + + with patch.object(em_mod, "get_pipeline", return_value=pipeline): + result = em_mod.em_Dedup({"user_id": "u1", "extracted": extracted}) + + assert result is extracted + + def test_em_reconcile_memories_reconciles_fact_and_episodic(self): + pipeline = MagicMock() + pipeline.reconcile_memories.side_effect = [ + {"kept": 2, "merged": 1, "contradicted": 0}, + {"kept": 1, "merged": 0, "contradicted": 0}, + ] + + with ( + patch.object(em_mod, "get_pipeline", return_value=pipeline), + patch("azure.cosmos.agent_memory.thresholds.get_dedup_pool_size", return_value=17), + ): + result = em_mod.em_ReconcileMemories({"user_id": "u1"}) + + # full_rebuild defaults False when the change-feed didn't request a backstop. + assert pipeline.reconcile_memories.call_args_list == [ + call(user_id="u1", n=17, memory_type="fact", full_rebuild=False), + call(user_id="u1", n=17, memory_type="episodic", full_rebuild=False), + ] + assert result == { + "fact": {"kept": 2, "merged": 1, "contradicted": 0}, + "episodic": {"kept": 1, "merged": 0, "contradicted": 0}, + } + + def test_em_reconcile_memories_forwards_full_rebuild(self): + pipeline = MagicMock() + pipeline.reconcile_memories.side_effect = [{"kept": 0}, {"kept": 0}] + with ( + patch.object(em_mod, "get_pipeline", return_value=pipeline), + patch("azure.cosmos.agent_memory.thresholds.get_dedup_pool_size", return_value=17), + ): + em_mod.em_ReconcileMemories({"user_id": "u1", "full_rebuild": True}) + + assert pipeline.reconcile_memories.call_args_list == [ + call(user_id="u1", n=17, memory_type="fact", full_rebuild=True), + call(user_id="u1", n=17, memory_type="episodic", full_rebuild=True), + ] + + def test_em_advance_extract_watermark_stamps_counter(self): + import asyncio + + from shared import cosmos_clients, counters + + container = MagicMock() + with ( + patch.object(cosmos_clients, "get_counter_container_async", new=AsyncMock(return_value=container)), + patch.object(counters, "advance_extract_watermark", new=AsyncMock()) as advance, + ): + result = asyncio.run( + em_mod.em_AdvanceExtractWatermark({"user_id": "u1", "thread_id": "t1", "count": 9}) + ) + + assert result is True + advance.assert_awaited_once_with(container, "thread:u1:t1", "u1", "t1", 9) + + class TestUserSummaryOrchestrator: def _orchestrator(self): return _user_function(us_mod.UserSummaryOrchestrator) diff --git a/tests/unit/processors/test_inprocess.py b/tests/unit/processors/test_inprocess.py index be8d841..28ddc2f 100644 --- a/tests/unit/processors/test_inprocess.py +++ b/tests/unit/processors/test_inprocess.py @@ -16,20 +16,23 @@ def test_process_thread_calls_summarize_extract_reconcile_in_order(): proc = InProcessProcessor(pipeline=pipeline) result = proc.process_thread(user_id="u1", thread_id="t1", turns=[]) - # Order of calls: summary -> extract -> reconcile + # Order of calls: summary -> extract -> reconcile (fact) -> reconcile (episodic) method_order = [c[0] for c in pipeline.method_calls] assert method_order == [ "generate_thread_summary", "extract_memories", "reconcile_memories", + "reconcile_memories", ] pipeline.generate_thread_summary.assert_called_once_with("u1", "t1") pipeline.extract_memories.assert_called_once_with("u1", "t1") - pipeline.reconcile_memories.assert_called_once_with("u1", 50) + assert pipeline.reconcile_memories.call_count == 2 + assert pipeline.reconcile_memories.call_args_list[0].kwargs["memory_type"] == "fact" + assert pipeline.reconcile_memories.call_args_list[1].kwargs["memory_type"] == "episodic" assert isinstance(result, ProcessThreadResult) assert result.thread_summary == {"id": "summary_u_t", "type": "thread_summary"} - assert result.reconciled_count == 3 + assert result.reconciled_count == 6 assert result.elapsed_ms >= 0 @@ -65,6 +68,20 @@ def test_generate_user_summary_no_summaries(): assert res.summary is None +def test_process_extract_memories_passes_recent_k_and_filters_to_ints(): + pipeline = MagicMock() + pipeline.extract_memories.return_value = { + "fact_count": 2, + "non_int_field": "skip me", + } + + proc = InProcessProcessor(pipeline=pipeline) + result = proc.process_extract_memories(user_id="u", thread_id="t", recent_k=3) + + pipeline.extract_memories.assert_called_once_with("u", "t", recent_k=3) + assert result == {"fact_count": 2} + + def test_close_is_noop(): proc = InProcessProcessor(pipeline=MagicMock()) assert proc.close() is None diff --git a/tests/unit/processors/test_protocol_satisfaction.py b/tests/unit/processors/test_protocol_satisfaction.py index dc63c98..6869f00 100644 --- a/tests/unit/processors/test_protocol_satisfaction.py +++ b/tests/unit/processors/test_protocol_satisfaction.py @@ -32,6 +32,7 @@ def process_extract_memories( *, user_id: str, thread_id: str, + recent_k: Optional[int] = None, ) -> dict[str, int]: return {} diff --git a/tests/unit/services/test_chaos_extract_persist.py b/tests/unit/services/test_chaos_extract_persist.py index d12e61b..82c61d0 100644 --- a/tests/unit/services/test_chaos_extract_persist.py +++ b/tests/unit/services/test_chaos_extract_persist.py @@ -11,6 +11,26 @@ from azure.cosmos.agent_memory.services.pipeline import PipelineService, _StoreContainerAdapter +@pytest.fixture(autouse=True) +def _pin_legacy_extract_dedup(monkeypatch): + monkeypatch.setattr( + "azure.cosmos.agent_memory.thresholds.get_dedup_vector_enabled", + lambda: False, + ) + monkeypatch.setattr( + "azure.cosmos.agent_memory.thresholds.get_dedup_context_vector_enabled", + lambda: False, + ) + monkeypatch.setattr( + "azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_vector_enabled", + lambda: False, + ) + monkeypatch.setattr( + "azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_context_vector_enabled", + lambda: False, + ) + + class _FlakyContainer: def __init__(self): self.docs: dict[str, dict[str, Any]] = {} diff --git a/tests/unit/services/test_dedup_vector.py b/tests/unit/services/test_dedup_vector.py new file mode 100644 index 0000000..a44f83a --- /dev/null +++ b/tests/unit/services/test_dedup_vector.py @@ -0,0 +1,383 @@ +from __future__ import annotations + +import json +from typing import Any +from unittest.mock import MagicMock + +from azure.cosmos.agent_memory.services.pipeline import PipelineService + + +def _make_pipeline() -> PipelineService: + p = PipelineService.__new__(PipelineService) + p._memories_container = MagicMock() + p._container = p._memories_container + p._embeddings = MagicMock() + p._embed_batch = MagicMock() + p._embed_one = MagicMock(return_value=[1.0]) + p._run_prompty = MagicMock( + return_value=json.dumps({"duplicate_groups": [], "contradicted_pairs": [], "kept_ids": []}) + ) + p._upsert_memory = MagicMock(side_effect=lambda doc: doc) + p._mark_superseded = MagicMock(return_value=True) + return p + + +def _doc(mid: str, content: str, memory_type: str = "fact", **extra: Any) -> dict[str, Any]: + tags = extra.pop("tags", [f"sys:{memory_type}"]) + metadata = extra.pop("metadata", {"category": "preference"} if memory_type == "fact" else { + "scope_type": "project", + "scope_value": "demo", + "lesson": content, + "outcome_valence": "neutral", + }) + return { + "id": mid, + "user_id": "u1", + "thread_id": "t1", + "type": memory_type, + "role": "system", + "content": content, + "content_hash": mid, + "confidence": 0.8, + "salience": 0.7, + "tags": tags, + "metadata": metadata, + "prompt_id": "extract_memories.prompty", + "prompt_version": "v1", + "created_at": "2025-01-01T00:00:00+00:00", + "updated_at": "2025-01-01T00:00:00+00:00", + **extra, + } + + +def test_vector_distance_function_reads_container_policy() -> None: + # The distance function comes from the container's vector embedding policy + # (read once, cached), NOT an env var. + p = _make_pipeline() + p._memories_container.read.return_value = { + "vectorEmbeddingPolicy": {"vectorEmbeddings": [{"path": "/embedding", "distanceFunction": "euclidean"}]} + } + assert p._vector_distance_function() == "euclidean" + # Cached: a later policy change is not re-read within the instance's lifetime. + p._memories_container.read.return_value = { + "vectorEmbeddingPolicy": {"vectorEmbeddings": [{"path": "/embedding", "distanceFunction": "cosine"}]} + } + assert p._vector_distance_function() == "euclidean" + assert p._memories_container.read.call_count == 1 + + +def test_distance_function_not_cached_on_read_failure() -> None: + # A transient container.read() failure must NOT poison the cache: it returns an + # uncached cosine default so the next call self-heals to the real (euclidean) + # policy. Caching cosine here would silently mis-handle a euclidean container. + p = _make_pipeline() + euclid = {"vectorEmbeddingPolicy": {"vectorEmbeddings": [{"path": "/embedding", "distanceFunction": "euclidean"}]}} + p._memories_container.read = MagicMock(side_effect=[RuntimeError("429 throttled"), euclid]) + + # First call: transient failure -> cosine, but NOT cached. + assert p._vector_distance_function() == "cosine" + assert getattr(p, "_distance_function_cache", None) is None + + # Second call: read succeeds -> real euclidean policy, now cached. + assert p._vector_distance_function() == "euclidean" + assert p._distance_function_cache == "euclidean" + + +def test_vector_candidates_orders_nearest_first_by_distance_function() -> None: + # Parity with async: ORDER BY direction follows the container distanceFunction. + p = _make_pipeline() + captured: dict[str, str] = {} + + def query_items(*, query: str, parameters, **kwargs): + del parameters, kwargs + captured["query"] = query + return iter( + [ + {"id": "near", "content": "a", "type": "fact", "score": 0.95}, + {"id": "far", "content": "b", "type": "fact", "score": 0.10}, + ] + ) + + p._memories_container.query_items.side_effect = query_items + + p._distance_function_cache = "cosine" + out = p._vector_candidates( + user_id="u1", embedding=[1.0, 0.0], memory_type="fact", top_k=2, exclude_ids=set() + ) + # Cosmos rejects an explicit ASC/DESC on ORDER BY VectorDistance(); it orders + # most-similar-first server-side. Direction-awareness lives in the Python sort. + assert "ORDER BY VectorDistance(c.embedding, @vec)" in captured["query"] + assert "VectorDistance(c.embedding, @vec) DESC" not in captured["query"] + assert "VectorDistance(c.embedding, @vec) ASC" not in captured["query"] + assert [c["id"] for c in out] == ["near", "far"] + + p._distance_function_cache = "euclidean" + out = p._vector_candidates( + user_id="u1", embedding=[1.0, 0.0], memory_type="fact", top_k=2, exclude_ids=set() + ) + assert "VectorDistance(c.embedding, @vec) ASC" not in captured["query"] + # euclidean: lower distance = more similar, so 0.10 ("far" label) sorts first. + assert [c["id"] for c in out] == ["far", "near"] + + +def test_dedup_extracted_vector_ladder_and_intra_batch() -> None: + p = _make_pipeline() + p._embed_batch.return_value = [[1.0, 0.0], [0.0, 1.0], [1.0, 0.0], [0.99, 0.01]] + p._vector_candidates = MagicMock( + side_effect=[ + [{"id": "existing-high", "content": "same", "type": "fact", "score": 0.99}], + [{"id": "existing-mid", "content": "near", "type": "fact", "score": 0.85}], + [{"id": "existing-low", "content": "far", "type": "fact", "score": 0.20}], + [{"id": "existing-low-2", "content": "far", "type": "fact", "score": 0.10}], + ] + ) + extracted = { + "facts": [ + _doc("f-high", "drop against existing"), + _doc("f-mid", "tag against existing"), + _doc("f-clean", "keep clean"), + _doc("f-intra", "near-dup in batch is now kept (deferred to reconcile)"), + ], + "episodic": [], + "updates": [], + } + + out = p.dedup_extracted_memories("u1", extracted) + + # Intra-batch new-vs-new dedup was removed: f-intra is compared only against + # persisted memories (Cosmos) -> novel here -> kept; reconcile catches any + # same-batch near-dups later. + assert [doc["id"] for doc in out["facts"]] == ["f-mid", "f-clean", "f-intra"] + assert out["facts"][0]["tags"][-1] == "sys:dup-candidate" + assert out["facts"][0]["metadata"]["dup_of"] == "existing-mid" + assert out["facts"][0]["metadata"]["dup_score"] == 0.85 + assert "sys:dup-candidate" not in out["facts"][1]["tags"] + assert all("embedding" in doc for doc in out["facts"]) + assert out["updates"][-1]["vector_dedup_skipped"] == 1 + assert out["updates"][-1]["dup_candidates_tagged"] == 1 + + +def test_candidate_mode_clears_tags_only_for_survivors() -> None: + # Latent-bug regression: a source consumed (superseded) by a merge must NOT be + # re-upserted by tag-clearing, which would resurrect it without superseded_by. + p = _make_pipeline() + f1 = _doc("f1", "a", tags=["sys:fact", "sys:dup-candidate"]) + f2 = _doc("f2", "b", tags=["sys:fact", "sys:dup-candidate"]) + f3 = _doc("f3", "c", tags=["sys:fact", "sys:dup-candidate"]) + p._build_candidate_clusters = MagicMock(return_value=([[f1, f2, f3]], 3, [f1, f2, f3])) + p._reconcile_pool = MagicMock(return_value=({"kept": 1, "merged": 2, "contradicted": 0}, {"f1", "f2"})) + cleared: list[str] = [] + p._clear_dup_candidate_tags = MagicMock(side_effect=lambda docs: cleared.extend(d["id"] for d in docs)) + p._emit_reconcile_outcome = MagicMock() + + result = p._reconcile_candidate_mode("u1", n=50, memory_type="fact", started_at=0.0) + + assert cleared == ["f3"] + assert result == { + "kept": 1, + "merged": 2, + "contradicted": 0, + "reconcile_clusters_sent": 1, + "reconcile_llm_calls_saved": 2, + } + + +def test_candidate_mode_clears_tags_on_orphan_seeds() -> None: + # Seeds tagged sys:dup-candidate that never join a cluster (no near-duplicate) + # must have their stale tag cleared so future sweeps don't re-scan them forever. + p = _make_pipeline() + orphan = _doc("orphan", "lonely", tags=["sys:fact", "sys:dup-candidate"]) + c1 = _doc("c1", "a", tags=["sys:fact", "sys:dup-candidate"]) + c2 = _doc("c2", "b", tags=["sys:fact", "sys:dup-candidate"]) + p._build_candidate_clusters = MagicMock(return_value=([[c1, c2]], 3, [orphan, c1, c2])) + p._reconcile_pool = MagicMock(return_value=({"kept": 2, "merged": 0, "contradicted": 0}, set())) + cleared: list[str] = [] + p._clear_dup_candidate_tags = MagicMock(side_effect=lambda docs: cleared.extend(d["id"] for d in docs)) + p._emit_reconcile_outcome = MagicMock() + + p._reconcile_candidate_mode("u1", n=50, memory_type="fact", started_at=0.0) + + # Cluster survivors (c1, c2) plus the orphan all get cleared. + assert set(cleared) == {"c1", "c2", "orphan"} + + +def test_euclidean_disables_near_exact_autodrop() -> None: + # On euclidean distance the cosine-calibrated DEDUP_SIM_HIGH auto-drop is + # disabled: a near-identical existing memory must NOT silently drop the new + # one. It falls through to borderline tagging (LLM reconcile adjudicates). + p = _make_pipeline() + p._vector_distance_function = MagicMock(return_value="euclidean") + p._embed_batch.return_value = [[1.0, 0.0]] + # euclidean score = distance; 0.05 is "near-exact" (would drop under cosine rules). + p._vector_candidates = MagicMock( + return_value=[{"id": "existing", "content": "same", "type": "fact", "score": 0.05}] + ) + extracted = {"facts": [_doc("f-new", "near identical")], "episodic": [], "updates": []} + + out = p.dedup_extracted_memories("u1", extracted) + + # Not dropped — kept and tagged for LLM reconcile instead. + assert [doc["id"] for doc in out["facts"]] == ["f-new"] + assert out["facts"][0]["tags"][-1] == "sys:dup-candidate" + assert out["updates"][-1]["vector_dedup_skipped"] == 0 + assert out["updates"][-1]["dup_candidates_tagged"] == 1 + + +def test_candidate_mode_has_no_inmemory_backstop() -> None: + # The periodic full-pool backstop is no longer driven by an in-memory sweep + # counter (unreliable on the FA per-worker singleton). Candidate mode does + # ONLY clustering now; the full-pool pass is requested by the caller via + # full_rebuild on a persisted-counter cadence. + p = _make_pipeline() + assert not hasattr(p, "_next_reconcile_sweep") + p._build_candidate_clusters = MagicMock(return_value=([], 0, [])) + p._active_memories_for_reconcile = MagicMock() + p._emit_reconcile_outcome = MagicMock() + + # Many sweeps in a row never escalate to a full-pool pass on their own. + for _ in range(30): + p._reconcile_candidate_mode("u1", n=50, memory_type="fact", started_at=0.0) + p._build_candidate_clusters.assert_called() + p._active_memories_for_reconcile.assert_not_called() + + +def test_full_rebuild_bypasses_candidate_mode(monkeypatch) -> None: + # Public reconcile(full_rebuild=True) must take the full-pool single-LLM-pass + # path even under candidate mode, so it catches dissimilar contradictions. + monkeypatch.setattr("azure.cosmos.agent_memory.thresholds.get_dedup_reconcile_mode", lambda: "candidate") + p = _make_pipeline() + pool = [_doc("f1", "User is vegetarian"), _doc("f2", "User loves steak")] + p._active_memories_for_reconcile = MagicMock(return_value=pool) + p._reconcile_pool = MagicMock(return_value=({"kept": 1, "merged": 0, "contradicted": 1}, set())) + p._reconcile_candidate_mode = MagicMock() + p._emit_reconcile_outcome = MagicMock() + + result = p.reconcile_memories("u1", n=50, memory_type="fact", full_rebuild=True) + + p._reconcile_candidate_mode.assert_not_called() + p._reconcile_pool.assert_called_once_with("u1", "fact", pool) + assert result["contradicted"] == 1 + + +def test_full_rebuild_clears_survivor_tags(monkeypatch) -> None: + # full_rebuild full-pool path must clear sys:dup-candidate on survivors so it + # doesn't leave stale tags/metadata on user-visible memories. + monkeypatch.setattr("azure.cosmos.agent_memory.thresholds.get_dedup_reconcile_mode", lambda: "candidate") + p = _make_pipeline() + pool = [ + _doc("f1", "a", tags=["sys:fact", "sys:dup-candidate"]), + _doc("f2", "b", tags=["sys:fact", "sys:dup-candidate"]), + ] + p._active_memories_for_reconcile = MagicMock(return_value=pool) + # f1 consumed (superseded), f2 survives. + p._reconcile_pool = MagicMock(return_value=({"kept": 1, "merged": 1, "contradicted": 0}, {"f1"})) + cleared: list[str] = [] + p._clear_dup_candidate_tags = MagicMock(side_effect=lambda docs: cleared.extend(d["id"] for d in docs)) + p._emit_reconcile_outcome = MagicMock() + + p.reconcile_memories("u1", n=50, memory_type="fact", full_rebuild=True) + + # Survivor f2 cleared; consumed f1 not re-upserted (would resurrect it). + assert cleared == ["f2"] + + +def test_sweep_survives_one_cluster_failure() -> None: + # A truncated/malformed LLM response on one cluster must not abort the sweep: + # remaining clusters still reconcile, orphan clearing still runs, and the failed + # cluster's tags are RETAINED (not cleared) so it retries next sweep. + p = _make_pipeline() + c1 = [ + _doc("a1", "x", tags=["sys:fact", "sys:dup-candidate"]), + _doc("a2", "y", tags=["sys:fact", "sys:dup-candidate"]), + ] + c2 = [ + _doc("b1", "p", tags=["sys:fact", "sys:dup-candidate"]), + _doc("b2", "q", tags=["sys:fact", "sys:dup-candidate"]), + ] + orphan = _doc("o1", "lonely", tags=["sys:fact", "sys:dup-candidate"]) + seeds = [c1[0], c1[1], c2[0], c2[1], orphan] + p._build_candidate_clusters = MagicMock(return_value=([c1, c2], 5, seeds)) + p._reconcile_pool = MagicMock( + side_effect=[RuntimeError("truncated LLM response"), ({"kept": 2, "merged": 0, "contradicted": 0}, set())] + ) + cleared: list[str] = [] + p._clear_dup_candidate_tags = MagicMock(side_effect=lambda docs: cleared.extend(d["id"] for d in docs)) + p._emit_reconcile_outcome = MagicMock() + + result = p._reconcile_candidate_mode("u1", n=50, memory_type="fact", started_at=0.0) + + # c1 failed -> its tags retained (not cleared) for retry; c2 survivors + orphan cleared. + assert "a1" not in cleared and "a2" not in cleared + assert {"b1", "b2", "o1"} <= set(cleared) + assert result["reconcile_clusters_sent"] == 2 + assert result["kept"] == 2 + + +def test_dedup_extracted_flag_off_is_noop(monkeypatch) -> None: + monkeypatch.setattr("azure.cosmos.agent_memory.thresholds.get_dedup_vector_enabled", lambda: False) + p = _make_pipeline() + extracted = {"facts": [_doc("f1", "content")], "episodic": [], "updates": []} + + out = p.dedup_extracted_memories("u1", extracted) + + assert out is extracted + p._embed_batch.assert_not_called() + + +def test_reconcile_memory_type_routing_episodic_and_procedural(monkeypatch) -> None: + monkeypatch.setattr("azure.cosmos.agent_memory.thresholds.get_dedup_reconcile_mode", lambda: "full_pool") + p = _make_pipeline() + episodes = [_doc("e1", "episode one", "episodic"), _doc("e2", "episode two", "episodic")] + p._memories_container.query_items.return_value = iter(episodes) + p._run_prompty.return_value = json.dumps({"duplicate_groups": [], "kept_ids": ["e1", "e2"]}) + + result = p.reconcile_memories("u1", memory_type="episodic") + + assert result["contradicted"] == 0 + assert p._run_prompty.call_args.args[0] == "dedup_episodic.prompty" + assert "episodics_text" in p._run_prompty.call_args.kwargs["inputs"] + + p._run_prompty.reset_mock() + assert p.reconcile_memories("u1", memory_type="procedural")["reconcile_clusters_sent"] == 0 + p._run_prompty.assert_not_called() + + +def test_candidate_mode_connected_components() -> None: + p = _make_pipeline() + seed = _doc( + "f1", + "User likes coffee", + embedding=[1.0, 0.0], + tags=["sys:fact", "sys:dup-candidate"], + metadata={"category": "preference", "dup_of": "f2", "dup_score": 0.9}, + ) + neighbor = _doc("f2", "User loves coffee", embedding=[0.95, 0.05]) + + def query_items(*, query: str, parameters: list[dict[str, Any]], **kwargs: Any): + del kwargs + params = {p["name"]: p["value"] for p in parameters} + if "ARRAY_CONTAINS" in query: + return iter([seed]) + ids = {value for name, value in params.items() if name.startswith("@id")} + if ids: + return iter([doc for doc in (seed, neighbor) if doc["id"] in ids]) + return iter([seed, neighbor]) + + p._memories_container.query_items.side_effect = query_items + p._vector_candidates = MagicMock( + return_value=[{"id": "f2", "content": neighbor["content"], "type": "fact", "score": 0.9}] + ) + p._run_prompty.return_value = json.dumps( + { + "duplicate_groups": [{"merged_content": "User likes coffee.", "source_ids": ["f1", "f2"]}], + "contradicted_pairs": [], + "kept_ids": [], + } + ) + + result = p.reconcile_memories("u1", memory_type="fact") + + assert result["reconcile_clusters_sent"] == 1 + assert result["merged"] == 2 + p._run_prompty.assert_called_once() + assert p._run_prompty.call_args.args[0] == "dedup.prompty" diff --git a/tests/unit/services/test_extract_dry.py b/tests/unit/services/test_extract_dry.py index 32b4390..c5c9adc 100644 --- a/tests/unit/services/test_extract_dry.py +++ b/tests/unit/services/test_extract_dry.py @@ -2,6 +2,7 @@ import json from typing import Any +from unittest.mock import AsyncMock import pytest @@ -14,10 +15,12 @@ class _SyncChat: def __init__(self, responses: list[dict[str, Any]]): self.responses = list(responses) self.calls = 0 + self.messages: list[list[dict[str, Any]]] = [] def generate(self, messages: list[dict[str, Any]], **opts: Any) -> str: - del messages, opts + del opts self.calls += 1 + self.messages.append(messages) return json.dumps(self.responses.pop(0)) @@ -58,6 +61,8 @@ async def generate(self, text: str) -> list[float]: class _Store: def __init__(self, docs: list[dict[str, Any]]): self.docs = [dict(doc) for doc in docs] + self.search_calls: list[dict[str, Any]] = [] + self.search_results: list[dict[str, Any]] = [] def query(self, sql: str, parameters=None, partition_key=None, cross_partition: bool = False): del partition_key, cross_partition @@ -101,6 +106,26 @@ def mark_superseded(self, old_doc: dict[str, Any], superseder_id: str, *, reason del old_doc, superseder_id, reason return True + def search( + self, + *, + search_terms: str, + user_id: str, + memory_types: list[str], + top_k: int, + include_superseded: bool = False, + ) -> list[dict[str, Any]]: + self.search_calls.append( + { + "search_terms": search_terms, + "user_id": user_id, + "memory_types": memory_types, + "top_k": top_k, + "include_superseded": include_superseded, + } + ) + return [dict(doc) for doc in self.search_results] + class _AsyncStore(_Store): async def query(self, sql: str, parameters=None, partition_key=None, cross_partition: bool = False): @@ -221,53 +246,19 @@ def test_extract_memories_dry_is_byte_deterministic_for_same_llm_response() -> N ) -@pytest.mark.asyncio -async def test_async_extract_memories_dry_shape_is_small_and_has_no_embeddings() -> None: - chat = _AsyncChat([_response()]) - embeddings = _AsyncEmbeddings() - memories_store = _AsyncStore([]) - turns_store = _AsyncStore([_turn(i) for i in range(50)]) - service = AsyncPipelineService( - memories_store, - chat, - embeddings, - containers=_async_containers_for_store(memories_store, turns_store=turns_store), - ) - - output = await service.extract_memories_dry("u1", "t1") - - assert set(output) == {"facts", "episodic", "updates", "processed_turn_docs"} - assert len(json.dumps(output)) < 32 * 1024 - assert all("embedding" not in doc for docs in (output["facts"], output["episodic"]) for doc in docs) - assert embeddings.calls == [] - - -@pytest.mark.asyncio -async def test_async_extract_memories_dry_is_byte_deterministic_for_same_llm_response() -> None: - store = _AsyncStore([]) - turns_store = _AsyncStore([_turn(1)]) - service = AsyncPipelineService( - store, - _AsyncChat([_response(), _response()]), - _AsyncEmbeddings(), - containers=_async_containers_for_store(store, turns_store=turns_store), - ) - - first = await service.extract_memories_dry("u1", "t1") - second = await service.extract_memories_dry("u1", "t1") - - assert json.dumps(first, sort_keys=True, separators=(",", ":")) == json.dumps( - second, sort_keys=True, separators=(",", ":") - ) - - -def test_dry_returns_processed_turn_docs_for_watermarking() -> None: - """``extract_memories_dry`` must surface the turn docs it processed so - ``persist_extracted_memories`` can stamp ``extracted_at`` on them and - the next extraction call doesn't reprocess them.""" +def test_extract_memories_dry_stage1_searches_user_turn_text_by_default() -> None: chat = _SyncChat([_response()]) memories_store = _Store([]) - turns_store = _Store([_turn(i) for i in range(3)]) + memories_store.search_results = [ + { + "id": "memory-hybrid", + "content": "Existing hybrid memory from search.", + "type": "fact", + "salience": 0.7, + } + ] + turns = [_turn(1), _turn(2)] + turns_store = _Store(turns) service = PipelineService( memories_store, chat, @@ -275,410 +266,213 @@ def test_dry_returns_processed_turn_docs_for_watermarking() -> None: containers=_containers_for_store(memories_store, turns_store=turns_store), ) - output = service.extract_memories_dry("u1", "t1") + service.extract_memories_dry("u1", "t1") - assert "processed_turn_docs" in output - assert {d["id"] for d in output["processed_turn_docs"]} == {"turn-0", "turn-1", "turn-2"} + assert memories_store.search_calls == [ + { + "search_terms": "\n".join(turn["content"] for turn in turns), + "user_id": "u1", + "memory_types": ["fact"], + "top_k": 10, + "include_superseded": False, + } + ] + assert "Existing hybrid memory from search." in json.dumps(chat.messages) -def test_dry_alone_does_not_mark_turns_as_extracted() -> None: - """A dry run is read-only: it must NOT stamp ``extracted_at`` on any - turn (the wet ``extract_memories`` orchestrator handles marking only - after a successful persist).""" - chat = _SyncChat([_response()]) +def test_extract_memories_dry_stage1_falls_back_to_transcript_without_user_turns() -> None: memories_store = _Store([]) - turns_store = _Store([_turn(i) for i in range(3)]) + turns = [ + { + **_turn(1), + "id": "assistant-turn-1", + "role": "assistant", + "content": "Assistant response with no user-role content.", + } + ] service = PipelineService( memories_store, - chat, + _SyncChat([_response()]), _SyncEmbeddings(), - containers=_containers_for_store(memories_store, turns_store=turns_store), + containers=_containers_for_store(memories_store, turns_store=_Store(turns)), ) service.extract_memories_dry("u1", "t1") - for doc in turns_store.docs: - assert "extracted_at" not in doc or doc["extracted_at"] is None + assert memories_store.search_calls == [ + { + "search_terms": service._build_transcript(turns), + "user_id": "u1", + "memory_types": ["fact"], + "top_k": 10, + "include_superseded": False, + } + ] -def test_extract_memories_marks_turns_after_successful_persist() -> None: - """The wet ``extract_memories`` must stamp ``extracted_at`` on each - turn it processed. Without this, the next extraction call re-loads - the same turns and the LLM re-decides UPDATE/CONTRADICT — which is - the runaway-extraction bug this fix is designed to prevent.""" +def test_extract_memories_dry_stage1_legacy_context_does_not_call_search(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("azure.cosmos.agent_memory.thresholds.get_dedup_context_vector_enabled", lambda: False) chat = _SyncChat([_response()]) - memories_store = _Store([]) - turns_store = _Store([_turn(i) for i in range(3)]) - service = PipelineService( - memories_store, - chat, - _SyncEmbeddings(), - containers=_containers_for_store(memories_store, turns_store=turns_store), + memories_store = _Store( + [ + { + "id": "legacy-memory", + "user_id": "u1", + "type": "fact", + "content": "Existing legacy memory from load.", + "salience": 0.6, + } + ] ) - - service.extract_memories("u1", "t1") - - marked_turns = [doc for doc in turns_store.docs if doc.get("extracted_at")] - assert len(marked_turns) == 3 - marked_ids = {doc["id"] for doc in marked_turns} - assert marked_ids == {"turn-0", "turn-1", "turn-2"} - - -def test_second_extract_call_does_not_reprocess_already_extracted_turns() -> None: - """End-to-end watermarking proof: after a first ``extract_memories`` - marks the turns, a second call with no NEW turns must produce zero - work — no second LLM call, no second persist. This is the property - that prevents reversed-supersede / hallucinated-meta-fact bugs.""" - chat = _SyncChat([_response(), _response()]) # second response should never be consumed - memories_store = _Store([]) - turns_store = _Store([_turn(i) for i in range(3)]) service = PipelineService( memories_store, chat, _SyncEmbeddings(), - containers=_containers_for_store(memories_store, turns_store=turns_store), + containers=_containers_for_store(memories_store, turns_store=_Store([_turn(1)])), ) - service.extract_memories("u1", "t1") - calls_after_first = chat.calls + service.extract_memories_dry("u1", "t1") - # Second invocation with no new turns: watermarked turns are filtered - # out by the query and the dry early-returns with empty items. - service.extract_memories("u1", "t1") - assert chat.calls == calls_after_first + assert memories_store.search_calls == [] + assert "Existing legacy memory from load." in json.dumps(chat.messages) @pytest.mark.asyncio -async def test_async_extract_memories_marks_turns_after_successful_persist() -> None: +async def test_async_extract_memories_dry_shape_is_small_and_has_no_embeddings(monkeypatch) -> None: + monkeypatch.setattr( + "azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_context_vector_enabled", lambda: False + ) chat = _AsyncChat([_response()]) + embeddings = _AsyncEmbeddings() memories_store = _AsyncStore([]) - turns_store = _AsyncStore([_turn(i) for i in range(3)]) + turns_store = _AsyncStore([_turn(i) for i in range(50)]) service = AsyncPipelineService( memories_store, chat, - _AsyncEmbeddings(), + embeddings, containers=_async_containers_for_store(memories_store, turns_store=turns_store), ) - await service.extract_memories("u1", "t1") - - marked_turns = [doc for doc in turns_store.docs if doc.get("extracted_at")] - assert len(marked_turns) == 3 - - -# --------------------------------------------------------------------------- -# Grounding-check regression tests -# -# These tests pin down the two known LLM extraction-time failure modes that -# previously corrupted the fact store and required a wheel hotfix: -# -# 1. The LLM synthesizes an ADD by paraphrase-merging 2+ existing facts -# (e.g. "user eats meat" + "user loves steak" → "user loves steak, -# indicating they eat meat") even though the new user turn said nothing -# on the topic. -# 2. The LLM emits a second invented CONTRADICT fact alongside the literal -# user statement (e.g. user says "I love steak"; LLM emits both -# "loves steak" AND "user eats meat" — the second is a phantom -# explicit-negation that polluted the store with claims the user -# didn't make). -# -# The fix is a prompt change that forbids both patterns. Because we can't -# directly test prompt-following at the unit-test level (no real LLM in -# the test loop), we test the structural safety net: -# ``check_extracted_fact_grounding`` logs a WARNING when these patterns -# slip through. The four scenarios below pair a "buggy" LLM response with -# a "clean" one for each pattern, asserting the WARNING fires (or not) -# accordingly. If a future change ever regresses the prompt and the LLM -# starts emitting these patterns, the WARNING in production telemetry -# becomes the visible signal. -# --------------------------------------------------------------------------- - - -def _existing_fact(fid: str, content: str) -> dict[str, Any]: - """Build a minimal existing-fact doc shaped like what - ``_load_existing_memories`` returns from Cosmos.""" - return { - "id": fid, - "user_id": "u1", - "type": "fact", - "content": content, - "content_hash": fid, - "salience": 0.8, - "confidence": 0.9, - "metadata": {"category": "preference"}, - "tags": ["sys:fact"], - } - - -def _moderate_hotels_turn() -> dict[str, Any]: - return { - "id": "turn-new", - "user_id": "u1", - "thread_id": "t1", - "role": "user", - "type": "turn", - "content": "Normally, I prefer moderate hotels.", - "created_at": "2026-06-02T19:00:00+00:00", - } - + output = await service.extract_memories_dry("u1", "t1") -def _steak_seafood_turn() -> dict[str, Any]: - return { - "id": "turn-new", - "user_id": "u1", - "thread_id": "t1", - "role": "user", - "type": "turn", - "content": "Actually, I love steak and seafood.", - "created_at": "2026-06-02T18:00:00+00:00", - } + assert set(output) == {"facts", "episodic", "updates", "processed_turn_docs"} + assert len(json.dumps(output)) < 32 * 1024 + assert all("embedding" not in doc for docs in (output["facts"], output["episodic"]) for doc in docs) + assert embeddings.calls == [] -def test_grounding_check_warns_when_add_synthesizes_from_multiple_existing_facts(caplog) -> None: - """Scenario 1 (buggy): the LLM emits a synthesized ADD whose tokens come - from 2+ existing facts but not from the new user turn. The grounding - check must emit a WARNING naming the offending fact.""" - existing = [ - _existing_fact("fact_meat", "The user eats meat."), - _existing_fact("fact_steak", "The user loves steak and seafood."), - ] - buggy_response = { - "facts": [ - { - "text": "The user normally prefers moderate hotels.", - "action": "ADD", - "category": "preference", - "confidence": 0.9, - "salience": 0.7, - }, - { - # synthesized — tokens come from existing fact_steak + fact_meat, - # not from the new "moderate hotels" turn - "text": "The user loves steak and seafood, indicating they eat meat.", - "action": "ADD", - "category": "preference", - "confidence": 0.9, - "salience": 0.7, - }, - ], - "episodic": [], - } - chat = _SyncChat([buggy_response]) - memories_store = _Store(existing) - turns_store = _Store([_moderate_hotels_turn()]) - service = PipelineService( - memories_store, - chat, - _SyncEmbeddings(), - containers=_containers_for_store(memories_store, turns_store=turns_store), +@pytest.mark.asyncio +async def test_async_extract_memories_dry_is_byte_deterministic_for_same_llm_response(monkeypatch) -> None: + monkeypatch.setattr( + "azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_context_vector_enabled", lambda: False + ) + store = _AsyncStore([]) + turns_store = _AsyncStore([_turn(1)]) + service = AsyncPipelineService( + store, + _AsyncChat([_response(), _response()]), + _AsyncEmbeddings(), + containers=_async_containers_for_store(store, turns_store=turns_store), ) - with caplog.at_level("WARNING", logger="azure.cosmos.agent_memory.pipeline"): - service.extract_memories_dry("u1", "t1") + first = await service.extract_memories_dry("u1", "t1") + second = await service.extract_memories_dry("u1", "t1") - synthesis_warnings = [ - rec - for rec in caplog.records - if "synthesized from" in rec.getMessage() and "steak and seafood, indicating they eat meat" in rec.getMessage() - ] - assert synthesis_warnings, ( - f"expected a WARNING flagging the synthesized fact; got: {[rec.getMessage() for rec in caplog.records]}" - ) - # The grounded "moderate hotels" fact must NOT trigger a warning. - assert not any( - "moderate hotels" in rec.getMessage() and "synthesized from" in rec.getMessage() for rec in caplog.records + assert json.dumps(first, sort_keys=True, separators=(",", ":")) == json.dumps( + second, sort_keys=True, separators=(",", ":") ) -def test_grounding_check_silent_when_add_is_grounded_in_user_turn(caplog) -> None: - """Scenario 2 (clean): with the same existing-facts context, a single - grounded ADD (only the new "moderate hotels" claim) must NOT trigger - any synthesis WARNING. This is the post-fix expected behaviour.""" - existing = [ - _existing_fact("fact_meat", "The user eats meat."), - _existing_fact("fact_steak", "The user loves steak and seafood."), - ] - clean_response = { - "facts": [ +@pytest.mark.asyncio +async def test_async_extract_memories_dry_stage1_searches_user_turn_text_by_default() -> None: + store = _AsyncStore([]) + store.search = AsyncMock( + return_value=[ { - "text": "The user normally prefers moderate hotels.", - "action": "ADD", - "category": "preference", - "confidence": 0.9, + "id": "fact-1", + "content": "The user prefers dark mode.", + "type": "fact", "salience": 0.7, } - ], - "episodic": [], - } - chat = _SyncChat([clean_response]) - memories_store = _Store(existing) - turns_store = _Store([_moderate_hotels_turn()]) - service = PipelineService( - memories_store, - chat, - _SyncEmbeddings(), - containers=_containers_for_store(memories_store, turns_store=turns_store), + ] ) - - with caplog.at_level("WARNING", logger="azure.cosmos.agent_memory.pipeline"): - service.extract_memories_dry("u1", "t1") - - grounding_warnings = [ - rec - for rec in caplog.records - if "synthesized from" in rec.getMessage() or "phantom-negation/restatement" in rec.getMessage() - ] - assert grounding_warnings == [], ( - f"clean output must not emit grounding warnings; got: {[rec.getMessage() for rec in grounding_warnings]}" + turns_store = _AsyncStore([_turn(1)]) + service = AsyncPipelineService( + store, + _AsyncChat([_response()]), + _AsyncEmbeddings(), + containers=_async_containers_for_store(store, turns_store=turns_store), ) + service._vector_candidates = AsyncMock(return_value=[]) + turns = [_turn(1)] + await service.extract_memories_dry("u1", "t1") -def test_grounding_check_warns_on_phantom_explicit_negation_fact(caplog) -> None: - """Scenario 3 (buggy): user said "Actually, I love steak and seafood"; - LLM emits both a literal-paraphrase fact AND a phantom "user eats meat" - fact (an invented explicit-negation of the prior vegetarian fact). - The phantom fact's tokens are NOT in the user turn and they overlap - a single existing fact ("does not eat meat") — the single-contributor - branch of the grounding heuristic must fire a WARNING.""" - existing = [_existing_fact("fact_veg", "The user does not eat meat.")] - buggy_response = { - "facts": [ - { - # legitimate literal paraphrase of the user turn - "text": "The user loves steak and seafood.", - "action": "CONTRADICT", - "supersedes_id": "fact_veg", - "category": "preference", - "confidence": 0.95, - "salience": 0.8, - }, - { - # phantom — user never said this; tokens come from existing fact_veg - "text": "The user eats meat.", - "action": "CONTRADICT", - "supersedes_id": "fact_veg", - "category": "preference", - "confidence": 0.95, - "salience": 0.8, - }, - ], - "episodic": [], - } - chat = _SyncChat([buggy_response]) - memories_store = _Store(existing) - turns_store = _Store([_steak_seafood_turn()]) - service = PipelineService( - memories_store, - chat, - _SyncEmbeddings(), - containers=_containers_for_store(memories_store, turns_store=turns_store), + store.search.assert_awaited_once_with( + search_terms="\n".join(turn["content"] for turn in turns), + user_id="u1", + memory_types=["fact"], + top_k=10, ) + service._vector_candidates.assert_not_awaited() - with caplog.at_level("WARNING", logger="azure.cosmos.agent_memory.pipeline"): - service.extract_memories_dry("u1", "t1") - phantom_warnings = [ - rec for rec in caplog.records if "phantom-negation" in rec.getMessage() and "eats meat" in rec.getMessage() +@pytest.mark.asyncio +async def test_async_extract_memories_dry_stage1_falls_back_to_transcript_without_user_turns() -> None: + store = _AsyncStore([]) + store.search = AsyncMock(return_value=[]) + turns = [ + { + **_turn(1), + "id": "assistant-turn-1", + "role": "assistant", + "content": "Assistant response with no user-role content.", + } ] - assert phantom_warnings, ( - f"expected a WARNING flagging the phantom-negation fact; got: {[rec.getMessage() for rec in caplog.records]}" - ) - # The legitimate "loves steak and seafood" fact must NOT trigger a warning; - # its tokens are grounded in the user turn. - assert not any( - "loves steak and seafood" in rec.getMessage() - and ("phantom-negation" in rec.getMessage() or "synthesized from" in rec.getMessage()) - for rec in caplog.records - ) - - -def test_grounding_check_silent_on_clean_implicit_contradict(caplog) -> None: - """Scenario 4 (clean): the post-fix expected behaviour for an implicit - contradiction — ONE fact with literal user text and a CONTRADICT - supersedes_id. No phantom-negation fact, no WARNING.""" - existing = [_existing_fact("fact_veg", "The user does not eat meat.")] - clean_response = { - "facts": [ - { - "text": "The user loves steak and seafood.", - "action": "CONTRADICT", - "supersedes_id": "fact_veg", - "category": "preference", - "confidence": 0.95, - "salience": 0.8, - } - ], - "episodic": [], - } - chat = _SyncChat([clean_response]) - memories_store = _Store(existing) - turns_store = _Store([_steak_seafood_turn()]) - service = PipelineService( - memories_store, - chat, - _SyncEmbeddings(), - containers=_containers_for_store(memories_store, turns_store=turns_store), + turns_store = _AsyncStore(turns) + service = AsyncPipelineService( + store, + _AsyncChat([_response()]), + _AsyncEmbeddings(), + containers=_async_containers_for_store(store, turns_store=turns_store), ) + transcript = service._build_transcript(turns) - with caplog.at_level("WARNING", logger="azure.cosmos.agent_memory.pipeline"): - service.extract_memories_dry("u1", "t1") + await service.extract_memories_dry("u1", "t1") - grounding_warnings = [ - rec - for rec in caplog.records - if "synthesized from" in rec.getMessage() or "phantom-negation/restatement" in rec.getMessage() - ] - assert grounding_warnings == [], ( - "clean implicit-contradict must not emit grounding warnings; got: " - f"{[rec.getMessage() for rec in grounding_warnings]}" + store.search.assert_awaited_once_with( + search_terms=transcript, + user_id="u1", + memory_types=["fact"], + top_k=10, ) @pytest.mark.asyncio -async def test_async_grounding_check_warns_on_synthesis(caplog) -> None: - """Async-path mirror of scenario 1: confirms the grounding heuristic - is wired into both sync and async extract pipelines.""" - existing = [ - _existing_fact("fact_meat", "The user eats meat."), - _existing_fact("fact_steak", "The user loves steak and seafood."), - ] - buggy_response = { - "facts": [ - { - "text": "The user normally prefers moderate hotels.", - "action": "ADD", - "category": "preference", - "confidence": 0.9, - "salience": 0.7, - }, +async def test_async_extract_memories_dry_stage1_legacy_path_when_context_vector_unset(monkeypatch) -> None: + monkeypatch.setattr( + "azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_context_vector_enabled", lambda: False + ) + store = _AsyncStore( + [ { - "text": "The user loves steak and seafood, indicating they eat meat.", - "action": "ADD", - "category": "preference", - "confidence": 0.9, - "salience": 0.7, - }, - ], - "episodic": [], - } - chat = _AsyncChat([buggy_response]) - memories_store = _AsyncStore(existing) - turns_store = _AsyncStore([_moderate_hotels_turn()]) + "id": "fact-1", + "user_id": "u1", + "type": "fact", + "content": "Existing fact.", + "content_hash": "hash-1", + } + ] + ) + store.search = AsyncMock(return_value=[]) + turns_store = _AsyncStore([_turn(1)]) service = AsyncPipelineService( - memories_store, - chat, + store, + _AsyncChat([_response()]), _AsyncEmbeddings(), - containers=_async_containers_for_store(memories_store, turns_store=turns_store), + containers=_async_containers_for_store(store, turns_store=turns_store), ) - with caplog.at_level("WARNING", logger="azure.cosmos.agent_memory.pipeline.aio"): - await service.extract_memories_dry("u1", "t1") + await service.extract_memories_dry("u1", "t1") - synthesis_warnings = [ - rec - for rec in caplog.records - if "synthesized from" in rec.getMessage() and "steak and seafood, indicating they eat meat" in rec.getMessage() - ] - assert synthesis_warnings, ( - f"expected an async WARNING flagging the synthesized fact; got: {[rec.getMessage() for rec in caplog.records]}" - ) + store.search.assert_not_awaited() diff --git a/tests/unit/services/test_pipeline_service.py b/tests/unit/services/test_pipeline_service.py index 587f680..27518b5 100644 --- a/tests/unit/services/test_pipeline_service.py +++ b/tests/unit/services/test_pipeline_service.py @@ -10,6 +10,22 @@ from azure.cosmos.agent_memory.services.pipeline import PipelineService, _StoreContainerAdapter +@pytest.fixture(autouse=True) +def _pin_legacy_dedup_paths(monkeypatch): + monkeypatch.setattr( + "azure.cosmos.agent_memory.thresholds.get_dedup_reconcile_mode", + lambda: "full_pool", + ) + monkeypatch.setattr( + "azure.cosmos.agent_memory.thresholds.get_dedup_vector_enabled", + lambda: False, + ) + monkeypatch.setattr( + "azure.cosmos.agent_memory.thresholds.get_dedup_context_vector_enabled", + lambda: False, + ) + + class FakeLLMService: """Test helper exposing chat_client + embeddings_client pair. @@ -203,7 +219,6 @@ def test_extract_memories_happy_path_writes_fact_and_episodic() -> None: { "scope_type": "project", "scope_value": "CI", - "text": "Stabilized flaky CI tests by adding retries.", "situation": "CI tests flaked intermittently", "action_taken": "Added retries", "outcome": "Tests stabilized", @@ -226,12 +241,14 @@ def test_extract_memories_happy_path_writes_fact_and_episodic() -> None: assert llm.embed_calls == [ [ "The user prefers dark mode.", - "Stabilized flaky CI tests by adding retries.", + "CI tests flaked intermittently → Added retries → Tests stabilized", ] ] -def test_extract_memories_contradict_supersedes_existing_fact() -> None: +def test_extract_memories_creates_new_fact_without_superseding() -> None: + # Extract no longer detects updates/contradictions — it just adds new facts. + # Contradiction resolution happens later in reconcile, not at extract time. old = _fact("old_fact", "The user prefers light mode.") store = FakeStore([old]) turns_store = FakeStore([_turn("Actually, I prefer dark mode now.")]) @@ -241,8 +258,6 @@ def test_extract_memories_contradict_supersedes_existing_fact() -> None: "facts": [ { "text": "The user prefers dark mode.", - "action": "CONTRADICT", - "supersedes_id": "old_fact", "category": "preference", } ] @@ -253,12 +268,10 @@ def test_extract_memories_contradict_supersedes_existing_fact() -> None: result = _pipeline(store, llm, turns_store=turns_store).extract_memories("u1", "t1") assert result["fact_count"] == 1 - assert result["contradicted_count"] == 1 - assert store.supersede_calls[0][2] == "contradict" + assert result["contradicted_count"] == 0 + assert store.supersede_calls == [] old_doc = next(doc for doc in store.docs if doc["id"] == "old_fact") - assert old_doc["superseded_by"] == store.upserts[0]["id"] - assert old_doc["supersede_reason"] == "contradict" - assert old_doc["superseded_at"] + assert "superseded_by" not in old_doc def test_synthesize_procedural_produces_procedural_memory() -> None: diff --git a/tests/unit/store/test_memory_store.py b/tests/unit/store/test_memory_store.py index 45a5117..e4371e1 100644 --- a/tests/unit/store/test_memory_store.py +++ b/tests/unit/store/test_memory_store.py @@ -273,6 +273,63 @@ def test_search_adds_created_time_range_filters(): assert params["@created_before"] == "2026-03-01T00:00:00+00:00" +def test_search_uses_keyword_params_for_hybrid_sql(): + memories = MagicMock() + memories.query_items.return_value = [] + embeddings = MagicMock() + embeddings.generate.return_value = [0.1, 0.2] + store = MemoryStore(containers=_containers(memories=memories), embeddings_client=embeddings) + + store.search("weather in Seattle", user_id="u1") + + call_kwargs = memories.query_items.call_args.kwargs + assert "ORDER BY RANK RRF" in call_kwargs["query"] + assert "FullTextScore(c.content, @kw0, @kw1)" in call_kwargs["query"] + params = _params_by_name(call_kwargs) + assert params["@kw0"] == "weather" + assert params["@kw1"] == "seattle" + + +def test_search_all_stopwords_falls_back_to_vector_only(): + memories = MagicMock() + memories.query_items.return_value = [] + embeddings = MagicMock() + embeddings.generate.return_value = [0.1, 0.2] + store = MemoryStore(containers=_containers(memories=memories), embeddings_client=embeddings) + + store.search("what is the", user_id="u1") + + call_kwargs = memories.query_items.call_args.kwargs + assert "ORDER BY VectorDistance" in call_kwargs["query"] + assert "RANK RRF" not in call_kwargs["query"] + assert not any(name.startswith("@kw") for name in _params_by_name(call_kwargs)) + + +def test_search_episodic_forwards_search_options(): + store = MemoryStore(containers=_containers()) + store.search = MagicMock(return_value=[]) + + store.search_episodic("u1", "weather") + + store.search.assert_called_once_with( + search_terms="weather", + user_id="u1", + memory_types=["episodic"], + top_k=5, + min_salience=None, + include_superseded=False, + ) + + +def test_build_episodic_context_forwards_search_options(): + store = MemoryStore(containers=_containers()) + store.search_episodic = MagicMock(return_value=[]) + + assert store.build_episodic_context("u1", "weather") == "" + + store.search_episodic.assert_called_once_with("u1", "weather", top_k=3) + + def test_add_cosmos_routes_by_type(): turns = MagicMock() memories = MagicMock() diff --git a/tests/unit/test_auto_trigger.py b/tests/unit/test_auto_trigger.py index 03c5420..4bb73e7 100644 --- a/tests/unit/test_auto_trigger.py +++ b/tests/unit/test_auto_trigger.py @@ -10,10 +10,52 @@ from unittest.mock import MagicMock, patch +import pytest +from azure.cosmos.exceptions import CosmosResourceNotFoundError + +from azure.cosmos.agent_memory import _counters from azure.cosmos.agent_memory.cosmos_memory_client import CosmosMemoryClient from azure.cosmos.agent_memory.processors import DurableFunctionProcessor, InProcessProcessor +class _FakeCounterContainer: + """Minimal in-memory stand-in for the Cosmos counter container. + + Exercises the REAL increment / watermark-read / watermark-advance helpers + end-to-end (no mocking of counter math), so a watermark/recent_k regression + can't slip through behind constant mocks. + """ + + def __init__(self) -> None: + self.store: dict[str, dict] = {} + self._etag = 0 + + def read_item(self, *, item, partition_key): + if item not in self.store: + raise CosmosResourceNotFoundError(message="404") + return dict(self.store[item]) + + def create_item(self, *, body): + self._etag += 1 + body = dict(body) + body["_etag"] = f"e{self._etag}" + self.store[body["id"]] = body + return dict(body) + + def upsert_item(self, *, body, **_kwargs): + self._etag += 1 + body = dict(body) + body["_etag"] = f"e{self._etag}" + self.store[body["id"]] = body + return dict(body) + + def patch_item(self, *, item, partition_key, patch_operations): + doc = self.store.setdefault(item, {"id": item}) + for op in patch_operations: + doc[op["path"].lstrip("/")] = op["value"] + return dict(doc) + + def _connected(processor=None) -> CosmosMemoryClient: client = CosmosMemoryClient(use_default_credential=False, processor=processor) client._memories_container_client = MagicMock() @@ -43,7 +85,7 @@ def test_push_to_cosmos_fires_inprocess_trigger_per_turn(monkeypatch): client.add_local(user_id="u1", role="user", thread_id="t1", content="hi") client.push_to_cosmos() - pipeline.extract_memories.assert_called_once_with("u1", "t1") + pipeline.extract_memories.assert_called_once_with("u1", "t1", recent_k=1) def test_push_to_cosmos_durable_does_not_fire_trigger(monkeypatch): @@ -143,10 +185,178 @@ def test_extract_fires_independently_of_summary(self, monkeypatch): client.add_local(user_id="u1", role="user", thread_id="t1", content="hi") client.push_to_cosmos() - processor.process_extract_memories.assert_called_once_with(user_id="u1", thread_id="t1") + processor.process_extract_memories.assert_called_once_with(user_id="u1", thread_id="t1", recent_k=1) processor.process_thread_summary.assert_not_called() processor.process_user_summary.assert_not_called() + @pytest.mark.parametrize( + ("n_facts", "batch_count", "counter_result", "watermark", "expected_recent_k"), + [ + (1, 1, (0, 1), None, 1), + (1, 3, (0, 3), None, 3), + (5, 1, (4, 5), None, 5), + (1, 1, (5, 10), 5, 5), + # Large backlog is NOT capped: recent_k spans every turn since the + # watermark (newest-recent_k slice covers exactly those), so the + # watermark can advance to new_count with no stranded turns. + (1, 1, (98, 100), 0, 100), + # BOOTSTRAP regression: no watermark yet but the counter is already + # ahead of this batch (earlier extracts failed). base=0 so recent_k = + # new_count (30) covers ALL turns — the old fallback max(n_facts, + # batch_count) would return 2 and strand turns 1-28 forever. + (1, 2, (20, 30), None, 30), + ], + ) + def test_extract_recent_k_uses_watermark_then_falls_back( + self, + monkeypatch, + n_facts, + batch_count, + counter_result, + watermark, + expected_recent_k, + ): + monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", str(n_facts)) + monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") + monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") + + processor = InProcessProcessor(pipeline=MagicMock()) + processor.process_extract_memories = MagicMock(return_value={}) + + client = _connected(processor=processor) + client._counter_container_client = MagicMock() + + with patch( + "azure.cosmos.agent_memory._counters.increment_counter_sync", + return_value=counter_result, + ), patch( + "azure.cosmos.agent_memory._counters.read_extract_watermark_sync", + return_value=watermark, + ), patch( + "azure.cosmos.agent_memory._counters.advance_extract_watermark_sync", + ) as advance: + for i in range(batch_count): + client.add_local(user_id="u1", role="user", thread_id="t1", content=f"hi {i}") + client.push_to_cosmos() + + processor.process_extract_memories.assert_called_once_with( + user_id="u1", + thread_id="t1", + recent_k=expected_recent_k, + ) + advance.assert_called_once() + + def test_watermark_round_trip_fail_then_succeed_no_strand(self, monkeypatch): + """End-to-end round-trip against a REAL in-memory counter (no constant + mocks): a thread's first extract fails, the second succeeds, and the + second must cover EVERY turn so far — not just its own batch — so turns + from the failed batch are never stranded. This is the bootstrap case the + constant-mock tests could not catch.""" + monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") + monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") + monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") + + counter = _FakeCounterContainer() + recorded_recent_k: list[int] = [] + + def extract(*, user_id, thread_id, recent_k): + recorded_recent_k.append(recent_k) + if len(recorded_recent_k) == 1: + raise RuntimeError("transient LLM outage") # first extract fails + return {} + + processor = InProcessProcessor(pipeline=MagicMock()) + processor.process_extract_memories = MagicMock(side_effect=extract) + client = _connected(processor=processor) + client._counter_container_client = counter + + # Batch 1: 10 turns -> counter 0->10, extract FAILS (watermark not advanced). + for i in range(10): + client.add_local(user_id="u1", role="user", thread_id="t1", content=f"a{i}") + client.push_to_cosmos() + + # Batch 2: 10 turns -> counter 10->20, extract SUCCEEDS. + for i in range(10): + client.add_local(user_id="u1", role="user", thread_id="t1", content=f"b{i}") + client.push_to_cosmos() + + # First fired with 10 (all turns so far); second with 20 (ALL turns, since + # the failed first extract left the watermark unset) — NOT 10. + assert recorded_recent_k == [10, 20] + # Watermark now seeded at the full count after the successful extract. + cid = _counters.thread_counter_id("u1", "t1") + assert _counters.read_extract_watermark_sync(counter, cid, "u1", "t1") == 20 + + def test_watermark_not_advanced_when_extract_fails(self, monkeypatch): + """advance-on-success: a failing extract must NOT move the watermark, so + the skipped turns are retried next sweep; failure is stamped instead.""" + monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") + monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") + monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") + + processor = InProcessProcessor(pipeline=MagicMock()) + processor.process_extract_memories = MagicMock(side_effect=RuntimeError("llm down")) + + client = _connected(processor=processor) + client._counter_container_client = MagicMock() + + with patch( + "azure.cosmos.agent_memory._counters.increment_counter_sync", + return_value=(0, 1), + ), patch( + "azure.cosmos.agent_memory._counters.read_extract_watermark_sync", + return_value=None, + ), patch( + "azure.cosmos.agent_memory._counters.advance_extract_watermark_sync", + ) as advance, patch( + "azure.cosmos.agent_memory._counters.stamp_failure_sync", + ) as stamp: + client.add_local(user_id="u1", role="user", thread_id="t1", content="hi") + client.push_to_cosmos() + + advance.assert_not_called() + stamp.assert_called_once() + + def test_reconcile_full_rebuild_on_persisted_counter_cadence(self, monkeypatch): + """Symmetry with the durable backend: the in-process auto-trigger requests + a full-pool reconcile (full_rebuild=True) on a PERSISTED-counter cadence — + every DEDUP_FULL_RECLUSTER_EVERY_N-th reconcile — not via an in-memory + per-instance sweep counter. Here that's every 2 turns.""" + monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") + monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") + monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") + monkeypatch.setenv("DEDUP_EVERY_N", "1") + monkeypatch.setattr( + "azure.cosmos.agent_memory.thresholds.get_dedup_full_recluster_every_n", lambda: 2 + ) + + rebuilds: list[bool] = [] + processor = InProcessProcessor(pipeline=MagicMock()) + processor.process_extract_memories = MagicMock(return_value={}) + processor.synthesize_procedural = MagicMock(return_value=None) + processor.process_reconcile = MagicMock( + side_effect=lambda *, user_id, full_rebuild=False: rebuilds.append(full_rebuild) + ) + + client = _connected(processor=processor) + client._counter_container_client = MagicMock() + + with patch( + "azure.cosmos.agent_memory._counters.increment_counter_sync", + side_effect=[(0, 1), (1, 2)], + ), patch( + "azure.cosmos.agent_memory._counters.read_extract_watermark_sync", + return_value=None, + ), patch( + "azure.cosmos.agent_memory._counters.advance_extract_watermark_sync", + ): + client.add_local(user_id="u1", role="user", thread_id="t1", content="a") + client.push_to_cosmos() # counter 0->1: reconcile, full crosses 2? no + client.add_local(user_id="u1", role="user", thread_id="t1", content="b") + client.push_to_cosmos() # counter 1->2: full backstop threshold (2) crossed + + assert rebuilds == [False, True] + def test_summary_fires_independently_when_threshold_crossed(self, monkeypatch): """N_summary=10 boundary fires summary; N_facts=0 prevents extract.""" monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "0") diff --git a/tests/unit/test_cosmos_memory_client.py b/tests/unit/test_cosmos_memory_client.py index c04b000..bae9f12 100644 --- a/tests/unit/test_cosmos_memory_client.py +++ b/tests/unit/test_cosmos_memory_client.py @@ -453,31 +453,31 @@ def test_create_memory_store_defaults_to_serverless(self): mock_lease_container, ] - mem = _make_client(cosmos_throughput_mode="serverless") - - with patch.dict("os.environ", {"COSMOS_DB_AUTOSCALE_MAX_RU": "not-an-int"}, clear=False): - with patch.dict( - "sys.modules", - { - "azure.cosmos": MagicMock( - CosmosClient=mock_cosmos_cls, - PartitionKey=MagicMock(), - ThroughputProperties=MagicMock(), - ), - }, - ): - mem.create_memory_store( - endpoint="https://fake.documents.azure.com:443/", - credential="fake-key", - throughput_mode="serverless", - ) + # serverless mode ignores autoscale config entirely, even an invalid value. + mem = _make_client(cosmos_throughput_mode="serverless", cosmos_autoscale_max_ru="not-an-int") + + with patch.dict( + "sys.modules", + { + "azure.cosmos": MagicMock( + CosmosClient=mock_cosmos_cls, + PartitionKey=MagicMock(), + ThroughputProperties=MagicMock(), + ), + }, + ): + mem.create_memory_store( + endpoint="https://fake.documents.azure.com:443/", + credential="fake-key", + throughput_mode="serverless", + ) for call in mock_db.create_container_if_not_exists.call_args_list: assert "offer_throughput" not in call.kwargs - def test_constructor_ignores_invalid_autoscale_env_in_serverless_mode(self): - with patch.dict("os.environ", {"COSMOS_DB_AUTOSCALE_MAX_RU": "not-an-int"}, clear=False): - mem = _make_client(cosmos_throughput_mode="serverless") + def test_constructor_ignores_autoscale_in_serverless_mode(self): + # Even an invalid autoscale value is ignored in serverless mode. + mem = _make_client(cosmos_throughput_mode="serverless", cosmos_autoscale_max_ru="not-an-int") assert mem._cosmos_autoscale_max_ru is None @@ -797,7 +797,7 @@ def test_search_cosmos(self): assert "VectorDistance" in call_kwargs["query"] assert len(result) == 1 - def test_search_hybrid(self): + def test_search_hybrid_uses_keyword_params(self): mem, container = _connected_client() container.query_items.return_value = [_make_doc()] @@ -806,14 +806,79 @@ def test_search_hybrid(self): mem.search_cosmos( search_terms="weather Seattle", - hybrid_search=True, top_k=5, ) call_kwargs = container.query_items.call_args.kwargs query = call_kwargs["query"] + params = {p["name"]: p["value"] for p in call_kwargs["parameters"]} assert "RANK RRF" in query - assert "FullTextScore" in query + assert "FullTextScore(c.content, @kw0, @kw1)" in query + assert params["@kw0"] == "weather" + assert params["@kw1"] == "seattle" + + def test_search_cosmos_forwards_search_options_to_store(self): + mem, _ = _connected_client() + store = MagicMock() + store._containers = mem._containers + store._embeddings_client = mem._embeddings_client + store.search.return_value = [] + mem._store = store + + mem.search_cosmos(search_terms="weather") + + store.search.assert_called_once_with( + search_terms="weather", + memory_id=None, + user_id=None, + role=None, + memory_types=None, + thread_id=None, + top_k=5, + tags_all=None, + tags_any=None, + exclude_tags=None, + include_superseded=False, + min_salience=None, + min_confidence=None, + created_after=None, + created_before=None, + ) + + def test_search_episodic_memories_forwards_search_options(self): + mem, _ = _connected_client() + store = MagicMock() + store._containers = mem._containers + store._embeddings_client = mem._embeddings_client + store.search_episodic.return_value = [] + mem._store = store + + mem.search_episodic_memories("u1", "weather") + + store.search_episodic.assert_called_once_with( + user_id="u1", + search_terms="weather", + top_k=5, + min_salience=None, + include_superseded=False, + ) + + def test_build_episodic_context_forwards_search_options(self): + mem, _ = _connected_client() + store = MagicMock() + store._containers = mem._containers + store._embeddings_client = mem._embeddings_client + store.build_episodic_context.return_value = "context" + mem._store = store + + result = mem.build_episodic_context("u1", "weather") + + assert result == "context" + store.build_episodic_context.assert_called_once_with( + user_id="u1", + query="weather", + top_k=3, + ) def test_search_turns(self): mem, container = _connected_client() diff --git a/tests/unit/test_pipeline_confidence.py b/tests/unit/test_pipeline_confidence.py index fce8308..3dcf20e 100644 --- a/tests/unit/test_pipeline_confidence.py +++ b/tests/unit/test_pipeline_confidence.py @@ -13,6 +13,18 @@ from azure.cosmos.agent_memory.store import MemoryStore +@pytest.fixture(autouse=True) +def _pin_legacy_extract_dedup(monkeypatch): + monkeypatch.setattr( + "azure.cosmos.agent_memory.thresholds.get_dedup_vector_enabled", + lambda: False, + ) + monkeypatch.setattr( + "azure.cosmos.agent_memory.thresholds.get_dedup_context_vector_enabled", + lambda: False, + ) + + def _make_pipeline(llm_response: dict): turns_container = MagicMock() memories_container = MagicMock() @@ -139,7 +151,6 @@ def test_extract_episodic_carries_confidence(): { "scope_type": "project", "scope_value": "CI revamp", - "text": "Set up CI by adding Ruff — faster lint times.", "situation": "Setup CI", "action_taken": "Added Ruff", "outcome": "Faster lint", @@ -286,17 +297,11 @@ def test_thread_ids_does_not_appear_in_query_or_parameters(self): # --------------------------------------------------------------------------- -def test_extract_drops_episodic_missing_text(caplog): - """An episodic with no ``text`` is dropped and surfaced via the return value. +def test_extract_scoped_intent_without_outcome_stores_correctly(caplog): + """An episodic with only scope fields (no situation/action/outcome) is kept. - Previously the pipeline synthesized boilerplate content like - ``"For the user's Paris trip, intent recorded."`` which was - semantically empty for embedding/recall. The fix is to require - the LLM to emit ``text`` (same field facts use) — if it doesn't, - drop the record so we don't poison the recall index. The drop is - logged at ERROR (it's data loss) and surfaced via the - ``dropped_episodic_count`` field on the return dict so callers - can monitor LLM-extraction compliance over time. + The doc must use the deterministic fallback content string, expose the + scope fields at the top level, and not emit a "dropping malformed" warning. """ pipeline, upserted = _make_pipeline( { @@ -306,34 +311,36 @@ def test_extract_drops_episodic_missing_text(caplog): "scope_value": "Paris", "confidence": 0.95, "salience": 0.8, - "tags": ["topic:travel", "topic:hotels"], } ] } ) - with caplog.at_level("ERROR", logger="azure.cosmos.agent_memory.pipeline"): - result = pipeline.extract_memories("u1", "t1") + with caplog.at_level("WARNING", logger="azure.cosmos.agent_memory.pipeline"): + pipeline.extract_memories("u1", "t1") eps = [d for d in upserted if d["type"] == "episodic"] - assert eps == [] - assert result["episodic_count"] == 0 - assert result["dropped_episodic_count"] == 1 - msgs = [rec.getMessage() for rec in caplog.records] - assert any("empty/missing text field" in m for m in msgs) - assert any("reason=missing_text" in m for m in msgs) - # Bumped from WARNING → ERROR because dropping == data loss. - assert any(rec.levelname == "ERROR" and "empty/missing text field" in rec.getMessage() for rec in caplog.records) + assert len(eps) == 1 + ep = eps[0] + assert ep["scope_type"] == "trip" + assert ep["scope_value"] == "Paris" + assert ep["metadata"]["scope_type"] == "trip" + assert ep["metadata"]["scope_value"] == "Paris" + assert ep["metadata"]["situation"] is None + assert ep["metadata"]["action_taken"] is None + assert ep["metadata"]["outcome"] is None + assert ep["content"] == "For the user's Paris trip, intent recorded." + assert ep["confidence"] == pytest.approx(0.95) + assert not any("dropping malformed episodic" in rec.getMessage() for rec in caplog.records) -def test_extract_past_event_episodic_uses_text_and_keeps_chain_in_metadata(): +def test_extract_past_event_episodic_uses_arrow_form_and_keeps_scope(): pipeline, upserted = _make_pipeline( { "episodic": [ { "scope_type": "project", "scope_value": "Acme revamp", - "text": "Migrated Acme DB by running the script — all rows migrated cleanly.", "situation": "Migrated DB", "action_taken": "Ran the script", "outcome": "All rows migrated", @@ -352,7 +359,7 @@ def test_extract_past_event_episodic_uses_text_and_keeps_chain_in_metadata(): pipeline.extract_memories("u1", "t1") [ep] = [d for d in upserted if d["type"] == "episodic"] - assert ep["content"] == "Migrated Acme DB by running the script — all rows migrated cleanly." + assert ep["content"] == "Migrated DB → Ran the script → All rows migrated" assert ep["scope_type"] == "project" assert ep["scope_value"] == "Acme revamp" md = ep["metadata"] @@ -366,13 +373,12 @@ def test_extract_past_event_episodic_uses_text_and_keeps_chain_in_metadata(): assert "topic:db" in ep["tags"] -def test_extract_episodic_uses_text_directly_no_synthesis(): - """The LLM-written ``text`` is the embedded ``content``, verbatim. +def test_extract_episodic_falls_back_to_arrow_form_when_summary_field_present(): + """The schema dropped ``summary``; pipeline now always uses arrow form. - The pipeline must NOT synthesize content from the s/a/o chain or - from scope fields — that's how we ended up with useless boilerplate - before the fix. Whatever the LLM emits in ``text`` is what gets - embedded. + Even if a non-strict LLM smuggles a ``summary`` field through, the + pipeline ignores it and builds content from + ``situation → action_taken → outcome``. """ pipeline, upserted = _make_pipeline( { @@ -380,7 +386,7 @@ def test_extract_episodic_uses_text_directly_no_synthesis(): { "scope_type": "trip", "scope_value": "Paris", - "text": "User wants luxury hotels for the Paris trip.", + "summary": "User wants luxury hotels for the Paris trip.", "situation": "Planning Paris trip", "action_taken": "Said luxury", "outcome": "Pending", @@ -392,106 +398,7 @@ def test_extract_episodic_uses_text_directly_no_synthesis(): pipeline.extract_memories("u1", "t1") [ep] = [d for d in upserted if d["type"] == "episodic"] - assert ep["content"] == "User wants luxury hotels for the Paris trip." - assert ep["metadata"]["situation"] == "Planning Paris trip" - assert ep["metadata"]["action_taken"] == "Said luxury" - assert ep["metadata"]["outcome"] == "Pending" - - -def test_extract_episodic_uses_text_alone_for_planned_intent(): - """Planned/in-flight episodics carry their meaning entirely in ``text``. - - This is the headline bug-1 scenario from the workshop: the LLM (correctly - following the prompt) emits only scope_type/scope_value/text for a - planned trip, and the pipeline must embed the text, not boilerplate. - """ - pipeline, upserted = _make_pipeline( - { - "episodic": [ - { - "scope_type": "trip", - "scope_value": "Tokyo", - "text": ("Planning a Tokyo trip with vegetarian and wheelchair-accessible-restaurant constraints."), - "confidence": 0.95, - "salience": 0.85, - "tags": ["topic:travel", "topic:accessibility"], - } - ] - } - ) - - pipeline.extract_memories("u1", "t1") - - [ep] = [d for d in upserted if d["type"] == "episodic"] - assert ep["content"] == ("Planning a Tokyo trip with vegetarian and wheelchair-accessible-restaurant constraints.") - assert ep["metadata"]["situation"] is None - assert ep["metadata"]["action_taken"] is None - assert ep["metadata"]["outcome"] is None - - -def test_extract_episodic_strips_whitespace_from_text(): - pipeline, upserted = _make_pipeline( - { - "episodic": [ - { - "scope_type": "trip", - "scope_value": "Paris", - "text": " Planning a Paris trip. ", - } - ] - } - ) - pipeline.extract_memories("u1", "t1") - [ep] = [d for d in upserted if d["type"] == "episodic"] - assert ep["content"] == "Planning a Paris trip." - - -def test_extract_compound_statement_yields_facts_across_categories(): - """Bug-2 scenario: a single user turn that combines preference + requirement - must produce two facts, not one merged "restaurant preferences" fact. - - Drives the prompt's tightened consolidation rule. We're mocking the LLM - response here so this is really a regression guard on the pipeline plumbing - (the prompt change is what makes a real LLM produce this shape). - """ - pipeline, upserted = _make_pipeline( - { - "facts": [ - { - "text": "The user does not eat meat.", - "category": "preference", - "subject": "user", - "predicate": "dietary_restriction", - "object": "no meat", - "confidence": 1.0, - "salience": 0.9, - "tags": ["topic:diet"], - "action": "ADD", - "supersedes_id": None, - }, - { - "text": "The user requires wheelchair-accessible restaurants.", - "category": "requirement", - "subject": "user", - "predicate": "accessibility_requirement", - "object": "wheelchair-accessible restaurants", - "confidence": 1.0, - "salience": 0.95, - "tags": ["topic:accessibility"], - "action": "ADD", - "supersedes_id": None, - }, - ] - } - ) - pipeline.extract_memories("u1", "t1") - - facts = [d for d in upserted if d["type"] == "fact"] - assert len(facts) == 2 - by_category = {f["metadata"]["category"]: f for f in facts} - assert set(by_category) == {"preference", "requirement"} - assert by_category["preference"]["content"] == "The user does not eat meat." - assert by_category["requirement"]["content"] == "The user requires wheelchair-accessible restaurants." + assert ep["content"] == "Planning Paris trip → Said luxury → Pending" def test_extract_drops_episodic_missing_scope_type(caplog): @@ -509,13 +416,10 @@ def test_extract_drops_episodic_missing_scope_type(caplog): ) with caplog.at_level("WARNING", logger="azure.cosmos.agent_memory.pipeline"): - result = pipeline.extract_memories("u1", "t1") + pipeline.extract_memories("u1", "t1") assert not any(d["type"] == "episodic" for d in upserted) assert any("dropping malformed episodic" in rec.getMessage() for rec in caplog.records) - assert any("reason=malformed_scope" in rec.getMessage() for rec in caplog.records) - # Malformed-scope drops also count toward the dropped_episodic_count signal. - assert result["dropped_episodic_count"] == 1 def test_extract_drops_episodic_missing_scope_value(caplog): @@ -578,7 +482,6 @@ def test_extract_strips_whitespace_from_scope_fields(): { "scope_type": " trip ", "scope_value": " Paris ", - "text": "Planning a Paris trip.", "confidence": 0.9, } ] @@ -590,4 +493,43 @@ def test_extract_strips_whitespace_from_scope_fields(): [ep] = [d for d in upserted if d["type"] == "episodic"] assert ep["scope_type"] == "trip" assert ep["scope_value"] == "Paris" - assert ep["content"] == "Planning a Paris trip." + assert ep["content"] == "For the user's Paris trip, intent recorded." + + +def test_extract_compound_statement_yields_facts_across_categories(): + """A single user turn that combines preference + requirement must produce two + facts, not one merged fact. Regression guard on the pipeline plumbing.""" + pipeline, upserted = _make_pipeline( + { + "facts": [ + { + "text": "The user does not eat meat.", + "category": "preference", + "subject": "user", + "predicate": "dietary_restriction", + "object": "no meat", + "confidence": 1.0, + "salience": 0.9, + "tags": ["topic:diet"], + }, + { + "text": "The user requires wheelchair-accessible restaurants.", + "category": "requirement", + "subject": "user", + "predicate": "accessibility_requirement", + "object": "wheelchair-accessible restaurants", + "confidence": 1.0, + "salience": 0.95, + "tags": ["topic:accessibility"], + }, + ] + } + ) + pipeline.extract_memories("u1", "t1") + + facts = [d for d in upserted if d["type"] == "fact"] + assert len(facts) == 2 + by_category = {f["metadata"]["category"]: f for f in facts} + assert set(by_category) == {"preference", "requirement"} + assert by_category["preference"]["content"] == "The user does not eat meat." + assert by_category["requirement"]["content"] == "The user requires wheelchair-accessible restaurants." diff --git a/tests/unit/test_procedural_synthesis.py b/tests/unit/test_procedural_synthesis.py index 889e3ca..5930104 100644 --- a/tests/unit/test_procedural_synthesis.py +++ b/tests/unit/test_procedural_synthesis.py @@ -15,6 +15,18 @@ from azure.cosmos.agent_memory.store import MemoryStore +@pytest.fixture(autouse=True) +def _pin_legacy_extract_dedup(monkeypatch): + monkeypatch.setattr( + "azure.cosmos.agent_memory.thresholds.get_dedup_vector_enabled", + lambda: False, + ) + monkeypatch.setattr( + "azure.cosmos.agent_memory.thresholds.get_dedup_context_vector_enabled", + lambda: False, + ) + + def _assert_iso8601(text: str) -> None: assert text datetime.fromisoformat(text) diff --git a/tests/unit/test_process_now.py b/tests/unit/test_process_now.py index d81a0f7..3d71c56 100644 --- a/tests/unit/test_process_now.py +++ b/tests/unit/test_process_now.py @@ -2,7 +2,7 @@ from __future__ import annotations -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, call, patch import pytest @@ -50,7 +50,13 @@ def test_process_now_with_inprocess_invokes_full_pipeline(): assert isinstance(client._processor, InProcessProcessor) pipeline.generate_thread_summary.assert_called_once_with("u1", "t1") pipeline.extract_memories.assert_called_once_with("u1", "t1") - pipeline.reconcile_memories.assert_called_once_with("u1", 50) + assert pipeline.reconcile_memories.call_count == 2 + pipeline.reconcile_memories.assert_has_calls( + [ + call("u1", n=50, memory_type="fact", full_rebuild=False), + call("u1", n=50, memory_type="episodic", full_rebuild=False), + ] + ) pipeline.synthesize_procedural.assert_called_once_with(user_id="u1", force=False) pipeline.generate_user_summary.assert_called_once_with("u1", None) assert result.procedural == {"id": "proc1", "type": "procedural"} diff --git a/tests/unit/test_reconcile.py b/tests/unit/test_reconcile.py index be3f63b..f0a1c32 100644 --- a/tests/unit/test_reconcile.py +++ b/tests/unit/test_reconcile.py @@ -29,6 +29,23 @@ from azure.cosmos.agent_memory.services.pipeline import PipelineService +@pytest.fixture(autouse=True) +def _pin_legacy_dedup_paths(monkeypatch): + """These tests cover the pre-candidate reconcile/extract code paths.""" + monkeypatch.setattr( + "azure.cosmos.agent_memory.thresholds.get_dedup_reconcile_mode", + lambda: "full_pool", + ) + monkeypatch.setattr( + "azure.cosmos.agent_memory.thresholds.get_dedup_vector_enabled", + lambda: False, + ) + monkeypatch.setattr( + "azure.cosmos.agent_memory.thresholds.get_dedup_context_vector_enabled", + lambda: False, + ) + + def _make_pipeline() -> PipelineService: p = PipelineService.__new__(PipelineService) p._embeddings = MagicMock() @@ -649,269 +666,6 @@ def test_fact_not_dropped_when_only_procedural_has_same_hash(self): assert fact_docs[0]["content"] == text -class TestEpisodicReconciliation: - """Episodic memories use scope as identity: the deterministic ID is - seeded only on ``(user_id, scope_type, scope_value)``. Any re-emission - for the same scope (paraphrased intent, added detail, reversed intent) - collides on upsert and replaces the prior record. The LLM does NOT - make ADD/UPDATE/CONTRADICT decisions for episodics — the scope IS the - identity. Distinct events under the same umbrella belong under - distinct ``scope_value`` strings (e.g. "Tokyo trip" vs - "Tokyo lost-wallet incident"). - """ - - def _build(self) -> PipelineService: - p = PipelineService.__new__(PipelineService) - p._embeddings = MagicMock() - p._embeddings.generate.return_value = [0.1] * 8 - p._embeddings.generate_batch.return_value = [[0.1] * 8] - p._container = MagicMock() - p._memories_container = p._container - p._turns_container = p._container - p._summaries_container = p._container - p._chat = MagicMock() - p._upsert_memory = MagicMock(side_effect=lambda doc: doc) - p._create_memory = MagicMock(side_effect=lambda doc: doc) - p._mark_superseded = MagicMock(return_value=True) - return p - - def _turns(self) -> list[dict]: - return [ - { - "id": "turn-1", - "role": "user", - "content": "x", - "type": "turn", - "created_at": "2024-01-01T00:00:00+00:00", - } - ] - - def _episodic_payload(self, **overrides) -> dict: - payload = { - "scope_type": "trip", - "scope_value": "Tokyo", - "text": "Planning a Tokyo trip with a luxury hotel preference.", - "situation": None, - "action_taken": None, - "outcome": None, - "outcome_valence": None, - "reasoning": None, - "lesson": None, - "domain": "travel", - "confidence": 0.95, - "salience": 0.8, - "tags": ["topic:travel"], - } - payload.update(overrides) - return payload - - def test_existing_episodics_are_rendered_into_prompt_inputs(self): - """The extractor must pass ``existing_episodics`` to the LLM, grouped - by ``(scope_type, scope_value)``. Without this, the model has no - context for refining or reversing the existing intent for that - scope when it emits the merged text.""" - p = self._build() - existing_text = "Planning a Tokyo trip with a luxury hotel preference." - existing_ep = { - "id": "ep_existing", - "type": "episodic", - "content": existing_text, - "content_hash": compute_content_hash(existing_text), - "thread_id": "__episodic__", - "salience": 0.8, - "metadata": {"scope_type": "trip", "scope_value": "Tokyo"}, - } - p._container.query_items.return_value = iter(self._turns()) - # Two queries are issued (one for facts, one for episodics) so each - # type gets its own 100-row budget — return [] for facts, the - # existing episodic for the episodic call. - p._load_existing_memories = MagicMock( - side_effect=lambda user_id, memory_types, **kw: [existing_ep] if memory_types == ["episodic"] else [] - ) - p._run_prompty = MagicMock(return_value=json.dumps({"facts": [], "episodic": [], "unclassified": []})) - - p.extract_memories("u1", "t1") - - # Two separate calls — one per type, each with its own budget. - load_calls = [c.args for c in p._load_existing_memories.call_args_list] - assert ("u1", ["fact"]) in load_calls - assert ("u1", ["episodic"]) in load_calls - call_kwargs = p._run_prompty.call_args.kwargs - inputs = call_kwargs["inputs"] - assert "existing_episodics" in inputs - rendered = inputs["existing_episodics"] - assert "trip = Tokyo" in rendered - assert "ep_existing" in rendered - assert existing_text in rendered - - def test_same_scope_episodics_collide_on_deterministic_id(self): - """Two episodics with the same (scope_type, scope_value) but - different ``text`` MUST produce the same deterministic ID so that - the second write overwrites the first via upsert. This is the - core mechanism that prevents near-duplicate episodic storage when - a recent-turn re-extraction window paraphrases the same intent. - """ - p = self._build() - p._container.query_items.return_value = iter(self._turns()) - p._load_existing_memories = MagicMock(return_value=[]) - # LLM emits two episodics under the SAME scope but with paraphrased - # text — this is the exact failure mode the user reported. - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "facts": [], - "episodic": [ - self._episodic_payload(text="Planning a Tokyo trip with a luxury hotel preference."), - self._episodic_payload(text="Planning a Tokyo trip with a preference for luxury hotels."), - ], - "unclassified": [], - } - ) - ) - - p.extract_memories("u1", "t1") - - upsert_calls = [c.args[0] for c in p._upsert_memory.call_args_list if c.args[0].get("type") == "episodic"] - # Both episodics flow through upsert (the persist path branches on - # type=episodic) and they MUST share the same det_id — that's what - # makes the second a Cosmos upsert that replaces the first. - assert len(upsert_calls) == 2 - assert upsert_calls[0]["id"] == upsert_calls[1]["id"] - # And neither went through create_item (which would 409 on the - # second write and silently lose the new richer text). - episodic_creates = [c for c in p._create_memory.call_args_list if c.args[0].get("type") == "episodic"] - assert episodic_creates == [] - - def test_different_scope_values_produce_different_ids(self): - """Two episodics with the same scope_type but different - scope_value (e.g. distinct trips, or distinct incidents within a - trip) MUST produce different deterministic IDs so they coexist. - """ - p = self._build() - p._container.query_items.return_value = iter(self._turns()) - p._load_existing_memories = MagicMock(return_value=[]) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "facts": [], - "episodic": [ - self._episodic_payload(scope_value="Tokyo", text="Trip A."), - self._episodic_payload(scope_value="Paris", text="Trip B."), - ], - "unclassified": [], - } - ) - ) - - p.extract_memories("u1", "t1") - - upsert_calls = [c.args[0] for c in p._upsert_memory.call_args_list if c.args[0].get("type") == "episodic"] - assert len(upsert_calls) == 2 - assert upsert_calls[0]["id"] != upsert_calls[1]["id"] - - def test_episodic_and_fact_with_same_content_do_not_collide(self): - """An episodic's deterministic ID is seeded on scope; a fact's is - seeded on content_hash. Even if their content text matches - verbatim, the IDs live in disjoint namespaces (ep_ vs fact_ prefix - plus different seeds) so both records persist.""" - p = self._build() - text = "Planning a Tokyo trip with a luxury hotel preference." - existing = [ - { - "id": "fact_existing", - "type": "fact", - "content": text, - "content_hash": compute_content_hash(text), - "thread_id": "t1", - "tags": ["sys:fact"], - } - ] - p._container.query_items.return_value = iter(self._turns()) - p._load_existing_memories = MagicMock(return_value=existing) - p._run_prompty = MagicMock( - return_value=json.dumps({"facts": [], "episodic": [self._episodic_payload(text=text)], "unclassified": []}) - ) - - out = p.extract_memories("u1", "t1") - - assert out["episodic_count"] == 1 - upsert_calls = [c.args[0] for c in p._upsert_memory.call_args_list if c.args[0].get("type") == "episodic"] - assert len(upsert_calls) == 1 - - def test_episodic_uses_sentinel_thread_id_for_partition_routing(self): - """Auto-extracted episodics MUST be persisted under the sentinel - ``thread_id="__episodic__"`` regardless of which thread emitted them. - - The memories container is partitioned hierarchically on - ``(user_id, thread_id)`` and Cosmos ``id`` uniqueness is per-partition - — so a deterministic ID seeded only on scope is only meaningful if - every episodic for that scope lands in the SAME partition. Writing - the live thread_id splits identical-scope episodics across two - partitions and breaks upsert dedup across threads. The originating - thread is preserved on ``metadata.originating_thread_id`` for audit. - """ - p = self._build() - p._container.query_items.return_value = iter(self._turns()) - p._load_existing_memories = MagicMock(return_value=[]) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "facts": [], - "episodic": [self._episodic_payload(text="Trip planning.")], - "unclassified": [], - } - ) - ) - - p.extract_memories("u1", "thread-alpha") - - upsert_calls = [c.args[0] for c in p._upsert_memory.call_args_list if c.args[0].get("type") == "episodic"] - assert len(upsert_calls) == 1 - doc = upsert_calls[0] - assert doc["thread_id"] == "__episodic__" - assert doc["metadata"]["originating_thread_id"] == "thread-alpha" - - def test_same_scope_episodics_collide_across_different_threads(self): - """The cross-thread regression test for the per-partition id bug. - - Same user, same scope ``(trip, Tokyo)``, but emitted from two - different threads (``thread-alpha`` and ``thread-beta``). Both - writes MUST produce the same det_id AND the same persisted - thread_id (the sentinel) so they land in one partition and the - second upsert replaces the first. Without the sentinel, the docs - live in two different partitions and you'd see duplicate - episodics for the same intent — exactly the bug the deterministic - ID was meant to prevent. - """ - p = self._build() - # Fresh iterator per call — both extract_memories calls need turns. - p._container.query_items.side_effect = lambda *a, **kw: iter(self._turns()) - p._load_existing_memories = MagicMock(return_value=[]) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "facts": [], - "episodic": [self._episodic_payload(text="Tokyo luxury hotel intent.")], - "unclassified": [], - } - ) - ) - - p.extract_memories("u1", "thread-alpha") - p.extract_memories("u1", "thread-beta") - - upsert_calls = [c.args[0] for c in p._upsert_memory.call_args_list if c.args[0].get("type") == "episodic"] - assert len(upsert_calls) == 2 - # Same det_id — scope is identity. - assert upsert_calls[0]["id"] == upsert_calls[1]["id"] - # Same partition (sentinel thread_id) — so the second upsert replaces - # the first instead of creating a duplicate in a sibling partition. - assert upsert_calls[0]["thread_id"] == upsert_calls[1]["thread_id"] == "__episodic__" - # Originating thread preserved on each metadata for audit. - assert upsert_calls[0]["metadata"]["originating_thread_id"] == "thread-alpha" - assert upsert_calls[1]["metadata"]["originating_thread_id"] == "thread-beta" - - class TestExtractEarlyReturnShape: """The no-memories early-return must include every key the success path returns; otherwise callers using ``result["exact_dedup_skipped"]`` @@ -931,7 +685,6 @@ def test_empty_thread_returns_full_dict_shape(self): "unclassified_count", "updated_count", "exact_dedup_skipped", - "dropped_episodic_count", ): assert key in out, f"missing key: {key}" assert out[key] == 0 @@ -1586,81 +1339,10 @@ def capture_prompty(name, inputs): assert "None" not in text -class TestExtractUpdateSupersedeReason: - """Extract-time UPDATE actions stamp ``supersede_reason="update"``, - distinct from reconcile-time ``"duplicate"`` (paraphrase merge) and - ``"contradict"`` (semantic conflict). The extract prompt defines - UPDATE as "contradicts or refines an existing memory" — labelling - these as ``"duplicate"`` makes audit trails ambiguous.""" - - def _build(self) -> PipelineService: - p = PipelineService.__new__(PipelineService) - p._embeddings = MagicMock() - p._embeddings.generate.return_value = [[0.1] * 8] - p._upsert_memory = MagicMock(side_effect=lambda doc: doc) - p._mark_superseded = MagicMock(return_value=True) - p._container = MagicMock() - p._memories_container = p._container - p._turns_container = p._container - p._summaries_container = p._container - p._chat = MagicMock() - p._load_existing_memories = MagicMock( - return_value=[ - { - "id": "fact_old", - "type": "fact", - "content": "User likes coffee", - "content_hash": "h_old", - } - ] - ) - p._container.read_item = MagicMock( - return_value={"id": "fact_old", "type": "fact", "content": "User likes coffee"} - ) - p._container.query_items.return_value = iter( - [ - { - "id": "turn-1", - "role": "user", - "content": "I love tea now", - "type": "turn", - "created_at": "2024-01-01T00:00:00+00:00", - } - ] - ) - return p - - def test_fact_update_uses_reason_update_not_duplicate(self): - p = self._build() - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "facts": [ - { - "text": "User now prefers tea over coffee", - "confidence": 0.9, - "salience": 0.7, - "action": "UPDATE", - "supersedes_id": "fact_old", - "tags": ["sys:fact"], - } - ], - "procedural": [], - "episodic": [], - } - ) - ) - p.extract_memories("u1", "t1") - assert p._mark_superseded.called - call_kwargs = p._mark_superseded.call_args.kwargs - assert call_kwargs.get("reason") == "update" - - class TestExtractUpdateSelfCollapseGuard: - """When an LLM emits ``UPDATE`` whose new content hashes to the same - deterministic id as the target (paraphrase-equivalent text), the - upsert would overwrite the audit metadata that ``_mark_superseded`` - just stamped on the target. Treat as a no-op.""" + """Procedural synthesis self-collapse guard: when synthesis would emit a + proc id identical to the existing one, treat as a no-op. (Fact extract-time + UPDATE was removed — facts/contradictions are reconciled, not extract-tagged.)""" def _build(self) -> PipelineService: p = PipelineService.__new__(PipelineService) @@ -1676,52 +1358,6 @@ def _build(self) -> PipelineService: p._load_existing_memories = MagicMock(return_value=[]) return p - def test_fact_update_with_self_referential_id_is_skipped(self): - from azure.cosmos.agent_memory._utils import compute_content_hash - from azure.cosmos.agent_memory.services.pipeline import _ID_SEED_SEP - - p = self._build() - text = "User likes tea" - seed = _ID_SEED_SEP.join(("u1", "t1", compute_content_hash(text))) - det_id = f"fact_{hashlib.sha256(seed.encode()).hexdigest()[:32]}" - - p._container.query_items.return_value = iter( - [ - { - "id": "turn-1", - "role": "user", - "content": "tea", - "type": "turn", - "created_at": "2024-01-01T00:00:00+00:00", - } - ] - ) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "facts": [ - { - "text": text, - "confidence": 0.9, - "salience": 0.6, - "action": "UPDATE", - "supersedes_id": det_id, - "tags": ["sys:fact"], - } - ], - "procedural": [], - "episodic": [], - } - ) - ) - - out = p.extract_memories("u1", "t1") - - assert p._mark_superseded.call_count == 0 - fact_upserts = [c for c in p._upsert_memory.call_args_list if c.args[0].get("type") == "fact"] - assert fact_upserts == [] - assert out["fact_count"] == 0 - def test_procedural_update_with_self_referential_id_is_skipped(self): from azure.cosmos.agent_memory._utils import compute_content_hash from azure.cosmos.agent_memory.services.pipeline import _ID_SEED_SEP diff --git a/tests/unit/test_thresholds.py b/tests/unit/test_thresholds.py index 1add114..d7f9d3d 100644 --- a/tests/unit/test_thresholds.py +++ b/tests/unit/test_thresholds.py @@ -2,6 +2,7 @@ import pytest +from azure.cosmos.agent_memory import thresholds from azure.cosmos.agent_memory.thresholds import ( DEFAULT_ENABLE_TURN_EMBEDDINGS, DEFAULT_TTL_BY_TYPE, @@ -50,3 +51,149 @@ def test_enable_turn_embeddings_truthy_values(monkeypatch, raw) -> None: def test_enable_turn_embeddings_falsy_values(monkeypatch, raw) -> None: monkeypatch.setenv("ENABLE_TURN_EMBEDDINGS", raw) assert get_enable_turn_embeddings() is False + + +@pytest.mark.parametrize( + ("env_name", "getter_name", "expected"), + [ + ("FACT_EXTRACTION_EVERY_N", "get_fact_extraction_every_n", 1), + ("THREAD_SUMMARY_EVERY_N", "get_thread_summary_every_n", 10), + ("USER_SUMMARY_EVERY_N", "get_user_summary_every_n", 20), + ("DEDUP_EVERY_N", "get_dedup_every_n", 5), + ("DEDUP_POOL_SIZE", "get_dedup_pool_size", 50), + ("PROCEDURAL_SYNTHESIS_AUTO", "get_procedural_synthesis_auto", True), + ], +) +def test_env_config_getters_defaults( + monkeypatch: pytest.MonkeyPatch, + env_name: str, + getter_name: str, + expected: object, +) -> None: + monkeypatch.delenv(env_name, raising=False) + + assert getattr(thresholds, getter_name)() == expected + + +@pytest.mark.parametrize( + ("env_name", "getter_name", "raw", "expected"), + [ + ("FACT_EXTRACTION_EVERY_N", "get_fact_extraction_every_n", "2", 2), + ("THREAD_SUMMARY_EVERY_N", "get_thread_summary_every_n", "11", 11), + ("USER_SUMMARY_EVERY_N", "get_user_summary_every_n", "21", 21), + ("DEDUP_EVERY_N", "get_dedup_every_n", "3", 3), + ("DEDUP_POOL_SIZE", "get_dedup_pool_size", "75", 75), + ("PROCEDURAL_SYNTHESIS_AUTO", "get_procedural_synthesis_auto", "false", False), + ], +) +def test_env_config_getters_parse_env( + monkeypatch: pytest.MonkeyPatch, + env_name: str, + getter_name: str, + raw: str, + expected: object, +) -> None: + monkeypatch.setenv(env_name, raw) + + assert getattr(thresholds, getter_name)() == expected + + +@pytest.mark.parametrize( + ("env_name", "getter_name", "expected"), + [ + ("FACT_EXTRACTION_EVERY_N", "get_fact_extraction_every_n", 1), + ("THREAD_SUMMARY_EVERY_N", "get_thread_summary_every_n", 10), + ("USER_SUMMARY_EVERY_N", "get_user_summary_every_n", 20), + ("DEDUP_EVERY_N", "get_dedup_every_n", 5), + ("DEDUP_POOL_SIZE", "get_dedup_pool_size", 50), + ], +) +def test_int_getters_reject_negative( + monkeypatch: pytest.MonkeyPatch, + env_name: str, + getter_name: str, + expected: int, +) -> None: + monkeypatch.setenv(env_name, "-1") + + assert getattr(thresholds, getter_name)() == expected + + +@pytest.mark.parametrize( + ("env_name", "getter_name", "expected"), + [ + ("FACT_EXTRACTION_EVERY_N", "get_fact_extraction_every_n", 1), + ("THREAD_SUMMARY_EVERY_N", "get_thread_summary_every_n", 10), + ("USER_SUMMARY_EVERY_N", "get_user_summary_every_n", 20), + ("DEDUP_EVERY_N", "get_dedup_every_n", 5), + ("DEDUP_POOL_SIZE", "get_dedup_pool_size", 50), + ], +) +def test_int_getters_invalid_use_default( + monkeypatch: pytest.MonkeyPatch, + env_name: str, + getter_name: str, + expected: int, +) -> None: + monkeypatch.setenv(env_name, "bogus") + + assert getattr(thresholds, getter_name)() == expected + + +def test_dedup_pool_size_clamps_high(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DEDUP_POOL_SIZE", "501") + + assert thresholds.get_dedup_pool_size() == 500 + + +def test_dedup_pool_size_rejects_zero(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DEDUP_POOL_SIZE", "0") + + assert thresholds.get_dedup_pool_size() == 50 + + +def test_procedural_synthesis_auto_invalid_uses_default(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PROCEDURAL_SYNTHESIS_AUTO", "bogus") + + assert thresholds.get_procedural_synthesis_auto() is True + + +def test_processor_owner_defaults_to_none(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("MEMORY_PROCESSOR_OWNER", raising=False) + + assert thresholds.get_processor_owner() is None + + +@pytest.mark.parametrize("raw", ["inprocess", "durable", "INPROCESS", "DURABLE"]) +def test_processor_owner_parse_env(monkeypatch: pytest.MonkeyPatch, raw: str) -> None: + monkeypatch.setenv("MEMORY_PROCESSOR_OWNER", raw) + + assert thresholds.get_processor_owner() == raw.lower() + + +def test_processor_owner_invalid_uses_default(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("MEMORY_PROCESSOR_OWNER", "bogus") + + assert thresholds.get_processor_owner() is None + + +def test_internalized_getters_return_fixed_constants_and_ignore_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DEDUP_CONTEXT_VECTOR_ENABLED", "false") + monkeypatch.setenv("DEDUP_CONTEXT_TOPK", "7") + monkeypatch.setenv("DEDUP_VECTOR_ENABLED", "false") + monkeypatch.setenv("DEDUP_SIM_HIGH", "0.50") + monkeypatch.setenv("DEDUP_SIM_LOW", "0.40") + monkeypatch.setenv("DEDUP_CANDIDATE_TOPK", "8") + monkeypatch.setenv("DEDUP_RECONCILE_MODE", "full_pool") + monkeypatch.setenv("DEDUP_CLUSTER_SIM", "0.10") + monkeypatch.setenv("DEDUP_FULL_RECLUSTER_EVERY_N", "4") + + assert thresholds.get_dedup_context_vector_enabled() is True + assert thresholds.get_dedup_context_topk() == 10 + assert thresholds.get_dedup_vector_enabled() is True + assert thresholds.get_dedup_sim_high() == 0.97 + assert thresholds.get_dedup_sim_low() == 0.80 + assert thresholds.get_dedup_candidate_topk() == 10 + assert thresholds.get_dedup_reconcile_mode() == "candidate" + assert thresholds.get_dedup_cluster_sim() == 0.60 + assert thresholds.get_dedup_full_recluster_every_n() == 12 diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 0749a0a..6226937 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -4,6 +4,7 @@ from azure.cosmos.agent_memory._utils import ( DEFAULT_TTL_BY_TYPE, + MAX_FULLTEXT_TERMS, _build_container_kwargs, _container_policies, _make_memory, @@ -12,7 +13,11 @@ _resolve_full_text_language, _resolve_vector_index_type, compute_content_hash, + distance_function_from_container_properties, + extract_keywords, normalize_ai_foundry_endpoint, + vector_order_direction, + vector_similarity_at_least, ) from azure.cosmos.agent_memory.exceptions import ConfigurationError, ValidationError @@ -192,65 +197,132 @@ def test_make_memory_invalid_type(): _make_memory(user_id="u1", role="user", content="test", memory_type="invalid") -def test_resolve_embedding_data_type_defaults(monkeypatch): - monkeypatch.delenv("AI_FOUNDRY_EMBEDDING_DATA_TYPE", raising=False) +def test_resolve_embedding_data_type_defaults(): assert _resolve_embedding_data_type(None) == "float32" -def test_resolve_embedding_data_type_from_env(monkeypatch): - monkeypatch.setenv("AI_FOUNDRY_EMBEDDING_DATA_TYPE", "int8") - assert _resolve_embedding_data_type(None) == "int8" - - -def test_resolve_embedding_data_type_explicit_overrides_env(monkeypatch): - monkeypatch.setenv("AI_FOUNDRY_EMBEDDING_DATA_TYPE", "int8") +def test_resolve_embedding_data_type_explicit(): + assert _resolve_embedding_data_type("int8") == "int8" assert _resolve_embedding_data_type("uint8") == "uint8" -def test_resolve_embedding_data_type_invalid_raises(monkeypatch): - monkeypatch.setenv("AI_FOUNDRY_EMBEDDING_DATA_TYPE", "bogus") +def test_resolve_embedding_data_type_invalid_raises(): with pytest.raises(ConfigurationError): - _resolve_embedding_data_type(None) + _resolve_embedding_data_type("bogus") -def test_resolve_distance_function_defaults(monkeypatch): - monkeypatch.delenv("AI_FOUNDRY_EMBEDDING_DISTANCE_FUNCTION", raising=False) +def test_resolve_distance_function_defaults(): assert _resolve_distance_function(None) == "cosine" -def test_resolve_distance_function_from_env(monkeypatch): - monkeypatch.setenv("AI_FOUNDRY_EMBEDDING_DISTANCE_FUNCTION", "dotproduct") - assert _resolve_distance_function(None) == "dotproduct" +def test_resolve_distance_function_explicit(): + assert _resolve_distance_function("dotproduct") == "dotproduct" + assert _resolve_distance_function("euclidean") == "euclidean" -def test_resolve_distance_function_invalid_raises(monkeypatch): - monkeypatch.setenv("AI_FOUNDRY_EMBEDDING_DISTANCE_FUNCTION", "manhattan") +def test_resolve_distance_function_invalid_raises(): with pytest.raises(ConfigurationError): - _resolve_distance_function(None) + _resolve_distance_function("manhattan") -def test_resolve_vector_index_type_defaults(monkeypatch): - monkeypatch.delenv("AI_FOUNDRY_EMBEDDING_VECTOR_INDEX_TYPE", raising=False) - assert _resolve_vector_index_type(None) == "quantizedFlat" +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_vector_similarity_at_least_cosine_and_dotproduct(): + # Higher score is more similar; threshold is a floor. + for fn in ("cosine", "dotproduct"): + assert vector_similarity_at_least(0.97, 0.97, fn) is True + assert vector_similarity_at_least(0.99, 0.97, fn) is True + assert vector_similarity_at_least(0.80, 0.97, fn) is False -def test_resolve_vector_index_type_from_env(monkeypatch): - monkeypatch.setenv("AI_FOUNDRY_EMBEDDING_VECTOR_INDEX_TYPE", "quantizedFlat") +def test_vector_similarity_at_least_euclidean_inverts(): + # Lower distance is more similar; threshold is a ceiling. + assert vector_similarity_at_least(0.10, 0.20, "euclidean") is True + assert vector_similarity_at_least(0.20, 0.20, "euclidean") is True + assert vector_similarity_at_least(0.50, 0.20, "euclidean") is False + + +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") == [ + "user", + "loves", + "hiking", + "trails", + ] + assert extract_keywords("") == [] + assert extract_keywords("the and of a an") == [] # all stopwords + + +def test_extract_keywords_capped_at_cosmos_fulltext_limit(): + # Cosmos FullTextScore rejects >30 terms; extraction must cap at exactly 30 so + # the hybrid query is always valid even for long multi-turn context strings. + text = " ".join(f"term{i}" for i in range(100)) + kws = extract_keywords(text) + assert len(kws) == MAX_FULLTEXT_TERMS == 30 + assert kws == [f"term{i}" for i in range(30)] + + +def test_resolve_vector_index_type_defaults(): assert _resolve_vector_index_type(None) == "quantizedFlat" -def test_resolve_vector_index_type_explicit_overrides_env(monkeypatch): - monkeypatch.setenv("AI_FOUNDRY_EMBEDDING_VECTOR_INDEX_TYPE", "quantizedFlat") +def test_resolve_vector_index_type_explicit(): assert _resolve_vector_index_type("flat") == "flat" -def test_resolve_vector_index_type_invalid_raises(monkeypatch): - monkeypatch.setenv("AI_FOUNDRY_EMBEDDING_VECTOR_INDEX_TYPE", "hnsw") +def test_resolve_vector_index_type_invalid_raises(): with pytest.raises(ConfigurationError): - _resolve_vector_index_type(None) + _resolve_vector_index_type("hnsw") -def test_container_policies_defaults_to_diskann(): +def test_container_policies_defaults_vector_index_type(): _, indexing_policy, _ = _container_policies( embedding_dimensions=1536, embedding_data_type="float32", @@ -266,19 +338,17 @@ def test_container_policies_uses_supplied_vector_index_type(): embedding_data_type="float32", distance_function="cosine", full_text_language="en-US", - vector_index_type="quantizedFlat", + vector_index_type="diskANN", ) - assert indexing_policy["vectorIndexes"] == [{"path": "/embedding", "type": "quantizedFlat"}] + assert indexing_policy["vectorIndexes"] == [{"path": "/embedding", "type": "diskANN"}] -def test_resolve_full_text_language_defaults(monkeypatch): - monkeypatch.delenv("COSMOS_DB_FULL_TEXT_LANGUAGE", raising=False) +def test_resolve_full_text_language_defaults(): assert _resolve_full_text_language(None) == "en-US" -def test_resolve_full_text_language_from_env(monkeypatch): - monkeypatch.setenv("COSMOS_DB_FULL_TEXT_LANGUAGE", "fr-FR") - assert _resolve_full_text_language(None) == "fr-FR" +def test_resolve_full_text_language_explicit(): + assert _resolve_full_text_language("fr-FR") == "fr-FR" # --------------------------------------------------------------------------- From 2cf898833bec58e2302a6d9cbfe8d4da9788e9ab Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Wed, 1 Jul 2026 09:40:21 -0700 Subject: [PATCH 2/8] fixing builds --- .../agent_memory/aio/services/pipeline.py | 8 +- .../cosmos/agent_memory/services/pipeline.py | 53 +++++------ azure/cosmos/agent_memory/thresholds.py | 14 +-- .../orchestrators/extract_memories.py | 11 ++- function_app/triggers/change_feed.py | 4 +- tests/integration/test_async_full_pipeline.py | 9 +- tests/integration/test_full_pipeline.py | 6 +- .../aio/services/test_dedup_vector_async.py | 16 +--- tests/unit/aio/test_auto_trigger.py | 92 ++++++++++--------- tests/unit/function_app/test_orchestrators.py | 4 +- tests/unit/services/test_dedup_vector.py | 25 ++--- tests/unit/test_auto_trigger.py | 73 ++++++++------- 12 files changed, 151 insertions(+), 164 deletions(-) diff --git a/azure/cosmos/agent_memory/aio/services/pipeline.py b/azure/cosmos/agent_memory/aio/services/pipeline.py index 4e5de1e..ba8798f 100644 --- a/azure/cosmos/agent_memory/aio/services/pipeline.py +++ b/azure/cosmos/agent_memory/aio/services/pipeline.py @@ -530,9 +530,7 @@ async def extract_memories_dry( m["content_hash"] for m in existing_for_hash if m.get("type") == "fact" and m.get("content_hash") } if get_dedup_context_vector_enabled(): - user_turns_text = "\n".join( - str(it.get("content", "")) for it in items if it.get("role") == "user" - ).strip() + user_turns_text = "\n".join(str(it.get("content", "")) for it in items if it.get("role") == "user").strip() context_query = user_turns_text or transcript existing = await self._store.search( search_terms=context_query, @@ -1778,9 +1776,7 @@ async def reconcile_memories( # "loves steak") — which candidate clustering, keyed on near-duplicate # similarity, would never group. Automatic sweeps use cheap candidate mode. if get_dedup_reconcile_mode() == "candidate" and not full_rebuild: - return await self._reconcile_candidate_mode( - user_id, n=n, memory_type=memory_type, started_at=started_at - ) + return await self._reconcile_candidate_mode(user_id, n=n, memory_type=memory_type, started_at=started_at) facts = await self._active_memories_for_reconcile(user_id, memory_type, n) result, consumed = await self._reconcile_pool(user_id, memory_type, facts) diff --git a/azure/cosmos/agent_memory/services/pipeline.py b/azure/cosmos/agent_memory/services/pipeline.py index b216816..a44f87f 100644 --- a/azure/cosmos/agent_memory/services/pipeline.py +++ b/azure/cosmos/agent_memory/services/pipeline.py @@ -414,9 +414,7 @@ def _query_active_memories( ) ) - def _load_memories_by_ids( - self, user_id: str, memory_type: str, ids: Iterable[str] - ) -> list[dict[str, Any]]: + def _load_memories_by_ids(self, user_id: str, memory_type: str, ids: Iterable[str]) -> list[dict[str, Any]]: ids = [mid for mid in dict.fromkeys(ids) if mid] if not ids: return [] @@ -600,9 +598,7 @@ def extract_memories_dry( transcript = self._build_transcript(items) existing = existing_for_hashes if threshold_config.get_dedup_context_vector_enabled(): - user_turns_text = "\n".join( - str(it.get("content", "")) for it in items if it.get("role") == "user" - ).strip() + user_turns_text = "\n".join(str(it.get("content", "")) for it in items if it.get("role") == "user").strip() context_query = user_turns_text or transcript existing = self._store.search( search_terms=context_query, @@ -1640,8 +1636,7 @@ def _build_candidate_clusters( candidate_ids = [ row["id"] for row in candidates - if row.get("id") - and vector_similarity_at_least(row.get("score", 0.0), cluster_sim, distance_function) + if row.get("id") and vector_similarity_at_least(row.get("score", 0.0), cluster_sim, distance_function) ] for cid in candidate_ids: edge_pairs.add(tuple(sorted((seed_id, cid)))) @@ -1698,9 +1693,7 @@ def _build_candidate_clusters( clusters.append([nodes_by_id[cid] for cid in component]) return clusters, len(nodes_by_id), seeds - def _reconcile_candidate_mode( - self, user_id: str, *, n: int, memory_type: str, started_at: float - ) -> dict[str, int]: + def _reconcile_candidate_mode(self, user_id: str, *, n: int, memory_type: str, started_at: float) -> dict[str, int]: # Candidate clustering only. The periodic full-pool backstop that catches # dissimilar-embedding contradictions ("vegetarian" vs "loves steak") is # driven by the caller via ``full_rebuild`` on a PERSISTED-counter cadence @@ -1810,9 +1803,7 @@ def reconcile_memories( # "loves steak") — which candidate clustering, keyed on near-duplicate # similarity, would never group. Automatic sweeps use cheap candidate mode. if threshold_config.get_dedup_reconcile_mode() == "candidate" and not full_rebuild: - return self._reconcile_candidate_mode( - user_id, n=n, memory_type=memory_type, started_at=started_at - ) + return self._reconcile_candidate_mode(user_id, n=n, memory_type=memory_type, started_at=started_at) facts = self._active_memories_for_reconcile(user_id, memory_type, n) result, consumed = self._reconcile_pool(user_id, memory_type, facts) @@ -1828,9 +1819,7 @@ def reconcile_memories( ) return result - def _active_memories_for_reconcile( - self, user_id: str, memory_type: str, n: int - ) -> list[dict[str, Any]]: + def _active_memories_for_reconcile(self, user_id: str, memory_type: str, n: int) -> list[dict[str, Any]]: # ---- Load up to N most recent active memories ---- # ORDER BY c.created_at DESC keeps the TOP cap deterministic across # physical partitions and matches the dedup prompt's tiebreaker @@ -2058,21 +2047,21 @@ def _reconcile_pool( try: merged_payload: dict[str, Any] = { - "id": merged_id, - "user_id": user_id, - "role": "system", - "type": memory_type, - "content": merged_content, - "thread_id": recent_thread_id or f"__reconciled__:{user_id}", - "confidence": confidence_val if confidence_val is not None else 0.5, - "salience": salience_val if salience_val is not None else 0.5, - "supersedes_ids": merged_supersedes, - "source_memory_ids": merged_source_memory_ids, - "tags": merged_tags, - "content_hash": merged_content_hash, - **self._prompt_lineage(prompt_filename), - "created_at": datetime.now(timezone.utc).isoformat(), - "updated_at": datetime.now(timezone.utc).isoformat(), + "id": merged_id, + "user_id": user_id, + "role": "system", + "type": memory_type, + "content": merged_content, + "thread_id": recent_thread_id or f"__reconciled__:{user_id}", + "confidence": confidence_val if confidence_val is not None else 0.5, + "salience": salience_val if salience_val is not None else 0.5, + "supersedes_ids": merged_supersedes, + "source_memory_ids": merged_source_memory_ids, + "tags": merged_tags, + "content_hash": merged_content_hash, + **self._prompt_lineage(prompt_filename), + "created_at": datetime.now(timezone.utc).isoformat(), + "updated_at": datetime.now(timezone.utc).isoformat(), } if memory_type == "fact": merged_payload["metadata"] = { diff --git a/azure/cosmos/agent_memory/thresholds.py b/azure/cosmos/agent_memory/thresholds.py index ca4c508..f25030c 100644 --- a/azure/cosmos/agent_memory/thresholds.py +++ b/azure/cosmos/agent_memory/thresholds.py @@ -35,15 +35,15 @@ # needs to become operator-facing we add the env plumbing back deliberately. # The dedup + hybrid-search features ship ON via these values. # --------------------------------------------------------------------------- -DEDUP_CONTEXT_VECTOR_ENABLED = True # Stage-1 relevance-ranked extraction context +DEDUP_CONTEXT_VECTOR_ENABLED = True # Stage-1 relevance-ranked extraction context DEDUP_CONTEXT_TOPK = 10 -DEDUP_VECTOR_ENABLED = True # Stage-3 vector near-dup ladder -DEDUP_SIM_HIGH = 0.97 # >= -> auto-skip near-exact -DEDUP_SIM_LOW = 0.80 # < -> novel; between -> tag candidate +DEDUP_VECTOR_ENABLED = True # Stage-3 vector near-dup ladder +DEDUP_SIM_HIGH = 0.97 # >= -> auto-skip near-exact +DEDUP_SIM_LOW = 0.80 # < -> novel; between -> tag candidate DEDUP_CANDIDATE_TOPK = 10 -DEDUP_RECONCILE_MODE = "candidate" # clustered candidate reconcile (vs legacy full_pool) -DEDUP_CLUSTER_SIM = 0.60 # Stage-5 clustering edge threshold -DEDUP_FULL_RECLUSTER_EVERY_N = 12 # full re-cluster safety net cadence +DEDUP_RECONCILE_MODE = "candidate" # clustered candidate reconcile (vs legacy full_pool) +DEDUP_CLUSTER_SIM = 0.60 # Stage-5 clustering edge threshold +DEDUP_FULL_RECLUSTER_EVERY_N = 12 # full re-cluster safety net cadence DEFAULT_TTL_BY_TYPE: dict[str, int] = { "turn": 2_592_000, diff --git a/function_app/orchestrators/extract_memories.py b/function_app/orchestrators/extract_memories.py index ec859bc..db56c24 100644 --- a/function_app/orchestrators/extract_memories.py +++ b/function_app/orchestrators/extract_memories.py @@ -127,10 +127,13 @@ def em_Extract(payload: dict) -> dict: @bp.activity_trigger(input_name="payload") def em_Dedup(payload: dict) -> dict: """vector-floor dedup ladder (gated; passthrough when disabled).""" - return get_pipeline().dedup_extracted_memories( - user_id=payload["user_id"], - extracted=payload["extracted"], - ) or payload["extracted"] + return ( + get_pipeline().dedup_extracted_memories( + user_id=payload["user_id"], + extracted=payload["extracted"], + ) + or payload["extracted"] + ) @bp.activity_trigger(input_name="payload") diff --git a/function_app/triggers/change_feed.py b/function_app/triggers/change_feed.py index 4630a46..6b7ad9d 100644 --- a/function_app/triggers/change_feed.py +++ b/function_app/triggers/change_feed.py @@ -220,9 +220,7 @@ async def process_changefeed_batch( # Persisted-counter backstop: every n_full_turns turns, force a # full-pool reconcile so dissimilar-embedding contradictions are # caught reliably on FA (not gated by the in-memory sweep counter). - should_full_reconcile = bool( - n_full_turns > 0 and crosses_threshold(old_count, new_count, n_full_turns) - ) + should_full_reconcile = bool(n_full_turns > 0 and crosses_threshold(old_count, new_count, n_full_turns)) watermark = await read_extract_watermark(counter_container, cid, user_id, thread_id) # Not capped: new_count - watermark is exactly the unextracted backlog # and the orchestrator advances the watermark to new_count, so capping diff --git a/tests/integration/test_async_full_pipeline.py b/tests/integration/test_async_full_pipeline.py index 88f66c4..a47804e 100644 --- a/tests/integration/test_async_full_pipeline.py +++ b/tests/integration/test_async_full_pipeline.py @@ -131,10 +131,7 @@ async def _async_seed_fact_with_embedding( sync helper). Retries through transient embedding-service blips so the extract-time vector floor always has a neighbour; skips honestly if the embedding service is genuinely unavailable.""" - check = ( - "SELECT c.id FROM c WHERE c.user_id = @uid AND c.content = @content " - "AND IS_DEFINED(c.embedding)" - ) + check = "SELECT c.id FROM c WHERE c.user_id = @uid AND c.content = @content AND IS_DEFINED(c.embedding)" params = [{"name": "@uid", "value": user_id}, {"name": "@content", "value": content}] for _ in range(retries): await mem.add_cosmos( @@ -145,9 +142,7 @@ async def _async_seed_fact_with_embedding( thread_id=thread_id, salience=0.7, ) - embedded = [ - doc async for doc in mem._memories_container_client.query_items(query=check, parameters=params) - ] + embedded = [doc async for doc in mem._memories_container_client.query_items(query=check, parameters=params)] if embedded: return await asyncio.sleep(1) diff --git a/tests/integration/test_full_pipeline.py b/tests/integration/test_full_pipeline.py index dcb217e..ed6be30 100644 --- a/tests/integration/test_full_pipeline.py +++ b/tests/integration/test_full_pipeline.py @@ -138,10 +138,7 @@ def _seed_fact_with_embedding( neighbour to match. Retry until an embedded copy exists (indexing is fast — the doc is vector-searchable within ~2s), and skip honestly if the embedding service is genuinely unavailable rather than reporting a false failure.""" - check = ( - "SELECT c.id FROM c WHERE c.user_id = @uid AND c.content = @content " - "AND IS_DEFINED(c.embedding)" - ) + check = "SELECT c.id FROM c WHERE c.user_id = @uid AND c.content = @content AND IS_DEFINED(c.embedding)" params = [{"name": "@uid", "value": user_id}, {"name": "@content", "value": content}] for _ in range(retries): mem.add_cosmos( @@ -624,4 +621,3 @@ def test_dedup_extracted_memories_flags_near_duplicate_of_stored_fact( ) finally: _cleanup(agent_memory, unique_user_id) - diff --git a/tests/unit/aio/services/test_dedup_vector_async.py b/tests/unit/aio/services/test_dedup_vector_async.py index be62553..e20a206 100644 --- a/tests/unit/aio/services/test_dedup_vector_async.py +++ b/tests/unit/aio/services/test_dedup_vector_async.py @@ -94,9 +94,7 @@ async def fake_query_items(_container, *, query, parameters): p._query_items = AsyncMock(side_effect=fake_query_items) p._distance_function_cache = "cosine" - out = await p._vector_candidates( - user_id="u1", embedding=[1.0, 0.0], memory_type="fact", top_k=2, exclude_ids=set() - ) + out = await p._vector_candidates(user_id="u1", embedding=[1.0, 0.0], memory_type="fact", top_k=2, exclude_ids=set()) # Cosmos rejects an explicit ASC/DESC on ORDER BY VectorDistance(); it orders # most-similar-first server-side. Direction-awareness lives in the Python sort. assert "ORDER BY VectorDistance(c.embedding, @vec)" in captured["query"] @@ -105,9 +103,7 @@ async def fake_query_items(_container, *, query, parameters): assert [c["id"] for c in out] == ["near", "far"] p._distance_function_cache = "euclidean" - out = await p._vector_candidates( - user_id="u1", embedding=[1.0, 0.0], memory_type="fact", top_k=2, exclude_ids=set() - ) + out = await p._vector_candidates(user_id="u1", embedding=[1.0, 0.0], memory_type="fact", top_k=2, exclude_ids=set()) assert "VectorDistance(c.embedding, @vec) ASC" not in captured["query"] # euclidean: lower distance = more similar, so 0.10 ("far" label) sorts first. # euclidean: lower distance = more similar, so 0.10 ("far" label) sorts first. @@ -202,9 +198,7 @@ async def clear(docs): @pytest.mark.asyncio async def test_full_rebuild_clears_survivor_tags(monkeypatch): # Async mirror: full_rebuild full-pool path clears survivor dup-candidate tags. - monkeypatch.setattr( - "azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_reconcile_mode", lambda: "candidate" - ) + monkeypatch.setattr("azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_reconcile_mode", lambda: "candidate") p = _service() pool = [ _fact("f1", "a", tags=["sys:fact", "sys:dup-candidate"]), @@ -333,9 +327,7 @@ async def test_dedup_skips_underspecified_doc_verbatim(): # Parity with sync: a doc with no/unknown type is passed through untouched # and never runs vector dedup (async previously defaulted type to the bucket). p = _service() - p._vector_candidates = AsyncMock( - return_value=[{"id": "x", "content": "c", "type": "fact", "score": 0.99}] - ) + p._vector_candidates = AsyncMock(return_value=[{"id": "x", "content": "c", "type": "fact", "score": 0.99}]) doc = _fact("f1", "content") doc.pop("type") doc.pop("embedding", None) diff --git a/tests/unit/aio/test_auto_trigger.py b/tests/unit/aio/test_auto_trigger.py index adb9f88..0c68e28 100644 --- a/tests/unit/aio/test_auto_trigger.py +++ b/tests/unit/aio/test_auto_trigger.py @@ -52,7 +52,6 @@ async def patch_item(self, *, item, partition_key, patch_operations): return dict(doc) - class TestAsyncAutoTriggerNonBlocking: @pytest.mark.asyncio async def test_push_to_cosmos_does_not_await_auto_trigger(self, monkeypatch): @@ -155,15 +154,13 @@ async def fake_upsert(body): @pytest.mark.parametrize( ("counter_result", "watermark", "expected_recent_k"), [ - ((5, 10), 5, 5), # backlog = new - watermark - ((0, 1), 1, 1), # new == watermark -> floored to 1 - ((98, 100), 0, 100), # large backlog is NOT capped + ((5, 10), 5, 5), # backlog = new - watermark + ((0, 1), 1, 1), # new == watermark -> floored to 1 + ((98, 100), 0, 100), # large backlog is NOT capped ((20, 30), None, 30), # BOOTSTRAP: no watermark -> base=0 -> recent_k = new_count ], ) - async def test_extract_recent_k_uses_watermark( - self, monkeypatch, counter_result, watermark, expected_recent_k - ): + async def test_extract_recent_k_uses_watermark(self, monkeypatch, counter_result, watermark, expected_recent_k): monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") @@ -177,16 +174,20 @@ async def test_extract_recent_k_uses_watermark( client._summaries_container_client = client._memories_container_client client._counter_container_client = MagicMock() - with patch( - "azure.cosmos.agent_memory._counters.increment_counter_async", - return_value=counter_result, - ), patch( - "azure.cosmos.agent_memory._counters.read_extract_watermark_async", - new=AsyncMock(return_value=watermark), - ), patch( - "azure.cosmos.agent_memory._counters.advance_extract_watermark_async", - new=AsyncMock(), - ) as advance: + with ( + patch( + "azure.cosmos.agent_memory._counters.increment_counter_async", + return_value=counter_result, + ), + patch( + "azure.cosmos.agent_memory._counters.read_extract_watermark_async", + new=AsyncMock(return_value=watermark), + ), + patch( + "azure.cosmos.agent_memory._counters.advance_extract_watermark_async", + new=AsyncMock(), + ) as advance, + ): client.add_local(user_id="u1", role="user", thread_id="t1", content="hi") await client.push_to_cosmos() await asyncio.gather(*list(client._background_tasks), return_exceptions=True) @@ -211,19 +212,24 @@ async def test_watermark_not_advanced_when_extract_fails(self, monkeypatch): client._summaries_container_client = client._memories_container_client client._counter_container_client = MagicMock() - with patch( - "azure.cosmos.agent_memory._counters.increment_counter_async", - return_value=(0, 1), - ), patch( - "azure.cosmos.agent_memory._counters.read_extract_watermark_async", - new=AsyncMock(return_value=None), - ), patch( - "azure.cosmos.agent_memory._counters.advance_extract_watermark_async", - new=AsyncMock(), - ) as advance, patch( - "azure.cosmos.agent_memory._counters.stamp_failure_async", - new=AsyncMock(), - ) as stamp: + with ( + patch( + "azure.cosmos.agent_memory._counters.increment_counter_async", + return_value=(0, 1), + ), + patch( + "azure.cosmos.agent_memory._counters.read_extract_watermark_async", + new=AsyncMock(return_value=None), + ), + patch( + "azure.cosmos.agent_memory._counters.advance_extract_watermark_async", + new=AsyncMock(), + ) as advance, + patch( + "azure.cosmos.agent_memory._counters.stamp_failure_async", + new=AsyncMock(), + ) as stamp, + ): client.add_local(user_id="u1", role="user", thread_id="t1", content="hi") await client.push_to_cosmos() await asyncio.gather(*list(client._background_tasks), return_exceptions=True) @@ -280,9 +286,7 @@ async def test_reconcile_full_rebuild_on_persisted_counter_cadence(self, monkeyp monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") monkeypatch.setenv("DEDUP_EVERY_N", "1") - monkeypatch.setattr( - "azure.cosmos.agent_memory.thresholds.get_dedup_full_recluster_every_n", lambda: 2 - ) + monkeypatch.setattr("azure.cosmos.agent_memory.thresholds.get_dedup_full_recluster_every_n", lambda: 2) rebuilds: list[bool] = [] processor = AsyncInProcessProcessor(pipeline=MagicMock()) @@ -298,15 +302,19 @@ async def test_reconcile_full_rebuild_on_persisted_counter_cadence(self, monkeyp client._summaries_container_client = client._memories_container_client client._counter_container_client = MagicMock() - with patch( - "azure.cosmos.agent_memory._counters.increment_counter_async", - new=AsyncMock(side_effect=[(0, 1), (1, 2)]), - ), patch( - "azure.cosmos.agent_memory._counters.read_extract_watermark_async", - new=AsyncMock(return_value=None), - ), patch( - "azure.cosmos.agent_memory._counters.advance_extract_watermark_async", - new=AsyncMock(), + with ( + patch( + "azure.cosmos.agent_memory._counters.increment_counter_async", + new=AsyncMock(side_effect=[(0, 1), (1, 2)]), + ), + patch( + "azure.cosmos.agent_memory._counters.read_extract_watermark_async", + new=AsyncMock(return_value=None), + ), + patch( + "azure.cosmos.agent_memory._counters.advance_extract_watermark_async", + new=AsyncMock(), + ), ): client.add_local(user_id="u1", role="user", thread_id="t1", content="a") await client.push_to_cosmos() diff --git a/tests/unit/function_app/test_orchestrators.py b/tests/unit/function_app/test_orchestrators.py index 9fc629f..808e9e1 100644 --- a/tests/unit/function_app/test_orchestrators.py +++ b/tests/unit/function_app/test_orchestrators.py @@ -462,9 +462,7 @@ def test_em_advance_extract_watermark_stamps_counter(self): patch.object(cosmos_clients, "get_counter_container_async", new=AsyncMock(return_value=container)), patch.object(counters, "advance_extract_watermark", new=AsyncMock()) as advance, ): - result = asyncio.run( - em_mod.em_AdvanceExtractWatermark({"user_id": "u1", "thread_id": "t1", "count": 9}) - ) + result = asyncio.run(em_mod.em_AdvanceExtractWatermark({"user_id": "u1", "thread_id": "t1", "count": 9})) assert result is True advance.assert_awaited_once_with(container, "thread:u1:t1", "u1", "t1", 9) diff --git a/tests/unit/services/test_dedup_vector.py b/tests/unit/services/test_dedup_vector.py index a44f83a..6b4c4bd 100644 --- a/tests/unit/services/test_dedup_vector.py +++ b/tests/unit/services/test_dedup_vector.py @@ -24,12 +24,17 @@ def _make_pipeline() -> PipelineService: def _doc(mid: str, content: str, memory_type: str = "fact", **extra: Any) -> dict[str, Any]: tags = extra.pop("tags", [f"sys:{memory_type}"]) - metadata = extra.pop("metadata", {"category": "preference"} if memory_type == "fact" else { - "scope_type": "project", - "scope_value": "demo", - "lesson": content, - "outcome_valence": "neutral", - }) + metadata = extra.pop( + "metadata", + {"category": "preference"} + if memory_type == "fact" + else { + "scope_type": "project", + "scope_value": "demo", + "lesson": content, + "outcome_valence": "neutral", + }, + ) return { "id": mid, "user_id": "u1", @@ -101,9 +106,7 @@ def query_items(*, query: str, parameters, **kwargs): p._memories_container.query_items.side_effect = query_items p._distance_function_cache = "cosine" - out = p._vector_candidates( - user_id="u1", embedding=[1.0, 0.0], memory_type="fact", top_k=2, exclude_ids=set() - ) + out = p._vector_candidates(user_id="u1", embedding=[1.0, 0.0], memory_type="fact", top_k=2, exclude_ids=set()) # Cosmos rejects an explicit ASC/DESC on ORDER BY VectorDistance(); it orders # most-similar-first server-side. Direction-awareness lives in the Python sort. assert "ORDER BY VectorDistance(c.embedding, @vec)" in captured["query"] @@ -112,9 +115,7 @@ def query_items(*, query: str, parameters, **kwargs): assert [c["id"] for c in out] == ["near", "far"] p._distance_function_cache = "euclidean" - out = p._vector_candidates( - user_id="u1", embedding=[1.0, 0.0], memory_type="fact", top_k=2, exclude_ids=set() - ) + out = p._vector_candidates(user_id="u1", embedding=[1.0, 0.0], memory_type="fact", top_k=2, exclude_ids=set()) assert "VectorDistance(c.embedding, @vec) ASC" not in captured["query"] # euclidean: lower distance = more similar, so 0.10 ("far" label) sorts first. assert [c["id"] for c in out] == ["far", "near"] diff --git a/tests/unit/test_auto_trigger.py b/tests/unit/test_auto_trigger.py index 4bb73e7..2f23420 100644 --- a/tests/unit/test_auto_trigger.py +++ b/tests/unit/test_auto_trigger.py @@ -226,15 +226,19 @@ def test_extract_recent_k_uses_watermark_then_falls_back( client = _connected(processor=processor) client._counter_container_client = MagicMock() - with patch( - "azure.cosmos.agent_memory._counters.increment_counter_sync", - return_value=counter_result, - ), patch( - "azure.cosmos.agent_memory._counters.read_extract_watermark_sync", - return_value=watermark, - ), patch( - "azure.cosmos.agent_memory._counters.advance_extract_watermark_sync", - ) as advance: + with ( + patch( + "azure.cosmos.agent_memory._counters.increment_counter_sync", + return_value=counter_result, + ), + patch( + "azure.cosmos.agent_memory._counters.read_extract_watermark_sync", + return_value=watermark, + ), + patch( + "azure.cosmos.agent_memory._counters.advance_extract_watermark_sync", + ) as advance, + ): for i in range(batch_count): client.add_local(user_id="u1", role="user", thread_id="t1", content=f"hi {i}") client.push_to_cosmos() @@ -300,17 +304,22 @@ def test_watermark_not_advanced_when_extract_fails(self, monkeypatch): client = _connected(processor=processor) client._counter_container_client = MagicMock() - with patch( - "azure.cosmos.agent_memory._counters.increment_counter_sync", - return_value=(0, 1), - ), patch( - "azure.cosmos.agent_memory._counters.read_extract_watermark_sync", - return_value=None, - ), patch( - "azure.cosmos.agent_memory._counters.advance_extract_watermark_sync", - ) as advance, patch( - "azure.cosmos.agent_memory._counters.stamp_failure_sync", - ) as stamp: + with ( + patch( + "azure.cosmos.agent_memory._counters.increment_counter_sync", + return_value=(0, 1), + ), + patch( + "azure.cosmos.agent_memory._counters.read_extract_watermark_sync", + return_value=None, + ), + patch( + "azure.cosmos.agent_memory._counters.advance_extract_watermark_sync", + ) as advance, + patch( + "azure.cosmos.agent_memory._counters.stamp_failure_sync", + ) as stamp, + ): client.add_local(user_id="u1", role="user", thread_id="t1", content="hi") client.push_to_cosmos() @@ -326,9 +335,7 @@ def test_reconcile_full_rebuild_on_persisted_counter_cadence(self, monkeypatch): monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") monkeypatch.setenv("DEDUP_EVERY_N", "1") - monkeypatch.setattr( - "azure.cosmos.agent_memory.thresholds.get_dedup_full_recluster_every_n", lambda: 2 - ) + monkeypatch.setattr("azure.cosmos.agent_memory.thresholds.get_dedup_full_recluster_every_n", lambda: 2) rebuilds: list[bool] = [] processor = InProcessProcessor(pipeline=MagicMock()) @@ -341,14 +348,18 @@ def test_reconcile_full_rebuild_on_persisted_counter_cadence(self, monkeypatch): client = _connected(processor=processor) client._counter_container_client = MagicMock() - with patch( - "azure.cosmos.agent_memory._counters.increment_counter_sync", - side_effect=[(0, 1), (1, 2)], - ), patch( - "azure.cosmos.agent_memory._counters.read_extract_watermark_sync", - return_value=None, - ), patch( - "azure.cosmos.agent_memory._counters.advance_extract_watermark_sync", + with ( + patch( + "azure.cosmos.agent_memory._counters.increment_counter_sync", + side_effect=[(0, 1), (1, 2)], + ), + patch( + "azure.cosmos.agent_memory._counters.read_extract_watermark_sync", + return_value=None, + ), + patch( + "azure.cosmos.agent_memory._counters.advance_extract_watermark_sync", + ), ): client.add_local(user_id="u1", role="user", thread_id="t1", content="a") client.push_to_cosmos() # counter 0->1: reconcile, full crosses 2? no From 8d3c0fd9d3db08a2edc18dabaeeb01d5b57bf8db Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Wed, 1 Jul 2026 18:24:05 -0700 Subject: [PATCH 3/8] Fixing embedding token limits --- .../cosmos/agent_memory/_embedding_tokens.py | 104 +++++++++++++++ azure/cosmos/agent_memory/aio/embeddings.py | 2 + .../agent_memory/aio/services/pipeline.py | 20 ++- azure/cosmos/agent_memory/embeddings.py | 2 + .../cosmos/agent_memory/services/pipeline.py | 20 ++- pyproject.toml | 1 + tests/unit/services/test_extract_dry.py | 77 +++++++++++ tests/unit/test_embedding_tokens.py | 123 ++++++++++++++++++ 8 files changed, 337 insertions(+), 12 deletions(-) create mode 100644 azure/cosmos/agent_memory/_embedding_tokens.py create mode 100644 tests/unit/test_embedding_tokens.py diff --git a/azure/cosmos/agent_memory/_embedding_tokens.py b/azure/cosmos/agent_memory/_embedding_tokens.py new file mode 100644 index 0000000..6917118 --- /dev/null +++ b/azure/cosmos/agent_memory/_embedding_tokens.py @@ -0,0 +1,104 @@ +"""Token-budget helpers for embedding inputs. + +Azure OpenAI embedding models (``text-embedding-3-*``) reject any single +input that exceeds their 8192-token context window. Callers in this toolkit +embed user-supplied and document-derived text of unbounded size — the +dedup-context query in extraction, ``search_turns`` queries, and stored turn +vectors when ``enable_turn_embeddings`` is on — so an oversized string would +otherwise crash the embed call or, at the write sites that swallow embed +errors, silently persist a record with no vector. + +Every embedding input is therefore capped to a safe token budget before the +API call. Truncation keeps the *head* of the text; preserving the tail of a +large document is a chunking concern handled at the ingestion layer, not +here (the embedder contract is one string -> one vector, never pooled). +""" + +from __future__ import annotations + +from functools import lru_cache +from typing import Any + +from azure.cosmos.agent_memory.logging import get_logger + +logger = get_logger(__name__) + +# ``text-embedding-3-*`` accept up to 8192 tokens per input. Leave headroom +# below the hard cap for tokenizer drift across model/deployment names and +# any special tokens the service may add server-side. +EMBEDDING_MAX_INPUT_TOKENS = 8000 + +# Defensive fallback used only when tiktoken cannot be imported (it is a +# declared dependency). English text averages ~4 characters per token. +_FALLBACK_CHARS_PER_TOKEN = 4 + + +@lru_cache(maxsize=8) +def _encoding_for_model(model: str) -> Any: + """Return a cached tiktoken encoding for *model*. + + Falls back to the ``cl100k_base`` encoding used by ``text-embedding-3-*`` + when the model name is unknown, and to ``None`` when tiktoken itself is + unavailable so the caller can apply a character-based estimate. + """ + try: + import tiktoken + except Exception: # pragma: no cover - tiktoken is a declared dependency + return None + try: + return tiktoken.encoding_for_model(model) + except Exception: + try: + return tiktoken.get_encoding("cl100k_base") + except Exception: # pragma: no cover - defensive + return None + + +def truncate_text_to_token_budget( + text: str, + model: str, + max_tokens: int = EMBEDDING_MAX_INPUT_TOKENS, +) -> str: + """Truncate *text* so it encodes to at most *max_tokens* tokens. + + Returns *text* unchanged when it is already within budget. Truncation + keeps the leading tokens and logs a warning so oversized inputs remain + observable — that warning is also the signal for where ingestion-time + chunking is needed. + """ + if not text: + return text + + # Fast path: every token spans at least one character, so any string + # whose character length is within the token budget cannot exceed it. + # This skips tokenization for the overwhelming majority of inputs + # (short turns and queries). + if len(text) <= max_tokens: + return text + + encoding = _encoding_for_model(model) + if encoding is None: + max_chars = max_tokens * _FALLBACK_CHARS_PER_TOKEN + if len(text) <= max_chars: + return text + logger.warning( + "Embedding input exceeds ~%d tokens; tiktoken unavailable, " + "truncating to %d chars (model=%s, original chars=%d)", + max_tokens, + max_chars, + model, + len(text), + ) + return text[:max_chars] + + tokens = encoding.encode(text) + if len(tokens) <= max_tokens: + return text + truncated = encoding.decode(tokens[:max_tokens]) + logger.warning( + "Embedding input truncated from %d to %d tokens (model=%s) to stay within the embedding context window", + len(tokens), + max_tokens, + model, + ) + return truncated diff --git a/azure/cosmos/agent_memory/aio/embeddings.py b/azure/cosmos/agent_memory/aio/embeddings.py index 503de8d..da29cbd 100644 --- a/azure/cosmos/agent_memory/aio/embeddings.py +++ b/azure/cosmos/agent_memory/aio/embeddings.py @@ -10,6 +10,7 @@ import asyncio from typing import Any +from azure.cosmos.agent_memory._embedding_tokens import truncate_text_to_token_budget from azure.cosmos.agent_memory.chat import resolve_api_version from azure.cosmos.agent_memory.exceptions import ConfigurationError from azure.cosmos.agent_memory.logging import get_logger @@ -94,6 +95,7 @@ def _ensure_client(self) -> Any: def _build_kwargs(self, input_: str | list[str]) -> dict[str, Any]: texts = [input_] if isinstance(input_, str) else input_ + texts = [truncate_text_to_token_budget(t, self._model) for t in texts] logger.debug( "Embedding request: model=%s, dimensions=%s, texts=%d", self._model, diff --git a/azure/cosmos/agent_memory/aio/services/pipeline.py b/azure/cosmos/agent_memory/aio/services/pipeline.py index ba8798f..83b6d43 100644 --- a/azure/cosmos/agent_memory/aio/services/pipeline.py +++ b/azure/cosmos/agent_memory/aio/services/pipeline.py @@ -532,12 +532,20 @@ async def extract_memories_dry( if get_dedup_context_vector_enabled(): user_turns_text = "\n".join(str(it.get("content", "")) for it in items if it.get("role") == "user").strip() context_query = user_turns_text or transcript - existing = await self._store.search( - search_terms=context_query, - user_id=user_id, - memory_types=["fact"], - top_k=get_dedup_context_topk(), - ) + try: + existing = await self._store.search( + search_terms=context_query, + user_id=user_id, + memory_types=["fact"], + top_k=get_dedup_context_topk(), + ) + except Exception as exc: # noqa: BLE001 + logger.warning( + "extract_memories_dry dedup-context vector search failed (%s); " + "falling back to hash-based existing memories", + exc, + ) + existing = existing_for_hash else: existing = existing_for_hash if existing: diff --git a/azure/cosmos/agent_memory/embeddings.py b/azure/cosmos/agent_memory/embeddings.py index 493c25f..1b7b4dd 100644 --- a/azure/cosmos/agent_memory/embeddings.py +++ b/azure/cosmos/agent_memory/embeddings.py @@ -11,6 +11,7 @@ from azure.cosmos.agent_memory.chat import resolve_api_version from azure.cosmos.agent_memory.logging import get_logger +from ._embedding_tokens import truncate_text_to_token_budget from .exceptions import ConfigurationError logger = get_logger(__name__) @@ -94,6 +95,7 @@ def _ensure_client(self) -> Any: def _build_kwargs(self, input_: str | list[str]) -> dict[str, Any]: texts = [input_] if isinstance(input_, str) else input_ + texts = [truncate_text_to_token_budget(t, self._model) for t in texts] logger.debug( "Embedding request: model=%s, dimensions=%s, texts=%d", self._model, diff --git a/azure/cosmos/agent_memory/services/pipeline.py b/azure/cosmos/agent_memory/services/pipeline.py index a44f87f..5b9616b 100644 --- a/azure/cosmos/agent_memory/services/pipeline.py +++ b/azure/cosmos/agent_memory/services/pipeline.py @@ -600,12 +600,20 @@ def extract_memories_dry( if threshold_config.get_dedup_context_vector_enabled(): user_turns_text = "\n".join(str(it.get("content", "")) for it in items if it.get("role") == "user").strip() context_query = user_turns_text or transcript - existing = self._store.search( - search_terms=context_query, - user_id=user_id, - memory_types=["fact"], - top_k=threshold_config.get_dedup_context_topk(), - ) + try: + existing = self._store.search( + search_terms=context_query, + user_id=user_id, + memory_types=["fact"], + top_k=threshold_config.get_dedup_context_topk(), + ) + except Exception as exc: # noqa: BLE001 + logger.warning( + "extract_memories_dry dedup-context vector search failed (%s); " + "falling back to hash-based existing memories", + exc, + ) + existing = existing_for_hashes if existing: existing_text = "\n".join( f"- [ID: {mem['id']}] {mem.get('content', '')} " diff --git a/pyproject.toml b/pyproject.toml index 2ead3b0..f6c7a56 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,7 @@ dependencies = [ "azure-identity>=1.20", "openai>=1.60", "aiohttp>=3.10", + "tiktoken>=0.7", "pydantic>=2.10", "prompty>=2.0.0a9", "jinja2>=3.1.4", diff --git a/tests/unit/services/test_extract_dry.py b/tests/unit/services/test_extract_dry.py index c5c9adc..635bcf9 100644 --- a/tests/unit/services/test_extract_dry.py +++ b/tests/unit/services/test_extract_dry.py @@ -476,3 +476,80 @@ async def test_async_extract_memories_dry_stage1_legacy_path_when_context_vector await service.extract_memories_dry("u1", "t1") store.search.assert_not_awaited() + + +def test_extract_memories_dry_stage1_search_failure_falls_back_to_hash_memories() -> None: + """A failing dedup-context vector search must not abort extraction; it + falls back to the hash-loaded existing memories (option 3 resilience).""" + chat = _SyncChat([_response()]) + memories_store = _Store( + [ + { + "id": "hash-memory", + "user_id": "u1", + "type": "fact", + "content": "Existing hash-based memory from load.", + "content_hash": "h1", + "salience": 0.6, + } + ] + ) + + def _boom(**_kwargs): + raise RuntimeError("vector search down") + + memories_store.search = _boom + service = PipelineService( + memories_store, + chat, + _SyncEmbeddings(), + containers=_containers_for_store(memories_store, turns_store=_Store([_turn(1)])), + ) + + # Must not raise despite the Stage-1 vector search failing. + output = service.extract_memories_dry("u1", "t1") + + assert output["facts"] # extraction proceeded + # Fallback surfaced the hash-loaded memory into the extraction prompt. + assert "Existing hash-based memory from load." in json.dumps(chat.messages) + + +@pytest.mark.asyncio +async def test_async_extract_memories_dry_stage1_search_failure_falls_back_to_hash_memories() -> None: + class _RecordingAsyncChat(_AsyncChat): + def __init__(self, responses): + super().__init__(responses) + self.messages: list = [] + + async def generate(self, messages, **opts): + self.messages.append(messages) + del opts + self.calls += 1 + return json.dumps(self.responses.pop(0)) + + chat = _RecordingAsyncChat([_response()]) + store = _AsyncStore( + [ + { + "id": "hash-memory", + "user_id": "u1", + "type": "fact", + "content": "Existing hash-based memory from load.", + "content_hash": "h1", + "salience": 0.6, + } + ] + ) + store.search = AsyncMock(side_effect=RuntimeError("vector search down")) + service = AsyncPipelineService( + store, + chat, + _AsyncEmbeddings(), + containers=_async_containers_for_store(store, turns_store=_AsyncStore([_turn(1)])), + ) + + output = await service.extract_memories_dry("u1", "t1") + + assert output["facts"] # extraction proceeded + store.search.assert_awaited_once() + assert "Existing hash-based memory from load." in json.dumps(chat.messages) diff --git a/tests/unit/test_embedding_tokens.py b/tests/unit/test_embedding_tokens.py new file mode 100644 index 0000000..6641e07 --- /dev/null +++ b/tests/unit/test_embedding_tokens.py @@ -0,0 +1,123 @@ +"""Unit tests for embedding token-budget truncation (option 2 safety net).""" + +from __future__ import annotations + +import tiktoken + +from azure.cosmos.agent_memory._embedding_tokens import ( + EMBEDDING_MAX_INPUT_TOKENS, + truncate_text_to_token_budget, +) +from azure.cosmos.agent_memory.aio.embeddings import AsyncEmbeddingsClient +from azure.cosmos.agent_memory.embeddings import EmbeddingsClient + +_MODEL = "text-embedding-3-large" + + +def _token_count(text: str) -> int: + return len(tiktoken.get_encoding("cl100k_base").encode(text)) + + +# --------------------------------------------------------------------------- +# truncate_text_to_token_budget() +# --------------------------------------------------------------------------- + + +class TestTruncateHelper: + def test_empty_string_unchanged(self) -> None: + assert truncate_text_to_token_budget("", _MODEL) == "" + + def test_short_text_unchanged(self) -> None: + text = "The user prefers dark mode." + assert truncate_text_to_token_budget(text, _MODEL) is text + + def test_text_within_budget_unchanged(self) -> None: + # ~1000 tokens, comfortably under the 8000 budget. + text = "word " * 1000 + assert truncate_text_to_token_budget(text, _MODEL) is text + + def test_oversized_text_truncated_within_budget(self) -> None: + # ``supercalifragilistic`` tokenizes to several tokens each, so this + # is well over the 8000-token budget. + text = "supercalifragilistic " * 3000 + assert _token_count(text) > EMBEDDING_MAX_INPUT_TOKENS + + result = truncate_text_to_token_budget(text, _MODEL) + + assert _token_count(result) <= EMBEDDING_MAX_INPUT_TOKENS + assert len(result) < len(text) + + def test_truncation_preserves_head(self) -> None: + head = "UNIQUE_HEAD_MARKER the user asked about billing. " + text = head + ("filler token stream " * 5000) + + result = truncate_text_to_token_budget(text, _MODEL) + + assert result.startswith("UNIQUE_HEAD_MARKER") + + def test_custom_max_tokens(self) -> None: + text = "word " * 1000 + result = truncate_text_to_token_budget(text, _MODEL, max_tokens=50) + assert _token_count(result) <= 50 + + def test_unknown_model_falls_back_to_cl100k(self) -> None: + text = "supercalifragilistic " * 3000 + # An unknown deployment name must still truncate (via cl100k_base). + result = truncate_text_to_token_budget(text, "some-custom-deployment") + assert _token_count(result) <= EMBEDDING_MAX_INPUT_TOKENS + + def test_char_fallback_when_tiktoken_unavailable(self, monkeypatch) -> None: + # Simulate tiktoken being unavailable so the char-based estimate runs. + monkeypatch.setattr( + "azure.cosmos.agent_memory._embedding_tokens._encoding_for_model", + lambda model: None, + ) + # 4 chars/token fallback → budget * 4 chars. Use a single long token-free + # string longer than that to force truncation. + text = "a" * (EMBEDDING_MAX_INPUT_TOKENS * 4 + 1000) + result = truncate_text_to_token_budget(text, _MODEL) + assert len(result) == EMBEDDING_MAX_INPUT_TOKENS * 4 + + def test_char_fallback_leaves_short_text_untouched(self, monkeypatch) -> None: + monkeypatch.setattr( + "azure.cosmos.agent_memory._embedding_tokens._encoding_for_model", + lambda model: None, + ) + # Longer than the token budget in chars, but within the char budget. + text = "a" * (EMBEDDING_MAX_INPUT_TOKENS + 100) + assert truncate_text_to_token_budget(text, _MODEL) is text + + +# --------------------------------------------------------------------------- +# Client wiring: _build_kwargs applies truncation on every embed path +# --------------------------------------------------------------------------- + + +def _oversized() -> str: + return "supercalifragilistic " * 3000 + + +class TestSyncClientWiring: + def test_build_kwargs_truncates_single_input(self) -> None: + client = EmbeddingsClient(endpoint="https://x", api_key="k", model=_MODEL) + kwargs = client._build_kwargs(_oversized()) + assert _token_count(kwargs["input"][0]) <= EMBEDDING_MAX_INPUT_TOKENS + + def test_build_kwargs_truncates_each_item_in_batch(self) -> None: + client = EmbeddingsClient(endpoint="https://x", api_key="k", model=_MODEL) + kwargs = client._build_kwargs([_oversized(), "small text", _oversized()]) + assert all(_token_count(t) <= EMBEDDING_MAX_INPUT_TOKENS for t in kwargs["input"]) + assert kwargs["input"][1] == "small text" + + +class TestAsyncClientWiring: + def test_build_kwargs_truncates_single_input(self) -> None: + client = AsyncEmbeddingsClient(endpoint="https://x", api_key="k", model=_MODEL) + kwargs = client._build_kwargs(_oversized()) + assert _token_count(kwargs["input"][0]) <= EMBEDDING_MAX_INPUT_TOKENS + + def test_build_kwargs_truncates_each_item_in_batch(self) -> None: + client = AsyncEmbeddingsClient(endpoint="https://x", api_key="k", model=_MODEL) + kwargs = client._build_kwargs([_oversized(), "small text"]) + assert all(_token_count(t) <= EMBEDDING_MAX_INPUT_TOKENS for t in kwargs["input"]) + assert kwargs["input"][1] == "small text" From 72149bc068182d385d5831f7d51e3175bb672abb Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Thu, 2 Jul 2026 15:10:44 -0700 Subject: [PATCH 4/8] Fixing schema parsing --- .../prompts/extract_memories.prompty | 10 +-- .../services/_pipeline_helpers.py | 14 +++- tests/unit/services/test_parse_llm_json.py | 80 +++++++++++++++++++ 3 files changed, 98 insertions(+), 6 deletions(-) create mode 100644 tests/unit/services/test_parse_llm_json.py diff --git a/azure/cosmos/agent_memory/prompts/extract_memories.prompty b/azure/cosmos/agent_memory/prompts/extract_memories.prompty index 723b458..af1b25a 100644 --- a/azure/cosmos/agent_memory/prompts/extract_memories.prompty +++ b/azure/cosmos/agent_memory/prompts/extract_memories.prompty @@ -19,23 +19,23 @@ inputs: --- system: -You are a precision memory extraction system. Your task is to read a conversation thread and extract structured memories worth retaining for future reference. You extract two distinct types of memory, each serving a different purpose. +You are a precision memory extraction system. Your task is to read the provided content — a conversation thread, and/or documents or reference material the user has shared — and extract structured memories worth retaining for future reference. You extract two distinct types of memory, each serving a different purpose. ## Memory Type Taxonomy -1. **Facts** — Declarative knowledge: concrete, verifiable statements about the user, their environment, decisions, or requirements. These are things that *are* true. +1. **Facts** — Declarative knowledge: concrete, verifiable statements about the user, their environment, decisions, or requirements — or salient factual information from the content the user is working with (documents, excerpts, or domain material they provide). These are things that *are* true. 2. **Episodic** — Past experiences: specific situations the user encountered, what they tried, and what happened. These are things that *happened*. 3. **Unclassified** — Use sparingly, only when an item is clearly worth retaining but you cannot confidently decide between fact / episodic. The pipeline will store these as facts with a `sys:unclassified` tag for later audit. Prefer this over forcing a wrong classification. -Every memory must be explicitly grounded in the conversation — never inferred, assumed, or speculated. +Every memory must be explicitly grounded in the provided content — never inferred, assumed, or speculated. ## Speaker Discrimination — Where Memories May Come From The transcript below is line-tagged: `[user]:` lines are the human's own words, `[agent]:` lines are the agent's response. These two sources are NOT interchangeable. -- **Facts about the user may ONLY come from `[user]:` lines.** The agent may restate, paraphrase, confirm, or make assumptions on the user's behalf ("Got it, you don't eat meat", "I assume you want a luxury hotel") — those are the agent's response, NOT the user's assertion. Never treat agent text as a new source of user facts. If a fact appears only in `[agent]:` content and is not asserted by the user, do not extract it. +- **Facts may ONLY come from `[user]:` lines.** The agent may restate, paraphrase, confirm, or make assumptions on the user's behalf ("Got it, you don't eat meat", "I assume you want a luxury hotel") — those are the agent's response, NOT the user's assertion. Never treat agent text as a new source of facts. If a fact appears only in `[agent]:` content and is not present in the user's lines, do not extract it. - **Episodic memories may use both speakers' content** — the user's stated intent or scope is the anchor (and must be present in `[user]:`), but the agent's content may help fill in the `action_taken` or `outcome` of a `situation → action_taken → outcome` arc when the agent carried out the action on the user's behalf. -- The agent's general world-knowledge answers (e.g. "Python 3.13 was released in October 2024") are never user facts and never user episodic memories. +- Unsolicited world-knowledge the agent volunteers in `[agent]:` lines (e.g. "Python 3.13 was released in October 2024") is the agent's answer, not a memory. However, factual and domain content the user provides, shares, or asks to remember — including documents, excerpts, and reference material in `[user]:` lines — IS extractable. ## Confidence Scoring diff --git a/azure/cosmos/agent_memory/services/_pipeline_helpers.py b/azure/cosmos/agent_memory/services/_pipeline_helpers.py index e94ce91..74b13d6 100644 --- a/azure/cosmos/agent_memory/services/_pipeline_helpers.py +++ b/azure/cosmos/agent_memory/services/_pipeline_helpers.py @@ -17,6 +17,9 @@ from typing import Any, Iterable, Optional from azure.cosmos.agent_memory.exceptions import LLMError +from azure.cosmos.agent_memory.logging import get_logger + +logger = get_logger(__name__) # Separator for deterministic id seeds. Using NUL ensures user_id / # thread_id values can never collide with literal section markers @@ -463,7 +466,7 @@ def parse_llm_json(text: str | None) -> dict[str, Any]: cleaned = cleaned[:-3] cleaned = cleaned.strip() try: - return json.loads(cleaned) + obj, end = json.JSONDecoder().raw_decode(cleaned) except json.JSONDecodeError as exc: preview = (text or "")[:200].replace("\n", " ") if _looks_truncated(cleaned, exc): @@ -475,6 +478,15 @@ def parse_llm_json(text: str | None) -> dict[str, Any]: f"recent_k, or split oversized turns). Decode error: {exc}. preview={preview!r}" ) from exc raise LLMError(f"LLM returned invalid JSON (preview={preview!r}): {exc}") from exc + trailing = cleaned[end:].strip() + if trailing: + logger.warning( + "LLM response had %d chars of extra data after the first JSON object; using the " + "first object and ignoring the remainder (trailing_preview=%r)", + len(trailing), + trailing[:120].replace("\n", " "), + ) + return obj def _looks_truncated(cleaned: str, exc: json.JSONDecodeError) -> bool: diff --git a/tests/unit/services/test_parse_llm_json.py b/tests/unit/services/test_parse_llm_json.py new file mode 100644 index 0000000..9352c75 --- /dev/null +++ b/tests/unit/services/test_parse_llm_json.py @@ -0,0 +1,80 @@ +"""Unit tests for parse_llm_json — resilient LLM JSON decoding (Issue A).""" + +from __future__ import annotations + +import logging + +import pytest + +from azure.cosmos.agent_memory.exceptions import LLMError +from azure.cosmos.agent_memory.services._pipeline_helpers import parse_llm_json + +_HELPER_LOGGER = "azure.cosmos.agent_memory.services._pipeline_helpers" + +# The exact string gpt-5.4 produced on the Ruler dataset: a valid object +# followed by a concatenated duplicate of itself. +_DOUBLED = '{"facts":[],"episodic":[],"unclassified":[]} {"facts":[],"episodic":[],"unclassified":[]}' + + +class TestHappyPath: + def test_plain_object(self) -> None: + assert parse_llm_json('{"facts": [], "episodic": []}') == {"facts": [], "episodic": []} + + def test_markdown_fenced_object(self) -> None: + assert parse_llm_json('```json\n{"facts": [1]}\n```') == {"facts": [1]} + + def test_object_with_surrounding_whitespace(self) -> None: + assert parse_llm_json(' \n {"a": 1}\n ') == {"a": 1} + + +class TestDoubledAndTrailing: + def test_doubled_object_returns_first(self) -> None: + # Previously raised "Extra data" and lost the whole extraction. + assert parse_llm_json(_DOUBLED) == {"facts": [], "episodic": [], "unclassified": []} + + def test_object_then_garbage_is_salvaged(self) -> None: + assert parse_llm_json('{"facts": [{"text": "x"}]} ') == {"facts": [{"text": "x"}]} + + def test_fenced_object_with_trailing_duplicate(self) -> None: + assert parse_llm_json('```json\n{"a": 1}\n``` {"a": 1}') == {"a": 1} + + def test_trailing_data_emits_warning(self, caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger=_HELPER_LOGGER): + parse_llm_json(_DOUBLED) + assert any("extra data after the first JSON object" in r.message for r in caplog.records) + + def test_clean_object_emits_no_warning(self, caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger=_HELPER_LOGGER): + parse_llm_json('{"facts": []}') + assert not any("extra data" in r.message for r in caplog.records) + + +class TestTruncationDetectionPreserved: + def test_unbalanced_braces_flagged_truncated(self) -> None: + with pytest.raises(LLMError) as exc: + parse_llm_json('{"facts": [{"text": "the user prefers') + assert "TRUNCATED" in str(exc.value) + + def test_unterminated_string_flagged_truncated(self) -> None: + with pytest.raises(LLMError) as exc: + parse_llm_json('{"facts": [], "note": "unterminated') + assert "TRUNCATED" in str(exc.value) + + +class TestGenuineErrors: + def test_none_raises(self) -> None: + with pytest.raises(LLMError) as exc: + parse_llm_json(None) + assert "no content" in str(exc.value) + + def test_empty_string_raises_invalid_not_truncated(self) -> None: + with pytest.raises(LLMError) as exc: + parse_llm_json("") + assert "invalid JSON" in str(exc.value) + assert "TRUNCATED" not in str(exc.value) + + def test_non_json_raises_invalid_not_truncated(self) -> None: + with pytest.raises(LLMError) as exc: + parse_llm_json("I could not find any memories.") + assert "invalid JSON" in str(exc.value) + assert "TRUNCATED" not in str(exc.value) From 8cd0275d887e816453bcf40386563c67535caf65 Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Fri, 10 Jul 2026 20:48:42 -0700 Subject: [PATCH 5/8] improving extraction prompt --- Samples/Advanced/advanced_memory_lifecycle.py | 8 +- Samples/Advanced/advanced_search_patterns.py | 22 +- .../Processing/processing_fact_extraction.py | 2 +- .../Processing/processing_thread_summary.py | 4 +- Samples/Processing/processing_user_profile.py | 14 +- Samples/Scenarios/scenario_chat_memory.py | 10 +- Samples/Scenarios/scenario_counter_tuning.py | 6 +- .../Scenarios/scenario_customer_support.py | 16 +- .../scenario_memory_reconciliation.py | 6 +- Samples/Scenarios/scenario_multi_agent.py | 24 +- Samples/Scenarios/scenario_rag_with_memory.py | 18 +- .../Scenarios/scenario_remote_processor.py | 10 +- .../scenario_remote_processor_async.py | 2 +- .../scenario_tagging_and_filtering.py | 24 +- .../cosmos/agent_memory/_base/base_client.py | 4 +- azure/cosmos/agent_memory/_counters.py | 90 +- .../cosmos/agent_memory/_embedding_tokens.py | 21 +- azure/cosmos/agent_memory/_utils.py | 6 +- azure/cosmos/agent_memory/aio/auto_trigger.py | 46 +- azure/cosmos/agent_memory/aio/chat.py | 6 +- .../agent_memory/aio/cosmos_memory_client.py | 33 +- azure/cosmos/agent_memory/aio/embeddings.py | 2 +- .../agent_memory/aio/processors/base.py | 1 - .../agent_memory/aio/processors/durable.py | 2 +- .../agent_memory/aio/processors/inprocess.py | 20 +- .../agent_memory/aio/services/pipeline.py | 1063 ++++------------ .../agent_memory/aio/store/memory_store.py | 62 +- azure/cosmos/agent_memory/auto_trigger.py | 49 +- azure/cosmos/agent_memory/chat.py | 10 +- .../agent_memory/cosmos_memory_client.py | 33 +- azure/cosmos/agent_memory/embeddings.py | 2 +- azure/cosmos/agent_memory/models.py | 11 +- azure/cosmos/agent_memory/processors/base.py | 1 - .../cosmos/agent_memory/processors/durable.py | 2 +- .../agent_memory/processors/inprocess.py | 21 +- azure/cosmos/agent_memory/prompts/_schemas.py | 63 +- .../cosmos/agent_memory/prompts/dedup.prompty | 82 +- .../prompts/extract_memories.prompty | 221 ++-- .../services/_pipeline_helpers.py | 108 +- .../cosmos/agent_memory/services/pipeline.py | 1110 +++++------------ .../agent_memory/store/_search_helpers.py | 5 +- .../cosmos/agent_memory/store/memory_store.py | 63 +- azure/cosmos/agent_memory/thresholds.py | 73 +- .../orchestrators/extract_memories.py | 13 +- function_app/shared/config.py | 12 +- function_app/shared/cosmos_clients.py | 4 +- function_app/shared/counters.py | 18 +- function_app/shared/pipeline_factory.py | 2 +- function_app/triggers/change_feed.py | 25 +- pyproject.toml | 2 +- tests/conftest.py | 2 +- tests/integration/test_async_full_pipeline.py | 6 +- .../test_changefeed_integration.py | 2 +- tests/integration/test_full_pipeline.py | 10 +- .../integration/test_processor_integration.py | 2 +- .../test_processor_integration_async.py | 4 +- tests/unit/aio/processors/test_durable.py | 2 +- tests/unit/aio/processors/test_inprocess.py | 4 +- .../aio/services/test_dedup_vector_async.py | 396 ++---- .../unit/aio/store/test_get_memory_history.py | 69 + tests/unit/aio/test_auto_trigger.py | 188 +-- tests/unit/aio/test_cosmos_memory_client.py | 2 +- tests/unit/aio/test_embeddings.py | 12 +- tests/unit/aio/test_procedural_synthesis.py | 2 +- tests/unit/aio/test_process_now.py | 6 +- tests/unit/aio/test_reconcile_telemetry.py | 4 +- tests/unit/conftest.py | 4 +- tests/unit/function_app/conftest.py | 2 +- tests/unit/function_app/test_change_feed.py | 58 +- tests/unit/function_app/test_counters.py | 14 +- tests/unit/function_app/test_orchestrators.py | 23 +- tests/unit/function_app/test_retry.py | 4 +- tests/unit/processors/test_durable.py | 2 +- tests/unit/processors/test_inprocess.py | 2 +- .../test_schema_prompty_conformance.py | 6 +- .../services/test_chaos_extract_persist.py | 8 - tests/unit/services/test_dedup_vector.py | 339 +++-- tests/unit/services/test_extract_dry.py | 303 ++--- .../unit/services/test_extraction_batching.py | 58 + tests/unit/services/test_parse_llm_json.py | 2 +- tests/unit/services/test_pipeline_service.py | 24 +- .../unit/services/test_transcript_metadata.py | 6 +- tests/unit/store/test_get_memory_history.py | 150 +++ tests/unit/store/test_memory_store.py | 2 +- tests/unit/test_async_hygiene.py | 6 +- tests/unit/test_auto_trigger.py | 166 +-- tests/unit/test_cosmos_memory_client.py | 14 +- tests/unit/test_counters.py | 6 +- tests/unit/test_embeddings.py | 4 +- tests/unit/test_memory_type_multi.py | 4 +- tests/unit/test_models.py | 2 +- tests/unit/test_pipeline_confidence.py | 33 - tests/unit/test_procedural_synthesis.py | 15 +- tests/unit/test_process_now.py | 14 +- tests/unit/test_reconcile.py | 722 +---------- tests/unit/test_thresholds.py | 18 +- 96 files changed, 1972 insertions(+), 4202 deletions(-) create mode 100644 tests/unit/aio/store/test_get_memory_history.py create mode 100644 tests/unit/services/test_extraction_batching.py create mode 100644 tests/unit/store/test_get_memory_history.py diff --git a/Samples/Advanced/advanced_memory_lifecycle.py b/Samples/Advanced/advanced_memory_lifecycle.py index c45541f..47ec5a5 100644 --- a/Samples/Advanced/advanced_memory_lifecycle.py +++ b/Samples/Advanced/advanced_memory_lifecycle.py @@ -1,10 +1,10 @@ -"""Advanced memory lifecycle — create → process → archive (delete raw turns). +"""Advanced memory lifecycle - create → process → archive (delete raw turns). Demonstrates the typical long-term-memory flow: 1. Add raw conversation turns 2. Extract structured memories (facts / procedural / episodic) 3. Generate a thread summary - 4. Delete raw turns — keeping only the compact derived memories + 4. Delete raw turns - keeping only the compact derived memories Uses the in-process ProcessingPipeline (same code as the Azure Function change-feed trigger). No Function deployment required. @@ -72,7 +72,7 @@ def main() -> None: ("user", "Hybrid. I want to use embeddings on book descriptions and reviews."), ("agent", "Cosmos DB for NoSQL with the vector index works well for that."), ("user", "I prefer Python and want to use FastAPI for the API layer."), - ("agent", "FastAPI is a great choice — fast, type-safe, async-native."), + ("agent", "FastAPI is a great choice - fast, type-safe, async-native."), ("user", "Last quarter I tried doing this with Pinecone and the costs blew up."), ]: mem.add_cosmos(user_id=user_id, role=role, content=content, thread_id=thread_id) @@ -97,7 +97,7 @@ def main() -> None: deleted += 1 print(f" deleted {deleted} raw turn(s)") - _header(6, "Final inventory — only compact long-term memory remains") + _header(6, "Final inventory - only compact long-term memory remains") _print_memories(mem, user_id, thread_id) _header(7, "Search still works against the archived knowledge") diff --git a/Samples/Advanced/advanced_search_patterns.py b/Samples/Advanced/advanced_search_patterns.py index dff042d..92c0e82 100644 --- a/Samples/Advanced/advanced_search_patterns.py +++ b/Samples/Advanced/advanced_search_patterns.py @@ -1,5 +1,5 @@ """ -Advanced Search Patterns — Agent Memory Toolkit +Advanced Search Patterns - Agent Memory Toolkit Demonstrates vector, hybrid, and filtered search patterns using Cosmos DB with AI Foundry embeddings. @@ -89,7 +89,7 @@ def seed_memories(mem: CosmosMemoryClient, user_id: str, thread_id: str) -> None # --------------------------------------------------------------------------- def vector_search(mem: CosmosMemoryClient, user_id: str) -> None: - """Pattern 1 — Semantic-style query (natural language, low keyword overlap).""" + """Pattern 1 - Semantic-style query (natural language, low keyword overlap).""" print_header("1. Semantic Search (natural-language query)") print(" Query: 'outdoor activities'") print(" Hybrid ranking leans on embedding similarity when there are few exact") @@ -104,11 +104,11 @@ def vector_search(mem: CosmosMemoryClient, user_id: str) -> None: def hybrid_search(mem: CosmosMemoryClient, user_id: str) -> None: - """Pattern 2 — Hybrid search (vector + full-text) is the default.""" + """Pattern 2 - Hybrid search (vector + full-text) is the default.""" print_header("2. Hybrid Search (vector + full-text)") print(" Query: 'hiking trails Pacific Northwest'") print(" Every search_cosmos call fuses embedding similarity with BM25 keyword") - print(" matching automatically — no flag required.\n") + print(" matching automatically - no flag required.\n") results = mem.search_cosmos( search_terms="hiking trails Pacific Northwest", @@ -119,8 +119,8 @@ def hybrid_search(mem: CosmosMemoryClient, user_id: str) -> None: def filtered_by_role(mem: CosmosMemoryClient, user_id: str) -> None: - """Pattern 3 — Filtered search: only user messages.""" - print_header("3. Filtered Search — by role ('user')") + """Pattern 3 - Filtered search: only user messages.""" + print_header("3. Filtered Search - by role ('user')") print(" Query: 'preferences'") print(" Restricts results to a specific conversation role.\n") @@ -134,8 +134,8 @@ def filtered_by_role(mem: CosmosMemoryClient, user_id: str) -> None: def filtered_by_memory_type(mem: CosmosMemoryClient, user_id: str) -> None: - """Pattern 4 — Filtered search: only 'fact' memories.""" - print_header("4. Filtered Search — by memory_type ('fact')") + """Pattern 4 - Filtered search: only 'fact' memories.""" + print_header("4. Filtered Search - by memory_type ('fact')") print(" Query: 'food preferences'") print(" Narrows results to a specific memory category.\n") @@ -149,8 +149,8 @@ def filtered_by_memory_type(mem: CosmosMemoryClient, user_id: str) -> None: def filtered_by_thread(mem: CosmosMemoryClient, user_id: str, thread_id: str) -> None: - """Pattern 5 — Filtered search: scoped to a single thread.""" - print_header("5. Filtered Search — by thread_id") + """Pattern 5 - Filtered search: scoped to a single thread.""" + print_header("5. Filtered Search - by thread_id") print(f" Query: 'activities' | thread: {thread_id[:8]}…") print(" Limits results to a specific conversation thread.\n") @@ -164,7 +164,7 @@ def filtered_by_thread(mem: CosmosMemoryClient, user_id: str, thread_id: str) -> def top_k_tuning(mem: CosmosMemoryClient, user_id: str) -> None: - """Pattern 6 — top-k tuning comparison.""" + """Pattern 6 - top-k tuning comparison.""" print_header("6. Top-K Tuning Comparison") print(" Query: 'hobbies and interests'") print(" Demonstrates how top_k affects the breadth of results.\n") diff --git a/Samples/Processing/processing_fact_extraction.py b/Samples/Processing/processing_fact_extraction.py index 36a3010..819e5ba 100644 --- a/Samples/Processing/processing_fact_extraction.py +++ b/Samples/Processing/processing_fact_extraction.py @@ -60,7 +60,7 @@ def main() -> None: ("agent", "8 years of Python experience is impressive!"), ("user", "I'm currently working on a project involving large language models and RAG."), ("agent", "That's a great area! LLMs combined with RAG can unlock powerful applications."), - ("user", "Last spring I went hiking on Mount Rainier with my dog — best trip ever."), + ("user", "Last spring I went hiking on Mount Rainier with my dog - best trip ever."), ("user", "When debugging an LLM, always check the prompt first then the response format."), ] diff --git a/Samples/Processing/processing_thread_summary.py b/Samples/Processing/processing_thread_summary.py index 89561e4..afd02c4 100644 --- a/Samples/Processing/processing_thread_summary.py +++ b/Samples/Processing/processing_thread_summary.py @@ -1,7 +1,7 @@ """Demonstrate per-thread summarisation (incremental updates included). ``CosmosMemoryClient.generate_thread_summary(...)`` runs the same in-process -ProcessingPipeline that the change-feed Azure Function uses — no Function +ProcessingPipeline that the change-feed Azure Function uses - no Function deployment is required for this sample. Required env vars (.env supported): @@ -73,7 +73,7 @@ def main() -> None: _banner("STEP 3 – more turns (incremental update path)") follow_up = [ ("user", "What about food? I'm a vegetarian."), - ("agent", "Japan has wonderful vegetarian options — try shojin-ryori (Buddhist temple cuisine), " + ("agent", "Japan has wonderful vegetarian options - try shojin-ryori (Buddhist temple cuisine), " "tofu specialties, and tempura vegetables."), ("user", "Are there any vegetarian restaurants you'd recommend in Kyoto?"), ("agent", "Shigetsu inside Tenryu-ji temple is famous for its shojin-ryori meals."), diff --git a/Samples/Processing/processing_user_profile.py b/Samples/Processing/processing_user_profile.py index de4bb9b..ecb35db 100644 --- a/Samples/Processing/processing_user_profile.py +++ b/Samples/Processing/processing_user_profile.py @@ -1,7 +1,7 @@ """Cross-thread user profile generation. Runs ``CosmosMemoryClient.generate_user_summary(...)`` which synthesises -a single user profile from per-thread summaries and extracted facts — +a single user profile from per-thread summaries and extracted facts - again, all in-process via the shared ProcessingPipeline. Required env vars (.env supported): @@ -55,7 +55,7 @@ def main() -> None: t1 = str(uuid.uuid4()) for role, content in [ ("user", "What's a good pasta recipe for a weeknight dinner?"), - ("agent", "Try aglio e olio — pasta, garlic, olive oil, chilli flakes, parsley."), + ("agent", "Try aglio e olio - pasta, garlic, olive oil, chilli flakes, parsley."), ("user", "Sounds great. I love simple Italian food, especially fresh ingredients."), ("agent", "Italian cuisine emphasises quality ingredients prepared simply."), ]: @@ -68,9 +68,9 @@ def main() -> None: _banner("Thread 2 – travel") t2 = str(uuid.uuid4()) for role, content in [ - ("user", "I want to visit Italy next year — Rome, Florence, Tuscany."), - ("agent", "Great itinerary! Tuscany is amazing in autumn — wine harvest season."), - ("user", "Perfect — I love wine, especially Chianti and Brunello."), + ("user", "I want to visit Italy next year - Rome, Florence, Tuscany."), + ("agent", "Great itinerary! Tuscany is amazing in autumn - wine harvest season."), + ("user", "Perfect - I love wine, especially Chianti and Brunello."), ("agent", "Brunello di Montalcino producers offer wonderful cellar tours."), ]: mem.add_cosmos(user_id=user_id, role=role, content=content, thread_id=t2) @@ -85,7 +85,7 @@ def main() -> None: ("user", "What's the best way to deploy a Python FastAPI app on Azure?"), ("agent", "Azure Container Apps is a great fit for FastAPI."), ("user", "Cool. I'm a Python engineer building AI tooling."), - ("agent", "Azure has excellent AI services — AI Foundry, AI Search, Cosmos DB for vectors."), + ("agent", "Azure has excellent AI services - AI Foundry, AI Search, Cosmos DB for vectors."), ]: mem.add_cosmos(user_id=user_id, role=role, content=content, thread_id=t3) @@ -93,7 +93,7 @@ def main() -> None: mem.generate_thread_summary(user_id=user_id, thread_id=t3) # ---------- generate user summary across all threads ---------- - _banner("User profile — synthesis across all 3 threads") + _banner("User profile - synthesis across all 3 threads") doc = mem.generate_user_summary(user_id=user_id, thread_ids=[t1, t2, t3]) print(f"id : {doc['id']}") print(f"source_count : {doc.get('metadata', {}).get('source_count')}") diff --git a/Samples/Scenarios/scenario_chat_memory.py b/Samples/Scenarios/scenario_chat_memory.py index 3a51c09..33680e9 100644 --- a/Samples/Scenarios/scenario_chat_memory.py +++ b/Samples/Scenarios/scenario_chat_memory.py @@ -97,7 +97,7 @@ def print_search_results(results: list[dict]) -> None: ("user", "Does caffeine really affect sleep that much?"), ( "agent", - "Yes — caffeine has a half-life of about 5-6 hours, so a coffee at " + "Yes - caffeine has a half-life of about 5-6 hours, so a coffee at " "3 PM still has half its caffeine at 9 PM. Limit caffeine to mornings.", ), ] @@ -169,14 +169,14 @@ def main() -> None: time.sleep(1) # ── Simulate reconnect: retrieve full thread ────────────── - banner("Reconnect — Retrieve travel thread") + banner("Reconnect - Retrieve travel thread") print(f" Fetching thread_id = {thread_1}\n") thread = mem.get_thread(thread_id=thread_1, user_id=user_id) print_thread(thread) print(f"\n Retrieved {len(thread)} turns from Cosmos DB.") # ── Cross-session search: recall relevant context ───────── - banner("Cross-session search — 'vegetarian food recommendations'") + banner("Cross-session search - 'vegetarian food recommendations'") results = mem.search_cosmos( search_terms="vegetarian food recommendations", user_id=user_id, @@ -184,7 +184,7 @@ def main() -> None: ) print_search_results(results) - banner("Cross-session search — 'caffeine and sleep'") + banner("Cross-session search - 'caffeine and sleep'") results = mem.search_cosmos( search_terms="caffeine and sleep", user_id=user_id, @@ -193,7 +193,7 @@ def main() -> None: print_search_results(results) # ── Show how a new session can build on old context ─────── - banner("New session — using recalled context") + banner("New session - using recalled context") thread_3 = str(uuid.uuid4()) print(f" thread_id = {thread_3}\n") diff --git a/Samples/Scenarios/scenario_counter_tuning.py b/Samples/Scenarios/scenario_counter_tuning.py index cd696a1..153d362 100644 --- a/Samples/Scenarios/scenario_counter_tuning.py +++ b/Samples/Scenarios/scenario_counter_tuning.py @@ -2,7 +2,7 @@ This sample demonstrates the *operational* surface of the :class:`DurableFunctionProcessor` deployment: the SDK never invokes an LLM -locally — instead, the sibling Azure Function app reads the Cosmos change +locally - instead, the sibling Azure Function app reads the Cosmos change feed and decides which orchestrator(s) to fire based on per-(user,thread) counters. @@ -72,7 +72,7 @@ def main() -> None: ("agent", "Great choice! Hoh Rainforest and Hurricane Ridge are must-sees."), ("user", "I'd like to camp 2 nights. Any permit guidance?"), ("agent", "You'll need a wilderness permit from recreation.gov for the Hoh."), - ("user", "Thanks — also, I'm vegetarian, please remember that."), + ("user", "Thanks - also, I'm vegetarian, please remember that."), ] print(f"Writing {len(transcript)} turns to Cosmos (thread={thread_id})...") for role, content in transcript: @@ -85,7 +85,7 @@ def main() -> None: ) # ------------------------------------------------------------------ - # In durable mode .process_now() is a no-op locally — the function app's + # In durable mode .process_now() is a no-op locally - the function app's # change-feed trigger has already begun work asynchronously. Use # process_now_and_wait() to *poll* for the summary doc; it's RU-costly so # use it only in demos/tests. diff --git a/Samples/Scenarios/scenario_customer_support.py b/Samples/Scenarios/scenario_customer_support.py index 1a5e459..3c7aa9c 100644 --- a/Samples/Scenarios/scenario_customer_support.py +++ b/Samples/Scenarios/scenario_customer_support.py @@ -1,10 +1,10 @@ -"""Customer-support scenario — build a unified customer profile across tickets. +"""Customer-support scenario - build a unified customer profile across tickets. Walks through three support tickets for the same customer; after each ticket we extract memories and update the user profile, so the agent can greet returning customers with personalised context. -In-process pipeline — no Azure Function deployment required. +In-process pipeline - no Azure Function deployment required. Required env vars (.env supported): @@ -47,14 +47,14 @@ def _process_ticket(mem: CosmosMemoryClient, user_id: str, ticket: str) -> None: def _run_ticket_1(mem: CosmosMemoryClient, user_id: str, ticket: str) -> None: - _banner(f"Ticket 1 ({ticket}) — initial complaint") + _banner(f"Ticket 1 ({ticket}) - initial complaint") _add_dialogue(mem, user_id, ticket, [ ("user", "Hi, I'm having issues with my Surface Pro 9. The battery drains in 3 hours."), ("agent", "I'm sorry to hear that. Could you share when you bought the device?"), ("user", "I bought it in March 2024 from the Microsoft Store online."), ("agent", "Thanks. I'll check warranty status. Could you also confirm your email?"), ("user", "It's alex.chen@example.com. The order number was MS-78234."), - ("agent", "Verified — the device is under warranty. I've initiated a battery diagnostic. " + ("agent", "Verified - the device is under warranty. I've initiated a battery diagnostic. " "Please run the Surface app diagnostic."), ("user", "OK, will do. By the way, I work as a software engineer so I rely on this device daily."), ]) @@ -62,21 +62,21 @@ def _run_ticket_1(mem: CosmosMemoryClient, user_id: str, ticket: str) -> None: def _run_ticket_2(mem: CosmosMemoryClient, user_id: str, ticket: str) -> None: - _banner(f"Ticket 2 ({ticket}) — accessory question") + _banner(f"Ticket 2 ({ticket}) - accessory question") _add_dialogue(mem, user_id, ticket, [ - ("user", "Hello again — Alex from the previous battery ticket. " + ("user", "Hello again - Alex from the previous battery ticket. " "The Surface app says battery is degraded."), ("agent", "Welcome back, Alex! Yes, that confirms the diagnostic. " "We'll ship a replacement battery service unit."), ("user", "Great. While we're talking, can you recommend a USB-C dock that works with the Surface Pro 9?"), ("agent", "The Surface Dock 2 or the Anker 778 are excellent choices for multi-monitor setups."), - ("user", "Multi-monitor is exactly what I need — I run 3 displays for development work."), + ("user", "Multi-monitor is exactly what I need - I run 3 displays for development work."), ]) _process_ticket(mem, user_id, ticket) def _run_ticket_3_greeting(mem: CosmosMemoryClient, user_id: str, ticket: str) -> None: - _banner(f"Ticket 3 ({ticket}) — personalised greeting using profile") + _banner(f"Ticket 3 ({ticket}) - personalised greeting using profile") # Build context from the user profile + extracted facts to personalise the greeting. profile = mem.get_user_summary(user_id) diff --git a/Samples/Scenarios/scenario_memory_reconciliation.py b/Samples/Scenarios/scenario_memory_reconciliation.py index 70d9930..b628efc 100644 --- a/Samples/Scenarios/scenario_memory_reconciliation.py +++ b/Samples/Scenarios/scenario_memory_reconciliation.py @@ -1,4 +1,4 @@ -"""Scenario: Memory reconciliation — duplicate merging and contradiction resolution. +"""Scenario: Memory reconciliation - duplicate merging and contradiction resolution. Demonstrates how `reconcile_memories` collapses paraphrased facts and resolves semantic contradictions in a single LLM pass: @@ -6,8 +6,8 @@ 1. Seed paraphrased facts about a user (will be merged). 2. Seed contradicting facts about the same subject (one will win, one will lose). 3. Run reconcile() and print the {kept, merged, contradicted} stats. -4. Show the live state — paraphrased duplicates collapsed, contradiction loser hidden. -5. Show the audit trail (include_superseded=True) — soft-deleted records carry +4. Show the live state - paraphrased duplicates collapsed, contradiction loser hidden. +5. Show the audit trail (include_superseded=True) - soft-deleted records carry supersede_reason, superseded_at, and superseded_by pointing at the survivor. Requirements: diff --git a/Samples/Scenarios/scenario_multi_agent.py b/Samples/Scenarios/scenario_multi_agent.py index e6222f6..2ea88d1 100644 --- a/Samples/Scenarios/scenario_multi_agent.py +++ b/Samples/Scenarios/scenario_multi_agent.py @@ -1,5 +1,5 @@ """ -Multi-Agent Collaboration — Agent Memory Toolkit +Multi-Agent Collaboration - Agent Memory Toolkit Demonstrates two agents (planner + researcher) sharing a single Cosmos DB thread so each can read the other's contributions via get_thread. @@ -46,7 +46,7 @@ def print_header(title: str) -> None: def print_thread(thread: list[dict], label: str = "Thread") -> None: """Pretty-print a thread, highlighting agent_id from metadata.""" - print(f"\n — {label} ({len(thread)} messages) —") + print(f"\n - {label} ({len(thread)} messages) -") for i, msg in enumerate(thread, 1): role = msg.get("role", "?") agent = msg.get("metadata", {}).get("agent_id", "n/a") @@ -74,7 +74,7 @@ def step1_user_question( mem: CosmosMemoryClient, user_id: str, thread_id: str, ) -> None: """User asks a complex, multi-part question.""" - print_header("Step 1 — User posts a complex question") + print_header("Step 1 - User posts a complex question") mem.add_cosmos( user_id=user_id, @@ -94,7 +94,7 @@ def step2_planner_creates_plan( mem: CosmosMemoryClient, user_id: str, thread_id: str, ) -> None: """Planner reads the user question and stores a structured plan.""" - print_header("Step 2 — Planner reads question & creates plan") + print_header("Step 2 - Planner reads question & creates plan") # Planner retrieves the thread to see the user's question thread = mem.get_thread(thread_id=thread_id, user_id=user_id) @@ -126,7 +126,7 @@ def step3_researcher_performs_research( mem: CosmosMemoryClient, user_id: str, thread_id: str, ) -> None: """Researcher reads the planner's plan, then stores research findings.""" - print_header("Step 3 — Researcher reads plan & stores findings") + print_header("Step 3 - Researcher reads plan & stores findings") # Researcher reads the full thread to find the plan thread = mem.get_thread(thread_id=thread_id, user_id=user_id) @@ -143,7 +143,7 @@ def step3_researcher_performs_research( findings = [ { "content": ( - "FINDING 1 — Manufacturing: EV battery production emits ~8 tonnes " + "FINDING 1 - Manufacturing: EV battery production emits ~8 tonnes " "CO₂ per 60 kWh pack. Fuel-cell stack production is lower per unit " "but hydrogen tanks require energy-intensive carbon-fibre." ), @@ -151,7 +151,7 @@ def step3_researcher_performs_research( }, { "content": ( - "FINDING 2 — Energy source: EVs charged on renewable grids achieve " + "FINDING 2 - Energy source: EVs charged on renewable grids achieve " "near-zero operational emissions. Green hydrogen (electrolysis) is " "promising but currently 3× less energy-efficient than direct " "grid-to-battery charging." @@ -160,7 +160,7 @@ def step3_researcher_performs_research( }, { "content": ( - "FINDING 3 — End-of-life: Li-ion battery recycling recovers ~95% " + "FINDING 3 - End-of-life: Li-ion battery recycling recovers ~95% " "of cobalt and nickel via hydrometallurgy. Fuel-cell membranes " "(PEM) contain platinum, which is recoverable but recycling " "infrastructure is nascent." @@ -185,7 +185,7 @@ def step4_planner_synthesises_answer( mem: CosmosMemoryClient, user_id: str, thread_id: str, ) -> None: """Planner reads the researcher's findings and generates a final answer.""" - print_header("Step 4 — Planner reads findings & produces final answer") + print_header("Step 4 - Planner reads findings & produces final answer") # Planner retrieves the full thread thread = mem.get_thread(thread_id=thread_id, user_id=user_id) @@ -221,7 +221,7 @@ def step5_review_shared_thread( mem: CosmosMemoryClient, user_id: str, thread_id: str, ) -> None: """Display the complete thread and per-agent filtered views.""" - print_header("Step 5 — Review the shared thread") + print_header("Step 5 - Review the shared thread") full_thread = mem.get_thread(thread_id=thread_id, user_id=user_id) print_thread(full_thread, label="Full shared thread") @@ -229,8 +229,8 @@ def step5_review_shared_thread( # Filtered views planner_only = filter_by_agent(full_thread, PLANNER) researcher_only = filter_by_agent(full_thread, RESEARCHER) - print_thread(planner_only, label=f"Filtered — {PLANNER} only") - print_thread(researcher_only, label=f"Filtered — {RESEARCHER} only") + print_thread(planner_only, label=f"Filtered - {PLANNER} only") + print_thread(researcher_only, label=f"Filtered - {RESEARCHER} only") turns = mem.get_thread(thread_id=thread_id, user_id=user_id) print_thread(turns, label="Thread turns") diff --git a/Samples/Scenarios/scenario_rag_with_memory.py b/Samples/Scenarios/scenario_rag_with_memory.py index 83bfaae..9adc8a0 100644 --- a/Samples/Scenarios/scenario_rag_with_memory.py +++ b/Samples/Scenarios/scenario_rag_with_memory.py @@ -117,7 +117,7 @@ def run_demo() -> None: thread_id = f"rag-session-{uuid.uuid4().hex[:8]}" # ------------------------------------------------------------------ - # Step 1 — Seed user facts and conversation history + # Step 1 - Seed user facts and conversation history # ------------------------------------------------------------------ _print_section("Step 1: Seeding user facts & conversation history") @@ -141,7 +141,7 @@ def run_demo() -> None: conversation = [ ("user", "I'm starting a new microservice and need a web framework."), ("agent", "Based on your Python preference, I'd suggest FastAPI or Flask."), - ("user", "I also need a database — something flexible for evolving schemas."), + ("user", "I also need a database - something flexible for evolving schemas."), ("agent", "Azure Cosmos DB with its NoSQL API would be a great fit."), ] for role, content in conversation: @@ -155,7 +155,7 @@ def run_demo() -> None: print(f" ✓ Stored {len(conversation)} conversation turns") # ------------------------------------------------------------------ - # Step 2 — New query arrives + # Step 2 - New query arrives # ------------------------------------------------------------------ _print_section("Step 2: New user query") @@ -163,7 +163,7 @@ def run_demo() -> None: print(f" User asks: \"{user_query}\"") # ------------------------------------------------------------------ - # Step 3 — Search memory for relevant context + # Step 3 - Search memory for relevant context # ------------------------------------------------------------------ _print_section("Step 3: Searching memory for relevant context") @@ -177,7 +177,7 @@ def run_demo() -> None: print(f" • [{m.get('memory_type')}] {m.get('content', '')[:80]}") # ------------------------------------------------------------------ - # Step 3b — Retrieve user summary (if one exists) + # Step 3b - Retrieve user summary (if one exists) # ------------------------------------------------------------------ _print_section("Step 3b: Retrieving user summary") @@ -189,7 +189,7 @@ def run_demo() -> None: print(" No summary available yet (generate one via generate_thread_summary).") # ------------------------------------------------------------------ - # Step 4 — Build augmented prompt + # Step 4 - Build augmented prompt # ------------------------------------------------------------------ _print_section("Step 4: Building augmented prompt") @@ -229,7 +229,7 @@ def run_demo() -> None: print(augmented_prompt) # ------------------------------------------------------------------ - # Step 5 — Show how memory personalises the response + # Step 5 - Show how memory personalises the response # ------------------------------------------------------------------ _print_section("Step 5: Personalised response (simulated)") @@ -239,7 +239,7 @@ def run_demo() -> None: **Web Framework → FastAPI** You already know Flask and expressed interest in trying FastAPI. It offers async support, automatic OpenAPI docs, and great - performance — ideal for the Azure-hosted SaaS microservice + performance - ideal for the Azure-hosted SaaS microservice you're building. **Database → Azure Cosmos DB (NoSQL API)** @@ -255,7 +255,7 @@ def run_demo() -> None: print(simulated_response) # ------------------------------------------------------------------ - # Cleanup — remove seeded data so the sample is re-runnable + # Cleanup - remove seeded data so the sample is re-runnable # ------------------------------------------------------------------ _print_section("Cleanup") diff --git a/Samples/Scenarios/scenario_remote_processor.py b/Samples/Scenarios/scenario_remote_processor.py index 23abf89..16f74ab 100644 --- a/Samples/Scenarios/scenario_remote_processor.py +++ b/Samples/Scenarios/scenario_remote_processor.py @@ -1,4 +1,4 @@ -"""Scenario: Remote processor — processing happens in the sibling Azure Function app. +"""Scenario: Remote processor - processing happens in the sibling Azure Function app. This sample shows how to wire :class:`DurableFunctionProcessor` so the SDK only writes raw turns to Cosmos DB, and a sibling Azure Function app (deployed via @@ -15,7 +15,7 @@ Behavior notes -------------- * :meth:`CosmosMemoryClient.process_now` is a debug-logged no-op when using the - durable processor — the function app owns processing. + durable processor - the function app owns processing. * :meth:`CosmosMemoryClient.process_now_and_wait` *polls* Cosmos for the resulting summary; that costs RUs so it's opt-in and intended for demos / tests. """ @@ -44,12 +44,12 @@ def main() -> None: user_id = "alice" thread_id = "thread-001" - # Write turns just like normal — the SDK never invokes the LLM here. + # Write turns just like normal - the SDK never invokes the LLM here. transcript = [ ("user", "Hi! I love Cosmos DB."), ("agent", "Cosmos DB is fantastic for low-latency global apps."), ("user", "Can it do vector search?"), - ("agent", "Yes — DiskANN indexes power semantic search natively."), + ("agent", "Yes - DiskANN indexes power semantic search natively."), ("user", "Great. What about hierarchical partition keys?"), ("agent", "HPK lets you co-locate related items for efficient queries."), ] @@ -67,7 +67,7 @@ def main() -> None: client.process_now(user_id=user_id, thread_id=thread_id) # Optional: block until the function app has produced a summary for this - # thread. Polls Cosmos every 0.5s until ``timeout`` — RU-costly, demo only. + # thread. Polls Cosmos every 0.5s until ``timeout`` - RU-costly, demo only. print("\nWaiting for the function app to produce a summary...") ok = client.process_now_and_wait(user_id=user_id, thread_id=thread_id, timeout=60.0) print(f"Summary available: {ok}") diff --git a/Samples/Scenarios/scenario_remote_processor_async.py b/Samples/Scenarios/scenario_remote_processor_async.py index 2c447e1..4ab9aca 100644 --- a/Samples/Scenarios/scenario_remote_processor_async.py +++ b/Samples/Scenarios/scenario_remote_processor_async.py @@ -40,7 +40,7 @@ async def main() -> None: ("user", "Hi! I love Cosmos DB."), ("agent", "Cosmos DB is fantastic for low-latency global apps."), ("user", "Can it do vector search?"), - ("agent", "Yes — DiskANN indexes power semantic search natively."), + ("agent", "Yes - DiskANN indexes power semantic search natively."), ("user", "Great. What about hierarchical partition keys?"), ("agent", "HPK lets you co-locate related items for efficient queries."), ] diff --git a/Samples/Scenarios/scenario_tagging_and_filtering.py b/Samples/Scenarios/scenario_tagging_and_filtering.py index e4401d8..5bdbbb7 100644 --- a/Samples/Scenarios/scenario_tagging_and_filtering.py +++ b/Samples/Scenarios/scenario_tagging_and_filtering.py @@ -4,9 +4,9 @@ feature. Tags are arbitrary string labels attached to a memory at write time; at read time you can filter retrieval with three composable predicates: -* ``tags_all=[...]`` — AND filter: every listed tag must be present -* ``tags_any=[...]`` — OR filter: at least one listed tag must be present -* ``exclude_tags=[...]`` — NOT filter: none of these tags may be present +* ``tags_all=[...]`` - AND filter: every listed tag must be present +* ``tags_any=[...]`` - OR filter: at least one listed tag must be present +* ``exclude_tags=[...]`` - NOT filter: none of these tags may be present Tags are stored in Cosmos as a JSON array on the memory document. They work across every memory type (``turn``, ``fact``, ``episodic``, ``procedural``, @@ -90,10 +90,10 @@ def main() -> None: print(f" + {mem_type:>10} tags={tags}") # ------------------------------------------------------------------ - # 2. AND filter: tags_all=[...] — every listed tag must be present + # 2. AND filter: tags_all=[...] - every listed tag must be present # Use get_memories() for filter-only retrieval (no vector search). # ------------------------------------------------------------------ - print_section("AND filter — tags_all=['workflow', 'important']") + print_section("AND filter - tags_all=['workflow', 'important']") results = client.get_memories( user_id=user_id, tags_all=["workflow", "important"], @@ -102,9 +102,9 @@ def main() -> None: # Expected: only the "Always confirm before purging" procedural memory. # ------------------------------------------------------------------ - # 3. OR filter: tags_any=[...] — any listed tag may match + # 3. OR filter: tags_any=[...] - any listed tag may match # ------------------------------------------------------------------ - print_section("OR filter — tags_any=['diet', 'travel']") + print_section("OR filter - tags_any=['diet', 'travel']") results = client.get_memories( user_id=user_id, tags_any=["diet", "travel"], @@ -113,23 +113,23 @@ def main() -> None: # Expected: the peanut allergy fact and the PyCon episodic memory. # ------------------------------------------------------------------ - # 4. NOT filter: exclude_tags=[...] — drop anything carrying these tags + # 4. NOT filter: exclude_tags=[...] - drop anything carrying these tags # ------------------------------------------------------------------ - print_section("NOT filter — exclude_tags=['ops'] (memory_types=['procedural'])") + print_section("NOT filter - exclude_tags=['ops'] (memory_types=['procedural'])") results = client.get_memories( user_id=user_id, memory_types=["procedural"], exclude_tags=["ops"], ) print_results(results) - # Expected: empty — both procedural memories carry the 'ops' tag. + # Expected: empty - both procedural memories carry the 'ops' tag. # ------------------------------------------------------------------ # 5. Combine tag filters with semantic search # search_cosmos() runs vector similarity against search_terms and # composes naturally with structural / tag filters. # ------------------------------------------------------------------ - print_section("Semantic + tag filter — search_terms='shipping a service' tags_all=['work']") + print_section("Semantic + tag filter - search_terms='shipping a service' tags_all=['work']") results = client.search_cosmos( search_terms="shipping a service", user_id=user_id, @@ -143,7 +143,7 @@ def main() -> None: # ------------------------------------------------------------------ # 6. Compose AND + NOT filters across types # ------------------------------------------------------------------ - print_section("AND + NOT — tags_all=['preference'] exclude_tags=['health']") + print_section("AND + NOT - tags_all=['preference'] exclude_tags=['health']") results = client.get_memories( user_id=user_id, tags_all=["preference"], diff --git a/azure/cosmos/agent_memory/_base/base_client.py b/azure/cosmos/agent_memory/_base/base_client.py index 73d662b..d21657a 100644 --- a/azure/cosmos/agent_memory/_base/base_client.py +++ b/azure/cosmos/agent_memory/_base/base_client.py @@ -276,7 +276,7 @@ def is_transient_tail_step_error(exc: BaseException) -> bool: become silent ``WARNING`` lines). Transient (swallow): - * ``LLMError`` — LLM-side defensive raises (no-choices, no-content). + * ``LLMError`` - LLM-side defensive raises (no-choices, no-content). * ``openai.RateLimitError`` / ``APITimeoutError`` / ``APIConnectionError``. * Any exception with ``status_code`` in :data:`_TRANSIENT_HTTP_STATUS_CODES` (covers ``CosmosHttpResponseError`` @@ -288,7 +288,7 @@ def is_transient_tail_step_error(exc: BaseException) -> bool: ``BadRequestError`` (status 400/401/403). * ``CosmosHttpResponseError`` with status 400/401/403/404/409. * Python builtins (``KeyError``, ``TypeError``, ``AttributeError``, - ``NameError`` …) — these are programmer bugs, not infra hiccups. + ``NameError`` …) - these are programmer bugs, not infra hiccups. """ from azure.cosmos.agent_memory.exceptions import LLMError diff --git a/azure/cosmos/agent_memory/_counters.py b/azure/cosmos/agent_memory/_counters.py index 65d888c..4f894af 100644 --- a/azure/cosmos/agent_memory/_counters.py +++ b/azure/cosmos/agent_memory/_counters.py @@ -6,18 +6,18 @@ Counter document shape:: - # thread-scoped — id = "thread:{user_id}:{thread_id}", PK = [user_id, thread_id] + # thread-scoped - id = "thread:{user_id}:{thread_id}", PK = [user_id, thread_id] { "id": ..., "user_id": ..., "thread_id": ..., "count": int, "last_batch_lsn": int|None, "last_batch_old_count": int, "last_failure_at": str|None, "last_failure_reason": str|None, "last_owner": str|None } - # user-scoped — id = "user:{user_id}", PK = [user_id, "__counters__"] + # user-scoped - id = "user:{user_id}", PK = [user_id, "__counters__"] { ... same fields ... } Unlike the FA-side helper, the SDK clients drive these counters without LSN replay protection because each ``push_to_cosmos()`` call is its own atomic -boundary — there is no change-feed redelivery to defend against. We +boundary - there is no change-feed redelivery to defend against. We **preserve any existing ``last_batch_lsn``** the FA may have written so the FA's monotonicity assumption stays valid even on shared deployments. @@ -79,7 +79,7 @@ def _maybe_warn_owner_mismatch(counter_id: str, existing_owner: Optional[str], o Operators run with either ``MEMORY_PROCESSOR_OWNER=inprocess`` or ``=durable``. If the same counter doc keeps getting touched by both - backends, that's a misconfiguration — both pipelines are extracting + backends, that's a misconfiguration - both pipelines are extracting against the same memories container. """ if not existing_owner or not observer_owner or existing_owner == observer_owner: @@ -90,7 +90,7 @@ def _maybe_warn_owner_mismatch(counter_id: str, existing_owner: Optional[str], o _warned_owner_mismatch.add(key) logger.warning( "Counter doc %s was last written by owner=%r but this process is " - "owner=%r — both backends appear to be processing the same container. " + "owner=%r - both backends appear to be processing the same container. " "Set MEMORY_PROCESSOR_OWNER consistently across all clients and the " "Function App to avoid double-extraction. Further mismatches for this " "counter will be logged at DEBUG level.", @@ -121,7 +121,7 @@ def increment_counter_sync( on shared deployments. Stamps ``last_owner`` (advisory) when *owner* is provided. - Returns ``(0, 0)`` and logs a warning if the container is unreachable — + Returns ``(0, 0)`` and logs a warning if the container is unreachable - auto-trigger failures must never block the user's primary write path. """ partition_key = [user_id, thread_id] @@ -170,7 +170,7 @@ def increment_counter_sync( if exc.status_code == 412 and attempt < MAX_RETRIES - 1: continue logger.warning( - "Counter increment failed counter_id=%s status=%s — auto-trigger skipped (increment dropped)", + "Counter increment failed counter_id=%s status=%s - auto-trigger skipped (increment dropped)", counter_id, exc.status_code, ) @@ -235,7 +235,7 @@ async def increment_counter_async( if exc.status_code == 412 and attempt < MAX_RETRIES - 1: continue logger.warning( - "Counter increment failed counter_id=%s status=%s — auto-trigger skipped", + "Counter increment failed counter_id=%s status=%s - auto-trigger skipped", counter_id, exc.status_code, ) @@ -307,7 +307,7 @@ def stamp_failure_sync( ) -> None: """Best-effort stamp ``last_failure_at`` / ``last_failure_reason`` on the counter doc. - Uses Cosmos ``patch_item`` so only the failure fields are touched — + Uses Cosmos ``patch_item`` so only the failure fields are touched - concurrent counter increments cannot lose updates here. Failures are logged and swallowed; we never want failure-stamping itself to break the user's write path. @@ -350,74 +350,6 @@ async def stamp_failure_async( logger.debug("stamp_failure_async failed counter_id=%s: %s", counter_id, exc) -def read_extract_watermark_sync( - container: Any, - counter_id: str, - user_id: str, - thread_id: str, -) -> Optional[int]: - """Return the count value at the last successful extract, or ``None``. - - The watermark lets recent_k cover every turn since the previous extract - succeeded, instead of just the current batch — so turns are never skipped - when extraction lags or transiently fails. Best-effort: returns ``None`` - on any read error so callers fall back to a batch-based recent_k. - """ - try: - doc = container.read_item(item=counter_id, partition_key=[user_id, thread_id]) - value = doc.get("last_extract_count") - return int(value) if value is not None else None - except Exception as exc: # pragma: no cover - best-effort - logger.debug("read_extract_watermark_sync failed counter_id=%s: %s", counter_id, exc) - return None - - -def advance_extract_watermark_sync( - container: Any, - counter_id: str, - user_id: str, - thread_id: str, - count: int, -) -> None: - """Stamp ``last_extract_count=count`` after a successful extract (sync).""" - patch_ops = [{"op": "add", "path": "/last_extract_count", "value": int(count)}] - try: - container.patch_item(item=counter_id, partition_key=[user_id, thread_id], patch_operations=patch_ops) - except Exception as exc: # pragma: no cover - best-effort - logger.debug("advance_extract_watermark_sync failed counter_id=%s: %s", counter_id, exc) - - -async def read_extract_watermark_async( - container: Any, - counter_id: str, - user_id: str, - thread_id: str, -) -> Optional[int]: - """Async version of :func:`read_extract_watermark_sync`.""" - try: - doc = await container.read_item(item=counter_id, partition_key=[user_id, thread_id]) - value = doc.get("last_extract_count") - return int(value) if value is not None else None - except Exception as exc: # pragma: no cover - best-effort - logger.debug("read_extract_watermark_async failed counter_id=%s: %s", counter_id, exc) - return None - - -async def advance_extract_watermark_async( - container: Any, - counter_id: str, - user_id: str, - thread_id: str, - count: int, -) -> None: - """Stamp ``last_extract_count=count`` after a successful extract (async).""" - patch_ops = [{"op": "add", "path": "/last_extract_count", "value": int(count)}] - try: - await container.patch_item(item=counter_id, partition_key=[user_id, thread_id], patch_operations=patch_ops) - except Exception as exc: # pragma: no cover - best-effort - logger.debug("advance_extract_watermark_async failed counter_id=%s: %s", counter_id, exc) - - __all__ = [ "USER_COUNTER_THREAD_ID", "thread_counter_id", @@ -427,8 +359,4 @@ async def advance_extract_watermark_async( "increment_counter_async", "stamp_failure_sync", "stamp_failure_async", - "read_extract_watermark_sync", - "advance_extract_watermark_sync", - "read_extract_watermark_async", - "advance_extract_watermark_async", ] diff --git a/azure/cosmos/agent_memory/_embedding_tokens.py b/azure/cosmos/agent_memory/_embedding_tokens.py index 6917118..1cc8e17 100644 --- a/azure/cosmos/agent_memory/_embedding_tokens.py +++ b/azure/cosmos/agent_memory/_embedding_tokens.py @@ -2,9 +2,9 @@ Azure OpenAI embedding models (``text-embedding-3-*``) reject any single input that exceeds their 8192-token context window. Callers in this toolkit -embed user-supplied and document-derived text of unbounded size — the +embed user-supplied and document-derived text of unbounded size - the dedup-context query in extraction, ``search_turns`` queries, and stored turn -vectors when ``enable_turn_embeddings`` is on — so an oversized string would +vectors when ``enable_turn_embeddings`` is on - so an oversized string would otherwise crash the embed call or, at the write sites that swallow embed errors, silently persist a record with no vector. @@ -54,6 +54,21 @@ def _encoding_for_model(model: str) -> Any: return None +def count_tokens(text: str, model: str = "gpt-4o") -> int: + """Return the token count of *text* for *model*. + + Uses tiktoken when available; falls back to a conservative character-based + estimate (~4 chars/token) so callers that batch by token budget still get a + usable (slightly over-counted) estimate when tiktoken is missing. + """ + if not text: + return 0 + encoding = _encoding_for_model(model) + if encoding is None: + return (len(text) // _FALLBACK_CHARS_PER_TOKEN) + 1 + return len(encoding.encode(text)) + + def truncate_text_to_token_budget( text: str, model: str, @@ -63,7 +78,7 @@ def truncate_text_to_token_budget( Returns *text* unchanged when it is already within budget. Truncation keeps the leading tokens and logs a warning so oversized inputs remain - observable — that warning is also the signal for where ingestion-time + observable - that warning is also the signal for where ingestion-time chunking is needed. """ if not text: diff --git a/azure/cosmos/agent_memory/_utils.py b/azure/cosmos/agent_memory/_utils.py index ac32f9e..9dbf059 100644 --- a/azure/cosmos/agent_memory/_utils.py +++ b/azure/cosmos/agent_memory/_utils.py @@ -77,7 +77,7 @@ def _normalize_for_hash(text: str) -> str: Deliberately conservative: lowercase, strip, and collapse internal runs of whitespace to a single space. Punctuation and word order still matter. - The point is to catch *identical* re-extractions cheaply — paraphrases + The point is to catch *identical* re-extractions cheaply - paraphrases are handled by the reconciliation LLM pass. """ return _WHITESPACE_RE.sub(" ", text.strip().lower()) @@ -313,7 +313,7 @@ def vector_autodrop_supported(distance_function: str) -> bool: The destructive ``DEDUP_SIM_HIGH`` auto-skip drops a new memory without an LLM check, relying on thresholds (~0.97) calibrated for cosine/dotproduct on normalized embeddings. Euclidean returns an *unbounded distance* (not a - [0,1] similarity), so those thresholds mis-fire — auto-drop is disabled for + [0,1] similarity), so those thresholds mis-fire - auto-drop is disabled for euclidean and the borderline tagging path (LLM-adjudicated) is used instead. """ return distance_function != "euclidean" @@ -602,7 +602,7 @@ def extract_keywords(text: Optional[str]) -> list[str]: Lowercases, tokenizes on alphanumeric runs (apostrophes/punctuation split into fragments that are themselves stopwords), removes stopwords, and de-duplicates while preserving first-seen order. The result is capped at ``MAX_FULLTEXT_TERMS`` - (30) — the hard limit on terms Azure Cosmos DB ``FullTextScore`` accepts — so the + (30) - the hard limit on terms Azure Cosmos DB ``FullTextScore`` accepts - so the hybrid search query is always valid. Returns ``[]`` when the text is empty or all stopwords, which the search layer treats as a signal to fall back to pure vector ranking. diff --git a/azure/cosmos/agent_memory/aio/auto_trigger.py b/azure/cosmos/agent_memory/aio/auto_trigger.py index 78f8716..60df3ae 100644 --- a/azure/cosmos/agent_memory/aio/auto_trigger.py +++ b/azure/cosmos/agent_memory/aio/auto_trigger.py @@ -64,10 +64,6 @@ async def maybe_trigger_steps( return n_dedup_turns = n_facts * n_dedup if n_facts > 0 and n_dedup > 0 else 0 - # Persisted-counter full-pool backstop cadence (durable-safe, mirrors the - # change-feed): every DEDUP_FULL_RECLUSTER_EVERY_N-th reconcile. - n_full_recluster = _threshold_int(thresholds, "get_dedup_full_recluster_every_n", "DEDUP_FULL_RECLUSTER_EVERY_N") - n_full_turns = n_dedup_turns * n_full_recluster if (n_dedup_turns > 0 and n_full_recluster > 0) else 0 user_batch_counts = await _trigger_thread_steps( processor, counter_container, @@ -75,7 +71,6 @@ async def maybe_trigger_steps( n_facts=n_facts, n_summary=n_summary, n_dedup_turns=n_dedup_turns, - n_full_turns=n_full_turns, thresholds=thresholds, ) await _trigger_user_steps(processor, counter_container, user_batch_counts, n_user=n_user) @@ -89,7 +84,6 @@ async def _trigger_thread_steps( n_facts: int, n_summary: int, n_dedup_turns: int, - n_full_turns: int, thresholds: Any = None, ) -> dict[str, int]: user_batch_counts: dict[str, int] = {} @@ -120,34 +114,11 @@ async def _trigger_thread_steps( fire_extract=n_facts > 0 and _counters.crosses_threshold(old_count, new_count, n_facts), fire_summary=n_summary > 0 and _counters.crosses_threshold(old_count, new_count, n_summary), fire_dedup=n_dedup_turns > 0 and _counters.crosses_threshold(old_count, new_count, n_dedup_turns), - fire_full_rebuild=n_full_turns > 0 and _counters.crosses_threshold(old_count, new_count, n_full_turns), thresholds=thresholds, ) return user_batch_counts -async def _watermark_recent_k( - counter_container: Any, - counter_id: str, - user_id: str, - thread_id: str, - *, - new_count: int, -) -> int: - """Async: recent_k covering every turn since the last successful extract. - - Not capped — ``new_count - watermark`` is exactly the unextracted backlog and - the newest-``recent_k`` slice covers precisely those turns, so the watermark - can advance to ``new_count`` with no stranded turns. **Bootstrap:** with no - watermark yet the base is ``0`` (``recent_k = new_count``), so turns added - during earlier failed extracts aren't stranded when the watermark first - advances to ``new_count``. - """ - watermark = await _counters.read_extract_watermark_async(counter_container, counter_id, user_id, thread_id) - base = watermark if watermark is not None else 0 - return max(new_count - base, 1) - - async def _fire_thread_steps( processor: AsyncInProcessProcessor, counter_container: Any, @@ -159,7 +130,6 @@ async def _fire_thread_steps( fire_extract: bool, fire_summary: bool, fire_dedup: bool, - fire_full_rebuild: bool = False, thresholds: Any = None, ) -> None: fire_procedural = fire_dedup and bool( @@ -171,20 +141,8 @@ async def _fire_thread_steps( ) ) if fire_extract: - recent_k = await _watermark_recent_k( - counter_container, - counter_id, - user_id, - thread_id, - new_count=new_count, - ) try: - await _call_async_compatible( - processor.process_extract_memories, user_id=user_id, thread_id=thread_id, recent_k=recent_k - ) - await _counters.advance_extract_watermark_async( - counter_container, counter_id, user_id, thread_id, new_count - ) + await _call_async_compatible(processor.process_extract_memories, user_id=user_id, thread_id=thread_id) except Exception as exc: logger.warning("Async auto-trigger process_extract_memories failed for %s/%s: %s", user_id, thread_id, exc) await _counters.stamp_failure_async( @@ -195,7 +153,7 @@ async def _fire_thread_steps( fire_dedup, "process_reconcile", processor.process_reconcile, - {"user_id": user_id, "full_rebuild": fire_full_rebuild}, + {"user_id": user_id}, ), ( fire_procedural, diff --git a/azure/cosmos/agent_memory/aio/chat.py b/azure/cosmos/agent_memory/aio/chat.py index 2f0edaf..d83d7d5 100644 --- a/azure/cosmos/agent_memory/aio/chat.py +++ b/azure/cosmos/agent_memory/aio/chat.py @@ -46,7 +46,7 @@ def _make_sync_token_provider_for_async(credential: Any, scope: str): awaitable yielding the bearer token. When the caller supplied a sync ``azure.identity.DefaultAzureCredential`` (the common case), we cannot use ``azure.identity.aio.get_bearer_token_provider`` because it - ``await``s the credential's ``get_token`` — which is not a coroutine on + ``await``s the credential's ``get_token`` - which is not a coroutine on sync credentials and would raise at runtime. Instead we wrap the sync ``get_token`` call in :func:`asyncio.to_thread` @@ -111,7 +111,7 @@ def _ensure_client(self) -> Any: ) # Detect sync vs async credential. Callers commonly pass a sync - # ``DefaultAzureCredential`` (from ``azure.identity``) — its + # ``DefaultAzureCredential`` (from ``azure.identity``) - its # ``get_token`` method is *not* awaitable and will hang the async # client if used with ``azure.identity.aio.get_bearer_token_provider``. # When the credential exposes an ``async def get_token`` we use the @@ -166,7 +166,7 @@ async def generate( exponential backoff using ``asyncio.sleep``. Any additional keyword arguments are forwarded directly to the OpenAI - client — typically sourced from a prompty file's ``model.parameters``. + client - typically sourced from a prompty file's ``model.parameters``. Raises ------ diff --git a/azure/cosmos/agent_memory/aio/cosmos_memory_client.py b/azure/cosmos/agent_memory/aio/cosmos_memory_client.py index af287aa..4f3e503 100644 --- a/azure/cosmos/agent_memory/aio/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/aio/cosmos_memory_client.py @@ -563,7 +563,7 @@ async def add_cosmos( For ``memory_type='turn'`` this also bumps the auto-trigger counter and schedules cadence-aware processing (extract / reconcile / procedural / - thread_summary / user_summary) as a background ``asyncio.Task`` — the + thread_summary / user_summary) as a background ``asyncio.Task`` - the same pattern :meth:`push_to_cosmos` uses for buffered turns. The await returns after the Cosmos write completes; cadence runs out-of-band so it does not block the caller. @@ -744,6 +744,31 @@ async def search_turns( created_before=created_before, ) + async def get_memory_history( + self, + memory_id: str, + user_id: str, + thread_id: Optional[str] = None, + *, + max_depth: int = 20, + ) -> list[dict[str, Any]]: + """Return a memory's superseded predecessors, most-recently-superseded first. + + AMT supersedes rather than deletes, so the full history of a fact - how a + preference, decision, or attribute changed over time - is preserved. This + walks the ``superseded_by`` chain backwards from *memory_id* and returns + every version it replaced (transitively), each carrying its + ``superseded_at`` / ``supersede_reason`` audit fields. Useful for + answering "what changed?" / "what was it before?" questions that a + current-value-only view cannot. + """ + return await self._get_store().get_memory_history( + memory_id, + user_id, + thread_id, + max_depth=max_depth, + ) + async def get_thread( self, thread_id: str, @@ -904,9 +929,7 @@ async def generate_user_summary( async def reconcile(self, user_id: str, n: Optional[int] = None) -> dict[str, int]: from azure.cosmos.agent_memory.thresholds import get_dedup_pool_size - return await self._get_pipeline().reconcile_memories( - user_id, n if n is not None else get_dedup_pool_size(), full_rebuild=True - ) + return await self._get_pipeline().reconcile_memories(user_id, n if n is not None else get_dedup_pool_size()) async def process_now(self, *, user_id: str, thread_id: str) -> "ProcessThreadResult": """Force the processor to run the full pipeline RIGHT NOW for one thread. @@ -921,7 +944,7 @@ async def process_now(self, *, user_id: str, thread_id: str) -> "ProcessThreadRe caught and logged as warnings so the per-thread work already persisted by the prior steps is not erased. Permanent failures (config bugs, auth errors, 4xx Cosmos errors, Python builtins like - ``KeyError`` / ``TypeError``) are re-raised — silencing them turns + ``KeyError`` / ``TypeError``) are re-raised - silencing them turns operational issues into invisible ``WARNING`` lines. """ _BaseMemoryClient._require_cosmos(self) diff --git a/azure/cosmos/agent_memory/aio/embeddings.py b/azure/cosmos/agent_memory/aio/embeddings.py index da29cbd..1811826 100644 --- a/azure/cosmos/agent_memory/aio/embeddings.py +++ b/azure/cosmos/agent_memory/aio/embeddings.py @@ -146,7 +146,7 @@ async def generate_batch( ConfigurationError If the endpoint or credentials are missing. openai.OpenAIError - Propagated from the SDK on API failure (no retry — see module + Propagated from the SDK on API failure (no retry - see module docstring). """ if not texts: diff --git a/azure/cosmos/agent_memory/aio/processors/base.py b/azure/cosmos/agent_memory/aio/processors/base.py index 421a5e5..aacc427 100644 --- a/azure/cosmos/agent_memory/aio/processors/base.py +++ b/azure/cosmos/agent_memory/aio/processors/base.py @@ -54,7 +54,6 @@ async def process_reconcile( self, *, user_id: str, - full_rebuild: bool = False, ) -> int: ... async def generate_user_summary( diff --git a/azure/cosmos/agent_memory/aio/processors/durable.py b/azure/cosmos/agent_memory/aio/processors/durable.py index 59fef03..ed38d8f 100644 --- a/azure/cosmos/agent_memory/aio/processors/durable.py +++ b/azure/cosmos/agent_memory/aio/processors/durable.py @@ -74,7 +74,7 @@ async def process_user_summary( ) return UserSummaryResult(summary=None) - async def process_reconcile(self, *, user_id: str, full_rebuild: bool = False) -> int: + async def process_reconcile(self, *, user_id: str) -> int: logger.debug( "AsyncDurableFunctionProcessor.process_reconcile no-op user_id=%s", user_id, diff --git a/azure/cosmos/agent_memory/aio/processors/inprocess.py b/azure/cosmos/agent_memory/aio/processors/inprocess.py index 62cac70..55cdf91 100644 --- a/azure/cosmos/agent_memory/aio/processors/inprocess.py +++ b/azure/cosmos/agent_memory/aio/processors/inprocess.py @@ -2,7 +2,7 @@ The underlying :class:`azure.cosmos.agent_memory.aio.services.pipeline.AsyncPipelineService` exposes native ``async def`` methods, so every call here is a direct -``await`` — no ``asyncio.to_thread`` adapter, no sync sub-clients. +``await`` - no ``asyncio.to_thread`` adapter, no sync sub-clients. """ from __future__ import annotations @@ -112,24 +112,20 @@ async def process_user_summary( summary = await self._pipeline.generate_user_summary(user_id, thread_ids) return UserSummaryResult(summary=summary if isinstance(summary, dict) else None) - async def process_reconcile(self, *, user_id: str, full_rebuild: bool = False) -> int: + async def process_reconcile(self, *, user_id: str) -> int: from ...thresholds import get_dedup_pool_size - return await self._reconcile_fact_and_episodic(user_id, get_dedup_pool_size(), full_rebuild=full_rebuild) + return await self._reconcile_fact_and_episodic(user_id, get_dedup_pool_size()) - async def _reconcile_fact_and_episodic(self, user_id: str, n: int, *, full_rebuild: bool = False) -> int: - """Reconcile facts and episodic memories; sum merged+contradicted counts. + async def _reconcile_fact_and_episodic(self, user_id: str, n: int) -> int: + """Reconcile facts and episodic memories; sum contradicted counts. SDK in-process processing reconciles both types (matching the Durable - backend) so episodic dups don't accrue forever. ``full_rebuild`` (set by - the auto-trigger on its persisted-counter full-recluster cadence) forces - the full-pool LLM pass that catches dissimilar-embedding contradictions. + backend) so contradictions are resolved for each. """ total = 0 for memory_type in ("fact", "episodic"): - reconciled = await self._pipeline.reconcile_memories( - user_id, n=n, memory_type=memory_type, full_rebuild=full_rebuild - ) + reconciled = await self._pipeline.reconcile_memories(user_id, n=n, memory_type=memory_type) total += self._extract_reconcile_count(reconciled) return total @@ -138,7 +134,7 @@ def _extract_reconcile_count(reconciled: Any) -> int: """Sum ``merged + contradicted`` from a ``reconcile_memories`` result. ``ProcessingPipeline.reconcile_memories`` returns a dict with - ``{"kept", "merged", "contradicted"}`` — both ``merged`` and + ``{"kept", "merged", "contradicted"}`` - both ``merged`` and ``contradicted`` represent facts that were consolidated or retired, so they contribute to the dedup-count metric. """ diff --git a/azure/cosmos/agent_memory/aio/services/pipeline.py b/azure/cosmos/agent_memory/aio/services/pipeline.py index 83b6d43..124c8cb 100644 --- a/azure/cosmos/agent_memory/aio/services/pipeline.py +++ b/azure/cosmos/agent_memory/aio/services/pipeline.py @@ -4,7 +4,7 @@ :class:`azure.cosmos.agent_memory.services.pipeline.PipelineService`. The two share all pure helpers via :mod:`azure.cosmos.agent_memory.services._pipeline_helpers`; only the IO call -sites differ — every Cosmos query, chat completion, and embedding call is +sites differ - every Cosmos query, chat completion, and embedding call is ``await``-ed against the async clients/store. """ @@ -55,12 +55,14 @@ VALID_VALENCES, PromptyLoader, _normalize_metadata_keys, + batch_turns_by_tokens, build_topic_tags, build_transcript, cap_structured_summary, chat_text, check_extracted_fact_grounding, coerce_valence, + is_retryable_llm_error, parse_llm_json, ) from azure.cosmos.agent_memory.services._pipeline_helpers import ( @@ -71,14 +73,9 @@ ) from azure.cosmos.agent_memory.store._search_helpers import top_literal from azure.cosmos.agent_memory.thresholds import ( - get_dedup_candidate_topk, - get_dedup_cluster_sim, - get_dedup_context_topk, - get_dedup_context_vector_enabled, - get_dedup_reconcile_mode, get_dedup_sim_high, - get_dedup_sim_low, get_dedup_vector_enabled, + get_extraction_batch_max_tokens, ) logger = get_logger("azure.cosmos.agent_memory.pipeline.aio") @@ -116,9 +113,6 @@ async def _collect_query(self, result: Any) -> list[dict[str, Any]]: async def query_items(self, **kwargs: Any) -> list[dict[str, Any]]: container = self._target_container() if container is not None and hasattr(container, "query_items"): - # Drop `enable_cross_partition_query` — async SDK leaks it to aiohttp - # (azure-cosmos 4.16.0 bug); SDK auto-detects cross-partition when - # partition_key is absent. kwargs.pop("enable_cross_partition_query", None) return await self._collect_query(container.query_items(**kwargs)) try: @@ -268,7 +262,7 @@ 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. + - 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). @@ -279,10 +273,6 @@ async def _vector_distance_function(self) -> str: try: props = await self._memories_container.read() except Exception: - # Transient read failure is indistinguishable from "no policy" once we - # drop to None — so DON'T cache. An uncached cosine default self-heals on - # the next call; caching it would pin cosine and silently mis-handle a - # euclidean container (cosine bands on euclidean distances → data loss). logger.debug( "vector dedup: could not read container vector policy; defaulting to cosine (not cached)", exc_info=True, @@ -368,7 +358,7 @@ def _prompt_lineage(self, filename: str) -> dict[str, str]: """Return ``{prompt_id, prompt_version}`` for stamping a doc. Safe no-op fallback (``prompt_version="v1"``) when the loader was - never initialised — happens in unit tests that build the service + never initialized - happens in unit tests that build the service via ``__new__`` to bypass real LLM/embedding clients. """ loader = getattr(self, "_prompty", None) @@ -410,7 +400,7 @@ async def _load_existing_memories( """Query active (non-superseded) memories for reconciliation context. Results are ordered by ``c._ts DESC`` so the most recently written - memories survive the cap — without ORDER BY, Cosmos returns rows + memories survive the cap - without ORDER BY, Cosmos returns rows in implementation-defined order and the dedup comparison set is non-deterministic. """ @@ -524,46 +514,56 @@ async def extract_memories_dry( logger.warning("extract_memories_dry no memories found user_id=%s thread_id=%s", user_id, thread_id) return {"facts": [], "episodic": [], "updates": [], "processed_turn_docs": []} - transcript = self._build_transcript(items) existing_for_hash = await self._load_existing_memories(user_id, ["fact"]) existing_fact_hashes: set[str] = { m["content_hash"] for m in existing_for_hash if m.get("type") == "fact" and m.get("content_hash") } - if get_dedup_context_vector_enabled(): - user_turns_text = "\n".join(str(it.get("content", "")) for it in items if it.get("role") == "user").strip() - context_query = user_turns_text or transcript + + # Token-bounded, per-batch extraction. Each batch is an independent LLM + # call, so a single poisoned turn fails only its own batch. Turns from + # succeeded and quarantined (non-retryable, e.g. content-filter) batches + # go into ``processed_turns`` and are stamped ``extracted_at`` by persist + # so they are never re-processed; turns from batches that fail with a + # *retryable* error are left un-stamped and retried on the next run. + batches = batch_turns_by_tokens(items, get_extraction_batch_max_tokens()) + facts: list[dict[str, Any]] = [] + episodic: list[dict[str, Any]] = [] + processed_turns: list[dict[str, Any]] = [] + deferred_turn_count = 0 + quarantined_turn_count = 0 + for batch in batches: + batch_transcript = self._build_transcript(batch) try: - existing = await self._store.search( - search_terms=context_query, - user_id=user_id, - memory_types=["fact"], - top_k=get_dedup_context_topk(), + response_text = await self._run_prompty( + "extract_memories.prompty", inputs={"transcript": batch_transcript} ) + parsed = self._parse_llm_json(response_text) + facts.extend(parsed.get("facts", [])) + episodic.extend(parsed.get("episodic", [])) + processed_turns.extend(batch) except Exception as exc: # noqa: BLE001 - logger.warning( - "extract_memories_dry dedup-context vector search failed (%s); " - "falling back to hash-based existing memories", - exc, - ) - existing = existing_for_hash - else: - existing = existing_for_hash - if existing: - existing_text = "\n".join( - f"- [ID: {mem['id']}] {mem.get('content', '')} " - f"(type={mem.get('type', 'fact')}, salience={mem.get('salience', 'N/A')})" - for mem in existing - ) - else: - existing_text = "(none)" - response_text = await self._run_prompty( - "extract_memories.prompty", - inputs={"existing_facts": existing_text, "transcript": transcript}, - ) - parsed = self._parse_llm_json(response_text) - facts = parsed.get("facts", []) - episodic = parsed.get("episodic", []) - unclassified = parsed.get("unclassified", []) + if is_retryable_llm_error(exc): + deferred_turn_count += len(batch) + logger.warning( + "extract_memories: deferring %d turns after retryable extraction error " + "(will retry next run) user_id=%s thread_id=%s err=%s", + len(batch), + user_id, + thread_id, + exc, + ) + else: + processed_turns.extend(batch) + quarantined_turn_count += len(batch) + logger.warning( + "extract_memories: quarantining %d turns after non-retryable extraction error " + "(e.g. content filter) - marking extracted so they do not re-poison future runs " + "user_id=%s thread_id=%s err=%s", + len(batch), + user_id, + thread_id, + exc, + ) doc_timestamp = self._stable_source_timestamp(items) fact_docs: list[dict[str, Any]] = [] @@ -604,10 +604,7 @@ async def extract_memories_dry( "confidence": 0.5 if confidence is None else confidence, **self._prompt_lineage("extract_memories.prompty"), "metadata": { - "category": fact.get("category") or "general", - "subject": fact.get("subject"), - "predicate": fact.get("predicate"), - "object": fact.get("object"), + "category": fact.get("category") or "other", "temporal_context": fact.get("temporal_context"), }, "salience": fact.get("salience") if fact.get("salience") is not None else 0.5, @@ -690,52 +687,23 @@ async def extract_memories_dry( } episodic_docs.append(self._validate_extracted_doc(doc)) - for item in unclassified: - text = item.get("text") - if not text: - continue - content_hash = compute_content_hash(text) - if content_hash in existing_fact_hashes: - logger.debug( - "extract_memories: skipping exact-dup unclassified hash=%s user_id=%s thread_id=%s", - content_hash, - user_id, - thread_id, - ) - exact_dedup_skipped += 1 - continue - seed = _ID_SEED_SEP.join((user_id, thread_id, content_hash)) - det_id = f"fact_{hashlib.sha256(seed.encode()).hexdigest()[:32]}" - topic_tags = build_topic_tags(item.get("tags", [])) - confidence = item.get("confidence") - doc = { - "id": det_id, - "user_id": user_id, - "thread_id": thread_id, - "role": "system", - "type": "fact", - "content": text, - "content_hash": content_hash, - "confidence": 0.5 if confidence is None else confidence, - **self._prompt_lineage("extract_memories.prompty"), - "metadata": {"category": "unclassified", "unclassified_reason": item.get("reason")}, - "salience": item.get("salience") if item.get("salience") is not None else 0.5, - "tags": ["sys:fact", "sys:auto-extracted", "sys:unclassified"] + topic_tags, - "created_at": doc_timestamp, - "updated_at": doc_timestamp, - } - fact_docs.append(self._validate_extracted_doc(doc)) - existing_fact_hashes.add(content_hash) - if exact_dedup_skipped: updates.append({"op": "stats", "exact_dedup_skipped": exact_dedup_skipped}) if dropped_episodic_count: updates.append({"op": "stats", "dropped_episodic_count": dropped_episodic_count}) + if deferred_turn_count or quarantined_turn_count: + updates.append( + { + "op": "stats", + "deferred_turn_count": deferred_turn_count, + "quarantined_turn_count": quarantined_turn_count, + } + ) check_extracted_fact_grounding( fact_docs, - items, - existing, + processed_turns, + existing_for_hash, user_id=user_id, thread_id=thread_id, logger=logger, @@ -745,7 +713,7 @@ async def extract_memories_dry( "facts": fact_docs, "episodic": episodic_docs, "updates": updates, - "processed_turn_docs": items, + "processed_turn_docs": processed_turns, } logger.info( "extract_memories_dry completed user_id=%s thread_id=%s fact_docs=%d episodic_docs=%d updates=%d", @@ -758,7 +726,9 @@ async def extract_memories_dry( return result async def dedup_extracted_memories(self, user_id: str, extracted: dict) -> dict: - """Apply gated vector-floor deduplication to extracted facts/episodes.""" + """Fold near-duplicate extracted docs into their existing canonical + memory *in place* (async mirror of the sync in-place dedup). + """ if not get_dedup_vector_enabled(): return extracted if not user_id: @@ -767,90 +737,147 @@ async def dedup_extracted_memories(self, user_id: str, extracted: dict) -> dict: raise ValidationError("extracted must be a dict") high = get_dedup_sim_high() - low = get_dedup_sim_low() - top_k = get_dedup_candidate_topk() distance_function = await self._vector_distance_function() - autodrop_ok = vector_autodrop_supported(distance_function) - if not autodrop_ok: + similarity_ok = vector_autodrop_supported(distance_function) + if not similarity_ok: self._warn_euclidean_autodrop_once(distance_function) + result = { "facts": [dict(doc) for doc in extracted.get("facts", [])], "episodic": [dict(doc) for doc in extracted.get("episodic", [])], "updates": [dict(op) for op in extracted.get("updates", [])], } + # Carry through any non-bucket keys (e.g. ``processed_turn_docs``) so this + # transform never silently drops caller state. + for _carry_key, _carry_value in extracted.items(): + if _carry_key not in result: + result[_carry_key] = _carry_value + docs = [doc for doc in result["facts"] + result["episodic"] if doc.get("content")] + # Similarity comparison is only meaningful for cosine/dotproduct; on a + # euclidean container we skip in-place folding and let everything ADD. + if not docs or not similarity_ok: + return result + missing_embeddings = [doc for doc in docs if not doc.get("embedding")] if missing_embeddings: embeddings = await self._embed_batch([str(doc["content"]) for doc in missing_embeddings]) for doc, embedding in zip(missing_embeddings, embeddings): doc["embedding"] = embedding - vector_dedup_skipped = 0 - dup_candidates_tagged = 0 - kept_ids: set[str] = set() - dropped_ids: set[str] = set() - filtered_by_key: dict[str, list[dict[str, Any]]] = {"facts": [], "episodic": []} - for key in ("facts", "episodic"): - for doc in result[key]: - if not doc.get("content"): - filtered_by_key[key].append(doc) - continue - doc_id = str(doc.get("id") or "") - memory_type = str(doc.get("type") or "") - if not doc_id or memory_type not in {"fact", "episodic"}: - # Parity with sync: under-specified docs (no id / unknown type) - # skip dedup and pass through verbatim. - filtered_by_key[key].append(doc) - continue - exclude_ids = kept_ids | dropped_ids | {doc_id, *(doc.get("supersedes_ids") or [])} - candidates = await self._vector_candidates( - user_id=user_id, - embedding=doc.get("embedding"), - memory_type=memory_type, - top_k=top_k, - exclude_ids=exclude_ids, - ) - best: dict[str, Any] | None = candidates[0] if candidates else None - score = float(best.get("score") or 0.0) if best else 0.0 - if best and autodrop_ok and vector_similarity_at_least(score, high, distance_function): - vector_dedup_skipped += 1 - dropped_ids.add(doc_id) - logger.info( - "dedup_extracted_memories: vector skip user_id=%s dropped=%r " - "surviving_id=%s surviving=%r score=%.4f", - user_id, - doc.get("content"), - best.get("id"), - best.get("content"), - score, - ) - continue - if best and vector_similarity_at_least(score, low, distance_function): - tags = list(doc.get("tags") or []) - if "sys:dup-candidate" not in tags: - tags.append("sys:dup-candidate") - doc["tags"] = tags - metadata = dict(doc.get("metadata") or {}) - metadata["dup_of"] = best.get("id") - metadata["dup_score"] = score - doc["metadata"] = metadata - dup_candidates_tagged += 1 - - kept_ids.add(doc_id) - filtered_by_key[key].append(doc) - - if vector_dedup_skipped or dup_candidates_tagged: - result["updates"].append( - { - "op": "stats", - "vector_dedup_skipped": vector_dedup_skipped, - "dup_candidates_tagged": dup_candidates_tagged, - } + inplace_updated = 0 + folded_ids: set[str] = set() + updated_target_ids: set[str] = set() + for doc in docs: + doc_id = str(doc.get("id") or "") + memory_type = str(doc.get("type") or "") + embedding = doc.get("embedding") or [] + if not doc_id or memory_type not in {"fact", "episodic"} or not embedding: + continue + + neighbor, score = await self._nearest_active_full( + user_id=user_id, + embedding=embedding, + memory_type=memory_type, + exclude_ids={doc_id} | set(doc.get("supersedes_ids") or []), ) - result["facts"] = filtered_by_key["facts"] - result["episodic"] = filtered_by_key["episodic"] + if not neighbor or not vector_similarity_at_least(score, high, distance_function): + continue # novel — leave in result for persist to ADD + + neighbor_id = str(neighbor.get("id") or "") + if not neighbor_id: + continue + if neighbor_id in updated_target_ids: + folded_ids.add(doc_id) + continue + if await self._apply_inplace_update(neighbor, doc): + updated_target_ids.add(neighbor_id) + inplace_updated += 1 + folded_ids.add(doc_id) + + if folded_ids: + for bucket in ("facts", "episodic"): + result[bucket] = [d for d in result[bucket] if str(d.get("id") or "") not in folded_ids] + if inplace_updated: + result["updates"].append({"op": "stats", "inplace_updated": inplace_updated}) return result + async def _nearest_active_full( + self, + *, + user_id: str, + embedding: list[float], + memory_type: str, + exclude_ids: set[str], + ) -> tuple[Optional[dict[str, Any]], float]: + """Async mirror: nearest active same-type memory returned as a *full* doc.""" + if not user_id or not embedding: + return None, 0.0 + query = ( + "SELECT TOP 5 c AS doc, VectorDistance(c.embedding, @vec) AS score " + "FROM c WHERE c.user_id = @user_id " + "AND c.type = @memory_type " + f"AND {_ACTIVE_DOC_FILTER} " + "AND IS_DEFINED(c.embedding) " + "ORDER BY VectorDistance(c.embedding, @vec)" + ) + try: + 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}, + ], + ) + except Exception as exc: # noqa: BLE001 + logger.warning("_nearest_active_full query failed user_id=%s err=%s", user_id, exc) + return None, 0.0 + for row in rows: + doc = row.get("doc") or {} + rid = str(doc.get("id") or "") + if rid and rid not in exclude_ids: + return doc, float(row.get("score") or 0.0) + return None, 0.0 + + async def _apply_inplace_update(self, neighbor: dict[str, Any], new_doc: dict[str, Any]) -> bool: + """Async mirror of the sync in-place refresh (recency-wins content+embedding).""" + try: + updated = dict(neighbor) + for sys_prop in ("_rid", "_self", "_etag", "_attachments", "_ts"): + updated.pop(sys_prop, None) + new_content = str(new_doc.get("content") or "") + updated["content"] = new_content + updated["content_hash"] = compute_content_hash(new_content) + if new_doc.get("embedding"): + updated["embedding"] = new_doc["embedding"] + updated["updated_at"] = datetime.now(timezone.utc).isoformat() + + new_sal = _max_or_none([neighbor.get("salience"), new_doc.get("salience")]) + if new_sal is not None: + updated["salience"] = new_sal + new_conf = _max_or_none([neighbor.get("confidence"), new_doc.get("confidence")]) + if new_conf is not None: + updated["confidence"] = new_conf + + merged_tags: list[str] = [] + for t in list(neighbor.get("tags") or []) + list(new_doc.get("tags") or []): + if t and t != "sys:dup-candidate" and t not in merged_tags: + merged_tags.append(t) + if merged_tags: + updated["tags"] = merged_tags + + await self._upsert_item(self._memories_container, body=updated) + return True + except Exception as exc: # noqa: BLE001 + logger.warning( + "in-place dedup update failed target_id=%s err=%s (keeping new doc as novel)", + neighbor.get("id"), + exc, + ) + return False + async def persist_extracted_memories( self, user_id: str, @@ -898,32 +925,16 @@ async def persist_extracted_memories( if op.get("op") == "stats": result["exact_dedup_skipped"] += int(op.get("exact_dedup_skipped") or 0) result["dropped_episodic_count"] += int(op.get("dropped_episodic_count") or 0) - if "vector_dedup_skipped" in op: - result["vector_dedup_skipped"] = result.get("vector_dedup_skipped", 0) + int( - op.get("vector_dedup_skipped") or 0 - ) - if "dup_candidates_tagged" in op: - result["dup_candidates_tagged"] = result.get("dup_candidates_tagged", 0) + int( - op.get("dup_candidates_tagged") or 0 - ) + if "inplace_updated" in op: + result["inplace_updated"] = result.get("inplace_updated", 0) + int(op.get("inplace_updated") or 0) logger.info("persist_extracted_memories completed user_id=%s counts=%s", user_id, result) - processed_turns = extracted.get("processed_turn_docs") or [] - if processed_turns: - marked = await self._mark_turns_extracted(processed_turns) - logger.info( - "persist_extracted_memories marked turns as extracted user_id=%s marked=%d/%d", - user_id, - marked, - len(processed_turns), - ) - return result async def _mark_turns_extracted(self, turn_docs: list[dict[str, Any]]) -> int: """Stamp ``extracted_at`` on each turn doc and upsert. Mirror of - the sync helper — per-turn failures are logged but never raise. + the sync helper - per-turn failures are logged but never raise. """ if not turn_docs: return 0 @@ -956,9 +967,33 @@ async def extract_memories( ) -> dict[str, int]: """Extract facts and episodic memories from a thread and persist them.""" extracted = await self.extract_memories_dry(user_id, thread_id, recent_k, turns=turns) + # Capture the processed turns from the DRY output as the single source of + # truth for stamping. Stamping happens here (not inside persist) so no + # intermediate transform (e.g. dedup) can drop ``processed_turn_docs`` + # and cause the same turns to be re-extracted forever. + processed_turns = extracted.get("processed_turn_docs") or [] if get_dedup_vector_enabled(): extracted = await self.dedup_extracted_memories(user_id, extracted) - return await self.persist_extracted_memories(user_id, extracted) + counts = await self.persist_extracted_memories(user_id, extracted) + if processed_turns: + marked = await self._mark_turns_extracted(processed_turns) + if marked < len(processed_turns): + logger.warning( + "extract_memories stamped only %d/%d processed turns as extracted user_id=%s " + "thread_id=%s (unstamped turns will be re-extracted next run)", + marked, + len(processed_turns), + user_id, + thread_id, + ) + else: + logger.info( + "extract_memories stamped %d processed turns as extracted user_id=%s thread_id=%s", + marked, + user_id, + thread_id, + ) + return counts async def synthesize_procedural( self, @@ -1067,38 +1102,12 @@ def _covered_by(prior: Optional[dict[str, Any]]) -> bool: if not current_source_ids: logger.info( - "synthesize_procedural skipping LLM user_id=%s — no behavioral facts or episodic lessons", + "synthesize_procedural skipping LLM user_id=%s - no behavioral facts or episodic lessons", user_id, ) return {"status": "unchanged", "procedural": prior_doc} - name_docs = await self._query_items( - self._memories_container, - query=( - "SELECT TOP 1 * FROM c WHERE c.user_id = @uid " - "AND c.type = @type " - f"AND {_ACTIVE_DOC_FILTER} " - "AND IS_DEFINED(c.metadata.category) " - "AND c.metadata.category = @category " - "AND IS_DEFINED(c.metadata.predicate) " - "AND c.metadata.predicate = @predicate " - "ORDER BY c._ts DESC" - ), - parameters=[ - {"name": "@uid", "value": user_id}, - {"name": "@type", "value": "fact"}, - {"name": "@category", "value": "biographical"}, - {"name": "@predicate", "value": "name"}, - ], - ) user_name = "the user" - if name_docs: - metadata = name_docs[0].get("metadata") or {} - name_candidate = metadata.get("object") - if not isinstance(name_candidate, str) or not name_candidate.strip(): - name_candidate = name_docs[0].get("content") - if isinstance(name_candidate, str) and name_candidate.strip(): - user_name = name_candidate.strip() def _render_bullets(values: list[str]) -> str: cleaned = [value.strip() for value in values if isinstance(value, str) and value.strip()] @@ -1117,7 +1126,7 @@ def _render_bullets(values: list[str]) -> str: # Retry loop: LLM call lives inside so that on a race-induced 409 # we (a) check whether the winner already covers our source set and # short-circuit if so, and (b) re-call the LLM with the winner as - # the new prior if not — keeping synthesized content monotonic in + # the new prior if not - keeping synthesized content monotonic in # source coverage, not just version number. written_doc: Optional[dict[str, Any]] = None for attempt in range(1, _PROCEDURAL_MAX_CREATE_ATTEMPTS + 1): @@ -1160,7 +1169,7 @@ def _render_bullets(values: list[str]) -> str: break except CosmosResourceExistsError: logger.info( - "synthesize_procedural id collision user_id=%s seq=%d attempt=%d/%d — re-reading", + "synthesize_procedural id collision user_id=%s seq=%d attempt=%d/%d - re-reading", user_id, new_seq, attempt, @@ -1557,204 +1566,15 @@ async def _load_memories_by_ids( parameters.extend({"name": f"@id{i}", "value": mid} for i, mid in enumerate(id_list)) return await self._query_items(self._memories_container, query=query, parameters=parameters) - async def _build_candidate_clusters( - self, - user_id: str, - memory_type: str, - n: int, - ) -> tuple[list[list[dict[str, Any]]], int, list[dict[str, Any]]]: - """Cluster dup-candidate seeds (+ vector neighbors) into connected components. - - Returns ``(clusters, node_count, seeds)`` where each cluster has >= 2 members, - ``node_count`` is the total distinct memories pulled into the graph (used to - report ``reconcile_llm_calls_saved``), and ``seeds`` is the tagged seed scan - (so the caller can clear stale tags on orphan seeds that never clustered). - The seed scan is bounded to ``n`` so a single cluster can never exceed the - reconcile prompt's pool cap. - """ - cluster_sim = get_dedup_cluster_sim() - top_k = get_dedup_candidate_topk() - distance_function = await self._vector_distance_function() - query = ( - f"SELECT TOP {top_literal(n, name='reconcile_memories.candidate_n')} * FROM c " - "WHERE c.user_id = @user_id " - "AND c.type = @memory_type " - "AND ARRAY_CONTAINS(c.tags, @tag) " - f"AND {_ACTIVE_DOC_FILTER} " - "ORDER BY c.created_at DESC" - ) - seeds = await self._query_items( - self._memories_container, - query=query, - parameters=[ - {"name": "@user_id", "value": user_id}, - {"name": "@memory_type", "value": memory_type}, - {"name": "@tag", "value": "sys:dup-candidate"}, - ], - ) - nodes_by_id: dict[str, dict[str, Any]] = {doc["id"]: doc for doc in seeds if doc.get("id")} - edges: set[tuple[str, str]] = set() - for seed in seeds: - sid = seed.get("id") - if not sid: - continue - dup_of = (seed.get("metadata") or {}).get("dup_of") if isinstance(seed.get("metadata"), dict) else None - if dup_of: - for doc in await self._load_memories_by_ids(user_id, memory_type, [dup_of]): - nodes_by_id[doc["id"]] = doc - edges.add(tuple(sorted((sid, doc["id"])))) - for cand in await self._vector_candidates( - user_id=user_id, - embedding=seed.get("embedding"), - memory_type=memory_type, - top_k=top_k, - exclude_ids={sid}, - ): - if vector_similarity_at_least(float(cand.get("score") or 0.0), cluster_sim, distance_function): - for doc in await self._load_memories_by_ids(user_id, memory_type, [cand.get("id")]): - nodes_by_id[doc["id"]] = doc - edges.add(tuple(sorted((sid, doc["id"])))) - node_ids = set(nodes_by_id) - for doc in list(nodes_by_id.values()): - did = doc.get("id") - if not did: - continue - for cand in await self._vector_candidates( - user_id=user_id, - embedding=doc.get("embedding"), - memory_type=memory_type, - top_k=top_k, - exclude_ids={did}, - ): - cid = cand.get("id") - if cid in node_ids and vector_similarity_at_least( - float(cand.get("score") or 0.0), cluster_sim, distance_function - ): - edges.add(tuple(sorted((did, cid)))) - - adjacency: dict[str, set[str]] = {node_id: set() for node_id in nodes_by_id} - for left, right in edges: - if left != right and left in adjacency and right in adjacency: - adjacency[left].add(right) - adjacency[right].add(left) - clusters: list[list[dict[str, Any]]] = [] - seen: set[str] = set() - for node_id in adjacency: - if node_id in seen: - continue - stack = [node_id] - component: list[str] = [] - seen.add(node_id) - while stack: - current = stack.pop() - component.append(current) - for nxt in adjacency[current]: - if nxt not in seen: - seen.add(nxt) - stack.append(nxt) - if len(component) >= 2: - # Cap cluster size at the reconcile pool limit: lowering the cluster - # threshold can chain many facts into one giant transitive component - # that would blow the prompt cap; keep the most-recent ``n``. - if len(component) > n: - component = component[:n] - clusters.append([nodes_by_id[cid] for cid in component]) - return clusters, len(nodes_by_id), seeds - - async def _reconcile_candidate_mode( - self, user_id: str, *, n: int, memory_type: str, started_at: float - ) -> dict[str, int]: - # Candidate clustering only. The periodic full-pool backstop that catches - # dissimilar-embedding contradictions ("vegetarian" vs "loves steak") is - # driven by the caller via ``full_rebuild`` on a PERSISTED-counter cadence - # (in-process auto-trigger + durable change-feed), not an in-memory sweep - # counter — the latter reset per worker/process and never fired reliably on - # the Function-App backend. - clusters, node_count, seeds = await self._build_candidate_clusters(user_id, memory_type, n) - aggregate = {"kept": 0, "merged": 0, "contradicted": 0} - clustered_ids: set[str] = set() - for cluster in clusters: - # Mark members as clustered BEFORE the LLM call so a failed cluster's - # seeds are not treated as orphans below (which would clear their tags - # and prevent a retry). A truncated/malformed LLM response on one - # cluster must not abort the sweep or starve the remaining clusters. - clustered_ids.update(doc["id"] for doc in cluster if doc.get("id")) - try: - counts, consumed = await self._reconcile_pool(user_id, memory_type, cluster) - except Exception as exc: - logger.warning( - "reconcile_memories: cluster reconcile failed user_id=%s memory_type=%s; " - "skipping cluster, tags retained for next sweep: %s", - user_id, - memory_type, - exc, - ) - continue - for key in aggregate: - aggregate[key] += int(counts.get(key, 0)) - # Clear dup-candidate tags only on survivors. Re-upserting a doc that - # was just superseded (duplicate source or contradiction loser) would - # resurrect it, since the in-memory cluster copy lacks superseded_by. - survivors = [doc for doc in cluster if doc.get("id") and doc["id"] not in consumed] - await self._clear_dup_candidate_tags(survivors) - # Orphan seeds: tagged dup-candidates that never joined a cluster have no - # near-duplicate, so clear the stale tag — otherwise every future sweep - # re-scans them as seeds and they accumulate forever. - orphan_seeds = [seed for seed in seeds if seed.get("id") and seed["id"] not in clustered_ids] - await self._clear_dup_candidate_tags(orphan_seeds) - aggregate["reconcile_clusters_sent"] = len(clusters) - aggregate["reconcile_llm_calls_saved"] = max(0, node_count - len(clusters)) - logger.info( - "reconcile_memories candidate completed user_id=%s memory_type=%s result=%s", - user_id, - memory_type, - aggregate, - ) - self._emit_reconcile_outcome( - started_at=started_at, - user_id=user_id, - candidates=node_count, - result=aggregate, - ) - return aggregate - - async def _clear_dup_candidate_tags(self, docs: Iterable[dict[str, Any]]) -> None: - for doc in docs: - tags = [tag for tag in (doc.get("tags") or []) if tag != "sys:dup-candidate"] - if tags == (doc.get("tags") or []): - continue - updated = dict(doc) - updated["tags"] = tags - metadata = dict(updated.get("metadata") or {}) - metadata.pop("dup_of", None) - metadata.pop("dup_score", None) - updated["metadata"] = metadata - updated["updated_at"] = datetime.now(timezone.utc).isoformat() - try: - await self._upsert_memory(updated) - except Exception: - logger.exception("reconcile_memories: failed to clear dup-candidate tag id=%s", doc.get("id")) + async def reconcile_memories(self, user_id: str, n: int = 50, *, memory_type: str = "fact") -> dict[str, int]: + """Resolve contradictions among a user's most-recent active memories. - async def reconcile_memories( - self, user_id: str, n: int = 50, *, memory_type: str = "fact", full_rebuild: bool = False - ) -> dict[str, int]: - """Reconcile a user's active facts in a single LLM pass. - - Loads the most recent ``n`` active (non-superseded) facts for - ``user_id``, asks the dedup prompt to classify them into - ``duplicate_groups``, ``contradicted_pairs``, and ``kept_ids``, then - applies both kinds of resolutions: - - * **Duplicates** — a fresh merged fact is upserted; every source is - soft-deleted with ``supersede_reason="duplicate"``. - * **Contradictions** — the loser is soft-deleted with - ``supersede_reason="contradict"`` and ``superseded_by`` set to - the winner. Dangling references are resolved transparently when a - contradicted id was just absorbed into a duplicate group. - - Returns ``{"kept": int, "merged": int, "contradicted": int}`` where - ``merged`` and ``contradicted`` count the *losers* that were - soft-deleted (duplicates and contradictions respectively). + Async mirror of the sync contradiction-only reconcile. Near-duplicate + paraphrases are folded in place at write time + (:meth:`dedup_extracted_memories`); this pass only supersedes the loser + of each ``contradicted_pairs`` entry — no clustering, no merged + documents, no re-merge churn. Episodic and procedural types are no-ops. + Returns ``{"kept", "merged", "contradicted"}`` with ``merged`` always 0. """ if not user_id: raise ValidationError("user_id is required") @@ -1764,34 +1584,16 @@ async def reconcile_memories( raise ValidationError(f"n must be <= 500 to bound prompt size and LLM cost, got {n}") if memory_type not in {"fact", "episodic", "procedural"}: raise ValidationError(f"memory_type must be one of fact, episodic, procedural, got {memory_type!r}") - if memory_type == "procedural": - result = { - "kept": 0, - "merged": 0, - "contradicted": 0, - "reconcile_clusters_sent": 0, - "reconcile_llm_calls_saved": 0, - } - logger.info("reconcile_memories procedural no-op user_id=%s result=%s", user_id, result) + if memory_type in {"episodic", "procedural"}: + result = {"kept": 0, "merged": 0, "contradicted": 0} + logger.info("reconcile_memories %s no-op user_id=%s result=%s", memory_type, user_id, result) return result started_at = time.monotonic() logger.info("reconcile_memories started user_id=%s n=%d memory_type=%s", user_id, n, memory_type) - # Explicit user-triggered reconcile (full_rebuild) always takes the - # full-pool single-LLM-pass path: it sees every active fact together, so it - # catches contradictions that aren't vector-similar (e.g. "vegetarian" vs - # "loves steak") — which candidate clustering, keyed on near-duplicate - # similarity, would never group. Automatic sweeps use cheap candidate mode. - if get_dedup_reconcile_mode() == "candidate" and not full_rebuild: - return await self._reconcile_candidate_mode(user_id, n=n, memory_type=memory_type, started_at=started_at) - facts = await self._active_memories_for_reconcile(user_id, memory_type, n) - result, consumed = await self._reconcile_pool(user_id, memory_type, facts) - # Clear dup-candidate tags on survivors so an explicit reconcile(full_rebuild=True) - # doesn't leave stale sys:dup-candidate/dup_of metadata on user-visible memories. - survivors = [doc for doc in facts if doc.get("id") and doc["id"] not in consumed] - await self._clear_dup_candidate_tags(survivors) + result = await self._reconcile_contradictions(user_id, memory_type, facts) self._emit_reconcile_outcome( started_at=started_at, user_id=user_id, @@ -1800,32 +1602,18 @@ async def reconcile_memories( ) return result - async def _reconcile_pool( + async def _reconcile_contradictions( self, user_id: str, memory_type: str, facts: list[dict[str, Any]] - ) -> tuple[dict[str, int], set[str]]: - """Reconcile an explicit pool of same-type memories in one LLM pass. - - Returns ``({"kept", "merged", "contradicted"}, consumed_ids)`` where - ``consumed_ids`` are the source/loser ids that were actually superseded, - so callers can skip them when clearing dup-candidate tags (re-upserting a - superseded source would resurrect it). Does not emit telemetry — the - caller owns the ``reconcile.outcome`` line. + ) -> dict[str, int]: + """Async mirror: resolve only ``contradicted_pairs`` within the pool. + + ``duplicate_groups`` are ignored (write-time in-place dedup handles + paraphrases), so no merged docs are minted and the pass is convergent. + Returns ``{"kept", "merged": 0, "contradicted"}``. """ if len(facts) <= 1: - logger.info( - "reconcile_memories: %d %s memories, nothing to reconcile", - len(facts), - memory_type, - ) - return {"kept": len(facts), "merged": 0, "contradicted": 0}, set() - - # ---- 2. Format the facts pool for the prompt ---- - # ``json.dumps`` escapes embedded quotes and pipes inside content so - # the visual grammar (`| Field:` separators, `""` quoting) - # stays unambiguous even on adversarial inputs like - # ``She said "hi" | weird``. IDs are kept raw because they're - # deterministic alphanumerics — quoting them risks the LLM copying - # the quotes back into ``source_ids``. + return {"kept": len(facts), "merged": 0, "contradicted": 0} + lines: list[str] = [] for i, cf in enumerate(facts, 1): content_quoted = json.dumps(cf.get("content", ""), ensure_ascii=False) @@ -1837,376 +1625,45 @@ async def _reconcile_pool( created_str = created_raw if created_raw else "N/A" lines.append( f"{i}. ID: {cf['id']} | Content: {content_quoted} | " - f"Confidence: {conf_str} | " - f"Salience: {sal_str} | " - f"Created: {created_str}" + f"Confidence: {conf_str} | Salience: {sal_str} | Created: {created_str}" ) facts_text = "\n".join(lines) - # ---- 3. Single LLM call over the entire pool ---- - # Polarity keyed on an explicit predicate so a future third type (e.g. - # procedural, were it ever routed here) can't silently diverge from sync. - is_episodic = memory_type == "episodic" - prompt_name = "dedup_episodic.prompty" if is_episodic else "dedup.prompty" - prompt_inputs = {"episodics_text": facts_text} if is_episodic else {"facts_text": facts_text} - response_text = await self._run_prompty(prompt_name, inputs=prompt_inputs) + response_text = await self._run_prompty("dedup.prompty", inputs={"facts_text": facts_text}) parsed = self._parse_llm_json(response_text) - - duplicate_groups = parsed.get("duplicate_groups", []) or [] - contradicted_pairs = [] if is_episodic else (parsed.get("contradicted_pairs", []) or []) - # ``kept_ids`` from the LLM is used below as a cross-check for - # accounting drift (hallucinated IDs, double-counting). The actual - # kept count is computed from facts minus consumed losers. - llm_kept_ids = list(parsed.get("kept_ids", []) or []) + contradicted_pairs = parsed.get("contradicted_pairs", []) or [] facts_by_id: dict[str, dict[str, Any]] = {f["id"]: f for f in facts} - - merged = 0 contradicted = 0 - # Tracks source_id -> merged_id rewrites so contradictions whose - # winner/loser landed in a duplicate group can be redirected to - # the surviving merged document. Only updated on *successful* - # supersede so stale redirects don't survive ETag races. - source_to_merged_id: dict[str, str] = {} - # Cache of merged docs we just upserted, keyed by merged_id. Lets - # the contradiction redirector reuse the in-memory dict instead of - # a cross-partition Cosmos round-trip for a doc we own. Also keeps - # the chain ETag-stable when the same merged doc absorbs both a - # duplicate group and a contradiction redirect in the same call. - merged_docs_by_id: dict[str, dict[str, Any]] = {} - # Set of source IDs that were *actually* superseded (counts toward - # ``merged``). Used by the kept-count cross-check below — earlier - # versions counted attempts and undercounted on ETag races. - consumed_source_ids: set[str] = set() - # Set of contradiction loser IDs that were *actually* superseded. consumed_loser_ids: set[str] = set() - # Original-pool winner IDs from successfully-applied contradictions. - # The LLM emits winners under ``contradicted_pairs``, never under - # ``kept_ids`` — so the kept-cross-check at the end must subtract - # them from the expected-kept set or every clean run looks like a - # mismatch. - contradiction_winner_ids_in_pool: set[str] = set() - - # ---- 4. Apply duplicate_groups FIRST ---- - for group in duplicate_groups: - source_ids = list(group.get("source_ids") or []) - merged_content = group.get("merged_content") - if not merged_content or not source_ids: - logger.debug( - "reconcile_memories: skipping malformed duplicate_group %r", - group, - ) - continue - - source_docs = [facts_by_id[sid] for sid in source_ids if sid in facts_by_id] - if not source_docs: - logger.debug( - "reconcile_memories: duplicate_group references unknown ids %r", - source_ids, - ) - continue - - # Filtered, hallucination-free view of the source ids that - # actually exist in the pool. Used both for ``supersedes_ids`` - # on the merged record and for the deterministic merged-id - # below so the merged doc faithfully represents reality. - valid_source_ids = [sid for sid in source_ids if sid in facts_by_id] - - if len(valid_source_ids) < 2: - logger.debug( - "reconcile_memories: skipping single-source duplicate_group %r", - source_ids, - ) - continue - - # Sort source_docs by Cosmos _ts DESC so the merged record's - # partition (thread_id) is picked deterministically from the - # newest source — independent of the LLM's source_ids order. - source_docs.sort(key=lambda d: d.get("_ts", 0), reverse=True) - - # Union tags across all source docs (preserve order, dedupe). - merged_tags: list[str] = [] - seen_tags: set[str] = set() - for src in source_docs: - for t in src.get("tags", []) or []: - if t == "sys:dup-candidate": - continue - if t not in seen_tags: - seen_tags.add(t) - merged_tags.append(t) - if not merged_tags: - merged_tags = [f"sys:{memory_type}"] - - # Union source_memory_ids across all source docs (provenance chain). - merged_source_memory_ids: list[str] = [] - seen_smi: set[str] = set() - for src in source_docs: - for smi in src.get("source_memory_ids", []) or []: - if smi not in seen_smi: - seen_smi.add(smi) - merged_source_memory_ids.append(smi) - - # Transitive supersedes_ids: include any prior chain hops the - # source docs already absorbed so the merged record carries - # the full provenance, not just the immediate parent layer. - merged_supersedes: list[str] = [] - seen_sup: set[str] = set() - for sid in valid_source_ids: - if sid not in seen_sup: - seen_sup.add(sid) - merged_supersedes.append(sid) - for src in source_docs: - for prior in src.get("supersedes_ids", []) or []: - if prior and prior not in seen_sup: - seen_sup.add(prior) - merged_supersedes.append(prior) - - # Newest source's thread_id wins (after _ts-desc sort above). - recent_thread_id = source_docs[0].get("thread_id", "") - - # If LLM omitted confidence/salience, returned a non-positive - # placeholder, returned a JSON ``true`` masquerading as numeric, - # or returned an out-of-range value (e.g. 1.05 — common when - # models confuse percent with [0,1]), fall back to max across - # the source docs. Out-of-range without a fallback would let - # ``MemoryRecord(...)`` raise on Pydantic validation and the - # blanket except below would silently drop the entire group. - llm_conf = group.get("confidence") - confidence_val = ( - float(llm_conf) - if _is_real_number(llm_conf) and 0 < llm_conf <= 1 - else _max_or_none(src.get("confidence") for src in source_docs) - ) - llm_sal = group.get("salience") - salience_val = ( - float(llm_sal) - if _is_real_number(llm_sal) and 0 < llm_sal <= 1 - else _max_or_none(src.get("salience") for src in source_docs) - ) - - # Deterministic merged id keyed on (user, "merged", content_hash) - # so re-running reconcile on the same merged content produces an - # idempotent upsert instead of a fresh UUID each cycle. Stable - # ids also keep the supersede chain shallow: a future paraphrase - # that gets folded into the same canonical merged content will - # see the same id rather than chaining through a new UUID. - merged_content_hash = compute_content_hash(merged_content) - merged_id_seed = _ID_SEED_SEP.join((user_id, "merged", merged_content_hash)) - merged_prefix = "ep_" if memory_type == "episodic" else "fact_" - merged_id = merged_prefix + hashlib.sha256(merged_id_seed.encode()).hexdigest()[:32] - try: - if memory_type == "episodic": - base_metadata = dict(source_docs[0].get("metadata") or {}) - base_metadata.update( - { - "lesson": merged_content, - "merged_via": "reconcile", - "merged_from_count": len(valid_source_ids), - } - ) - base_metadata.setdefault("scope_type", base_metadata.get("scope_type") or "general") - base_metadata.setdefault("scope_value", base_metadata.get("scope_value") or "general") - base_metadata.setdefault("outcome_valence", base_metadata.get("outcome_valence") or "neutral") - record_cls = EpisodicRecord - prompt_lineage = self._prompt_lineage("dedup_episodic.prompty") - metadata = base_metadata - else: - record_cls = FactRecord - prompt_lineage = self._prompt_lineage("dedup.prompty") - metadata = { - "category": "preference", - "merged_via": "reconcile", - "merged_from_count": len(valid_source_ids), - } - merged_record = construct_internal( - record_cls, - { - "id": merged_id, - "user_id": user_id, - "role": "system", - "type": memory_type, - "content": merged_content, - "thread_id": recent_thread_id or f"__reconciled__:{user_id}", - "confidence": confidence_val if confidence_val is not None else 0.5, - "salience": salience_val if salience_val is not None else 0.5, - "supersedes_ids": merged_supersedes, - "source_memory_ids": merged_source_memory_ids, - "tags": merged_tags, - "content_hash": merged_content_hash, - "metadata": metadata, - **prompt_lineage, - "created_at": datetime.now(timezone.utc).isoformat(), - "updated_at": datetime.now(timezone.utc).isoformat(), - }, - ) - except Exception: - logger.exception( - "reconcile_memories: failed to build merged record for group %r", - group, - ) - continue - - # Generate embedding for the merged content so retrieval can - # rank it against future queries from the moment it lands. - # If embedding fails, abort this duplicate group entirely: - # writing a merged doc with no embedding and then superseding - # the sources would create a search-index hole until the next - # reconcile retried. Better to leave the duplicates in place. - try: - merged_record.embedding = await self._embed_one(merged_content) - except Exception: - logger.exception( - "reconcile_memories: embedding failed for merged id=%s; " - "aborting duplicate group to avoid search-index hole", - merged_record.id, - ) - continue - - merged_doc = merged_record.to_doc() - try: - merged_doc = await self._upsert_memory(merged_doc) - except Exception: - logger.exception( - "reconcile_memories: upsert failed for merged id=%s; aborting duplicate group", - merged_record.id, - ) - continue - merged_docs_by_id[merged_record.id] = merged_doc - - group_supersede_count = 0 - for sid in valid_source_ids: - src_doc = facts_by_id.get(sid) - if src_doc is None: - # Defensive — already filtered above, kept for clarity. - continue - # Only update redirect/consumed-set on *successful* supersede. - # Losing the ETag race means another writer beat us; the - # source doc is still active from our perspective and should - # not be treated as consumed. - if await self._mark_superseded(src_doc, merged_record.id, reason="duplicate"): - merged += 1 - group_supersede_count += 1 - source_to_merged_id[sid] = merged_record.id - consumed_source_ids.add(sid) - - # If every supersede attempt for this group failed (typically - # an ETag race against a concurrent reconcile that already - # superseded the same sources to the *same* deterministic - # merged id), do NOT delete the merged doc. A delete here - # would orphan the sources whose ``superseded_by`` already - # points at this merged id — they'd become invisible to - # default reads (filter ``superseded_by IS NULL``) and to the - # reconcile pool, causing permanent data loss. The merged doc - # is idempotent (deterministic id), so leaving it in place is - # consistent with whatever the winning concurrent writer - # produced. - if group_supersede_count == 0: - logger.info( - "reconcile_memories: no sources superseded for merged id=%s " - "(likely ETag race with concurrent reconcile); leaving " - "merged doc in place — idempotent upsert is self-healing", - merged_record.id, - ) - - # ---- 5. Apply contradicted_pairs SECOND with dangling-id resolution ---- for pair in contradicted_pairs: winner_id = pair.get("winner_id") loser_id = pair.get("loser_id") - if not winner_id or not loser_id: - logger.debug( - "reconcile_memories: skipping malformed contradicted_pair %r", - pair, - ) + if not winner_id or not loser_id or winner_id == loser_id: continue - - # Redirect through any duplicate-merge that absorbed the id. - resolved_winner = source_to_merged_id.get(winner_id, winner_id) - resolved_loser_id = source_to_merged_id.get(loser_id, loser_id) - - # Validate the (resolved) winner. The LLM is instructed never to - # invent IDs — if it does, refuse to write a dangling - # ``superseded_by`` pointer that breaks the audit trail. - if resolved_winner not in facts_by_id and resolved_winner not in merged_docs_by_id: + if winner_id not in facts_by_id: logger.warning( - "reconcile_memories: hallucinated winner_id=%s (resolved=%s) " - "not in pool or merged set; skipping pair %r", + "reconcile_memories: hallucinated winner_id=%s not in pool; skipping pair %r", winner_id, - resolved_winner, pair, ) continue - - if resolved_winner == resolved_loser_id: - # Both sides collapsed into the same merged doc — the - # contradiction is moot. Drop it silently. - logger.debug( - "reconcile_memories: contradiction collapsed into duplicate group " - "(winner=%s loser=%s -> %s); skipping", - winner_id, - loser_id, - resolved_winner, - ) - continue - - loser_doc = facts_by_id.get(resolved_loser_id) - if loser_doc is None and resolved_loser_id != loser_id: - # The original loser was just merged. Reuse the in-memory - # merged doc so we skip a cross-partition re-fetch — we - # own the (user_id, thread_id) partition and just wrote - # it. This in-memory copy carries the ``_etag`` returned - # by ``_upsert_memory``'s captured upsert response, so - # the supersede below takes the ETag-protected - # ``replace_item`` branch — concurrency-safe against any - # other reconcile that may have touched the same merged - # id in parallel. - loser_doc = merged_docs_by_id.get(resolved_loser_id) - + loser_doc = facts_by_id.get(loser_id) if loser_doc is None: - logger.warning( - "reconcile_memories: loser doc not found for pair %r (resolved_loser=%s)", - pair, - resolved_loser_id, - ) continue - - if await self._mark_superseded(loser_doc, resolved_winner, reason="contradict"): + if await self._mark_superseded(loser_doc, winner_id, reason="contradict"): contradicted += 1 - # Track the *original* loser_id from the LLM so the kept - # cross-check below can reconcile against the input pool. - if loser_id in facts_by_id: - consumed_loser_ids.add(loser_id) - # If the winner is an original pool member (not a freshly - # minted merged doc), record it so the kept-cross-check - # doesn't flag a clean run. - if winner_id in facts_by_id: - contradiction_winner_ids_in_pool.add(winner_id) - - # The pipeline's "kept" semantic = facts that survive as live - # records in the pool. The LLM's ``kept_ids`` semantic = - # everything *not* mentioned in duplicate_groups or - # contradicted_pairs. They differ by exactly the contradiction - # winners (winners survive but are listed under contradicted_pairs). - consumed_ids = consumed_source_ids | consumed_loser_ids - kept_actual = {fid for fid in facts_by_id.keys() if fid not in consumed_ids} - kept = len(kept_actual) - # Cross-check: the LLM's kept_ids set should equal kept_actual - # minus the contradiction winners. Mismatch usually means the LLM - # hallucinated an id or double-counted a fact across categories. - expected_llm_kept = kept_actual - contradiction_winner_ids_in_pool - llm_kept_set = {kid for kid in llm_kept_ids if kid in facts_by_id} - if llm_kept_set != expected_llm_kept: - symdiff = sorted(llm_kept_set ^ expected_llm_kept)[:10] - logger.info( - "reconcile_memories: kept_ids mismatch (llm=%d valid=%d, expected=%d). " - "Likely a hallucinated or double-counted fact id. Sample diff (≤10): %s", - len(llm_kept_ids), - len(llm_kept_set), - len(expected_llm_kept), - symdiff, - ) - result = {"kept": kept, "merged": merged, "contradicted": contradicted} - logger.info("reconcile_memories completed: %s", result) - return result, consumed_ids + consumed_loser_ids.add(loser_id) + + kept = len([fid for fid in facts_by_id if fid not in consumed_loser_ids]) + result = {"kept": kept, "merged": 0, "contradicted": contradicted} + logger.info( + "reconcile_memories contradiction pass user_id=%s memory_type=%s result=%s", + user_id, + memory_type, + result, + ) + return result async def build_procedural_context(self, user_id: str) -> str: """Return the active synthesized procedural prompt for system injection.""" diff --git a/azure/cosmos/agent_memory/aio/store/memory_store.py b/azure/cosmos/agent_memory/aio/store/memory_store.py index b7aec2e..4d5942b 100644 --- a/azure/cosmos/agent_memory/aio/store/memory_store.py +++ b/azure/cosmos/agent_memory/aio/store/memory_store.py @@ -32,6 +32,7 @@ from azure.cosmos.agent_memory.logging import get_logger from azure.cosmos.agent_memory.models import MemoryRecord from azure.cosmos.agent_memory.store._search_helpers import ( + MEMORY_PROJECTION, add_salience_filter, add_tag_filters, build_search_sql, @@ -460,7 +461,7 @@ async def delete( raise MemoryNotFoundError(memory_id=memory_id, user_id=user_id, thread_id=thread_id) from exc except CosmosAccessConditionFailedError as exc: raise MemoryConflictError( - f"delete conflicted for memory_id={memory_id!r} — document was modified after the type check" + f"delete conflicted for memory_id={memory_id!r} - document was modified after the type check" ) from exc except Exception as exc: raise _wrap_cosmos_exception(exc, message=f"async delete failed for {memory_id}: {exc}") from exc @@ -866,6 +867,65 @@ async def search( partition_key=partition_key, ) + async def get_memory_history( + self, + memory_id: str, + user_id: str, + thread_id: Optional[str] = None, + *, + max_depth: int = 20, + ) -> list[dict[str, Any]]: + """Return a memory's superseded predecessors, most-recently-superseded first. + + When a fact is updated or + contradicted, the prior document is retained with ``superseded_by`` + pointing at its replacement (see :meth:`mark_superseded`). This walks + that chain backwards from *memory_id* so callers can reason about how a + fact evolved - knowledge updates, preference reversals, relocations - + instead of seeing only the current value. + + The document identified by *memory_id* is not itself included; the return + value is everything it superseded, transitively, bounded by *max_depth* + to guard against cycles. Each entry carries the ``superseded_at`` / + ``supersede_reason`` audit fields so callers can order and explain the + transitions. + """ + if not memory_id: + raise ValidationError("memory_id is required") + if not user_id: + raise ValidationError("user_id is required") + partition_key, _ = query_scope(user_id, thread_id) + history: list[dict[str, Any]] = [] + seen: set[str] = {memory_id} + frontier: list[str] = [memory_id] + for _ in range(max(1, max_depth)): + id_params = [{"name": f"@sid{i}", "value": sid} for i, sid in enumerate(frontier)] + placeholders = ", ".join(param["name"] for param in id_params) + parameters: list[dict[str, Any]] = [*id_params, {"name": "@user_id", "value": user_id}] + where = f"c.superseded_by IN ({placeholders}) AND c.user_id = @user_id" + if thread_id is not None: + where += " AND c.thread_id = @thread_id" + parameters.append({"name": "@thread_id", "value": thread_id}) + sql = f"SELECT {MEMORY_PROJECTION} FROM c WHERE {where}" + rows = await self.query( + sql, + parameters, + container_key=ContainerKey.MEMORIES, + partition_key=partition_key, + ) + frontier = [] + for doc in rows: + doc_id = doc.get("id") + if not doc_id or doc_id in seen: + continue + seen.add(doc_id) + history.append(doc) + frontier.append(doc_id) + if not frontier: + break + history.sort(key=lambda d: d.get("superseded_at") or d.get("created_at") or "", reverse=True) + return history + async def search_turns( self, search_terms: Optional[str] = None, diff --git a/azure/cosmos/agent_memory/auto_trigger.py b/azure/cosmos/agent_memory/auto_trigger.py index a494181..ee41322 100644 --- a/azure/cosmos/agent_memory/auto_trigger.py +++ b/azure/cosmos/agent_memory/auto_trigger.py @@ -108,12 +108,6 @@ def maybe_trigger_steps( return n_dedup_turns = n_facts * n_dedup if n_facts > 0 and n_dedup > 0 else 0 - # Full-pool reconcile backstop cadence, derived from the PERSISTED counter so - # it fires reliably regardless of process/worker lifetime (the old in-memory - # per-instance sweep counter reset and under-fired). Every - # DEDUP_FULL_RECLUSTER_EVERY_N-th reconcile = every n_dedup_turns * N turns. - n_full_recluster = _threshold_int(thresholds, "get_dedup_full_recluster_every_n", "DEDUP_FULL_RECLUSTER_EVERY_N") - n_full_turns = n_dedup_turns * n_full_recluster if (n_dedup_turns > 0 and n_full_recluster > 0) else 0 user_batch_counts = _trigger_thread_steps( processor, counter_container, @@ -121,7 +115,6 @@ def maybe_trigger_steps( n_facts=n_facts, n_summary=n_summary, n_dedup_turns=n_dedup_turns, - n_full_turns=n_full_turns, thresholds=thresholds, ) _trigger_user_steps(processor, counter_container, user_batch_counts, n_user=n_user) @@ -135,7 +128,6 @@ def _trigger_thread_steps( n_facts: int, n_summary: int, n_dedup_turns: int, - n_full_turns: int, thresholds: Any = None, ) -> dict[str, int]: user_batch_counts: dict[str, int] = {} @@ -166,39 +158,11 @@ def _trigger_thread_steps( fire_extract=n_facts > 0 and _counters.crosses_threshold(old_count, new_count, n_facts), fire_summary=n_summary > 0 and _counters.crosses_threshold(old_count, new_count, n_summary), fire_dedup=n_dedup_turns > 0 and _counters.crosses_threshold(old_count, new_count, n_dedup_turns), - fire_full_rebuild=n_full_turns > 0 and _counters.crosses_threshold(old_count, new_count, n_full_turns), thresholds=thresholds, ) return user_batch_counts -def _watermark_recent_k( - counter_container: Any, - counter_id: str, - user_id: str, - thread_id: str, - *, - new_count: int, -) -> int: - """recent_k covering every turn since the last successful extract. - - ``new_count - watermark`` is exactly the count of turns added since the last - successful extract, and the newest-``recent_k`` slice in the pipeline picks - precisely those turns — so the whole backlog is covered and the watermark can - safely advance to ``new_count``. **Not capped:** capping would extract only the - newest N and strand the oldest ``backlog - N`` turns. - - **Bootstrap:** before a thread's first successful extract there is no watermark, - so we treat it as ``0`` — ``recent_k = new_count`` covers every turn the thread - has so far. Using only the current batch size here would strand turns added - during earlier *failed* extracts, because the watermark still advances to - ``new_count`` on the first success. - """ - watermark = _counters.read_extract_watermark_sync(counter_container, counter_id, user_id, thread_id) - base = watermark if watermark is not None else 0 - return max(new_count - base, 1) - - def _fire_thread_steps( processor: InProcessProcessor, counter_container: Any, @@ -210,7 +174,6 @@ def _fire_thread_steps( fire_extract: bool, fire_summary: bool, fire_dedup: bool, - fire_full_rebuild: bool = False, thresholds: Any = None, ) -> None: fire_procedural = fire_dedup and bool( @@ -222,16 +185,8 @@ def _fire_thread_steps( ) ) if fire_extract: - recent_k = _watermark_recent_k( - counter_container, - counter_id, - user_id, - thread_id, - new_count=new_count, - ) try: - processor.process_extract_memories(user_id=user_id, thread_id=thread_id, recent_k=recent_k) - _counters.advance_extract_watermark_sync(counter_container, counter_id, user_id, thread_id, new_count) + processor.process_extract_memories(user_id=user_id, thread_id=thread_id) except Exception as exc: logger.warning("Auto-trigger process_extract_memories failed for %s/%s: %s", user_id, thread_id, exc) _counters.stamp_failure_sync( @@ -241,7 +196,7 @@ def _fire_thread_steps( ( fire_dedup, "process_reconcile", - lambda: processor.process_reconcile(user_id=user_id, full_rebuild=fire_full_rebuild), + lambda: processor.process_reconcile(user_id=user_id), ), ( fire_procedural, diff --git a/azure/cosmos/agent_memory/chat.py b/azure/cosmos/agent_memory/chat.py index e036c6c..6376469 100644 --- a/azure/cosmos/agent_memory/chat.py +++ b/azure/cosmos/agent_memory/chat.py @@ -156,12 +156,6 @@ def _build_kwargs( if response_format is not None: kwargs["response_format"] = response_format kwargs.update(extra) - # Force temperature=1.0 across all callers. Newer Azure OpenAI models - # (gpt-5.x family, o-series reasoning models) only accept the default - # value (1.0) and reject any other; older models (gpt-4o, gpt-4o-mini) - # accept 1.0 as a valid value. Hardcoding to 1.0 keeps behavior uniform - # across the deployment matrix and lets prompt engineering — not a - # sampling knob — be the sole control for output determinism. kwargs["temperature"] = 1.0 return kwargs @@ -180,7 +174,7 @@ def generate( exponential backoff. Any additional keyword arguments (e.g. ``top_p``, ``seed``) are - forwarded directly to ``client.chat.completions.create`` — this lets + forwarded directly to ``client.chat.completions.create`` - this lets callers pass through ``model.parameters`` from a prompty file without modification. @@ -236,7 +230,7 @@ def generate( except openai.APIError as exc: status = getattr(exc, "status_code", None) # Reasoning models (gpt-5, o-series) reject custom sampling - # parameters with 400. Strip the offending param and retry — + # parameters with 400. Strip the offending param and retry - # this does NOT consume a retry slot since it's a request-shape # repair, not a transient failure. bad_param = unsupported_param(exc) if status == 400 else None diff --git a/azure/cosmos/agent_memory/cosmos_memory_client.py b/azure/cosmos/agent_memory/cosmos_memory_client.py index 7b0349a..61d31d6 100644 --- a/azure/cosmos/agent_memory/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/cosmos_memory_client.py @@ -519,7 +519,7 @@ def add_cosmos( For ``memory_type='turn'`` this also bumps the auto-trigger counter and schedules cadence-aware processing (extract / reconcile / procedural / thread_summary / user_summary) via :func:`maybe_trigger_steps`, exactly - like :meth:`push_to_cosmos` does for buffered turns — so the + like :meth:`push_to_cosmos` does for buffered turns - so the ``FACT_EXTRACTION_EVERY_N`` / ``THREAD_SUMMARY_EVERY_N`` / ``USER_SUMMARY_EVERY_N`` / ``DEDUP_EVERY_N`` knobs apply uniformly whether the caller uses the buffer or writes through directly. @@ -707,6 +707,31 @@ def search_turns( created_before=created_before, ) + def get_memory_history( + self, + memory_id: str, + user_id: str, + thread_id: Optional[str] = None, + *, + max_depth: int = 20, + ) -> list[dict[str, Any]]: + """Return a memory's superseded predecessors, most-recently-superseded first. + + AMT supersedes rather than deletes, so the full history of a fact - how a + preference, decision, or attribute changed over time - is preserved. This + walks the ``superseded_by`` chain backwards from *memory_id* and returns + every version it replaced (transitively), each carrying its + ``superseded_at`` / ``supersede_reason`` audit fields. Useful for + answering "what changed?" / "what was it before?" questions that a + current-value-only view cannot. + """ + return self._get_store().get_memory_history( + memory_id, + user_id, + thread_id, + max_depth=max_depth, + ) + def get_thread( self, thread_id: str, @@ -884,9 +909,7 @@ def reconcile(self, user_id: str, n: Optional[int] = None) -> dict[str, int]: """Reconcile a user's facts via the contradiction-aware dedup pass.""" from .thresholds import get_dedup_pool_size - return self._get_pipeline().reconcile_memories( - user_id, n if n is not None else get_dedup_pool_size(), full_rebuild=True - ) + return self._get_pipeline().reconcile_memories(user_id, n if n is not None else get_dedup_pool_size()) def process_now(self, *, user_id: str, thread_id: str) -> "ProcessThreadResult": """Force the processor to run the full pipeline RIGHT NOW for one thread. @@ -901,7 +924,7 @@ def process_now(self, *, user_id: str, thread_id: str) -> "ProcessThreadResult": caught and logged as warnings so the per-thread work already persisted by the prior steps is not erased. Permanent failures (config bugs, auth errors, 4xx Cosmos errors, Python builtins like - ``KeyError`` / ``TypeError``) are re-raised — silencing them turns + ``KeyError`` / ``TypeError``) are re-raised - silencing them turns operational issues into invisible ``WARNING`` lines. """ self._require_cosmos() diff --git a/azure/cosmos/agent_memory/embeddings.py b/azure/cosmos/agent_memory/embeddings.py index 1b7b4dd..65f08ca 100644 --- a/azure/cosmos/agent_memory/embeddings.py +++ b/azure/cosmos/agent_memory/embeddings.py @@ -146,7 +146,7 @@ def generate_batch( ConfigurationError If the endpoint or credentials are missing. openai.OpenAIError - Propagated from the SDK on API failure (no retry — see module + Propagated from the SDK on API failure (no retry - see module docstring). """ if not texts: diff --git a/azure/cosmos/agent_memory/models.py b/azure/cosmos/agent_memory/models.py index 6ae3526..ac6c795 100644 --- a/azure/cosmos/agent_memory/models.py +++ b/azure/cosmos/agent_memory/models.py @@ -1,8 +1,8 @@ """Pydantic data models for the Agent Memory Toolkit. -Six concrete record types — :class:`TurnRecord`, :class:`ThreadSummaryRecord`, +Six concrete record types - :class:`TurnRecord`, :class:`ThreadSummaryRecord`, :class:`UserSummaryRecord`, :class:`FactRecord`, :class:`EpisodicRecord`, -:class:`ProceduralRecord` — share :class:`MemoryRecordBase` and serialize +:class:`ProceduralRecord` - share :class:`MemoryRecordBase` and serialize to/from Cosmos DB-compatible dicts. ``MemoryRecord`` is exported as the union type for return-signature use and @@ -331,7 +331,7 @@ def _strip_unset_optional(data: dict[str, Any]) -> dict[str, Any]: class TurnRecord(MemoryRecordBase): - """A single conversation turn — raw user/agent/tool/system message.""" + """A single conversation turn - raw user/agent/tool/system message.""" memory_type: Literal[MemoryType.turn] = Field( # type: ignore[assignment] alias="type", default=MemoryType.turn @@ -392,11 +392,8 @@ def _require_thread_ids_in_metadata(self) -> "UserSummaryRecord": _FACT_ALLOWED_CATEGORIES = { "preference", "requirement", - "decision", "biographical", - "temporal", - "relational", - "action_item", + "other", } diff --git a/azure/cosmos/agent_memory/processors/base.py b/azure/cosmos/agent_memory/processors/base.py index 2fbf9b2..164c948 100644 --- a/azure/cosmos/agent_memory/processors/base.py +++ b/azure/cosmos/agent_memory/processors/base.py @@ -100,7 +100,6 @@ def process_reconcile( self, *, user_id: str, - full_rebuild: bool = False, ) -> int: ... def generate_user_summary( diff --git a/azure/cosmos/agent_memory/processors/durable.py b/azure/cosmos/agent_memory/processors/durable.py index 9834dcf..3a2b580 100644 --- a/azure/cosmos/agent_memory/processors/durable.py +++ b/azure/cosmos/agent_memory/processors/durable.py @@ -80,7 +80,7 @@ def process_user_summary( ) return UserSummaryResult(summary=None) - def process_reconcile(self, *, user_id: str, full_rebuild: bool = False) -> int: + def process_reconcile(self, *, user_id: str) -> int: logger.debug( "DurableFunctionProcessor.process_reconcile no-op user_id=%s", user_id, diff --git a/azure/cosmos/agent_memory/processors/inprocess.py b/azure/cosmos/agent_memory/processors/inprocess.py index aac1760..734d6b0 100644 --- a/azure/cosmos/agent_memory/processors/inprocess.py +++ b/azure/cosmos/agent_memory/processors/inprocess.py @@ -125,30 +125,25 @@ def process_user_summary( summary = self._pipeline.generate_user_summary(user_id, thread_ids) return UserSummaryResult(summary=summary if isinstance(summary, dict) else None) - def process_reconcile(self, *, user_id: str, full_rebuild: bool = False) -> int: - """Run reconciliation standalone. Returns count of facts merged + contradicted. + def process_reconcile(self, *, user_id: str) -> int: + """Run reconciliation standalone. Returns count of facts contradicted. Pool size is read from ``DEDUP_POOL_SIZE`` (env-tunable, default 50, capped at 500) so the auto-trigger and the standalone path agree. - ``full_rebuild`` (set by the auto-trigger on its persisted-counter - full-recluster cadence) forces the full-pool LLM pass that catches - dissimilar-embedding contradictions. """ from ..thresholds import get_dedup_pool_size - return self._reconcile_fact_and_episodic(user_id, get_dedup_pool_size(), full_rebuild=full_rebuild) + return self._reconcile_fact_and_episodic(user_id, get_dedup_pool_size()) - def _reconcile_fact_and_episodic(self, user_id: str, n: int, *, full_rebuild: bool = False) -> int: - """Reconcile facts and episodic memories; sum merged+contradicted counts. + def _reconcile_fact_and_episodic(self, user_id: str, n: int) -> int: + """Reconcile facts and episodic memories; sum contradicted counts. SDK in-process processing reconciles both types (matching the Durable - backend) so episodic dups don't accrue forever. + backend) so contradictions are resolved for each. """ total = 0 for memory_type in ("fact", "episodic"): - reconciled = self._pipeline.reconcile_memories( - user_id, n=n, memory_type=memory_type, full_rebuild=full_rebuild - ) + reconciled = self._pipeline.reconcile_memories(user_id, n=n, memory_type=memory_type) total += self._extract_reconcile_count(reconciled) return total @@ -157,7 +152,7 @@ def _extract_reconcile_count(reconciled: Any) -> int: """Sum ``merged + contradicted`` from a ``reconcile_memories`` result. ``ProcessingPipeline.reconcile_memories`` returns a dict with - ``{"kept", "merged", "contradicted"}`` — both ``merged`` and + ``{"kept", "merged", "contradicted"}`` - both ``merged`` and ``contradicted`` represent facts that were consolidated or retired, so they contribute to the dedup-count metric. """ diff --git a/azure/cosmos/agent_memory/prompts/_schemas.py b/azure/cosmos/agent_memory/prompts/_schemas.py index 6c3306c..1d98908 100644 --- a/azure/cosmos/agent_memory/prompts/_schemas.py +++ b/azure/cosmos/agent_memory/prompts/_schemas.py @@ -13,7 +13,7 @@ There are no truly optional fields; "optional" is expressed with ``"type": ["string", "null"]`` so the field is always present but may be ``null``. -* ``additionalProperties: false`` on every object — the model cannot +* ``additionalProperties: false`` on every object - the model cannot invent extra keys (e.g. ``reasoning`` or ``confidence`` siblings to the real payload that gpt-5.x was leaking into json_object outputs). @@ -26,25 +26,11 @@ from typing import Any # --------------------------------------------------------------------------- -# dedup.prompty — reconcile a pool of active facts +# dedup.prompty - reconcile a pool of active facts # --------------------------------------------------------------------------- DEDUP_SCHEMA: dict[str, Any] = { "type": "object", "properties": { - "duplicate_groups": { - "type": "array", - "items": { - "type": "object", - "properties": { - "merged_content": {"type": "string"}, - "source_ids": {"type": "array", "items": {"type": "string"}}, - "confidence": {"type": ["number", "null"]}, - "salience": {"type": ["number", "null"]}, - }, - "required": ["merged_content", "source_ids", "confidence", "salience"], - "additionalProperties": False, - }, - }, "contradicted_pairs": { "type": "array", "items": { @@ -60,13 +46,13 @@ }, "kept_ids": {"type": "array", "items": {"type": "string"}}, }, - "required": ["duplicate_groups", "contradicted_pairs", "kept_ids"], + "required": ["contradicted_pairs", "kept_ids"], "additionalProperties": False, } # --------------------------------------------------------------------------- -# dedup_episodic.prompty — reconcile a pool of active episodic memories +# dedup_episodic.prompty - reconcile a pool of active episodic memories # (MERGE-ONLY: same-event duplicates collapse; no contradiction/deletion) # --------------------------------------------------------------------------- DEDUP_EPISODIC_SCHEMA: dict[str, Any] = { @@ -94,7 +80,7 @@ # --------------------------------------------------------------------------- -# extract_memories.prompty — extract facts + episodic + unclassified +# extract_memories.prompty - extract facts + episodic + unclassified # --------------------------------------------------------------------------- _FACT_ITEM = { "type": "object", @@ -105,16 +91,10 @@ "enum": [ "preference", "requirement", - "decision", "biographical", - "temporal", - "relational", - "action_item", + "other", ], }, - "subject": {"type": ["string", "null"]}, - "predicate": {"type": ["string", "null"]}, - "object": {"type": ["string", "null"]}, "confidence": {"type": "number"}, "salience": {"type": "number"}, "temporal_context": {"type": ["string", "null"]}, @@ -123,9 +103,6 @@ "required": [ "text", "category", - "subject", - "predicate", - "object", "confidence", "salience", "temporal_context", @@ -170,33 +147,19 @@ "additionalProperties": False, } -_UNCLASSIFIED_ITEM = { - "type": "object", - "properties": { - "text": {"type": "string"}, - "confidence": {"type": "number"}, - "salience": {"type": "number"}, - "tags": {"type": "array", "items": {"type": "string"}}, - "reason": {"type": "string"}, - }, - "required": ["text", "confidence", "salience", "tags", "reason"], - "additionalProperties": False, -} - EXTRACT_MEMORIES_SCHEMA: dict[str, Any] = { "type": "object", "properties": { "facts": {"type": "array", "items": _FACT_ITEM}, "episodic": {"type": "array", "items": _EPISODIC_ITEM}, - "unclassified": {"type": "array", "items": _UNCLASSIFIED_ITEM}, }, - "required": ["facts", "episodic", "unclassified"], + "required": ["facts", "episodic"], "additionalProperties": False, } # --------------------------------------------------------------------------- -# summarize.prompty — first-pass thread summary +# summarize.prompty - first-pass thread summary # # Mirrors the 6-field shape the prompty actually instructs the model to emit # (see ``summarize.prompty`` lines ~69-88). Strict mode would silently drop @@ -238,14 +201,14 @@ # --------------------------------------------------------------------------- -# summarize_update.prompty — incremental thread summary update +# summarize_update.prompty - incremental thread summary update # Same shape as the first-pass schema; both prompties emit the same payload. # --------------------------------------------------------------------------- SUMMARIZE_UPDATE_SCHEMA: dict[str, Any] = SUMMARIZE_SCHEMA # --------------------------------------------------------------------------- -# user_summary.prompty — first-pass user profile +# user_summary.prompty - first-pass user profile # # Mirrors the 8 sections the prompty body documents (see # ``user_summary.prompty`` lines ~35-86 and the JSON example block). Each @@ -280,14 +243,14 @@ # --------------------------------------------------------------------------- -# user_summary_update.prompty — incremental user profile update +# user_summary_update.prompty - incremental user profile update # Same shape as the first-pass schema. # --------------------------------------------------------------------------- USER_SUMMARY_UPDATE_SCHEMA: dict[str, Any] = USER_SUMMARY_SCHEMA # --------------------------------------------------------------------------- -# synthesize_procedural.prompty — agent self-improvement / procedural prompt +# synthesize_procedural.prompty - agent self-improvement / procedural prompt # --------------------------------------------------------------------------- SYNTHESIZE_PROCEDURAL_SCHEMA: dict[str, Any] = { "type": "object", @@ -300,7 +263,7 @@ # --------------------------------------------------------------------------- -# Registry — maps prompty filename → (schema_name, schema_dict) +# Registry - maps prompty filename → (schema_name, schema_dict) # --------------------------------------------------------------------------- PROMPTY_SCHEMAS: dict[str, tuple[str, dict[str, Any]]] = { "dedup.prompty": ("DedupOutput", DEDUP_SCHEMA), diff --git a/azure/cosmos/agent_memory/prompts/dedup.prompty b/azure/cosmos/agent_memory/prompts/dedup.prompty index d45bb60..21a80d9 100644 --- a/azure/cosmos/agent_memory/prompts/dedup.prompty +++ b/azure/cosmos/agent_memory/prompts/dedup.prompty @@ -1,7 +1,7 @@ --- name: dedup version: v1 -description: Reconcile a pool of active facts — collapse duplicates and resolve semantic contradictions in one pass. +description: Resolve semantic contradictions within a pool of active facts. model: apiType: chat options: @@ -16,36 +16,22 @@ inputs: --- system: -You are a precision fact-reconciliation system. You receive a pool of active facts (each with an ID, content, confidence, salience, and creation timestamp) and must classify them into three orthogonal buckets: duplicates, contradictions, and unique kept facts. +You are a precision fact-reconciliation system. You receive a pool of active facts (each with an ID, content, confidence, salience, and creation timestamp) and must find **contradictions** — pairs of facts that assert opposing claims about the same subject — and pick a winner and a loser for each. ## Your Goal Produce a clean reconciliation that: -1. Collapses paraphrases of the same claim into a single merged fact. -2. Resolves opposing claims about the same subject by picking a winner and a loser. -3. Leaves everything else untouched. +1. Resolves opposing claims about the same subject by picking a winner and a loser. +2. Leaves everything else untouched. -## Three Orthogonal Outcomes (mutually exclusive, exhaustive) +Near-duplicate paraphrases are NOT your concern here — they are folded together earlier, at write time. Do not merge, rewrite, or collapse facts. Your only job is to identify genuine contradictions. -Every input fact must end up in **at least one** of these three places: -- `duplicate_groups[*].source_ids` — the fact is a paraphrase of one or more other facts. -- `contradicted_pairs[*]` — the fact is either the `winner_id` or the `loser_id` of a contradiction. -- `kept_ids` — the fact is unique: neither paraphrased by nor in opposition to anything else. - -The only legal multi-membership is the dual case described in the **Interaction Rule** below: a duplicate-group source that is also a contradiction loser/winner. Otherwise, place each fact in exactly one bucket. No input fact may be omitted from the output. - -## What is a DUPLICATE +## Two Outcomes -Two or more facts that restate the **same claim about the same subject** in different words. - -Examples: -- "User prefers aisle seats" vs "User likes aisle seats when flying" → duplicate. -- "User works at Acme Corp" vs "User is a senior data engineer at Acme Corp" → duplicate (the more specific claim subsumes the less specific one — see "Partial overlap" below). +Every input fact ends up in one of two places: +- `contradicted_pairs[*]` — the fact is either the `winner_id` or the `loser_id` of a contradiction. +- `kept_ids` — the fact is not in opposition to anything else. -Resolution: emit one entry in `duplicate_groups` with: -- `merged_content` — a clean, self-contained restatement (third person, under 40 words for facts). Do NOT concatenate the originals; synthesize. -- `source_ids` — every original ID that participates in the duplicate group. -- `confidence` — the **max** confidence among the source facts. -- `salience` — the **max** salience among the source facts. +A fact that is merely a paraphrase of, or on the same topic as, another fact (but not contradicting it) belongs in `kept_ids`. ## What is a CONTRADICTION @@ -69,25 +55,19 @@ A contradiction is a pair. If three or more facts mutually contradict, emit pair ## What is KEPT -Facts that are unique — neither a paraphrase of, nor in opposition to, anything else in the pool. List their IDs in `kept_ids`. +Every fact that is not the loser or winner of a contradiction. List their IDs in `kept_ids`. Paraphrases and same-topic-but-different-claim facts go here — they are NOT contradictions. -## Interaction Rule (dangling references) +## Rules -When a fact participates in **both** a duplicate group **and** a contradiction, express the contradiction in terms of one of the original `source_ids` (any one is fine). The downstream pipeline will redirect the contradiction to the merged document automatically. +**Do NOT invent new IDs.** Only IDs that appear in the input pool are valid in `winner_id` or `loser_id`. -**Do NOT invent new IDs** for merged documents in the `winner_id` or `loser_id` fields. Only IDs that appear in the input pool are valid. +1. **Conservative bias.** If you cannot confidently classify two facts as contradictions, put them both in `kept_ids`. A false contradiction (dropping a fact the user never retracted) is worse than retaining both. -## Decision Guidelines +2. **Topic vs claim distinction.** Two facts on the same topic but making *different, compatible* claims are NOT a contradiction. Example: "User prefers dark mode" and "User uses VS Code" share a topic but do not oppose each other — both go in `kept_ids`. -1. **Conservative bias.** If you cannot confidently classify two facts as duplicates or as contradictions, put them both in `kept_ids`. Over-merging or false-contradicting is worse than retaining a redundancy. +3. **Opposition, not refinement.** A more specific restatement of the same claim ("works at Acme" → "senior engineer at Acme") is NOT a contradiction — both go in `kept_ids`. -2. **Topic vs claim distinction.** Two facts on the same topic but making different claims are NOT duplicates. Example: "User prefers dark mode" and "User uses VS Code" share the UX-tooling topic but assert different things — both go in `kept_ids`. - -3. **Partial overlap → duplicate via merging.** If one fact is a strict refinement of another (same subject, same direction of claim, more detail), treat them as a duplicate group and let `merged_content` carry the more specific claim. Example: "User works at Acme Corp" + "User is a senior data engineer at Acme Corp" → merged_content = "User is a senior data engineer at Acme Corp". - -4. **Don't drop information.** A merged_content must preserve every distinct piece of information from its sources. If you cannot preserve everything in one clean sentence, the facts probably aren't true duplicates — keep them separate. - -5. **Don't fabricate.** Never introduce facts, qualifiers, or entities that are not present in at least one source. +4. **Don't fabricate.** Never introduce facts, qualifiers, or entities that are not present in the pool. ## Input Format @@ -100,11 +80,8 @@ N. ID: | Content: "" | Confidence: 0.85 | Salience: 0.7 | Created: Some facts may show `Confidence: N/A`, `Salience: N/A`, or `Created: N/A`. -- Treat `N/A` confidence/salience as **unknown** — set those fields to `null` - in the `duplicate_groups` output (do NOT omit the keys); the pipeline will - fall back to `max(source.confidence)` / `max(source.salience)`. -- Treat `N/A` created_at as **older than any real timestamp** for - winner-selection purposes (a fact with a real date beats one with N/A). +- Treat `N/A` confidence/salience as **unknown** for tiebreaking purposes. +- Treat `N/A` created_at as **older than any real timestamp** for winner-selection purposes (a fact with a real date beats one with N/A). ## Output Format @@ -112,18 +89,12 @@ You must output ONLY valid JSON matching this exact schema. No preamble, no expl ```json { - "duplicate_groups": [ - {"merged_content": "", "source_ids": ["", ""], "confidence": 0.0, "salience": 0.0} - ], "contradicted_pairs": [ {"winner_id": "", "loser_id": "", "reason": ""} ], "kept_ids": ["", ""] } ``` - - -The `confidence` and `salience` values **must** be real numbers between 0 and 1 in the actual output — the `0.0` above is a structural placeholder, not a literal value to echo. Compute them as the maximum across source facts in the group. If you cannot determine confidence or salience, set the field to `null` (the pipeline will fall back to `max(source.*)` from the source records). **Do not omit the keys** — the response schema requires both fields to appear on every group, with either a number or `null` as the value. If a bucket is empty, emit it as an empty array (`[]`) rather than omitting the key. @@ -134,30 +105,23 @@ If a bucket is empty, emit it as an empty array (`[]`) rather than omitting the 1. ID: F1 | Content: "User is vegetarian" | Confidence: 0.9 | Salience: 0.9 | Created: 2024-01-01T00:00:00Z 2. ID: F2 | Content: "User loves a good ribeye steak" | Confidence: 0.95 | Salience: 0.9 | Created: 2024-01-09T00:00:00Z 3. ID: F3 | Content: "User prefers aisle seats on flights" | Confidence: 0.9 | Salience: 0.7 | Created: 2024-01-06T00:00:00Z -4. ID: F4 | Content: "User likes aisle seats when flying" | Confidence: 0.85 | Salience: 0.7 | Created: 2024-01-07T00:00:00Z -5. ID: F5 | Content: "User works at Acme Corp" | Confidence: 0.9 | Salience: 0.8 | Created: 2023-12-12T00:00:00Z -6. ID: F6 | Content: "User is a senior data engineer at Acme Corp" | Confidence: 0.9 | Salience: 0.8 | Created: 2024-01-10T00:00:00Z +4. ID: F4 | Content: "User is a senior data engineer at Acme Corp" | Confidence: 0.9 | Salience: 0.8 | Created: 2024-01-10T00:00:00Z ``` **Expected output:** ```json { - "duplicate_groups": [ - {"merged_content": "User prefers aisle seats on flights", "source_ids": ["F3", "F4"], "confidence": 0.9, "salience": 0.7}, - {"merged_content": "User is a senior data engineer at Acme Corp", "source_ids": ["F5", "F6"], "confidence": 0.9, "salience": 0.8} - ], "contradicted_pairs": [ {"winner_id": "F2", "loser_id": "F1", "reason": "F2 is more recent and asserts a contradictory dietary preference"} ], - "kept_ids": [] + "kept_ids": ["F3", "F4"] } ``` Notice: -- F3 and F4 are paraphrases of the same aisle-seat preference → one duplicate group, merged_content keeps the cleaner phrasing, confidence = max(0.9, 0.85) = 0.9, salience = max(0.7, 0.7) = 0.7. -- F5 and F6 are partial-overlap duplicates; the more specific claim becomes merged_content. - F1 and F2 directly contradict each other on diet; F2 wins because it is more recent. -- Every input ID (F1–F6) appears exactly once across the three buckets. `kept_ids` is empty here because nothing is left over. +- F3 and F4 make unique, non-opposing claims → both kept. +- Every input ID (F1–F4) appears exactly once across the two buckets. user: Reconcile the following pool of active facts: diff --git a/azure/cosmos/agent_memory/prompts/extract_memories.prompty b/azure/cosmos/agent_memory/prompts/extract_memories.prompty index af1b25a..1ce30a9 100644 --- a/azure/cosmos/agent_memory/prompts/extract_memories.prompty +++ b/azure/cosmos/agent_memory/prompts/extract_memories.prompty @@ -1,7 +1,7 @@ --- name: extract_memories version: v1 -description: Extract facts and episodic memories from a conversation, reconciled against existing memories. +description: Extract facts and episodic memories from a conversation or provided content. model: apiType: chat options: @@ -11,31 +11,33 @@ model: response_format: type: json_object inputs: - existing_facts: - type: string - default: '(none)' transcript: type: string --- system: -You are a precision memory extraction system. Your task is to read the provided content — a conversation thread, and/or documents or reference material the user has shared — and extract structured memories worth retaining for future reference. You extract two distinct types of memory, each serving a different purpose. +You are a precision memory extraction system. Your task is to read the provided content - a conversation thread, and/or documents or reference material the user has shared - and extract structured memories worth retaining for future reference. You extract two distinct types of memory, each serving a different purpose. ## Memory Type Taxonomy -1. **Facts** — Declarative knowledge: concrete, verifiable statements about the user, their environment, decisions, or requirements — or salient factual information from the content the user is working with (documents, excerpts, or domain material they provide). These are things that *are* true. -2. **Episodic** — Past experiences: specific situations the user encountered, what they tried, and what happened. These are things that *happened*. -3. **Unclassified** — Use sparingly, only when an item is clearly worth retaining but you cannot confidently decide between fact / episodic. The pipeline will store these as facts with a `sys:unclassified` tag for later audit. Prefer this over forcing a wrong classification. +1. **Facts** - Declarative knowledge: concrete, verifiable statements about the user, their environment, decisions, or requirements - or salient factual information from the content the user is working with (documents, excerpts, or domain material they provide). These are things that *are* true. +2. **Episodic** - Past experiences: specific situations the user encountered, what they tried, and what happened. These are things that *happened*. + +Every memory must be explicitly grounded in the provided content - never inferred, assumed, or speculated. -Every memory must be explicitly grounded in the provided content — never inferred, assumed, or speculated. +## Extraction Discipline - Read This First -## Speaker Discrimination — Where Memories May Come From +- **Extract only what was explicitly stated. Do not infer mental states.** Never manufacture a memory about what the user *thinks*, *feels*, *believes*, *wants*, *likes*, *enjoys*, *finds*, *prefers*, or *is worried/frustrated about* unless the user said so in those (or clearly equivalent) words. Turning "the API took 8 seconds" into "the user thinks the API is slow" is inference, not extraction - do not do it. +- **When the user states a concrete fact, record that fact once - and stop.** Do not also emit a second, softer "the user thinks/feels X" memory derived from it. Record each piece of information once; don't restate the same fact under a different framing. +- **Fewer, higher-value memories beat a long padded list.** If a passage is trivia, filler, or has no future-reference value, extract nothing from it. An empty array is a valid - and often the correct - answer for a given passage. You are not rewarded for volume. + +## Speaker Discrimination - Where Memories May Come From The transcript below is line-tagged: `[user]:` lines are the human's own words, `[agent]:` lines are the agent's response. These two sources are NOT interchangeable. -- **Facts may ONLY come from `[user]:` lines.** The agent may restate, paraphrase, confirm, or make assumptions on the user's behalf ("Got it, you don't eat meat", "I assume you want a luxury hotel") — those are the agent's response, NOT the user's assertion. Never treat agent text as a new source of facts. If a fact appears only in `[agent]:` content and is not present in the user's lines, do not extract it. -- **Episodic memories may use both speakers' content** — the user's stated intent or scope is the anchor (and must be present in `[user]:`), but the agent's content may help fill in the `action_taken` or `outcome` of a `situation → action_taken → outcome` arc when the agent carried out the action on the user's behalf. -- Unsolicited world-knowledge the agent volunteers in `[agent]:` lines (e.g. "Python 3.13 was released in October 2024") is the agent's answer, not a memory. However, factual and domain content the user provides, shares, or asks to remember — including documents, excerpts, and reference material in `[user]:` lines — IS extractable. +- **Facts may ONLY come from `[user]:` lines.** The agent may restate, paraphrase, confirm, or make assumptions on the user's behalf ("Got it, you don't eat meat", "I assume you want a luxury hotel") - those are the agent's response, NOT the user's assertion. Never treat agent text as a new source of facts. If a fact appears only in `[agent]:` content and is not present in the user's lines, do not extract it. +- **Episodic memories may use both speakers' content** - the user's stated intent or scope is the anchor (and must be present in `[user]:`), but the agent's content may help fill in the `action_taken` or `outcome` of a `situation → action_taken → outcome` arc when the agent carried out the action on the user's behalf. +- Unsolicited world-knowledge the agent volunteers in `[agent]:` lines (e.g. "Python 3.13 was released in October 2024") is the agent's answer, not a memory. However, factual and domain content the user provides, shares, or asks to remember - including documents, excerpts, and reference material in `[user]:` lines - IS extractable. ## Confidence Scoring @@ -45,14 +47,14 @@ Every extracted memory (in any bucket) must include a `confidence` field in `[0. |-------|---------| | 0.9–1.0 | Directly stated and unambiguous | | 0.7–0.9 | Clearly implied, no contradicting evidence | -| 0.5–0.7 | Inferred from context — plausible but not explicit | -| < 0.5 | Speculative — should probably be in `unclassified` instead | +| 0.5–0.7 | Inferred from context - plausible but not explicit | +| < 0.5 | Speculative - likely should not be extracted at all | Retrieval may filter by `min_confidence` so under-confident extractions get suppressed automatically. --- -## 1. Facts — Declarative Knowledge +## 1. Facts - Declarative Knowledge Extract concrete, factual statements that fall into these categories: @@ -60,47 +62,52 @@ Extract concrete, factual statements that fall into these categories: |----------|-------------|---------| | `preference` | Things the user likes, dislikes, or prefers | "User prefers dark mode in all IDEs" | | `requirement` | Explicit constraints or needs | "The project must comply with SOC 2" | -| `decision` | Conclusions reached or choices confirmed | "User chose PostgreSQL over MySQL for the backend" | | `biographical` | Stable personal or professional details | "User is a senior engineer at Contoso" | -| `temporal` | Time-bound facts (deadlines, dates, schedules) | "The MVP deadline is March 1st, 2025" | -| `relational` | Relationships between entities | "User reports to Sarah, the VP of Engineering" | -| `action_item` | Commitments, tasks, or follow-ups | "User will send the API spec by Friday" | +| `other` | Any factual claim that doesn't fit the above (decisions, events, dates, relationships, tasks, domain facts, etc.) | "User chose PostgreSQL over MySQL"; "The MVP deadline is March 1st, 2025" | + +Category is a lightweight hint, not a hard decision - when a fact doesn't clearly fit `preference`, +`requirement`, or `biographical`, use `other`. **Never withhold or drop a fact because it is hard to +categorize** - extracting the fact matters far more than labeling it. + +**Capturing time:** Whenever a fact involves any time reference - an exact date, a duration, a +frequency, a deadline, or a relative expression ("yesterday", "3 weeks ago", "last June", "every +Monday") - keep that time expression verbatim in the fact `text` AND record it in `temporal_context`. +Time-bound answers ("when did X happen?", "how long?") are impossible to recover later if the time +detail is dropped, so never omit or round it. ### Fact Formatting Rules -- Each fact must be self-contained and intelligible without context — no pronouns like "it" or "they" without antecedents +- Each fact must be self-contained and intelligible without context - no pronouns like "it" or "they" without antecedents - Write in third person ("The user...", "The project...") -- Keep each fact concise — under 40 words -- Consolidate closely related items into a single fact (e.g., multiple search results on the same topic) -- Only split into separate facts when the claims are about genuinely different topics -- Each fact will be stored as its own document with its own vector embedding — self-contained phrasing is essential for retrieval - -### Avoiding Duplicates -Before adding a fact, check it against the existing memories provided below. **Only emit facts that are genuinely new.** If a fact is already captured in existing memories with no meaningful change, simply omit it. Do not try to update, merge, or flag contradictions with existing facts — that is handled later by a separate reconciliation step. Your only job here is to extract new facts and episodic memories from the transcript. +- **Preserve every specific detail stated in the source: exact dates, times, durations, numbers, quantities, amounts, proper nouns, names, brands, product/model names, and named locations or events.** NEVER generalize a specific into a category - keep "June 3rd" (not "a date"), "7 days" (not "about a week"), "Fitbit Versa 3" (not "a smartwatch"), "$4,500" (not "some money"). A fact that drops the specific detail is useless for later recall - the specific IS the memory. +- **Group by topic - one contextually rich memory per distinct topic or event, not one per claim.** Capture the whole coherent picture (the core fact plus its immediate context, qualifiers, and every specific) in a single self-contained memory instead of shattering it into fragments. Split into separate memories only when the content spans genuinely distinct topics (e.g. career vs. family vs. a trip) or when one topic is so detail-dense that a single memory would run past ~3 sentences - then split along natural sub-groupings (e.g. one memory per enemy type in a stat block), never by shaving off individual details. Do NOT merge two unrelated topics to save space, and do NOT split one topic just to make memories smaller. +- **Capture transitions as one memory - the new state AND what it replaced.** When the user changes, switches, replaces, upgrades, or stops something in favor of something else, the link between old and new is essential context: record both together. "The user switched from almond milk to oat milk after developing an almond sensitivity" - not two disconnected facts. If the change is explicitly temporary or a trial ("for a month", "trying out"), capture that too. Keeping old→new in one memory prevents the two states from later surfacing as contradictory-looking facts. +- Keep phrasing tight, but never generalize away a concrete detail (exact date, number, duration, name, brand) to save words. This preserves detail *within* the facts you choose to keep - it is not a directive to maximize how many facts you extract. +- Each memory is stored as its own document with its own vector embedding, so keep every memory to a **single coherent topic** - blending two unrelated topics into one memory produces a muddy embedding that retrieves poorly for both. The target is rich-but-focused: all the specifics of one topic, nothing borrowed from another. --- -## 2. Episodic — Situated Memories +## 2. Episodic - Situated Memories Extract memories that are tied to a **specific situation, scope, or context** the user is in. Episodic covers three cases: -- **Planned / in-flight** — a stated intent or preference scoped to a particular trip, project, event, session, or other bounded context that has not yet finished. -- **Past with outcome** — a completed event the user described, following the `situation → action_taken → outcome` pattern. -- **Ongoing context** — a temporary state of affairs that will not persist as a standing fact (e.g. "right now I'm focused on X"). +- **Planned / in-flight** - a stated intent or preference scoped to a particular trip, project, event, session, or other bounded context that has not yet finished. +- **Past with outcome** - a completed event the user described, following the `situation → action_taken → outcome` pattern. +- **Ongoing context** - a temporary state of affairs that will not persist as a standing fact (e.g. "right now I'm focused on X"). ### Required Fields -- **scope_type** — short, free-form noun describing the kind of context (e.g. `trip`, `project`, `event`, `session`, `release`, `campaign`). Pick whatever vocabulary fits the user's domain. Do not invent a value if one is not implied — if no scope is present, the memory probably belongs in facts. -- **scope_value** — the specific instance of that scope (e.g. `Paris 2025`, `Acme revamp`, `Q3 launch`). +- **scope_type** - short, free-form noun describing the kind of context (e.g. `trip`, `project`, `event`, `session`, `release`, `campaign`). Pick whatever vocabulary fits the user's domain. Do not invent a value if one is not implied - if no scope is present, the memory probably belongs in facts. +- **scope_value** - the specific instance of that scope (e.g. `Paris 2025`, `Acme revamp`, `Q3 launch`). Both must be non-empty. ### Optional Fields (include only when applicable) -- **situation** — the context or problem faced (present for past/in-flight events) -- **action_taken** — what was specifically tried or done (present for past events) -- **outcome** — what actually happened as a result (present only when the event has concluded) -- **outcome_valence** — `positive` | `negative` | `mixed` | `neutral` (present only with `outcome`) -- **reasoning** — why it worked or failed -- **lesson** — a transferable takeaway -- **domain** — topic area +- **situation** - the context or problem faced (present for past/in-flight events) +- **action_taken** - what was specifically tried or done (present for past events) +- **outcome** - what actually happened as a result (present only when the event has concluded) +- **outcome_valence** - `positive` | `negative` | `mixed` | `neutral` (present only with `outcome`) +- **reasoning** - why it worked or failed +- **lesson** - a transferable takeaway +- **domain** - topic area For planned/in-flight or ongoing-context memories, leave `situation`, `action_taken`, `outcome`, `outcome_valence`, `reasoning`, and `lesson` as `null`. The scope fields alone carry the meaning. @@ -120,22 +127,7 @@ This produces: When in doubt: - If it's a **standing state** ("X is true", with no bounded context) → Fact - If it's a **story** ("We tried X and Y happened") → Episodic -- If it's a **state scoped to a particular context** — a trip, project, event, session, release, campaign, or any other bounded container — → Episodic, with `scope_type` and `scope_value` filled in. The classic trap is a preference that sounds general but is qualified by "for this trip / on this project / just for today" — those are episodic, not fact. - ---- - -## Existing Memories -The following facts already exist for this user. They are provided so you can avoid re-extracting facts that are already known. If a fact from the new transcript is already captured here, omit it. - -### Existing Memories Are Context, Not a Source - -The existing memories block is provided ONLY for grounding and for deciding whether a fact from the new transcript is already known and should be omitted. It is NOT part of the conversation transcript. - -- New memories must come exclusively from the new transcript below. -- Never extract a fact or episodic memory solely because it appears in existing memories. -- If information appears only in existing memories and not in the new transcript, do not emit it. - -{{existing_facts}} +- If it's a **state scoped to a particular context** - a trip, project, event, session, release, campaign, or any other bounded container - → Episodic, with `scope_type` and `scope_value` filled in. The classic trap is a preference that sounds general but is qualified by "for this trip / on this project / just for today" - those are episodic, not fact. --- @@ -155,10 +147,10 @@ Assign a salience score to every extracted memory using this scale: ## What to Exclude Do NOT extract: -- Opinions or speculation ("User thinks the API might be slow") +- **Inferred opinions, feelings, or mental states** - e.g. "User thinks the API might be slow", "User enjoys concise answers", "User is frustrated with the tool". Never write "the user thinks / feels / believes / wants / likes / enjoys / finds X" unless the user literally said it. Synthesizing the user's attitude from what they did or described is inference, not extraction. - Filler or pleasantries ("User said thanks") - Uncertain or hypothetical statements ("User mentioned they might switch tools") -- Redundant memories — if the same information appears multiple times, extract it only once +- Redundant memories - if the same information appears multiple times, extract it only once - Raw agent reasoning or intermediate steps that did not produce a confirmed outcome - Memories that are only meaningful within the context of this conversation and have no future reference value @@ -177,8 +169,8 @@ A fact is a standing claim that holds outside any specific context. The test: if **For Facts:** 1. Could someone act on or reference this fact without reading the original thread? 2. Is this fact stated explicitly, not inferred? -3. Is each fact truly atomic — one claim per entry? -4. Can any facts be merged because they describe variants of the same thing? +3. Is each memory scoped to a single coherent topic, with all of that topic's specifics captured together - not shattered into fragments, and not blended with a second topic? +4. Does each fact retain every specific detail (dates, numbers, durations, names, brands) from the source, with nothing generalized away? **For Episodic:** 1. Did this event actually happen, or is it hypothetical? @@ -199,33 +191,16 @@ A fact is a standing claim that holds outside any specific context. The test: if { "facts": [ { - "text": "The user's name is Alex.", + "text": "The user is Alex, a data engineer at Acme Corp.", "category": "biographical", - "subject": "user", - "predicate": "name", - "object": "Alex", "confidence": 1.0, "salience": 0.9, "temporal_context": null, - "tags": ["topic:identity"] - }, - { - "text": "Alex is a data engineer at Acme Corp.", - "category": "biographical", - "subject": "Alex", - "predicate": "role_at", - "object": "data engineer at Acme Corp", - "confidence": 1.0, - "salience": 0.8, - "temporal_context": null, "tags": ["topic:identity", "topic:career"] }, { - "text": "Alex's new ETL pipeline project has a deadline at end of Q2.", - "category": "temporal", - "subject": "ETL pipeline project", - "predicate": "deadline", - "object": "end of Q2", + "text": "Alex's team just kicked off a new ETL pipeline project with a deadline at end of Q2.", + "category": "other", "confidence": 0.95, "salience": 0.9, "temporal_context": "end of Q2", @@ -236,6 +211,8 @@ A fact is a standing claim that holds outside any specific context. The test: if } ``` +Name and role are one coherent biographical topic → a single rich memory, not two fragments. The ETL project is a distinct topic → its own memory. + ### Example 2: Troubleshooting experience **Conversation:** @@ -247,10 +224,7 @@ A fact is a standing claim that holds outside any specific context. The test: if "facts": [ { "text": "The user's team increased Kubernetes pod memory limits from 512MB to 1GB after OOM-killing issues.", - "category": "decision", - "subject": "user's team", - "predicate": "changed", - "object": "Kubernetes pod memory limits from 512MB to 1GB", + "category": "other", "confidence": 0.95, "salience": 0.7, "temporal_context": "last month", @@ -288,9 +262,6 @@ A fact is a standing claim that holds outside any specific context. The test: if { "text": "The user's analytics database runs on BigQuery.", "category": "biographical", - "subject": "user", - "predicate": "analytics_database", - "object": "BigQuery", "confidence": 0.95, "salience": 0.7, "temporal_context": null, @@ -299,9 +270,6 @@ A fact is a standing claim that holds outside any specific context. The test: if { "text": "The marketing team's budget is $50,000 for the current quarter.", "category": "requirement", - "subject": "marketing team", - "predicate": "budget", - "object": "$50,000 for the current quarter", "confidence": 0.95, "salience": 0.9, "temporal_context": "current quarter", @@ -318,7 +286,7 @@ A fact is a standing claim that holds outside any specific context. The test: if > User: "I usually prefer budget hotels." > User: "For this Paris trip, I want luxury accommodations." -The first statement is a standing preference and belongs in `facts`. The second is qualified by "for this Paris trip" — dropping the scope changes the meaning, so it belongs in `episodic` with the scope captured structurally. The two memories coexist; the trip-scoped intent never enters the fact pool and never collides with the standing preference. +The first statement is a standing preference and belongs in `facts`. The second is qualified by "for this Paris trip" - dropping the scope changes the meaning, so it belongs in `episodic` with the scope captured structurally. The two memories coexist; the trip-scoped intent never enters the fact pool and never collides with the standing preference. **Output:** ```json @@ -327,9 +295,6 @@ The first statement is a standing preference and belongs in `facts`. The second { "text": "The user usually prefers budget hotels.", "category": "preference", - "subject": "user", - "predicate": "hotel_preference", - "object": "budget hotels", "confidence": 0.95, "salience": 0.7, "temporal_context": null, @@ -355,24 +320,69 @@ The first statement is a standing preference and belongs in `facts`. The second } ``` +### Example 5: Discipline - extract only what was stated, infer nothing + +**Conversation:** +> User: "The dashboard took about 8 seconds to load the report. Anyway, the weather's been nice this week." + +Extract only the literal, user-stated fact. The load time is concrete and worth keeping. "The weather's been nice" is filler with no future-reference value → excluded. Do NOT manufacture inferred opinions such as "The user thinks the dashboard is slow" or "The user enjoys nice weather" - the user asserted neither. + +**Output:** +```json +{ + "facts": [ + { + "text": "The user's dashboard took about 8 seconds to load the report.", + "category": "other", + "confidence": 0.9, + "salience": 0.4, + "temporal_context": null, + "tags": ["topic:performance", "topic:dashboard"] + } + ], + "episodic": [] +} +``` + +### Example 6: Capturing a transition as a single memory + +**Conversation:** +> User: "I moved our CI from Jenkins to GitHub Actions last quarter - the Jenkins maintenance was eating too much of my time." + +The change is one coherent memory: it records the new state, what it replaced, and why. Do NOT split it into a disconnected "uses GitHub Actions" plus "used Jenkins" - that loses the relationship and later looks like two conflicting facts. + +**Output:** +```json +{ + "facts": [ + { + "text": "The user's team migrated CI from Jenkins to GitHub Actions last quarter because Jenkins maintenance was consuming too much of the user's time.", + "category": "other", + "confidence": 0.95, + "salience": 0.7, + "temporal_context": "last quarter", + "tags": ["topic:CI", "topic:tooling"] + } + ], + "episodic": [] +} +``` + --- ## Output Format -You must output ONLY valid JSON matching the schema below. No preamble, no explanation, no closing remarks — just the JSON object. Each array can be empty if no memories of that type are found. Omit any fact that is already captured in the existing memories. +You must output ONLY valid JSON matching the schema below. No preamble, no explanation, no closing remarks - just the JSON object. Each array can be empty if no memories of that type are found. ```json { "facts": [ { "text": "Self-contained fact for embedding", - "category": "preference|requirement|decision|biographical|temporal|relational|action_item", - "subject": "entity", - "predicate": "relationship", - "object": "value", + "category": "preference|requirement|biographical|other", "confidence": 0.95, "salience": 0.8, - "temporal_context": "optional time ref or null", + "temporal_context": "the exact time expression stated (e.g. 'June 3rd', '3 weeks ago', 'by Friday', 'every Monday') or null if none", "tags": ["topic:x"] } ], @@ -391,15 +401,6 @@ You must output ONLY valid JSON matching the schema below. No preamble, no expla "salience": 0.7, "tags": ["topic:x"] } - ], - "unclassified": [ - { - "text": "Self-contained statement worth retaining", - "confidence": 0.5, - "salience": 0.5, - "tags": ["topic:x"], - "reason": "why this could not be classified as fact/episodic" - } ] } ``` diff --git a/azure/cosmos/agent_memory/services/_pipeline_helpers.py b/azure/cosmos/agent_memory/services/_pipeline_helpers.py index 74b13d6..4bc36f9 100644 --- a/azure/cosmos/agent_memory/services/_pipeline_helpers.py +++ b/azure/cosmos/agent_memory/services/_pipeline_helpers.py @@ -16,11 +16,59 @@ from pathlib import Path from typing import Any, Iterable, Optional +from azure.cosmos.agent_memory._embedding_tokens import count_tokens from azure.cosmos.agent_memory.exceptions import LLMError from azure.cosmos.agent_memory.logging import get_logger logger = get_logger(__name__) +_NON_RETRYABLE_LLM_MARKERS = ( + "content_filter", + "content management policy", + "context_length_exceeded", + "maximum context length", +) + + +def is_retryable_llm_error(exc: BaseException) -> bool: + """Classify an extraction LLM failure as retryable (transient) or not.""" + text = str(exc).lower() + return not any(marker in text for marker in _NON_RETRYABLE_LLM_MARKERS) + + +def batch_turns_by_tokens( + items: list[dict[str, Any]], + max_tokens: int, + *, + model: str = "gpt-5.4", +) -> list[list[dict[str, Any]]]: + """Greedily pack ordered *items* into batches whose combined content stays + within *max_tokens*. + + Token-bounded batching keeps each extraction call small enough that (a) the + model can attend to every turn (more complete extraction - smaller windows + extract more faithfully) and (b) a single poisoned turn fails only its own + batch instead of the whole backlog. A turn larger than *max_tokens* on its + own becomes a singleton batch (never dropped). + """ + if not items: + return [] + batches: list[list[dict[str, Any]]] = [] + current: list[dict[str, Any]] = [] + current_tokens = 0 + for item in items: + item_tokens = count_tokens(str(item.get("content") or ""), model) + if current and current_tokens + item_tokens > max_tokens: + batches.append(current) + current = [] + current_tokens = 0 + current.append(item) + current_tokens += item_tokens + if current: + batches.append(current) + return batches + + # Separator for deterministic id seeds. Using NUL ensures user_id / # thread_id values can never collide with literal section markers # (e.g. a thread literally named ``"merged"`` cannot collide with the @@ -241,7 +289,7 @@ def build_transcript( If *True*, group messages under ``=== Thread ===`` headers. metadata_keys: Allow-list of metadata keys to surface in each transcript line. - Defaults to ``None`` (no metadata serialized — only ``[role]: + Defaults to ``None`` (no metadata serialized - only ``[role]: content`` lines). When provided, only the listed keys are emitted, in iteration order. Keys absent from a given turn's metadata are silently skipped. @@ -250,7 +298,7 @@ def build_transcript( ``TurnRecord.metadata`` that the extraction LLM should see (e.g. ``["agent_id", "timestamp"]``). Leaving it unset keeps free-form metadata blobs (raw tool calls, IDE schema, etc.) out of every - prompt — they're often 10-100x larger than the dialog itself and + prompt - they're often 10-100x larger than the dialog itself and dilute extraction quality. Accepts any iterable of strings except ``str`` itself (which would @@ -340,6 +388,46 @@ def build_transcript( "may", "might", "must", + "say", + "says", + "said", + "saying", + "tell", + "tells", + "told", + "ask", + "asks", + "asked", + "mention", + "mentions", + "mentioned", + "stated", + "noted", + "added", + "replied", + "want", + "wants", + "wanted", + "decide", + "decides", + "decided", + "propose", + "proposes", + "proposed", + "suggest", + "suggests", + "suggested", + "planned", + "choose", + "chooses", + "chose", + "like", + "likes", + "liked", + "later", + "then", + "also", + "again", } ) @@ -366,7 +454,7 @@ def check_extracted_fact_grounding( Catches two known LLM failure modes that previously corrupted the fact store: - 1. **Synthesis from existing facts** — the LLM emits an ADD whose content + 1. **Synthesis from existing facts** - the LLM emits an ADD whose content paraphrase-merges two or more existing facts (e.g. existing "user eats meat" + "user loves steak" → emitted "user loves steak, indicating they eat meat") even though the new user turn says nothing @@ -374,7 +462,7 @@ def check_extracted_fact_grounding( but the visible artefact is a chain of "duplicate" supersedes that the user never triggered. - 2. **Phantom explicit-negation** — the LLM emits a second CONTRADICT fact + 2. **Phantom explicit-negation** - the LLM emits a second CONTRADICT fact alongside the literal user statement (e.g. user says "I love steak and seafood"; LLM emits both "user loves steak and seafood" and an invented "user eats meat" CONTRADICT) when the supersedes_id on the literal fact @@ -386,8 +474,8 @@ def check_extracted_fact_grounding( facts → strong synthesis signal. If they come from a single existing fact with >=50%% overlap → weaker phantom-negation signal. - Logs a WARNING for each suspected fact. Does NOT drop facts — downstream - reconciliation remains the dedup authority — but the WARNING is the + Logs a WARNING for each suspected fact. Does NOT drop facts - downstream + reconciliation remains the dedup authority - but the WARNING is the deterministic test signal that catches regressions. """ if not fact_docs or not turn_items: @@ -421,7 +509,7 @@ def check_extracted_fact_grounding( if len(contributors) >= 2: logger.warning( "extract_memories: emitted fact appears synthesized from %d existing facts " - "(ungrounded in user turns) — extract should ground only in this turn's [user] lines. " + "(ungrounded in user turns) - extract should ground only in this turn's [user] lines. " "doc_id=%s content=%r ungrounded_tokens=%s contributor_ids=%s " "user_id=%s thread_id=%s", len(contributors), @@ -438,7 +526,7 @@ def check_extracted_fact_grounding( if overlap_ratio >= 0.5: logger.warning( "extract_memories: emitted fact has ungrounded tokens overlapping a single existing fact " - "(possible phantom-negation/restatement) — extract should ground only in this turn's " + "(possible phantom-negation/restatement) - extract should ground only in this turn's " "[user] lines. doc_id=%s content=%r ungrounded_tokens=%s overlap_existing_id=%s " "overlap_ratio=%.2f user_id=%s thread_id=%s", doc.get("id"), @@ -472,7 +560,7 @@ def parse_llm_json(text: str | None) -> dict[str, Any]: if _looks_truncated(cleaned, exc): raise LLMError( "LLM JSON output appears TRUNCATED (decode error at the very end of a " - f"{len(cleaned)}-char body — the model almost certainly hit its output-token " + f"{len(cleaned)}-char body - the model almost certainly hit its output-token " "cap mid-object). Increase 'maxOutputTokens' in the calling prompty, or reduce " "the amount of input per call (e.g. lower the fact-extraction batch size / " f"recent_k, or split oversized turns). Decode error: {exc}. preview={preview!r}" @@ -564,7 +652,7 @@ def prepare(self, filename: str, inputs: dict[str, Any]) -> tuple[list[dict[str, return messages, params -# Allowed values for the EpisodicRecord ``outcome_valence`` field — mirrors +# Allowed values for the EpisodicRecord ``outcome_valence`` field - mirrors # ``azure.cosmos.agent_memory.models._EPISODIC_ALLOWED_VALENCES`` but kept inline # to avoid an import cycle (helpers must not import models). VALID_VALENCES = frozenset({"positive", "negative", "neutral", "mixed"}) diff --git a/azure/cosmos/agent_memory/services/pipeline.py b/azure/cosmos/agent_memory/services/pipeline.py index 5b9616b..a5058e1 100644 --- a/azure/cosmos/agent_memory/services/pipeline.py +++ b/azure/cosmos/agent_memory/services/pipeline.py @@ -54,12 +54,14 @@ VALID_VALENCES, PromptyLoader, _normalize_metadata_keys, + batch_turns_by_tokens, build_topic_tags, build_transcript, cap_structured_summary, chat_text, check_extracted_fact_grounding, coerce_valence, + is_retryable_llm_error, parse_llm_json, ) from azure.cosmos.agent_memory.services._pipeline_helpers import ( @@ -210,7 +212,7 @@ def _prompt_lineage(self, filename: str) -> dict[str, str]: """Return ``{prompt_id, prompt_version}`` for stamping a doc. Safe no-op fallback (``prompt_version="v1"``) when the loader was - never initialised — happens in unit tests that build the service + never initialised - happens in unit tests that build the service via ``__new__`` to bypass real LLM/embedding clients. """ loader = getattr(self, "_prompty", None) @@ -252,7 +254,7 @@ def _load_existing_memories( """Query active (non-superseded) memories for reconciliation context. Results are ordered by ``c._ts DESC`` so the most recently written - memories survive the cap — without ORDER BY, Cosmos returns rows + memories survive the cap - without ORDER BY, Cosmos returns rows in implementation-defined order and the dedup comparison set is non-deterministic. """ @@ -283,7 +285,7 @@ def _load_existing_memories( def _vector_distance_function(self) -> str: """Return the container's configured Cosmos ``distanceFunction`` (cached). - Read from the container's vector embedding policy (``container.read()``) — + 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 @@ -296,7 +298,7 @@ def _vector_distance_function(self) -> str: 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 + # "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). @@ -438,23 +440,6 @@ def _load_memories_by_ids(self, user_id: str, memory_type: str, ids: Iterable[st ) ) - def _clear_dup_candidate_tags(self, docs: Iterable[dict[str, Any]]) -> None: - for doc in docs: - tags = [tag for tag in (doc.get("tags") or []) if tag != "sys:dup-candidate"] - if tags == (doc.get("tags") or []): - continue - updated = dict(doc) - updated["tags"] = tags - metadata = dict(updated.get("metadata") or {}) - metadata.pop("dup_of", None) - metadata.pop("dup_score", None) - updated["metadata"] = metadata - updated["updated_at"] = datetime.now(timezone.utc).isoformat() - try: - self._upsert_memory(updated) - except Exception: - logger.exception("reconcile_memories: failed to clear dup-candidate tag id=%s", doc.get("id")) - def _upsert_memory(self, doc: dict[str, Any]) -> dict[str, Any]: """Upsert a fact, episodic, or procedural document to the memories container.""" response = self._memories_container.upsert_item(body=doc) @@ -484,6 +469,7 @@ def _empty_extract_counts() -> dict[str, int]: "contradicted_count": 0, "exact_dedup_skipped": 0, "dropped_episodic_count": 0, + "inplace_updated": 0, } @staticmethod @@ -595,42 +581,51 @@ def extract_memories_dry( existing_fact_hashes: set[str] = { m["content_hash"] for m in existing_for_hashes if m.get("type") == "fact" and m.get("content_hash") } - transcript = self._build_transcript(items) - existing = existing_for_hashes - if threshold_config.get_dedup_context_vector_enabled(): - user_turns_text = "\n".join(str(it.get("content", "")) for it in items if it.get("role") == "user").strip() - context_query = user_turns_text or transcript + + # Token-bounded, per-batch extraction. Each batch is an independent LLM + # call, so (a) each stays small enough to extract faithfully and (b) a + # single poisoned turn fails only its own batch. Turns from succeeded and + # quarantined (non-retryable, e.g. content-filter) batches go into + # ``processed_turns`` and will be stamped ``extracted_at`` by persist so + # they are never re-processed; turns from batches that fail with a + # *retryable* error are left un-stamped and retried on the next run. + batches = batch_turns_by_tokens(items, threshold_config.get_extraction_batch_max_tokens()) + facts: list[dict[str, Any]] = [] + episodic: list[dict[str, Any]] = [] + processed_turns: list[dict[str, Any]] = [] + deferred_turn_count = 0 + quarantined_turn_count = 0 + for batch in batches: + batch_transcript = self._build_transcript(batch) try: - existing = self._store.search( - search_terms=context_query, - user_id=user_id, - memory_types=["fact"], - top_k=threshold_config.get_dedup_context_topk(), - ) + response_text = self._run_prompty("extract_memories.prompty", inputs={"transcript": batch_transcript}) + parsed = self._parse_llm_json(response_text) + facts.extend(parsed.get("facts", [])) + episodic.extend(parsed.get("episodic", [])) + processed_turns.extend(batch) except Exception as exc: # noqa: BLE001 - logger.warning( - "extract_memories_dry dedup-context vector search failed (%s); " - "falling back to hash-based existing memories", - exc, - ) - existing = existing_for_hashes - if existing: - existing_text = "\n".join( - f"- [ID: {mem['id']}] {mem.get('content', '')} " - f"(type={mem.get('type', 'fact')}, salience={mem.get('salience', 'N/A')})" - for mem in existing - ) - else: - existing_text = "(none)" - - response_text = self._run_prompty( - "extract_memories.prompty", - inputs={"existing_facts": existing_text, "transcript": transcript}, - ) - parsed = self._parse_llm_json(response_text) - facts = parsed.get("facts", []) - episodic = parsed.get("episodic", []) - unclassified = parsed.get("unclassified", []) + if is_retryable_llm_error(exc): + deferred_turn_count += len(batch) + logger.warning( + "extract_memories: deferring %d turns after retryable extraction error " + "(will retry next run) user_id=%s thread_id=%s err=%s", + len(batch), + user_id, + thread_id, + exc, + ) + else: + processed_turns.extend(batch) + quarantined_turn_count += len(batch) + logger.warning( + "extract_memories: quarantining %d turns after non-retryable extraction error " + "(e.g. content filter) - marking extracted so they do not re-poison future runs " + "user_id=%s thread_id=%s err=%s", + len(batch), + user_id, + thread_id, + exc, + ) doc_timestamp = self._stable_source_timestamp(items) fact_docs: list[dict[str, Any]] = [] @@ -671,10 +666,7 @@ def extract_memories_dry( "confidence": 0.5 if confidence is None else confidence, **self._prompt_lineage("extract_memories.prompty"), "metadata": { - "category": fact.get("category") or "general", - "subject": fact.get("subject"), - "predicate": fact.get("predicate"), - "object": fact.get("object"), + "category": fact.get("category") or "other", "temporal_context": fact.get("temporal_context"), }, "salience": fact.get("salience") if fact.get("salience") is not None else 0.5, @@ -757,52 +749,23 @@ def extract_memories_dry( } episodic_docs.append(self._validate_extracted_doc(doc)) - for item in unclassified: - text = item.get("text") - if not text: - continue - content_hash = compute_content_hash(text) - if content_hash in existing_fact_hashes: - logger.debug( - "extract_memories: skipping exact-dup unclassified hash=%s user_id=%s thread_id=%s", - content_hash, - user_id, - thread_id, - ) - exact_dedup_skipped += 1 - continue - seed = _ID_SEED_SEP.join((user_id, thread_id, content_hash)) - det_id = f"fact_{hashlib.sha256(seed.encode()).hexdigest()[:32]}" - topic_tags = build_topic_tags(item.get("tags", [])) - confidence = item.get("confidence") - doc = { - "id": det_id, - "user_id": user_id, - "thread_id": thread_id, - "role": "system", - "type": "fact", - "content": text, - "content_hash": content_hash, - "confidence": 0.5 if confidence is None else confidence, - **self._prompt_lineage("extract_memories.prompty"), - "metadata": {"category": "unclassified", "unclassified_reason": item.get("reason")}, - "salience": item.get("salience") if item.get("salience") is not None else 0.5, - "tags": ["sys:fact", "sys:auto-extracted", "sys:unclassified"] + topic_tags, - "created_at": doc_timestamp, - "updated_at": doc_timestamp, - } - fact_docs.append(self._validate_extracted_doc(doc)) - existing_fact_hashes.add(content_hash) - if exact_dedup_skipped: updates.append({"op": "stats", "exact_dedup_skipped": exact_dedup_skipped}) if dropped_episodic_count: updates.append({"op": "stats", "dropped_episodic_count": dropped_episodic_count}) + if deferred_turn_count or quarantined_turn_count: + updates.append( + { + "op": "stats", + "deferred_turn_count": deferred_turn_count, + "quarantined_turn_count": quarantined_turn_count, + } + ) check_extracted_fact_grounding( fact_docs, - items, - existing, + processed_turns, + existing_for_hashes, user_id=user_id, thread_id=thread_id, logger=logger, @@ -812,7 +775,7 @@ def extract_memories_dry( "facts": fact_docs, "episodic": episodic_docs, "updates": updates, - "processed_turn_docs": items, + "processed_turn_docs": processed_turns, } logger.info( "extract_memories_dry completed user_id=%s thread_id=%s fact_docs=%d episodic_docs=%d updates=%d", @@ -829,7 +792,22 @@ def dedup_extracted_memories( user_id: str, extracted: dict[str, list[dict[str, Any]]], ) -> dict[str, list[dict[str, Any]]]: - """Apply the gated Stage-3 vector dedup ladder to extracted docs.""" + """Fold near-duplicate extracted docs into their existing canonical + memory *in place*, instead of tagging them for an async merge sweep. + + For each newly extracted fact/episodic doc we find its single nearest + active same-type neighbor. If similarity is at/above ``SIM_HIGH`` the new + doc is a near-duplicate: we refresh the existing neighbor in place — + recency-wins content + embedding, unioned tags, max salience/confidence, + bumped ``updated_at`` — keeping its id, and drop the new doc so it is not + written as a fresh record. Everything below the threshold is novel and + flows through to ``persist_extracted_memories`` unchanged. + + This makes the write path convergent: a restatement updates one existing + document rather than minting a new one that a later reconcile sweep must + merge and supersede. There is no ``sys:dup-candidate`` tagging and no + clustering — reconcile only resolves contradictions. + """ if not threshold_config.get_dedup_vector_enabled(): return extracted if not user_id: @@ -838,22 +816,26 @@ def dedup_extracted_memories( raise ValidationError("extracted must be a dict") high = threshold_config.get_dedup_sim_high() - low = threshold_config.get_dedup_sim_low() - top_k = threshold_config.get_dedup_candidate_topk() distance_function = self._vector_distance_function() - autodrop_ok = vector_autodrop_supported(distance_function) - if not autodrop_ok: + similarity_ok = vector_autodrop_supported(distance_function) + if not similarity_ok: self._warn_euclidean_autodrop_once(distance_function) - vector_dedup_skipped = 0 - dup_candidates_tagged = 0 result: dict[str, list[dict[str, Any]]] = { "facts": [dict(doc) for doc in extracted.get("facts", [])], "episodic": [dict(doc) for doc in extracted.get("episodic", [])], "updates": [dict(op) for op in extracted.get("updates", [])], } + # Carry through any non-bucket keys (e.g. ``processed_turn_docs``) so this + # transform never silently drops caller state. + for _carry_key, _carry_value in extracted.items(): + if _carry_key not in result: + result[_carry_key] = _carry_value + docs = [doc for bucket in ("facts", "episodic") for doc in result[bucket] if doc.get("content")] - if not docs: + # Similarity comparison is only meaningful for cosine/dotproduct; on a + # euclidean container we skip in-place folding and let everything ADD. + if not docs or not similarity_ok: return result missing_embeddings = [doc for doc in docs if not doc.get("embedding")] @@ -862,65 +844,135 @@ def dedup_extracted_memories( for doc, embedding in zip(missing_embeddings, embeddings): doc["embedding"] = embedding - kept_ids: set[str] = set() - dropped_ids: set[str] = set() + inplace_updated = 0 + folded_ids: set[str] = set() + updated_target_ids: set[str] = set() for doc in docs: doc_id = str(doc.get("id") or "") memory_type = str(doc.get("type") or "") - if not doc_id or memory_type not in {"fact", "episodic"}: + embedding = doc.get("embedding") or [] + if not doc_id or memory_type not in {"fact", "episodic"} or not embedding: continue - best: dict[str, Any] | None = None - candidates = self._vector_candidates( + neighbor, score = self._nearest_active_full( user_id=user_id, - embedding=doc.get("embedding") or [], + embedding=embedding, memory_type=memory_type, - top_k=top_k, - exclude_ids=kept_ids | dropped_ids | {doc_id} | set(doc.get("supersedes_ids") or []), + exclude_ids={doc_id} | set(doc.get("supersedes_ids") or []), ) - if candidates: - best = candidates[0] + if not neighbor or not vector_similarity_at_least(score, high, distance_function): + continue # novel — leave in result for persist to ADD - score = float(best.get("score") or 0.0) if best else 0.0 - if best and autodrop_ok and vector_similarity_at_least(score, high, distance_function): - logger.info( - "vector dedup skipped new memory id=%s type=%s score=%.4f surviving_id=%s " - "new_content=%r surviving_content=%r", - doc_id, - memory_type, - score, - best.get("id"), - doc.get("content"), - best.get("content"), - ) - vector_dedup_skipped += 1 - dropped_ids.add(doc_id) + neighbor_id = str(neighbor.get("id") or "") + if not neighbor_id: continue + if neighbor_id in updated_target_ids: + # A prior near-dup in this batch already refreshed this target; + # drop this one too rather than re-writing the same document. + folded_ids.add(doc_id) + continue + if self._apply_inplace_update(neighbor, doc): + updated_target_ids.add(neighbor_id) + inplace_updated += 1 + folded_ids.add(doc_id) + # If the in-place update failed, leave the doc in result so persist + # ADDs it as a novel record — never silently lose an extraction. + + if folded_ids: + for bucket in ("facts", "episodic"): + result[bucket] = [doc for doc in result[bucket] if str(doc.get("id") or "") not in folded_ids] + if inplace_updated: + result["updates"].append({"op": "stats", "inplace_updated": inplace_updated}) + return result - if best and vector_similarity_at_least(score, low, distance_function): - tags = list(doc.get("tags") or []) - if "sys:dup-candidate" not in tags: - tags.append("sys:dup-candidate") - doc["tags"] = tags - metadata = dict(doc.get("metadata") or {}) - metadata["dup_of"] = best.get("id") - metadata["dup_score"] = score - doc["metadata"] = metadata - dup_candidates_tagged += 1 - - kept_ids.add(doc_id) - - for bucket in ("facts", "episodic"): - result[bucket] = [doc for doc in result[bucket] if doc.get("id") not in dropped_ids] - if vector_dedup_skipped or dup_candidates_tagged: - result["updates"].append( - { - "op": "stats", - "vector_dedup_skipped": vector_dedup_skipped, - "dup_candidates_tagged": dup_candidates_tagged, - } + def _nearest_active_full( + self, + *, + user_id: str, + embedding: list[float], + memory_type: str, + exclude_ids: set[str], + ) -> tuple[Optional[dict[str, Any]], float]: + """Return the single nearest active same-type memory as a *full* doc. + + Unlike ``_vector_candidates`` (which projects only id/content/score), + this returns the complete stored document so the caller can refresh it in + place. Pulls a small TOP-k and returns the first not in ``exclude_ids``. + """ + if not user_id or not embedding: + return None, 0.0 + query = ( + "SELECT TOP 5 c AS doc, VectorDistance(c.embedding, @vec) AS score " + "FROM c WHERE c.user_id = @user_id " + "AND c.type = @memory_type " + f"AND {_ACTIVE_DOC_FILTER} " + "AND IS_DEFINED(c.embedding) " + "ORDER BY VectorDistance(c.embedding, @vec)" + ) + try: + 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, + ) ) - return result + except Exception as exc: # noqa: BLE001 + logger.warning("_nearest_active_full query failed user_id=%s err=%s", user_id, exc) + return None, 0.0 + for row in rows: + doc = row.get("doc") or {} + rid = str(doc.get("id") or "") + if rid and rid not in exclude_ids: + return doc, float(row.get("score") or 0.0) + return None, 0.0 + + def _apply_inplace_update(self, neighbor: dict[str, Any], new_doc: dict[str, Any]) -> bool: + """Refresh an existing memory in place with a near-duplicate's content. + + Recency wins: the neighbor keeps its id / created_at / partition but takes + the new doc's content + embedding, unions tags, and takes the max + salience/confidence. Returns True on a successful upsert; False otherwise + (the caller then keeps the new doc as a novel ADD, so nothing is lost). + """ + try: + updated = dict(neighbor) + for sys_prop in ("_rid", "_self", "_etag", "_attachments", "_ts"): + updated.pop(sys_prop, None) + new_content = str(new_doc.get("content") or "") + updated["content"] = new_content + updated["content_hash"] = compute_content_hash(new_content) + if new_doc.get("embedding"): + updated["embedding"] = new_doc["embedding"] + updated["updated_at"] = datetime.now(timezone.utc).isoformat() + + new_sal = _max_or_none([neighbor.get("salience"), new_doc.get("salience")]) + if new_sal is not None: + updated["salience"] = new_sal + new_conf = _max_or_none([neighbor.get("confidence"), new_doc.get("confidence")]) + if new_conf is not None: + updated["confidence"] = new_conf + + merged_tags: list[str] = [] + for t in list(neighbor.get("tags") or []) + list(new_doc.get("tags") or []): + if t and t != "sys:dup-candidate" and t not in merged_tags: + merged_tags.append(t) + if merged_tags: + updated["tags"] = merged_tags + + self._memories_container.upsert_item(body=updated) + return True + except Exception as exc: # noqa: BLE001 + logger.warning( + "in-place dedup update failed target_id=%s err=%s (keeping new doc as novel)", + neighbor.get("id"), + exc, + ) + return False def persist_extracted_memories( self, @@ -974,22 +1026,12 @@ def persist_extracted_memories( if op.get("op") == "stats": result["exact_dedup_skipped"] += int(op.get("exact_dedup_skipped") or 0) result["dropped_episodic_count"] += int(op.get("dropped_episodic_count") or 0) - for key in ("vector_dedup_skipped", "dup_candidates_tagged"): + for key in ("inplace_updated",): if key in op: result[key] = result.get(key, 0) + int(op.get(key) or 0) logger.info("persist_extracted_memories completed user_id=%s counts=%s", user_id, result) - processed_turns = extracted.get("processed_turn_docs") or [] - if processed_turns: - marked = self._mark_turns_extracted(processed_turns) - logger.info( - "persist_extracted_memories marked turns as extracted user_id=%s marked=%d/%d", - user_id, - marked, - len(processed_turns), - ) - return result def _mark_turns_extracted(self, turn_docs: list[dict[str, Any]]) -> int: @@ -997,7 +1039,7 @@ def _mark_turns_extracted(self, turn_docs: list[dict[str, Any]]) -> int: We upsert the full doc (rather than patch) because the container adapter only exposes upsert. Per-turn failures are logged but do - not raise — the worst case is one turn gets re-extracted on the + not raise - the worst case is one turn gets re-extracted on the next call, which is bounded and recoverable. """ if not turn_docs: @@ -1031,9 +1073,33 @@ def extract_memories( ) -> dict[str, int]: """Extract facts and episodic memories from a thread and persist them.""" extracted = self.extract_memories_dry(user_id, thread_id, recent_k, turns=turns) + # Capture the processed turns from the DRY output as the single source of + # truth for stamping. Stamping happens here (not inside persist) so no + # intermediate transform (e.g. dedup) can drop ``processed_turn_docs`` + # and cause the same turns to be re-extracted forever. + processed_turns = extracted.get("processed_turn_docs") or [] if threshold_config.get_dedup_vector_enabled(): extracted = self.dedup_extracted_memories(user_id, extracted) - return self.persist_extracted_memories(user_id, extracted) + counts = self.persist_extracted_memories(user_id, extracted) + if processed_turns: + marked = self._mark_turns_extracted(processed_turns) + if marked < len(processed_turns): + logger.warning( + "extract_memories stamped only %d/%d processed turns as extracted user_id=%s " + "thread_id=%s (unstamped turns will be re-extracted next run)", + marked, + len(processed_turns), + user_id, + thread_id, + ) + else: + logger.info( + "extract_memories stamped %d processed turns as extracted user_id=%s thread_id=%s", + marked, + user_id, + thread_id, + ) + return counts def synthesize_procedural( self, @@ -1148,40 +1214,12 @@ def _covered_by(prior: Optional[dict[str, Any]]) -> bool: if not current_source_ids: logger.info( - "synthesize_procedural skipping LLM user_id=%s — no behavioral facts or episodic lessons", + "synthesize_procedural skipping LLM user_id=%s - no behavioral facts or episodic lessons", user_id, ) return {"status": "unchanged", "procedural": prior_doc} - name_docs = list( - self._memories_container.query_items( - query=( - "SELECT TOP 1 * FROM c WHERE c.user_id = @uid " - "AND c.type = @type " - f"AND {_ACTIVE_DOC_FILTER} " - "AND IS_DEFINED(c.metadata.category) " - "AND c.metadata.category = @category " - "AND IS_DEFINED(c.metadata.predicate) " - "AND c.metadata.predicate = @predicate " - "ORDER BY c._ts DESC" - ), - parameters=[ - {"name": "@uid", "value": user_id}, - {"name": "@type", "value": "fact"}, - {"name": "@category", "value": "biographical"}, - {"name": "@predicate", "value": "name"}, - ], - enable_cross_partition_query=True, - ) - ) user_name = "the user" - if name_docs: - metadata = name_docs[0].get("metadata") or {} - name_candidate = metadata.get("object") - if not isinstance(name_candidate, str) or not name_candidate.strip(): - name_candidate = name_docs[0].get("content") - if isinstance(name_candidate, str) and name_candidate.strip(): - user_name = name_candidate.strip() def _render_bullets(values: list[str]) -> str: cleaned = [value.strip() for value in values if isinstance(value, str) and value.strip()] @@ -1200,7 +1238,7 @@ def _render_bullets(values: list[str]) -> str: # Retry loop: LLM call lives inside so that on a race-induced 409 # we (a) check whether the winner already covers our source set and # short-circuit if so, and (b) re-call the LLM with the winner as - # the new prior if not — keeping synthesized content monotonic in + # the new prior if not - keeping synthesized content monotonic in # source coverage, not just version number. written_doc: Optional[dict[str, Any]] = None for attempt in range(1, _PROCEDURAL_MAX_CREATE_ATTEMPTS + 1): @@ -1243,7 +1281,7 @@ def _render_bullets(values: list[str]) -> str: break except CosmosResourceExistsError: logger.info( - "synthesize_procedural id collision user_id=%s seq=%d attempt=%d/%d — re-reading", + "synthesize_procedural id collision user_id=%s seq=%d attempt=%d/%d - re-reading", user_id, new_seq, attempt, @@ -1597,191 +1635,24 @@ def _emit_reconcile_outcome( }, ) - def _build_candidate_clusters( - self, user_id: str, memory_type: str, n: int - ) -> tuple[list[list[dict[str, Any]]], int, list[dict[str, Any]]]: - """Cluster dup-candidate seeds (+ vector neighbors) into connected components. - - Returns ``(clusters, node_count, seeds)`` where each cluster has >= 2 members, - ``node_count`` is the total distinct memories pulled into the graph (used to - report ``reconcile_llm_calls_saved``), and ``seeds`` is the tagged seed scan - (so the caller can clear stale tags on orphan seeds that never clustered). - The seed scan is bounded to ``n`` so a single cluster can never exceed the - reconcile prompt's pool cap. - """ - cluster_sim = threshold_config.get_dedup_cluster_sim() - top_k = threshold_config.get_dedup_candidate_topk() - distance_function = self._vector_distance_function() - seeds = self._query_active_memories(user_id, memory_type, limit=n, tagged_only=True) - nodes_by_id: dict[str, dict[str, Any]] = {doc["id"]: doc for doc in seeds if doc.get("id")} - edge_pairs: set[tuple[str, str]] = set() - - dup_of_ids = [ - (doc.get("metadata") or {}).get("dup_of") - for doc in seeds - if isinstance(doc.get("metadata"), dict) and (doc.get("metadata") or {}).get("dup_of") - ] - for doc in self._load_memories_by_ids(user_id, memory_type, dup_of_ids): - if doc.get("id"): - nodes_by_id[doc["id"]] = doc + def reconcile_memories(self, user_id: str, n: int = 50, *, memory_type: str = "fact") -> dict[str, int]: + """Resolve contradictions among a user's most-recent active memories. - for seed in seeds: - seed_id = seed.get("id") - if not seed_id: - continue - dup_of = (seed.get("metadata") or {}).get("dup_of") if isinstance(seed.get("metadata"), dict) else None - if dup_of: - edge_pairs.add(tuple(sorted((seed_id, dup_of)))) - if not seed.get("embedding"): - continue - candidates = self._vector_candidates( - user_id=user_id, - embedding=seed.get("embedding") or [], - memory_type=memory_type, - top_k=top_k, - exclude_ids={seed_id}, - ) - candidate_ids = [ - row["id"] - for row in candidates - if row.get("id") and vector_similarity_at_least(row.get("score", 0.0), cluster_sim, distance_function) - ] - for cid in candidate_ids: - edge_pairs.add(tuple(sorted((seed_id, cid)))) - for doc in self._load_memories_by_ids(user_id, memory_type, candidate_ids): - if doc.get("id"): - nodes_by_id[doc["id"]] = doc - - node_ids = list(nodes_by_id) - node_id_set = set(node_ids) - for node_id in node_ids: - node = nodes_by_id[node_id] - if not node.get("embedding"): - continue - for row in self._vector_candidates( - user_id=user_id, - embedding=node.get("embedding") or [], - memory_type=memory_type, - top_k=top_k, - exclude_ids={node_id}, - ): - neighbor_id = row.get("id") - if neighbor_id in node_id_set and vector_similarity_at_least( - float(row.get("score") or 0.0), cluster_sim, distance_function - ): - edge_pairs.add(tuple(sorted((node_id, neighbor_id)))) - - adjacency: dict[str, set[str]] = {node_id: set() for node_id in node_ids} - for left_id, right_id in edge_pairs: - if left_id != right_id and left_id in adjacency and right_id in adjacency: - adjacency[left_id].add(right_id) - adjacency[right_id].add(left_id) - - clusters: list[list[dict[str, Any]]] = [] - seen: set[str] = set() - for node_id in node_ids: - if node_id in seen: - continue - stack = [node_id] - component: list[str] = [] - seen.add(node_id) - while stack: - current = stack.pop() - component.append(current) - for neighbor in adjacency.get(current, set()): - if neighbor not in seen: - seen.add(neighbor) - stack.append(neighbor) - if len(component) >= 2: - # Cap cluster size at the reconcile pool limit: lowering the cluster - # threshold can chain many facts into one giant transitive component - # that would blow the prompt cap; keep the most-recent ``n``. - if len(component) > n: - component = component[:n] - clusters.append([nodes_by_id[cid] for cid in component]) - return clusters, len(nodes_by_id), seeds - - def _reconcile_candidate_mode(self, user_id: str, *, n: int, memory_type: str, started_at: float) -> dict[str, int]: - # Candidate clustering only. The periodic full-pool backstop that catches - # dissimilar-embedding contradictions ("vegetarian" vs "loves steak") is - # driven by the caller via ``full_rebuild`` on a PERSISTED-counter cadence - # (in-process auto-trigger + durable change-feed), not an in-memory sweep - # counter — the latter reset per worker/process and never fired reliably on - # the Function-App backend. - clusters, node_count, seeds = self._build_candidate_clusters(user_id, memory_type, n) - aggregate = {"kept": 0, "merged": 0, "contradicted": 0} - clustered_ids: set[str] = set() - for cluster in clusters: - # Mark members as clustered BEFORE the LLM call so a failed cluster's - # seeds are not treated as orphans below (which would clear their tags - # and prevent a retry). A truncated/malformed LLM response on one - # cluster must not abort the sweep or starve the remaining clusters. - clustered_ids.update(doc["id"] for doc in cluster if doc.get("id")) - try: - counts, consumed = self._reconcile_pool(user_id, memory_type, cluster) - except Exception as exc: - logger.warning( - "reconcile_memories: cluster reconcile failed user_id=%s memory_type=%s; " - "skipping cluster, tags retained for next sweep: %s", - user_id, - memory_type, - exc, - ) - continue - for key in aggregate: - aggregate[key] += int(counts.get(key, 0)) - # Clear dup-candidate tags only on survivors. Re-upserting a doc that - # was just superseded (duplicate source or contradiction loser) would - # resurrect it, since the in-memory cluster copy lacks superseded_by. - survivors = [doc for doc in cluster if doc.get("id") and doc["id"] not in consumed] - self._clear_dup_candidate_tags(survivors) - # Orphan seeds: tagged dup-candidates that never joined a cluster have no - # near-duplicate, so clear the stale tag — otherwise every future sweep - # re-scans them as seeds and they accumulate forever. - orphan_seeds = [seed for seed in seeds if seed.get("id") and seed["id"] not in clustered_ids] - self._clear_dup_candidate_tags(orphan_seeds) - aggregate["reconcile_clusters_sent"] = len(clusters) - aggregate["reconcile_llm_calls_saved"] = max(0, node_count - len(clusters)) - logger.info( - "reconcile_memories candidate completed user_id=%s memory_type=%s result=%s", - user_id, - memory_type, - aggregate, - ) - self._emit_reconcile_outcome( - started_at=started_at, - user_id=user_id, - candidates=node_count, - result=aggregate, - ) - return aggregate + Loads up to ``n`` active (non-superseded) ``memory_type`` records and + asks the dedup prompt to identify ``contradicted_pairs`` — opposing + claims about the same subject (e.g. "deadline March 1" vs "March 15"). + Each loser is soft-deleted with ``supersede_reason="contradict"`` and + ``superseded_by`` set to the winner. - def reconcile_memories( - self, user_id: str, n: int = 50, *, memory_type: str = "fact", full_rebuild: bool = False - ) -> dict[str, int]: - """Reconcile a user's active facts in a single LLM pass. - - Loads the most recent ``n`` active (non-superseded) facts for - ``user_id``, asks the dedup prompt to classify them into - ``duplicate_groups``, ``contradicted_pairs``, and ``kept_ids``, then - applies both kinds of resolutions: - - * **Duplicates** — a fresh merged fact is upserted; every source is - soft-deleted with ``supersede_reason="duplicate"``. - * **Contradictions** — the loser is soft-deleted with - ``supersede_reason="contradict"`` and ``superseded_by`` set to - the winner. Dangling references are resolved transparently when a - contradicted id was just absorbed into a duplicate group. - - ``full_rebuild`` (set by the explicit public ``reconcile()`` entrypoint) - makes candidate mode cluster the whole active pool instead of only - Stage-3-tagged seeds, so a user-triggered reconcile dedups every active - fact. Automatic background sweeps leave it ``False`` for the cheap - tagged-only path. - - Returns ``{"kept": int, "merged": int, "contradicted": int}`` where - ``merged`` and ``contradicted`` count the *losers* that were - soft-deleted (duplicates and contradictions respectively). + Near-duplicate *paraphrases* are no longer merged here: the write-time + in-place dedup (:meth:`dedup_extracted_memories`) folds restatements into + their canonical record before they land, so reconcile is a bounded, + convergent contradiction pass — no clustering, no synthesized merged + documents, no re-merge churn. Episodic and procedural types are no-ops + (episodic has no contradiction semantics; its near-dups fold at write + time). + + Returns ``{"kept", "merged", "contradicted"}``; ``merged`` is always 0. """ if not user_id: raise ValidationError("user_id is required") @@ -1791,34 +1662,16 @@ def reconcile_memories( raise ValidationError(f"n must be <= 500 to bound prompt size and LLM cost, got {n}") if memory_type not in {"fact", "episodic", "procedural"}: raise ValidationError(f"memory_type must be one of fact, episodic, procedural; got {memory_type!r}") - if memory_type == "procedural": - result = { - "kept": 0, - "merged": 0, - "contradicted": 0, - "reconcile_clusters_sent": 0, - "reconcile_llm_calls_saved": 0, - } - logger.info("reconcile_memories procedural no-op user_id=%s result=%s", user_id, result) + if memory_type in {"episodic", "procedural"}: + result = {"kept": 0, "merged": 0, "contradicted": 0} + logger.info("reconcile_memories %s no-op user_id=%s result=%s", memory_type, user_id, result) return result started_at = time.monotonic() logger.info("reconcile_memories started user_id=%s n=%d memory_type=%s", user_id, n, memory_type) - # Explicit user-triggered reconcile (full_rebuild) always takes the - # full-pool single-LLM-pass path: it sees every active fact together, so it - # catches contradictions that aren't vector-similar (e.g. "vegetarian" vs - # "loves steak") — which candidate clustering, keyed on near-duplicate - # similarity, would never group. Automatic sweeps use cheap candidate mode. - if threshold_config.get_dedup_reconcile_mode() == "candidate" and not full_rebuild: - return self._reconcile_candidate_mode(user_id, n=n, memory_type=memory_type, started_at=started_at) - facts = self._active_memories_for_reconcile(user_id, memory_type, n) - result, consumed = self._reconcile_pool(user_id, memory_type, facts) - # Clear dup-candidate tags on survivors so an explicit reconcile(full_rebuild=True) - # doesn't leave stale sys:dup-candidate/dup_of metadata on user-visible memories. - survivors = [doc for doc in facts if doc.get("id") and doc["id"] not in consumed] - self._clear_dup_candidate_tags(survivors) + result = self._reconcile_contradictions(user_id, memory_type, facts) self._emit_reconcile_outcome( started_at=started_at, user_id=user_id, @@ -1827,63 +1680,19 @@ def reconcile_memories( ) return result - def _active_memories_for_reconcile(self, user_id: str, memory_type: str, n: int) -> list[dict[str, Any]]: - # ---- Load up to N most recent active memories ---- - # ORDER BY c.created_at DESC keeps the TOP cap deterministic across - # physical partitions and matches the dedup prompt's tiebreaker - # ("more recent created_at first"). Cosmos's _ts is the last-write - # timestamp, which would diverge from created_at after any UPDATE. - query = ( - f"SELECT TOP {top_literal(n, name='reconcile_memories.n')} * FROM c " - "WHERE c.user_id = @user_id " - "AND c.type = @memory_type " - f"AND {_ACTIVE_DOC_FILTER} " - "ORDER BY c.created_at DESC" - ) - return list( - self._memories_container.query_items( - query=query, - parameters=[ - {"name": "@user_id", "value": user_id}, - {"name": "@memory_type", "value": memory_type}, - ], - enable_cross_partition_query=True, - ) - ) - - def _reconcile_pool( - self, user_id: str, memory_type: str, facts: list[dict[str, Any]] - ) -> tuple[dict[str, int], set[str]]: - """Reconcile an explicit pool of same-type memories in one LLM pass. + def _reconcile_contradictions(self, user_id: str, memory_type: str, facts: list[dict[str, Any]]) -> dict[str, int]: + """Resolve contradictions within an explicit pool of same-type memories. - Returns ``({"kept", "merged", "contradicted"}, consumed_ids)`` where - ``consumed_ids`` are the source/loser ids that were actually superseded, - so callers can skip them when clearing dup-candidate tags (re-upserting a - superseded source would resurrect it). Does not emit telemetry — the - caller owns the ``reconcile.outcome`` line. + Runs the dedup prompt over the pool and applies ONLY its + ``contradicted_pairs`` — the loser of each pair is soft-deleted with + ``superseded_by`` set to the winner. ``duplicate_groups`` in the response + are ignored (write-time in-place dedup already handles paraphrases), so + no merged documents are minted and the pass is convergent. Returns + ``{"kept", "merged": 0, "contradicted"}``. """ - # Polarity keyed on an explicit predicate so a future third type (e.g. - # procedural, were it ever routed here) can't silently diverge from async. - is_episodic = memory_type == "episodic" - prompt_filename = "dedup_episodic.prompty" if is_episodic else "dedup.prompty" - prompt_input_key = "episodics_text" if is_episodic else "facts_text" - allow_contradictions = not is_episodic - if len(facts) <= 1: - logger.info( - "reconcile_memories: %d %s memories, nothing to reconcile", - len(facts), - memory_type, - ) - return {"kept": len(facts), "merged": 0, "contradicted": 0}, set() - - # ---- 2. Format the facts pool for the prompt ---- - # ``json.dumps`` escapes embedded quotes and pipes inside content so - # the visual grammar (`| Field:` separators, `""` quoting) - # stays unambiguous even on adversarial inputs like - # ``She said "hi" | weird``. IDs are kept raw because they're - # deterministic alphanumerics — quoting them risks the LLM copying - # the quotes back into ``source_ids``. + return {"kept": len(facts), "merged": 0, "contradicted": 0} + lines: list[str] = [] for i, cf in enumerate(facts, 1): content_quoted = json.dumps(cf.get("content", ""), ensure_ascii=False) @@ -1895,372 +1704,69 @@ def _reconcile_pool( created_str = created_raw if created_raw else "N/A" lines.append( f"{i}. ID: {cf['id']} | Content: {content_quoted} | " - f"Confidence: {conf_str} | " - f"Salience: {sal_str} | " - f"Created: {created_str}" + f"Confidence: {conf_str} | Salience: {sal_str} | Created: {created_str}" ) facts_text = "\n".join(lines) - # ---- 3. Single LLM call over the entire pool ---- - response_text = self._run_prompty( - prompt_filename, - inputs={prompt_input_key: facts_text}, - ) + response_text = self._run_prompty("dedup.prompty", inputs={"facts_text": facts_text}) parsed = self._parse_llm_json(response_text) - - duplicate_groups = parsed.get("duplicate_groups", []) or [] - contradicted_pairs = (parsed.get("contradicted_pairs", []) or []) if allow_contradictions else [] - # ``kept_ids`` from the LLM is used below as a cross-check for - # accounting drift (hallucinated IDs, double-counting). The actual - # kept count is computed from facts minus consumed losers. - llm_kept_ids = list(parsed.get("kept_ids", []) or []) + contradicted_pairs = parsed.get("contradicted_pairs", []) or [] facts_by_id: dict[str, dict[str, Any]] = {f["id"]: f for f in facts} - - merged = 0 contradicted = 0 - # Tracks source_id -> merged_id rewrites so contradictions whose - # winner/loser landed in a duplicate group can be redirected to - # the surviving merged document. Only updated on *successful* - # supersede so stale redirects don't survive ETag races. - source_to_merged_id: dict[str, str] = {} - # Cache of merged docs we just upserted, keyed by merged_id. Lets - # the contradiction redirector reuse the in-memory dict instead of - # a cross-partition Cosmos round-trip for a doc we own. Also keeps - # the chain ETag-stable when the same merged doc absorbs both a - # duplicate group and a contradiction redirect in the same call. - merged_docs_by_id: dict[str, dict[str, Any]] = {} - # Set of source IDs that were *actually* superseded (counts toward - # ``merged``). Used by the kept-count cross-check below — earlier - # versions counted attempts and undercounted on ETag races. - consumed_source_ids: set[str] = set() - # Set of contradiction loser IDs that were *actually* superseded. consumed_loser_ids: set[str] = set() - # Original-pool winner IDs from successfully-applied contradictions. - # The LLM emits winners under ``contradicted_pairs``, never under - # ``kept_ids`` — so the kept-cross-check at the end must subtract - # them from the expected-kept set or every clean run looks like a - # mismatch. - contradiction_winner_ids_in_pool: set[str] = set() - - # ---- 4. Apply duplicate_groups FIRST ---- - for group in duplicate_groups: - source_ids = list(group.get("source_ids") or []) - merged_content = group.get("merged_content") - if not merged_content or not source_ids: - logger.debug( - "reconcile_memories: skipping malformed duplicate_group %r", - group, - ) - continue - - source_docs = [facts_by_id[sid] for sid in source_ids if sid in facts_by_id] - if not source_docs: - logger.debug( - "reconcile_memories: duplicate_group references unknown ids %r", - source_ids, - ) - continue - - # Filtered, hallucination-free view of the source ids that - # actually exist in the pool. Used both for ``supersedes_ids`` - # on the merged record and for the deterministic merged-id - # below so the merged doc faithfully represents reality. - valid_source_ids = [sid for sid in source_ids if sid in facts_by_id] - - if len(valid_source_ids) < 2: - logger.debug( - "reconcile_memories: skipping single-source duplicate_group %r", - source_ids, - ) - continue - - # Sort source_docs by Cosmos _ts DESC so the merged record's - # partition (thread_id) is picked deterministically from the - # newest source — independent of the LLM's source_ids order. - source_docs.sort(key=lambda d: d.get("_ts", 0), reverse=True) - - # Union tags across all source docs (preserve order, dedupe). - # Strip sys:dup-candidate so the merged doc isn't reborn as a - # permanent reconcile seed. - merged_tags: list[str] = [] - seen_tags: set[str] = set() - for src in source_docs: - for t in src.get("tags", []) or []: - if t == "sys:dup-candidate": - continue - if t not in seen_tags: - seen_tags.add(t) - merged_tags.append(t) - if not merged_tags: - merged_tags = [f"sys:{memory_type}"] - - # Union source_memory_ids across all source docs (provenance chain). - merged_source_memory_ids: list[str] = [] - seen_smi: set[str] = set() - for src in source_docs: - for smi in src.get("source_memory_ids", []) or []: - if smi not in seen_smi: - seen_smi.add(smi) - merged_source_memory_ids.append(smi) - - # Transitive supersedes_ids: include any prior chain hops the - # source docs already absorbed so the merged record carries - # the full provenance, not just the immediate parent layer. - merged_supersedes: list[str] = [] - seen_sup: set[str] = set() - for sid in valid_source_ids: - if sid not in seen_sup: - seen_sup.add(sid) - merged_supersedes.append(sid) - for src in source_docs: - for prior in src.get("supersedes_ids", []) or []: - if prior and prior not in seen_sup: - seen_sup.add(prior) - merged_supersedes.append(prior) - - # Newest source's thread_id wins (after _ts-desc sort above). - recent_thread_id = source_docs[0].get("thread_id", "") - - # If LLM omitted confidence/salience, returned a non-positive - # placeholder, returned a JSON ``true`` masquerading as numeric, - # or returned an out-of-range value (e.g. 1.05 — common when - # models confuse percent with [0,1]), fall back to max across - # the source docs. Out-of-range without a fallback would let - # ``MemoryRecord(...)`` raise on Pydantic validation and the - # blanket except below would silently drop the entire group. - llm_conf = group.get("confidence") - confidence_val = ( - float(llm_conf) - if _is_real_number(llm_conf) and 0 < llm_conf <= 1 - else _max_or_none(src.get("confidence") for src in source_docs) - ) - llm_sal = group.get("salience") - salience_val = ( - float(llm_sal) - if _is_real_number(llm_sal) and 0 < llm_sal <= 1 - else _max_or_none(src.get("salience") for src in source_docs) - ) - - # Deterministic merged id keyed on (user, "merged", content_hash) - # so re-running reconcile on the same merged content produces an - # idempotent upsert instead of a fresh UUID each cycle. Stable - # ids also keep the supersede chain shallow: a future paraphrase - # that gets folded into the same canonical merged content will - # see the same id rather than chaining through a new UUID. - merged_content_hash = compute_content_hash(merged_content) - merged_id_seed = _ID_SEED_SEP.join((user_id, "merged", merged_content_hash)) - id_prefix = "fact_" if memory_type == "fact" else "ep_" - merged_id = id_prefix + hashlib.sha256(merged_id_seed.encode()).hexdigest()[:32] - - try: - merged_payload: dict[str, Any] = { - "id": merged_id, - "user_id": user_id, - "role": "system", - "type": memory_type, - "content": merged_content, - "thread_id": recent_thread_id or f"__reconciled__:{user_id}", - "confidence": confidence_val if confidence_val is not None else 0.5, - "salience": salience_val if salience_val is not None else 0.5, - "supersedes_ids": merged_supersedes, - "source_memory_ids": merged_source_memory_ids, - "tags": merged_tags, - "content_hash": merged_content_hash, - **self._prompt_lineage(prompt_filename), - "created_at": datetime.now(timezone.utc).isoformat(), - "updated_at": datetime.now(timezone.utc).isoformat(), - } - if memory_type == "fact": - merged_payload["metadata"] = { - "category": "preference", - "merged_via": "reconcile", - "merged_from_count": len(valid_source_ids), - } - record_cls = FactRecord - else: - source_meta = dict(source_docs[0].get("metadata") or {}) - source_meta.update( - { - "lesson": merged_content, - "merged_via": "reconcile", - "merged_from_count": len(valid_source_ids), - } - ) - source_meta.setdefault("scope_type", source_meta.get("scope_type") or "general") - source_meta.setdefault("scope_value", source_meta.get("scope_value") or "general") - source_meta.setdefault("outcome_valence", source_meta.get("outcome_valence") or "neutral") - merged_payload["metadata"] = source_meta - record_cls = EpisodicRecord - merged_record = construct_internal(record_cls, merged_payload) - except Exception: - logger.exception( - "reconcile_memories: failed to build merged record for group %r", - group, - ) - continue - - # Generate embedding for the merged content so retrieval can - # rank it against future queries from the moment it lands. - # If embedding fails, abort this duplicate group entirely: - # writing a merged doc with no embedding and then superseding - # the sources would create a search-index hole until the next - # reconcile retried. Better to leave the duplicates in place. - try: - merged_record.embedding = self._embed_one(merged_content) - except Exception: - logger.exception( - "reconcile_memories: embedding failed for merged id=%s; " - "aborting duplicate group to avoid search-index hole", - merged_record.id, - ) - continue - - merged_doc = merged_record.to_doc() - try: - merged_doc = self._upsert_memory(merged_doc) - except Exception: - logger.exception( - "reconcile_memories: upsert failed for merged id=%s; aborting duplicate group", - merged_record.id, - ) - continue - merged_docs_by_id[merged_record.id] = merged_doc - - group_supersede_count = 0 - for sid in valid_source_ids: - src_doc = facts_by_id.get(sid) - if src_doc is None: - # Defensive — already filtered above, kept for clarity. - continue - # Only update redirect/consumed-set on *successful* supersede. - # Losing the ETag race means another writer beat us; the - # source doc is still active from our perspective and should - # not be treated as consumed. - if self._mark_superseded(src_doc, merged_record.id, reason="duplicate"): - merged += 1 - group_supersede_count += 1 - source_to_merged_id[sid] = merged_record.id - consumed_source_ids.add(sid) - - # If every supersede attempt for this group failed (typically - # an ETag race against a concurrent reconcile that already - # superseded the same sources to the *same* deterministic - # merged id), do NOT delete the merged doc. A delete here - # would orphan the sources whose ``superseded_by`` already - # points at this merged id — they'd become invisible to - # default reads (filter ``superseded_by IS NULL``) and to the - # reconcile pool, causing permanent data loss. The merged doc - # is idempotent (deterministic id), so leaving it in place is - # consistent with whatever the winning concurrent writer - # produced. - if group_supersede_count == 0: - logger.info( - "reconcile_memories: no sources superseded for merged id=%s " - "(likely ETag race with concurrent reconcile); leaving " - "merged doc in place — idempotent upsert is self-healing", - merged_record.id, - ) - - # ---- 5. Apply contradicted_pairs SECOND with dangling-id resolution ---- for pair in contradicted_pairs: winner_id = pair.get("winner_id") loser_id = pair.get("loser_id") - if not winner_id or not loser_id: - logger.debug( - "reconcile_memories: skipping malformed contradicted_pair %r", - pair, - ) + if not winner_id or not loser_id or winner_id == loser_id: continue - - # Redirect through any duplicate-merge that absorbed the id. - resolved_winner = source_to_merged_id.get(winner_id, winner_id) - resolved_loser_id = source_to_merged_id.get(loser_id, loser_id) - - # Validate the (resolved) winner. The LLM is instructed never to - # invent IDs — if it does, refuse to write a dangling - # ``superseded_by`` pointer that breaks the audit trail. - if resolved_winner not in facts_by_id and resolved_winner not in merged_docs_by_id: + if winner_id not in facts_by_id: logger.warning( - "reconcile_memories: hallucinated winner_id=%s (resolved=%s) " - "not in pool or merged set; skipping pair %r", + "reconcile_memories: hallucinated winner_id=%s not in pool; skipping pair %r", winner_id, - resolved_winner, pair, ) continue - - if resolved_winner == resolved_loser_id: - # Both sides collapsed into the same merged doc — the - # contradiction is moot. Drop it silently. - logger.debug( - "reconcile_memories: contradiction collapsed into duplicate group " - "(winner=%s loser=%s -> %s); skipping", - winner_id, - loser_id, - resolved_winner, - ) - continue - - loser_doc = facts_by_id.get(resolved_loser_id) - if loser_doc is None and resolved_loser_id != loser_id: - # The original loser was just merged. Reuse the in-memory - # merged doc so we skip a cross-partition re-fetch — we - # own the (user_id, thread_id) partition and just wrote - # it. This in-memory copy carries the ``_etag`` returned - # by ``_upsert_memory``'s captured upsert response, so - # the supersede below takes the ETag-protected - # ``replace_item`` branch — concurrency-safe against any - # other reconcile that may have touched the same merged - # id in parallel. - loser_doc = merged_docs_by_id.get(resolved_loser_id) - + loser_doc = facts_by_id.get(loser_id) if loser_doc is None: - logger.warning( - "reconcile_memories: loser doc not found for pair %r (resolved_loser=%s)", - pair, - resolved_loser_id, - ) continue - - if self._mark_superseded(loser_doc, resolved_winner, reason="contradict"): + if self._mark_superseded(loser_doc, winner_id, reason="contradict"): contradicted += 1 - # Track the *original* loser_id from the LLM so the kept - # cross-check below can reconcile against the input pool. - if loser_id in facts_by_id: - consumed_loser_ids.add(loser_id) - # If the winner is an original pool member (not a freshly - # minted merged doc), record it so the kept-cross-check - # doesn't flag a clean run. - if winner_id in facts_by_id: - contradiction_winner_ids_in_pool.add(winner_id) - - # The pipeline's "kept" semantic = facts that survive as live - # records in the pool. The LLM's ``kept_ids`` semantic = - # everything *not* mentioned in duplicate_groups or - # contradicted_pairs. They differ by exactly the contradiction - # winners (winners survive but are listed under contradicted_pairs). - consumed_ids = consumed_source_ids | consumed_loser_ids - kept_actual = {fid for fid in facts_by_id.keys() if fid not in consumed_ids} - kept = len(kept_actual) - # Cross-check: the LLM's kept_ids set should equal kept_actual - # minus the contradiction winners. Mismatch usually means the LLM - # hallucinated an id or double-counted a fact across categories. - expected_llm_kept = kept_actual - contradiction_winner_ids_in_pool - llm_kept_set = {kid for kid in llm_kept_ids if kid in facts_by_id} - if llm_kept_set != expected_llm_kept: - symdiff = sorted(llm_kept_set ^ expected_llm_kept)[:10] - logger.info( - "reconcile_memories: kept_ids mismatch (llm=%d valid=%d, expected=%d). " - "Likely a hallucinated or double-counted fact id. Sample diff (≤10): %s", - len(llm_kept_ids), - len(llm_kept_set), - len(expected_llm_kept), - symdiff, + consumed_loser_ids.add(loser_id) + + kept = len([fid for fid in facts_by_id if fid not in consumed_loser_ids]) + result = {"kept": kept, "merged": 0, "contradicted": contradicted} + logger.info( + "reconcile_memories contradiction pass user_id=%s memory_type=%s result=%s", + user_id, + memory_type, + result, + ) + return result + + def _active_memories_for_reconcile(self, user_id: str, memory_type: str, n: int) -> list[dict[str, Any]]: + # ---- Load up to N most recent active memories ---- + # ORDER BY c.created_at DESC keeps the TOP cap deterministic across + # physical partitions and matches the dedup prompt's tiebreaker + # ("more recent created_at first"). Cosmos's _ts is the last-write + # timestamp, which would diverge from created_at after any UPDATE. + query = ( + f"SELECT TOP {top_literal(n, name='reconcile_memories.n')} * FROM c " + "WHERE c.user_id = @user_id " + "AND c.type = @memory_type " + f"AND {_ACTIVE_DOC_FILTER} " + "ORDER BY c.created_at DESC" + ) + return list( + self._memories_container.query_items( + query=query, + parameters=[ + {"name": "@user_id", "value": user_id}, + {"name": "@memory_type", "value": memory_type}, + ], + enable_cross_partition_query=True, ) - result = {"kept": kept, "merged": merged, "contradicted": contradicted} - logger.info("reconcile_memories completed: %s", result) - return result, consumed_ids + ) def build_procedural_context(self, user_id: str) -> str: """Return the active synthesized procedural prompt for system injection.""" diff --git a/azure/cosmos/agent_memory/store/_search_helpers.py b/azure/cosmos/agent_memory/store/_search_helpers.py index 032965c..5e58269 100644 --- a/azure/cosmos/agent_memory/store/_search_helpers.py +++ b/azure/cosmos/agent_memory/store/_search_helpers.py @@ -15,7 +15,8 @@ MEMORY_PROJECTION = ( "c.id, c.user_id, c.thread_id, c.role, c.type, c.content, " - "c.metadata, c.created_at, c.tags, c.salience, c.confidence, c.superseded_by" + "c.metadata, c.created_at, c.tags, c.salience, c.confidence, " + "c.superseded_by, c.superseded_at, c.supersede_reason" ) @@ -71,7 +72,7 @@ def coerce_embedding(result: Any) -> list[float]: return result if isinstance(result, list) and not result: raise ConfigurationError( - "Embedder returned an empty vector — likely an upstream embedding failure", + "Embedder returned an empty vector - likely an upstream embedding failure", parameter="embeddings_client", ) raise ConfigurationError("Embedder must return list[float]", parameter="embeddings_client") diff --git a/azure/cosmos/agent_memory/store/memory_store.py b/azure/cosmos/agent_memory/store/memory_store.py index dd5d345..f67b532 100644 --- a/azure/cosmos/agent_memory/store/memory_store.py +++ b/azure/cosmos/agent_memory/store/memory_store.py @@ -30,6 +30,7 @@ from azure.cosmos.agent_memory.logging import get_logger from azure.cosmos.agent_memory.models import MemoryRecord from azure.cosmos.agent_memory.store._search_helpers import ( + MEMORY_PROJECTION, add_salience_filter, add_tag_filters, build_search_sql, @@ -491,7 +492,7 @@ def delete( raise MemoryNotFoundError(memory_id=memory_id, user_id=user_id, thread_id=thread_id) from exc except CosmosAccessConditionFailedError as exc: raise MemoryConflictError( - f"delete conflicted for memory_id={memory_id!r} — document was modified after the type check" + f"delete conflicted for memory_id={memory_id!r} - document was modified after the type check" ) from exc except Exception as exc: raise _wrap_cosmos_exception(exc, message=f"delete failed for {memory_id}: {exc}") from exc @@ -903,6 +904,66 @@ def search( cross_partition=cross_partition, ) + def get_memory_history( + self, + memory_id: str, + user_id: str, + thread_id: Optional[str] = None, + *, + max_depth: int = 20, + ) -> list[dict[str, Any]]: + """Return a memory's superseded predecessors, most-recently-superseded first. + + AMT supersedes rather than deletes: when a fact is updated or + contradicted, the prior document is retained with ``superseded_by`` + pointing at its replacement (see :meth:`mark_superseded`). This walks + that chain backwards from *memory_id* so callers can reason about how a + fact evolved - knowledge updates, preference reversals, relocations - + instead of seeing only the current value. + + The document identified by *memory_id* is not itself included; the return + value is everything it superseded, transitively, bounded by *max_depth* + to guard against cycles. Each entry carries the ``superseded_at`` / + ``supersede_reason`` audit fields so callers can order and explain the + transitions. + """ + if not memory_id: + raise ValidationError("memory_id is required") + if not user_id: + raise ValidationError("user_id is required") + partition_key, cross_partition = query_scope(user_id, thread_id) + history: list[dict[str, Any]] = [] + seen: set[str] = {memory_id} + frontier: list[str] = [memory_id] + for _ in range(max(1, max_depth)): + id_params = [{"name": f"@sid{i}", "value": sid} for i, sid in enumerate(frontier)] + placeholders = ", ".join(param["name"] for param in id_params) + parameters: list[dict[str, Any]] = [*id_params, {"name": "@user_id", "value": user_id}] + where = f"c.superseded_by IN ({placeholders}) AND c.user_id = @user_id" + if thread_id is not None: + where += " AND c.thread_id = @thread_id" + parameters.append({"name": "@thread_id", "value": thread_id}) + sql = f"SELECT {MEMORY_PROJECTION} FROM c WHERE {where}" + rows = self.query( + sql, + parameters, + container_key=ContainerKey.MEMORIES, + partition_key=partition_key, + cross_partition=cross_partition, + ) + frontier = [] + for doc in rows: + doc_id = doc.get("id") + if not doc_id or doc_id in seen: + continue + seen.add(doc_id) + history.append(doc) + frontier.append(doc_id) + if not frontier: + break + history.sort(key=lambda d: d.get("superseded_at") or d.get("created_at") or "", reverse=True) + return history + def search_turns( self, search_terms: Optional[str] = None, diff --git a/azure/cosmos/agent_memory/thresholds.py b/azure/cosmos/agent_memory/thresholds.py index f25030c..a1edc6a 100644 --- a/azure/cosmos/agent_memory/thresholds.py +++ b/azure/cosmos/agent_memory/thresholds.py @@ -18,7 +18,7 @@ DEFAULT_FACT_EXTRACTION_EVERY_N = 1 DEFAULT_THREAD_SUMMARY_EVERY_N = 10 DEFAULT_USER_SUMMARY_EVERY_N = 20 -# Dedup runs on its own cadence — every Nth extract (NOT every Nth turn), +# Dedup runs on its own cadence - every Nth extract (NOT every Nth turn), # because dedup is O(N²) over all active facts and dominates per-push cost # when FACT_EXTRACTION_EVERY_N=1. Default 5 = one dedup sweep per 5 extracts. # Set to 1 to dedup on every extract; set to 0 to disable entirely. @@ -29,21 +29,15 @@ DEFAULT_DEDUP_POOL_SIZE = 50 # --------------------------------------------------------------------------- -# INTERNAL dedup/search tuning — NOT customer-configurable. +# INTERNAL dedup/search tuning - NOT customer-configurable. # These ship as fixed feature constants (no env vars, not in any settings # template). They are maintainer-tunable here in code only; if a knob ever # needs to become operator-facing we add the env plumbing back deliberately. # The dedup + hybrid-search features ship ON via these values. # --------------------------------------------------------------------------- -DEDUP_CONTEXT_VECTOR_ENABLED = True # Stage-1 relevance-ranked extraction context -DEDUP_CONTEXT_TOPK = 10 -DEDUP_VECTOR_ENABLED = True # Stage-3 vector near-dup ladder -DEDUP_SIM_HIGH = 0.97 # >= -> auto-skip near-exact -DEDUP_SIM_LOW = 0.80 # < -> novel; between -> tag candidate -DEDUP_CANDIDATE_TOPK = 10 -DEDUP_RECONCILE_MODE = "candidate" # clustered candidate reconcile (vs legacy full_pool) -DEDUP_CLUSTER_SIM = 0.60 # Stage-5 clustering edge threshold -DEDUP_FULL_RECLUSTER_EVERY_N = 12 # full re-cluster safety net cadence +EXTRACTION_BATCH_MAX_TOKENS = 7000 +DEDUP_VECTOR_ENABLED = True # write-time in-place near-dup folding +DEDUP_SIM_HIGH = 0.97 # >= -> fold new memory into existing canonical in place DEFAULT_TTL_BY_TYPE: dict[str, int] = { "turn": 2_592_000, @@ -65,7 +59,7 @@ # vector index, so this only governs whether vectors are generated/searched. DEFAULT_ENABLE_TURN_EMBEDDINGS = False -# Owner exclusivity — declares which backend is authoritative for the shared +# Owner exclusivity - declares which backend is authoritative for the shared # memories + counter container. When set, the *other* backend skips its # auto-trigger and logs a loud warning. Default unset preserves today's # behavior (no enforcement) for backward compatibility. @@ -143,7 +137,7 @@ def get_dedup_pool_size() -> int: the pipeline; values above are clamped to 500 with a WARN.""" raw = _parse_threshold("DEDUP_POOL_SIZE", DEFAULT_DEDUP_POOL_SIZE) if raw == 0: - # 0 isn't meaningful for a pool size — fall back to default. + # 0 isn't meaningful for a pool size - fall back to default. logger.warning( "DEDUP_POOL_SIZE=0 is invalid for a pool size; using default %d", DEFAULT_DEDUP_POOL_SIZE, @@ -156,18 +150,13 @@ def get_dedup_pool_size() -> int: # --------------------------------------------------------------------------- -# Internal dedup/search feature accessors — return fixed constants (no env). +# Internal dedup/search feature accessors - return fixed constants (no env). # Kept as thin functions so call sites stay stable and the values can be # changed in one place; NOT customer-configurable. # --------------------------------------------------------------------------- -def get_dedup_context_vector_enabled() -> bool: - """Whether Stage-1 extraction context uses vector retrieval (internal; on).""" - return DEDUP_CONTEXT_VECTOR_ENABLED - - -def get_dedup_context_topk() -> int: - """Top-K memories to retrieve for Stage-1 extraction context (internal).""" - return DEDUP_CONTEXT_TOPK +def get_extraction_batch_max_tokens() -> int: + """Token budget per extraction batch (internal; see EXTRACTION_BATCH_MAX_TOKENS).""" + return EXTRACTION_BATCH_MAX_TOKENS def get_dedup_vector_enabled() -> bool: @@ -176,35 +165,11 @@ def get_dedup_vector_enabled() -> bool: def get_dedup_sim_high() -> float: - """Similarity at/above which near-exact memories are auto-skipped (internal).""" + """Similarity at/above which a new memory is folded into its existing + canonical record in place (internal).""" return DEDUP_SIM_HIGH -def get_dedup_sim_low() -> float: - """Similarity below which memories are treated as novel (internal).""" - return DEDUP_SIM_LOW - - -def get_dedup_candidate_topk() -> int: - """Top-K existing memories pulled per new memory in Stage-3 (internal).""" - return DEDUP_CANDIDATE_TOPK - - -def get_dedup_reconcile_mode() -> str: - """Reconcile mode: ``candidate`` clustering (internal; the shipped feature).""" - return DEDUP_RECONCILE_MODE - - -def get_dedup_cluster_sim() -> float: - """Similarity threshold for Stage-5 clustering edges (internal).""" - return DEDUP_CLUSTER_SIM - - -def get_dedup_full_recluster_every_n() -> int: - """Full re-cluster safety-net cadence, every Nth reconcile sweep (internal).""" - return DEDUP_FULL_RECLUSTER_EVERY_N - - def get_procedural_synthesis_auto() -> bool: """Whether procedural synthesis auto-fires after extract. @@ -231,7 +196,7 @@ def get_processor_owner() -> Optional[str]: """Return the configured ``MEMORY_PROCESSOR_OWNER`` or ``None``. Each side reads this to decide whether to run its auto-trigger. The - contract is **asymmetric** by design — there is no cross-process lock, + contract is **asymmetric** by design - there is no cross-process lock, so the two sides default differently to avoid double-firing: * **SDK** (in-process) fires when the value is ``None`` (unset) or @@ -247,7 +212,7 @@ def get_processor_owner() -> Optional[str]: . note:: This is **operator-configured exclusivity, not enforced**. Counter writes still stamp ``last_owner`` and emit a one-shot WARN when the - observed owner disagrees with the writer — treat that as a + observed owner disagrees with the writer - treat that as a configuration audit signal, not a guarantee. """ raw = os.environ.get("MEMORY_PROCESSOR_OWNER") @@ -281,15 +246,9 @@ def get_processor_owner() -> Optional[str]: "get_user_summary_every_n", "get_dedup_every_n", "get_dedup_pool_size", - "get_dedup_context_vector_enabled", - "get_dedup_context_topk", + "get_extraction_batch_max_tokens", "get_dedup_vector_enabled", "get_dedup_sim_high", - "get_dedup_sim_low", - "get_dedup_candidate_topk", - "get_dedup_reconcile_mode", - "get_dedup_cluster_sim", - "get_dedup_full_recluster_every_n", "get_procedural_synthesis_auto", "get_enable_turn_embeddings", "get_processor_owner", diff --git a/function_app/orchestrators/extract_memories.py b/function_app/orchestrators/extract_memories.py index db56c24..4952b11 100644 --- a/function_app/orchestrators/extract_memories.py +++ b/function_app/orchestrators/extract_memories.py @@ -34,7 +34,6 @@ def ExtractMemoriesOrchestrator(context: df.DurableOrchestrationContext): user_id = payload["user_id"] thread_id = payload["thread_id"] should_reconcile = bool(payload.get("reconcile", False)) - full_rebuild = bool(payload.get("full_rebuild", False)) recent_k = payload.get("recent_k") retry = default_retry_options() @@ -71,7 +70,7 @@ def ExtractMemoriesOrchestrator(context: df.DurableOrchestrationContext): reconciled = yield context.call_activity_with_retry( "em_ReconcileMemories", retry, - {"user_id": user_id, "full_rebuild": full_rebuild}, + {"user_id": user_id}, ) if config.get_procedural_synthesis_auto(): count = payload.get("count") @@ -172,16 +171,10 @@ def em_ReconcileMemories(payload: dict) -> dict: # GA keeps reconcile single-activity: its LLM dedup decisions and supersession # operations are larger/more coupled than the extract→dedup→persist split handled here. user_id = payload["user_id"] - # full_rebuild forces the full-pool single-LLM-pass path (catches dissimilar - # contradictions). The change-feed sets it on a persisted-counter cadence so it - # fires reliably on FA, where the in-memory candidate-mode sweep counter can't. - full_rebuild = bool(payload.get("full_rebuild", False)) pipeline = get_pipeline() from azure.cosmos.agent_memory.thresholds import get_dedup_pool_size n = get_dedup_pool_size() - facts = pipeline.reconcile_memories(user_id=user_id, n=n, memory_type="fact", full_rebuild=full_rebuild) or {} - episodic = ( - pipeline.reconcile_memories(user_id=user_id, n=n, memory_type="episodic", full_rebuild=full_rebuild) or {} - ) + facts = pipeline.reconcile_memories(user_id=user_id, n=n, memory_type="fact") or {} + episodic = pipeline.reconcile_memories(user_id=user_id, n=n, memory_type="episodic") or {} return {"fact": facts, "episodic": episodic} diff --git a/function_app/shared/config.py b/function_app/shared/config.py index da4f4fa..40a1b97 100644 --- a/function_app/shared/config.py +++ b/function_app/shared/config.py @@ -3,11 +3,11 @@ All knobs are read from environment variables / Azure Functions app settings. Defaults: -* ``FACT_EXTRACTION_EVERY_N`` — default 1 (per-turn extraction) -* ``THREAD_SUMMARY_EVERY_N`` — default 10 (rolling summary cadence) -* ``USER_SUMMARY_EVERY_N`` — default 20 -* ``PROCEDURAL_SYNTHESIS_AUTO`` — default true -* ``MAX_BATCH_SIZE`` — default 20 +* ``FACT_EXTRACTION_EVERY_N`` - default 1 (per-turn extraction) +* ``THREAD_SUMMARY_EVERY_N`` - default 10 (rolling summary cadence) +* ``USER_SUMMARY_EVERY_N`` - default 20 +* ``PROCEDURAL_SYNTHESIS_AUTO`` - default true +* ``MAX_BATCH_SIZE`` - default 20 The fact-extraction default of ``1`` - every new turn produces fresh facts. Operators can raise this for cost-sensitive @@ -27,7 +27,7 @@ contribution in one ETag-guarded write. Even if the resulting jump crosses several ``EVERY_N`` boundaries (e.g. ``EVERY_N=10`` and a batch of 100 turns advances a counter from 0 → 100), exactly **one** orchestrator is -started — keyed on the new counter value — and it folds in only the +started - keyed on the new counter value - and it folds in only the trailing ``MAX_BATCH_SIZE`` items. For high-throughput bursts this means summarization is coarser than ``EVERY_N`` alone implies. Operators who need finer per-boundary fan-out should lower ``MAX_BATCH_SIZE`` and raise diff --git a/function_app/shared/cosmos_clients.py b/function_app/shared/cosmos_clients.py index 825a2e2..ab56a35 100644 --- a/function_app/shared/cosmos_clients.py +++ b/function_app/shared/cosmos_clients.py @@ -10,7 +10,7 @@ * ``get_counter_container_async()`` returns an *async* AsyncContainerProxy used by the change-feed trigger to update counters. -Clients are lazily constructed and cached at module level — Azure Functions +Clients are lazily constructed and cached at module level - Azure Functions re-uses the same Python worker across invocations, so we want one client per worker instead of one per invocation. """ @@ -113,7 +113,7 @@ async def get_counter_container_async(): async def close_async_clients() -> None: """Close the cached async Cosmos client and credential. - Idempotent — safe to call multiple times. Exposed so tests and explicit + Idempotent - safe to call multiple times. Exposed so tests and explicit shutdown hooks can release the underlying ``aiohttp`` session cleanly. """ global _async_cosmos_client, _async_counter_container, _async_credential diff --git a/function_app/shared/counters.py b/function_app/shared/counters.py index 1adb0e5..be0d14c 100644 --- a/function_app/shared/counters.py +++ b/function_app/shared/counters.py @@ -6,11 +6,11 @@ Counter document shape:: - # thread-scoped — id = "thread:{user_id}:{thread_id}", PK = [user_id, thread_id] + # thread-scoped - id = "thread:{user_id}:{thread_id}", PK = [user_id, thread_id] { "id": ..., "user_id": ..., "thread_id": ..., "count": int, "last_batch_lsn": int|None, "last_batch_old_count": int } - # user-scoped — id = "user:{user_id}", PK = [user_id, "__counters__"] + # user-scoped - id = "user:{user_id}", PK = [user_id, "__counters__"] { "id": ..., "user_id": ..., "thread_id": "__counters__", "count": int, "last_batch_lsn": int|None, "last_batch_old_count": int } """ @@ -41,7 +41,7 @@ def _maybe_warn_owner_mismatch( ) -> None: """Log a one-shot WARN when the counter's previous writer differs from us. - Advisory-only — the FA still runs the orchestration. ``MEMORY_PROCESSOR_OWNER`` + Advisory-only - the FA still runs the orchestration. ``MEMORY_PROCESSOR_OWNER`` is operator-configured exclusivity, not a server-side lock; this just surfaces accidental double-deployment so it shows up in App Insights. """ @@ -116,7 +116,7 @@ async def increment_counter_by( old_count = existing_doc.get("count", 0) etag = existing_doc.get("_etag") except CosmosResourceNotFoundError: - pass # first time — will create + pass # first time - will create # Owner-mismatch detection (advisory only). if existing_doc is not None: @@ -131,7 +131,7 @@ async def increment_counter_by( # re-balance, host crash → checkpoint regression where another # batch landed in between) also short-circuit. For the exact # match we replay the cached result; for the strict-greater case - # we return (current, current) — no threshold crossing — because + # we return (current, current) - no threshold crossing - because # the batch's effect is already absorbed in a later state we have # no cached pre-batch value for. if ( @@ -179,7 +179,7 @@ async def increment_counter_by( new_doc["last_failure_reason"] = existing_doc.get("last_failure_reason") if "last_extract_count" in existing_doc: new_doc["last_extract_count"] = existing_doc.get("last_extract_count") - # Stamp the writing backend (advisory only — not enforced server-side). + # Stamp the writing backend (advisory only - not enforced server-side). if owner is not None: new_doc["last_owner"] = owner elif existing_doc is not None and "last_owner" in existing_doc: @@ -222,7 +222,7 @@ async def increment_counter_by( ) continue if exc.status_code == 412: - # Exhausted ETag retries — RAISE so the change-feed batch retries + # Exhausted ETag retries - RAISE so the change-feed batch retries # via at-least-once redelivery. The last_batch_lsn replay-protection # ensures the next attempt won't double-increment. Silently # returning (old, old) here would advance the lease without ever @@ -235,7 +235,7 @@ async def increment_counter_by( raise raise - # Should never reach here — the loop either returns or raises every iteration. + # Should never reach here - the loop either returns or raises every iteration. raise RuntimeError(f"increment_counter_by({counter_id}) exhausted MAX_RETRIES without resolution") @@ -255,7 +255,7 @@ def crosses_threshold(old_count: int, new_count: int, n: int) -> bool: Raises: ValueError: if ``n <= 0``. Callers should gate on ``n > 0`` instead of - relying on a "disabled" sentinel here — keeping this strict makes + relying on a "disabled" sentinel here - keeping this strict makes misuse loud. """ if n <= 0: diff --git a/function_app/shared/pipeline_factory.py b/function_app/shared/pipeline_factory.py index edf5c03..02d64fc 100644 --- a/function_app/shared/pipeline_factory.py +++ b/function_app/shared/pipeline_factory.py @@ -1,7 +1,7 @@ """Lazy PipelineService factory (MI auth, sync clients). The activities reuse :class:`azure.cosmos.agent_memory.services.pipeline.PipelineService` -verbatim — no business logic is duplicated in the function app. +verbatim - no business logic is duplicated in the function app. """ from __future__ import annotations diff --git a/function_app/triggers/change_feed.py b/function_app/triggers/change_feed.py index 6b7ad9d..80d806e 100644 --- a/function_app/triggers/change_feed.py +++ b/function_app/triggers/change_feed.py @@ -98,7 +98,7 @@ async def process_changefeed_batch( # * Function-App deployments: must set ``MEMORY_PROCESSOR_OWNER=durable`` # explicitly to enable the change-feed trigger. Without that opt-in # both backends would double-extract / double-dedup against the same - # writes — exactly the customer-day-one footgun this guard prevents. + # writes - exactly the customer-day-one footgun this guard prevents. from azure.cosmos.agent_memory.thresholds import ( PROCESSOR_OWNER_DURABLE, get_processor_owner, @@ -110,7 +110,7 @@ async def process_changefeed_batch( if not _warned_owner_skip: _warned_owner_skip = True logger.warning( - "Change-feed trigger skipping batch — MEMORY_PROCESSOR_OWNER is %r, " + "Change-feed trigger skipping batch - MEMORY_PROCESSOR_OWNER is %r, " "set it to 'durable' on the Function App to enable Durable orchestration. " "Further skipped batches will be logged at DEBUG level.", owner, @@ -129,16 +129,6 @@ async def process_changefeed_batch( # Reconcile fires every (n_facts * n_dedup) turns, matching the SDK # auto-trigger contract. Disabled when either knob is 0. n_dedup_turns = n_facts * n_dedup if (n_facts > 0 and n_dedup > 0) else 0 - # Full-pool reconcile backstop cadence. The candidate-mode "every Nth sweep" - # gate is an in-memory per-worker counter, which on the FA backend (per-worker - # singleton pipeline) resets on recycle/scale and never reaches N for a user — - # so the full-pool pass that catches dissimilar-embedding contradictions almost - # never fires. Derive it from the PERSISTED counter instead: every - # DEDUP_FULL_RECLUSTER_EVERY_N-th reconcile = every n_dedup_turns * N turns. - from azure.cosmos.agent_memory.thresholds import get_dedup_full_recluster_every_n - - n_full_recluster = get_dedup_full_recluster_every_n() - n_full_turns = n_dedup_turns * n_full_recluster if (n_dedup_turns > 0 and n_full_recluster > 0) else 0 if n_thread == 0 and n_facts == 0 and n_user == 0: return # all orchestrators disabled @@ -217,16 +207,12 @@ async def process_changefeed_batch( if n_facts > 0 and crosses_threshold(old_count, new_count, n_facts): instance_id = f"extract:{user_id}:{thread_id}:{new_count}" should_reconcile = bool(n_dedup_turns > 0 and crosses_threshold(old_count, new_count, n_dedup_turns)) - # Persisted-counter backstop: every n_full_turns turns, force a - # full-pool reconcile so dissimilar-embedding contradictions are - # caught reliably on FA (not gated by the in-memory sweep counter). - should_full_reconcile = bool(n_full_turns > 0 and crosses_threshold(old_count, new_count, n_full_turns)) watermark = await read_extract_watermark(counter_container, cid, user_id, thread_id) # Not capped: new_count - watermark is exactly the unextracted backlog # and the orchestrator advances the watermark to new_count, so capping # would strand the oldest turns (DEDUP_POOL_SIZE is the reconcile knob, # not the extraction window). Bootstrap: with no watermark yet, base=0 - # so recent_k = new_count covers every turn so far — using only this + # so recent_k = new_count covers every turn so far - using only this # batch (new_count - old_count) would strand turns added during earlier # failed extracts once the watermark first advances to new_count. base = watermark if watermark is not None else 0 @@ -240,7 +226,6 @@ async def process_changefeed_batch( "thread_id": thread_id, "count": new_count, "reconcile": should_reconcile, - "full_rebuild": should_full_reconcile, "recent_k": recent_k, }, orchestration_errors, @@ -301,7 +286,7 @@ async def _safe_start( # retries: it shows up as a 409 Conflict from the Durable client. # The Durable Python SDK currently raises bare ``Exception`` for # everything (no typed class for OrchestrationAlreadyExists), so - # we keep two signals — status_code first, then the canonical + # we keep two signals - status_code first, then the canonical # message string. We always log ``type(exc).__name__`` so any future # SDK switch to a typed exception shows up immediately in App # Insights instead of silently breaking string matching. @@ -347,6 +332,6 @@ async def _safe_start( ) @bp.durable_client_input(client_name="starter") async def on_memory_change(documents: func.DocumentList, starter) -> None: - """Change-feed trigger entry point — delegates to :func:`process_changefeed_batch`.""" + """Change-feed trigger entry point - delegates to :func:`process_changefeed_batch`.""" docs = [dict(doc) for doc in documents] await process_changefeed_batch(docs, starter) diff --git a/pyproject.toml b/pyproject.toml index f6c7a56..60f7ee8 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.0b7" description = "Store, retrieve, and transform AI agent memories backed by Azure Cosmos DB" readme = "README.md" license = {file = "LICENSE"} diff --git a/tests/conftest.py b/tests/conftest.py index 87206ed..d239fbb 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -131,7 +131,7 @@ def ai_foundry_endpoint(): @pytest.fixture(scope="session") def ai_foundry_api_key(): - """Azure OpenAI API key from env vars (optional — Entra ID is preferred).""" + """Azure OpenAI API key from env vars (optional - Entra ID is preferred).""" return os.environ.get("AI_FOUNDRY_API_KEY", "") diff --git a/tests/integration/test_async_full_pipeline.py b/tests/integration/test_async_full_pipeline.py index a47804e..2bb6347 100644 --- a/tests/integration/test_async_full_pipeline.py +++ b/tests/integration/test_async_full_pipeline.py @@ -3,7 +3,7 @@ The async pipeline mirrors every sync dedup/reconcile code path line-for-line (watermark, euclidean guard, ``exclude_ids`` parity, reconcile cadence, the vector-floor dedup ladder). The sync suite covers all of that against a live -backend, but the async client had **no** live coverage — its processor test is +backend, but the async client had **no** live coverage - its processor test is fully mocked. This module exercises the real async end-to-end flow (write turns → extract → reconcile → search) so the async mirror can't silently diverge from sync without a test failing. @@ -146,7 +146,7 @@ async def _async_seed_fact_with_embedding( if embedded: return await asyncio.sleep(1) - pytest.skip(f"embedding service unavailable — could not seed an embedded fact for {content!r}") + pytest.skip(f"embedding service unavailable - could not seed an embedded fact for {content!r}") async def _async_wait_vector_searchable( @@ -228,7 +228,7 @@ async def test_dedup_extracted_memories_flags_near_duplicate_of_stored_fact( ): """Async extract-time vector floor drops/tags a near-duplicate fact. - Parity check with the sync ``TestExtractTimeVectorDedup`` — guards the + Parity check with the sync ``TestExtractTimeVectorDedup`` - guards the async ``dedup_extracted_memories`` mirror (``_vector_candidates`` + similarity bands) against a live backend. Driven with a controlled near-duplicate (no LLM variance) so the assertion is deterministic. diff --git a/tests/integration/test_changefeed_integration.py b/tests/integration/test_changefeed_integration.py index d9075c8..f6a702d 100644 --- a/tests/integration/test_changefeed_integration.py +++ b/tests/integration/test_changefeed_integration.py @@ -12,7 +12,7 @@ A running Azure Functions host (deployed or local ``func start``) with the change-feed trigger configured is required, along with the ``turns``, ``counter``, and ``leases`` containers. Post container-split, the change-feed -trigger binds to the ``turns`` container — turn documents must be inserted +trigger binds to the ``turns`` container - turn documents must be inserted there, not into ``memories``. """ diff --git a/tests/integration/test_full_pipeline.py b/tests/integration/test_full_pipeline.py index ed6be30..2f9e8ac 100644 --- a/tests/integration/test_full_pipeline.py +++ b/tests/integration/test_full_pipeline.py @@ -5,7 +5,7 @@ episodic extraction, deduplication) inline, and reading back results via Cosmos DB queries and vector / hybrid search. -The Azure Function host is **not** required — the same ProcessingPipeline +The Azure Function host is **not** required - the same ProcessingPipeline that the change-feed trigger invokes is also exposed directly on ``CosmosMemoryClient`` (``extract_memories``, ``generate_thread_summary``, ``generate_user_summary``, ``reconcile``). @@ -135,7 +135,7 @@ def _seed_fact_with_embedding( ``add_cosmos`` generates the embedding synchronously; a transient embedding-service blip logs "proceeding without embedding" and stores the doc without a vector, which would leave the extract-time vector floor with no - neighbour to match. Retry until an embedded copy exists (indexing is fast — + neighbour to match. Retry until an embedded copy exists (indexing is fast - the doc is vector-searchable within ~2s), and skip honestly if the embedding service is genuinely unavailable rather than reporting a false failure.""" check = "SELECT c.id FROM c WHERE c.user_id = @uid AND c.content = @content AND IS_DEFINED(c.embedding)" @@ -157,7 +157,7 @@ def _seed_fact_with_embedding( if embedded: return time.sleep(1) - pytest.skip(f"embedding service unavailable — could not seed an embedded fact for {content!r}") + pytest.skip(f"embedding service unavailable - could not seed an embedded fact for {content!r}") def _wait_vector_searchable( @@ -199,7 +199,7 @@ def test_generate_thread_summary(self, agent_memory, unique_user_id, unique_thre ("user", "What are some good restaurants in Paris?"), ("agent", "Le Comptoir du Panthéon is a classic bistro in the 5th arrondissement."), ("user", "What kind of cuisine do they serve?"), - ("agent", "Traditional French bistro fare — confit de canard, steak frites, etc."), + ("agent", "Traditional French bistro fare - confit de canard, steak frites, etc."), ], ) time.sleep(1) @@ -585,7 +585,7 @@ def test_dedup_extracted_memories_flags_near_duplicate_of_stored_fact( ): try: # Seed a stored fact (embedded + vector-indexed) to dedup against. - # Concrete, minimally-reworded facts embed ~0.93-0.98 cosine — well + # Concrete, minimally-reworded facts embed ~0.93-0.98 cosine - well # inside the DEDUP_SIM_LOW (0.80) / DEDUP_SIM_HIGH (0.97) bands. _seed_fact_with_embedding( agent_memory, unique_user_id, unique_thread_id, "The user has a cat named Whiskers." diff --git a/tests/integration/test_processor_integration.py b/tests/integration/test_processor_integration.py index b5a0be2..5e83f02 100644 --- a/tests/integration/test_processor_integration.py +++ b/tests/integration/test_processor_integration.py @@ -1,7 +1,7 @@ """Integration-style tests for the processor protocol surface. Exercises ``CosmosMemoryClient.process_now`` / ``process_now_and_wait`` end-to-end with -a fully mocked Cosmos container — no live Azure calls — to validate that +a fully mocked Cosmos container - no live Azure calls - to validate that the SDK wires the active :class:`MemoryProcessor` correctly through the public API. """ diff --git a/tests/integration/test_processor_integration_async.py b/tests/integration/test_processor_integration_async.py index 5769c3e..e96b8d7 100644 --- a/tests/integration/test_processor_integration_async.py +++ b/tests/integration/test_processor_integration_async.py @@ -31,7 +31,7 @@ class TestAsyncInProcessProcessNowEndToEnd: async def test_process_now_invokes_pipeline_with_correct_args(self): pipeline = MagicMock() # ``AsyncInProcessProcessor`` awaits these pipeline methods, so the - # mocks must be ``AsyncMock`` — a plain ``MagicMock`` returns the dict + # mocks must be ``AsyncMock`` - a plain ``MagicMock`` returns the dict # synchronously, which ``await`` cannot consume. pipeline.generate_thread_summary = AsyncMock( return_value={ @@ -167,7 +167,7 @@ async def _no_sleep(*_a, **_k): monkeypatch.setattr("asyncio.sleep", _no_sleep) - # Tiny timeout — loop.time() advances normally, so this expires fast. + # Tiny timeout - loop.time() advances normally, so this expires fast. ok = await client.process_now_and_wait(user_id="u-to", thread_id="th-to", timeout=0.001) assert ok is False diff --git a/tests/unit/aio/processors/test_durable.py b/tests/unit/aio/processors/test_durable.py index bf89abe..9a61b97 100644 --- a/tests/unit/aio/processors/test_durable.py +++ b/tests/unit/aio/processors/test_durable.py @@ -1,4 +1,4 @@ -"""Tests for AsyncDurableFunctionProcessor — verifies all methods are no-ops.""" +"""Tests for AsyncDurableFunctionProcessor - verifies all methods are no-ops.""" from __future__ import annotations diff --git a/tests/unit/aio/processors/test_inprocess.py b/tests/unit/aio/processors/test_inprocess.py index ec22819..c933e7d 100644 --- a/tests/unit/aio/processors/test_inprocess.py +++ b/tests/unit/aio/processors/test_inprocess.py @@ -53,7 +53,7 @@ async def test_close_is_noop(): @pytest.mark.asyncio async def test_process_reconcile_invokes_pipeline_with_env_pool_size(monkeypatch): """Regression: this was completely broken (ModuleNotFoundError on - ``from ..thresholds``) — auto-trigger silently never reconciled in + ``from ..thresholds``) - auto-trigger silently never reconciled in async deployments. Verify the call now succeeds and forwards the env-tunable pool size from ``get_dedup_pool_size``.""" monkeypatch.setenv("DEDUP_POOL_SIZE", "37") @@ -68,12 +68,10 @@ async def test_process_reconcile_invokes_pipeline_with_env_pool_size(monkeypatch assert pipeline.reconcile_memories.await_args_list[0].kwargs == { "n": 37, "memory_type": "fact", - "full_rebuild": False, } assert pipeline.reconcile_memories.await_args_list[1].kwargs == { "n": 37, "memory_type": "episodic", - "full_rebuild": False, } assert count == 10 # (merged+contradicted) x2 types diff --git a/tests/unit/aio/services/test_dedup_vector_async.py b/tests/unit/aio/services/test_dedup_vector_async.py index e20a206..a193b78 100644 --- a/tests/unit/aio/services/test_dedup_vector_async.py +++ b/tests/unit/aio/services/test_dedup_vector_async.py @@ -111,328 +111,204 @@ async def fake_query_items(_container, *, query, parameters): @pytest.mark.asyncio -async def test_candidate_mode_clears_tags_only_for_survivors(): - # Latent-bug regression (async mirror): consumed sources must not be - # re-upserted by tag-clearing, which would resurrect them without superseded_by. +async def test_dedup_extracted_folds_near_dup_in_place_and_keeps_novel(): p = _service() - f1 = _fact("f1", "a", tags=["sys:fact", "sys:dup-candidate"]) - f2 = _fact("f2", "b", tags=["sys:fact", "sys:dup-candidate"]) - f3 = _fact("f3", "c", tags=["sys:fact", "sys:dup-candidate"]) - p._build_candidate_clusters = AsyncMock(return_value=([[f1, f2, f3]], 3, [f1, f2, f3])) - p._reconcile_pool = AsyncMock(return_value=({"kept": 1, "merged": 2, "contradicted": 0}, {"f1", "f2"})) - cleared: list[str] = [] - - async def clear(docs): - cleared.extend(d["id"] for d in docs) - - p._clear_dup_candidate_tags = AsyncMock(side_effect=clear) - p._emit_reconcile_outcome = MagicMock() - - result = await p._reconcile_candidate_mode("u1", n=50, memory_type="fact", started_at=0.0) - - assert cleared == ["f3"] - assert result == { - "kept": 1, - "merged": 2, - "contradicted": 0, - "reconcile_clusters_sent": 1, - "reconcile_llm_calls_saved": 2, + p._vector_distance_function = AsyncMock(return_value="cosine") + p._embed_batch.return_value = [[1.0, 0.0], [0.0, 1.0]] + p._nearest_active_full = AsyncMock( + side_effect=[ + ({"id": "existing-1", "content": "same", "type": "fact"}, 0.99), + (None, 0.0), + ] + ) + p._apply_inplace_update = AsyncMock(return_value=True) + extracted = { + "facts": [_fact("f-dup", "restatement"), _fact("f-novel", "brand new")], + "episodic": [], + "updates": [], } + out = await p.dedup_extracted_memories("u1", extracted) -@pytest.mark.asyncio -async def test_candidate_mode_clears_tags_on_orphan_seeds(): - # Async mirror: orphan dup-candidate seeds (no cluster) get their stale tag cleared. - p = _service() - orphan = _fact("orphan", "lonely", tags=["sys:fact", "sys:dup-candidate"]) - c1 = _fact("c1", "a", tags=["sys:fact", "sys:dup-candidate"]) - c2 = _fact("c2", "b", tags=["sys:fact", "sys:dup-candidate"]) - p._build_candidate_clusters = AsyncMock(return_value=([[c1, c2]], 3, [orphan, c1, c2])) - p._reconcile_pool = AsyncMock(return_value=({"kept": 2, "merged": 0, "contradicted": 0}, set())) - cleared: list[str] = [] - - async def clear(docs): - cleared.extend(d["id"] for d in docs) - - p._clear_dup_candidate_tags = AsyncMock(side_effect=clear) - p._emit_reconcile_outcome = MagicMock() - - await p._reconcile_candidate_mode("u1", n=50, memory_type="fact", started_at=0.0) - - assert set(cleared) == {"c1", "c2", "orphan"} + assert [doc["id"] for doc in out["facts"]] == ["f-novel"] + p._apply_inplace_update.assert_awaited_once() + target, new_doc = p._apply_inplace_update.call_args.args + assert target["id"] == "existing-1" + assert new_doc["id"] == "f-dup" + assert out["updates"][-1]["inplace_updated"] == 1 @pytest.mark.asyncio -async def test_sweep_survives_one_cluster_failure(): - # Async mirror: one failing cluster must not abort the sweep; failed cluster's - # tags are retained, remaining cluster + orphan are still cleared. +async def test_dedup_extracted_failed_inplace_update_keeps_new_doc(): p = _service() - c1 = [ - _fact("a1", "x", tags=["sys:fact", "sys:dup-candidate"]), - _fact("a2", "y", tags=["sys:fact", "sys:dup-candidate"]), - ] - c2 = [ - _fact("b1", "p", tags=["sys:fact", "sys:dup-candidate"]), - _fact("b2", "q", tags=["sys:fact", "sys:dup-candidate"]), - ] - orphan = _fact("o1", "lonely", tags=["sys:fact", "sys:dup-candidate"]) - p._build_candidate_clusters = AsyncMock(return_value=([c1, c2], 5, [*c1, *c2, orphan])) - p._reconcile_pool = AsyncMock( - side_effect=[RuntimeError("truncated LLM response"), ({"kept": 2, "merged": 0, "contradicted": 0}, set())] - ) - cleared: list[str] = [] - - async def clear(docs): - cleared.extend(d["id"] for d in docs) - - p._clear_dup_candidate_tags = AsyncMock(side_effect=clear) - p._emit_reconcile_outcome = MagicMock() + p._vector_distance_function = AsyncMock(return_value="cosine") + p._embed_batch.return_value = [[1.0, 0.0]] + p._nearest_active_full = AsyncMock(return_value=({"id": "existing-1", "content": "same", "type": "fact"}, 0.99)) + p._apply_inplace_update = AsyncMock(return_value=False) + extracted = {"facts": [_fact("f-dup", "restatement")], "episodic": [], "updates": []} - result = await p._reconcile_candidate_mode("u1", n=50, memory_type="fact", started_at=0.0) + out = await p.dedup_extracted_memories("u1", extracted) - assert "a1" not in cleared and "a2" not in cleared - assert {"b1", "b2", "o1"} <= set(cleared) - assert result["reconcile_clusters_sent"] == 2 + assert [doc["id"] for doc in out["facts"]] == ["f-dup"] + assert all(op.get("op") != "stats" or "inplace_updated" not in op for op in out["updates"]) @pytest.mark.asyncio -async def test_full_rebuild_clears_survivor_tags(monkeypatch): - # Async mirror: full_rebuild full-pool path clears survivor dup-candidate tags. - monkeypatch.setattr("azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_reconcile_mode", lambda: "candidate") +async def test_dedup_extracted_below_threshold_is_novel(): p = _service() - pool = [ - _fact("f1", "a", tags=["sys:fact", "sys:dup-candidate"]), - _fact("f2", "b", tags=["sys:fact", "sys:dup-candidate"]), - ] - p._active_memories_for_reconcile = AsyncMock(return_value=pool) - p._reconcile_pool = AsyncMock(return_value=({"kept": 1, "merged": 1, "contradicted": 0}, {"f1"})) - cleared: list[str] = [] - - async def clear(docs): - cleared.extend(d["id"] for d in docs) - - p._clear_dup_candidate_tags = AsyncMock(side_effect=clear) - p._emit_reconcile_outcome = MagicMock() + p._vector_distance_function = AsyncMock(return_value="cosine") + p._embed_batch.return_value = [[1.0, 0.0]] + p._nearest_active_full = AsyncMock(return_value=({"id": "existing-1", "content": "near", "type": "fact"}, 0.85)) + p._apply_inplace_update = AsyncMock(return_value=True) + extracted = {"facts": [_fact("f-new", "somewhat similar")], "episodic": [], "updates": []} - await p.reconcile_memories("u1", n=50, memory_type="fact", full_rebuild=True) + out = await p.dedup_extracted_memories("u1", extracted) - assert cleared == ["f2"] + assert [doc["id"] for doc in out["facts"]] == ["f-new"] + p._apply_inplace_update.assert_not_awaited() @pytest.mark.asyncio -async def test_dedup_extracted_memories_flag_off_is_noop(monkeypatch): - monkeypatch.setattr("azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_vector_enabled", lambda: False) +async def test_euclidean_disables_inplace_folding(): p = _service() - extracted = {"facts": [_fact("f1", "User likes tea")], "episodic": [], "updates": []} + p._vector_distance_function = AsyncMock(return_value="euclidean") + p._embed_batch.return_value = [[1.0, 0.0]] + p._nearest_active_full = AsyncMock() + p._apply_inplace_update = AsyncMock() + extracted = {"facts": [_fact("f-new", "near identical")], "episodic": [], "updates": []} out = await p.dedup_extracted_memories("u1", extracted) - assert out is extracted - p._embed_batch.assert_not_called() + assert [doc["id"] for doc in out["facts"]] == ["f-new"] + p._nearest_active_full.assert_not_awaited() + p._apply_inplace_update.assert_not_awaited() @pytest.mark.asyncio -async def test_euclidean_disables_near_exact_autodrop(): - # Async mirror: euclidean disables the cosine-calibrated near-exact auto-drop; - # the near-identical new memory is kept + tagged for LLM reconcile. +async def test_apply_inplace_update_recency_wins_and_unions(): p = _service() - p._vector_distance_function = AsyncMock(return_value="euclidean") - fact = _fact("f-new", "near identical") - fact.pop("embedding", None) - p._embed_batch.return_value = [[1.0, 0.0]] - p._vector_candidates = AsyncMock( - return_value=[{"id": "existing", "content": "same", "type": "fact", "score": 0.05}] + p._upsert_item = AsyncMock() + neighbor = _fact("existing-1", "old content", tags=["sys:fact", "topic:a"]) + neighbor["confidence"] = 0.6 + neighbor["salience"] = 0.5 + neighbor["updated_at"] = "2025-01-01T00:00:00+00:00" + neighbor["_etag"] = "etag-xyz" + new_doc = _fact( + "f-new", "new richer content", embedding=[0.5, 0.5], tags=["sys:fact", "topic:b", "sys:dup-candidate"] ) + new_doc["confidence"] = 0.9 + new_doc["salience"] = 0.8 - out = await p.dedup_extracted_memories("u1", {"facts": [fact], "episodic": [], "updates": []}) + ok = await p._apply_inplace_update(neighbor, new_doc) - assert [doc["id"] for doc in out["facts"]] == ["f-new"] - assert out["facts"][0]["tags"][-1] == "sys:dup-candidate" - assert out["updates"][-1]["vector_dedup_skipped"] == 0 - assert out["updates"][-1]["dup_candidates_tagged"] == 1 + assert ok is True + written = p._upsert_item.call_args.kwargs["body"] + assert written["id"] == "existing-1" + assert written["content"] == "new richer content" + assert written["embedding"] == [0.5, 0.5] + assert written["salience"] == 0.8 + assert written["confidence"] == 0.9 + assert "topic:a" in written["tags"] and "topic:b" in written["tags"] + assert "sys:dup-candidate" not in written["tags"] + assert "_etag" not in written @pytest.mark.asyncio -async def test_dedup_extracted_memories_passes_user_id_per_concurrent_call(): +async def test_nearest_active_full_returns_full_doc_and_skips_excluded(): p = _service() - seen: list[tuple[str, str]] = [] - - async def vector_candidates(**kwargs): - await asyncio.sleep(0) - exclude_ids = kwargs["exclude_ids"] - doc_id = next(iter(exclude_ids)) - seen.append((doc_id, kwargs["user_id"])) - return [] - - p._vector_candidates = AsyncMock(side_effect=vector_candidates) - user_a_doc = _fact("user-a-doc", "User A likes tea") - user_b_doc = _fact("user-b-doc", "User B likes coffee") - - await asyncio.gather( - p.dedup_extracted_memories("user-a", {"facts": [user_a_doc], "episodic": [], "updates": []}), - p.dedup_extracted_memories("user-b", {"facts": [user_b_doc], "episodic": [], "updates": []}), - ) + doc_a = _fact("a", "first") + doc_b = _fact("b", "second") - assert dict(seen) == {"user-a-doc": "user-a", "user-b-doc": "user-b"} + async def query_items(_container, *, query, parameters): + del query, parameters + return [{"doc": doc_a, "score": 0.99}, {"doc": doc_b, "score": 0.80}] + + p._query_items = AsyncMock(side_effect=query_items) + neighbor, score = await p._nearest_active_full( + user_id="u1", embedding=[1.0, 0.0], memory_type="fact", exclude_ids={"a"} + ) + assert neighbor["id"] == "b" + assert score == 0.80 @pytest.mark.asyncio -async def test_dedup_extracted_memories_vector_ladder_and_intra_batch(): +async def test_dedup_extracted_memories_flag_off_is_noop(monkeypatch): + monkeypatch.setattr("azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_vector_enabled", lambda: False) p = _service() - docs = [ - _fact("drop-existing", "User likes tea"), - _fact("tag-existing", "User likes coffee"), - _fact("clean", "User likes water"), - _fact("batch-keeper", "User likes green tea"), - _fact("drop-batch", "User likes green tea too"), - ] - for doc in docs: - doc.pop("embedding", None) - p._embed_batch.return_value = [ - [1.0, 0.0], - [0.0, 1.0], - [1.0, -1.0], - [0.7, 0.7], - [0.7, 0.7], - ] - p._vector_candidates = AsyncMock( - side_effect=[ - [{"id": "old-high", "content": "User likes tea.", "type": "fact", "score": 0.99}], - [{"id": "old-mid", "content": "User enjoys coffee.", "type": "fact", "score": 0.90}], - [{"id": "old-low", "content": "Unrelated", "type": "fact", "score": 0.20}], - [], - [], - ] - ) + extracted = {"facts": [_fact("f1", "content")], "episodic": [], "updates": []} - out = await p.dedup_extracted_memories("u1", {"facts": docs, "episodic": [], "updates": []}) + out = await p.dedup_extracted_memories("u1", extracted) - # Intra-batch new-vs-new dedup was removed: drop-batch (no existing match) - # is now kept; same-batch near-dups are deferred to reconcile. - ids = [doc["id"] for doc in out["facts"]] - assert ids == ["tag-existing", "clean", "batch-keeper", "drop-batch"] - assert all("embedding" in doc for doc in out["facts"]) - tagged = out["facts"][0] - assert "sys:dup-candidate" in tagged["tags"] - assert tagged["metadata"]["dup_of"] == "old-mid" - assert tagged["metadata"]["dup_score"] == 0.90 - assert "sys:dup-candidate" not in out["facts"][1]["tags"] - stats = out["updates"][-1] - assert stats["vector_dedup_skipped"] == 1 - assert stats["dup_candidates_tagged"] == 1 + assert out is extracted + p._embed_batch.assert_not_awaited() @pytest.mark.asyncio -async def test_dedup_skips_underspecified_doc_verbatim(): - # Parity with sync: a doc with no/unknown type is passed through untouched - # and never runs vector dedup (async previously defaulted type to the bucket). +async def test_dedup_extracted_memories_passes_user_id_per_concurrent_call(): + # Two concurrent dedup calls for different users must each query with their own + # user_id (no shared mutable state leaking one user's id into another's query). p = _service() - p._vector_candidates = AsyncMock(return_value=[{"id": "x", "content": "c", "type": "fact", "score": 0.99}]) - doc = _fact("f1", "content") - doc.pop("type") - doc.pop("embedding", None) - p._embed_batch.return_value = [[1.0, 0.0]] + p._vector_distance_function = AsyncMock(return_value="cosine") + seen_users: list[str] = [] + + async def nearest(*, user_id, embedding, memory_type, exclude_ids): + del embedding, memory_type, exclude_ids + seen_users.append(user_id) + return None, 0.0 + + p._nearest_active_full = AsyncMock(side_effect=nearest) - out = await p.dedup_extracted_memories("u1", {"facts": [doc], "episodic": [], "updates": []}) + async def run(uid): + p2 = _service() + p2._vector_distance_function = AsyncMock(return_value="cosine") + p2._nearest_active_full = AsyncMock(side_effect=nearest) + p2._embed_batch.return_value = [[1.0, 0.0]] + await p2.dedup_extracted_memories(uid, {"facts": [_fact("f", "c")], "episodic": [], "updates": []}) - assert [d["id"] for d in out["facts"]] == ["f1"] - assert "sys:dup-candidate" not in out["facts"][0].get("tags", []) - p._vector_candidates.assert_not_awaited() + await asyncio.gather(run("userA"), run("userB")) + assert set(seen_users) == {"userA", "userB"} @pytest.mark.asyncio -async def test_reconcile_memory_type_routes_episodic_merge_only(monkeypatch): - monkeypatch.setattr("azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_reconcile_mode", lambda: "full_pool") +async def test_reconcile_memory_type_routes_episodic_and_procedural_noop(): p = _service() - episodes = [_episode("ep_1", "CI failed then retries fixed it"), _episode("ep_2", "CI failed; retries fixed it")] - p._active_memories_for_reconcile = AsyncMock(return_value=episodes) - p._run_prompty = AsyncMock( - return_value=json.dumps( - { - "duplicate_groups": [ - { - "merged_content": "CI failed, then retries fixed it", - "source_ids": ["ep_1", "ep_2"], - "confidence": 0.9, - "salience": 0.8, - } - ], - "kept_ids": [], - "contradicted_pairs": [{"winner_id": "ep_1", "loser_id": "ep_2"}], - } - ) - ) + p._run_prompty = AsyncMock() - result = await p.reconcile_memories("u1", memory_type="episodic") + episodic_result = await p.reconcile_memories("u1", memory_type="episodic") + assert episodic_result == {"kept": 0, "merged": 0, "contradicted": 0} - assert result == {"kept": 0, "merged": 2, "contradicted": 0} - assert p._run_prompty.await_args.kwargs["inputs"].keys() == {"episodics_text"} - assert p._run_prompty.await_args.args[0] == "dedup_episodic.prompty" + procedural_result = await p.reconcile_memories("u1", memory_type="procedural") + assert procedural_result == {"kept": 0, "merged": 0, "contradicted": 0} + + p._run_prompty.assert_not_awaited() @pytest.mark.asyncio -async def test_candidate_reconcile_builds_connected_component(): +async def test_reconcile_fact_contradiction_only(): p = _service() - docs = { - "f1": _fact("f1", "User likes aisle seats", tags=["sys:fact", "sys:dup-candidate"]), - "f2": _fact("f2", "User prefers aisle seats"), - "f3": _fact("f3", "User enjoys aisle seats"), - } - - async def query_items(_container, **kwargs): - params = {p["name"]: p["value"] for p in kwargs.get("parameters", [])} - if params.get("@tag") == "sys:dup-candidate": - return [docs["f1"]] - ids = [value for name, value in params.items() if name.startswith("@id")] - return [docs[mid] for mid in ids if mid in docs] - - p._query_items = AsyncMock(side_effect=query_items) - p._vector_candidates = AsyncMock( - side_effect=[ - [ - {"id": "f2", "content": docs["f2"]["content"], "type": "fact", "score": 0.90}, - {"id": "f3", "content": docs["f3"]["content"], "type": "fact", "score": 0.88}, - ], - [ - {"id": "f2", "content": docs["f2"]["content"], "type": "fact", "score": 0.90}, - {"id": "f3", "content": docs["f3"]["content"], "type": "fact", "score": 0.88}, - ], - [ - {"id": "f1", "content": docs["f1"]["content"], "type": "fact", "score": 0.90}, - {"id": "f3", "content": docs["f3"]["content"], "type": "fact", "score": 0.89}, - ], - [ - {"id": "f1", "content": docs["f1"]["content"], "type": "fact", "score": 0.88}, - {"id": "f2", "content": docs["f2"]["content"], "type": "fact", "score": 0.89}, - ], - ] - ) + facts = [ + _fact("f1", "User's deadline is March 1"), + _fact("f2", "User's deadline is March 15"), + ] + facts[0]["created_at"] = "2024-01-01T00:00:00+00:00" + facts[1]["created_at"] = "2024-02-01T00:00:00+00:00" + p._active_memories_for_reconcile = AsyncMock(return_value=facts) p._run_prompty = AsyncMock( return_value=json.dumps( { - "duplicate_groups": [ - { - "merged_content": "User prefers aisle seats", - "source_ids": ["f1", "f2", "f3"], - "confidence": 0.9, - "salience": 0.8, - } - ], - "contradicted_pairs": [], - "kept_ids": [], + "duplicate_groups": [{"merged_content": "ignored", "source_ids": ["f1", "f2"]}], + "contradicted_pairs": [{"winner_id": "f2", "loser_id": "f1", "reason": "more recent"}], + "kept_ids": ["f2"], } ) ) - result = await p.reconcile_memories("u1") + result = await p.reconcile_memories("u1", memory_type="fact") - assert result == { - "kept": 0, - "merged": 3, - "contradicted": 0, - "reconcile_clusters_sent": 1, - "reconcile_llm_calls_saved": 2, - } - p._run_prompty.assert_awaited_once() - facts_text = p._run_prompty.await_args.kwargs["inputs"]["facts_text"] - assert all(fid in facts_text for fid in ["f1", "f2", "f3"]) + assert result == {"kept": 1, "merged": 0, "contradicted": 1} + p._upsert_memory.assert_not_awaited() + assert p._mark_superseded.await_count == 1 + assert p._mark_superseded.call_args.args[0]["id"] == "f1" + assert p._mark_superseded.call_args.args[1] == "f2" + assert p._mark_superseded.call_args.kwargs["reason"] == "contradict" + assert p._run_prompty.call_args.args[0] == "dedup.prompty" diff --git a/tests/unit/aio/store/test_get_memory_history.py b/tests/unit/aio/store/test_get_memory_history.py new file mode 100644 index 0000000..fe4527a --- /dev/null +++ b/tests/unit/aio/store/test_get_memory_history.py @@ -0,0 +1,69 @@ +"""Async parity tests for get_memory_history.""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest + +from azure.cosmos.agent_memory._container_routing import ContainerKey +from azure.cosmos.agent_memory.aio.store import AsyncMemoryStore +from azure.cosmos.agent_memory.exceptions import ValidationError + + +class AsyncIterator: + def __init__(self, items): + self._items = iter(items) + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self._items) + except StopIteration: + raise StopAsyncIteration + + +def _containers(*, memories=None): + return { + ContainerKey.TURNS: MagicMock(), + ContainerKey.MEMORIES: memories if memories is not None else MagicMock(), + ContainerKey.SUMMARIES: MagicMock(), + } + + +def _fact(fact_id, *, superseded_at=None, reason=None): + return { + "id": fact_id, + "user_id": "u1", + "thread_id": "t1", + "type": "fact", + "content": fact_id, + "metadata": {"category": "biographical"}, + "created_at": "2025-01-01T00:00:00+00:00", + "superseded_at": superseded_at, + "supersede_reason": reason, + } + + +async def test_async_multi_level_chain_newest_first(): + memories = MagicMock() + memories.query_items.side_effect = [ + AsyncIterator([_fact("v2", superseded_at="2025-03-01T00:00:00+00:00", reason="contradict")]), + AsyncIterator([_fact("v1", superseded_at="2025-02-01T00:00:00+00:00", reason="update")]), + AsyncIterator([]), + ] + store = AsyncMemoryStore(containers=_containers(memories=memories)) + + history = await store.get_memory_history("current", user_id="u1", thread_id="t1") + + assert [d["id"] for d in history] == ["v2", "v1"] + + +async def test_async_missing_ids_raise(): + store = AsyncMemoryStore(containers=_containers()) + with pytest.raises(ValidationError): + await store.get_memory_history("", user_id="u1") + with pytest.raises(ValidationError): + await store.get_memory_history("current", user_id="") diff --git a/tests/unit/aio/test_auto_trigger.py b/tests/unit/aio/test_auto_trigger.py index 0c68e28..bc218c9 100644 --- a/tests/unit/aio/test_auto_trigger.py +++ b/tests/unit/aio/test_auto_trigger.py @@ -13,7 +13,6 @@ import pytest from azure.cosmos.exceptions import CosmosResourceNotFoundError -from azure.cosmos.agent_memory import _counters from azure.cosmos.agent_memory.aio.cosmos_memory_client import AsyncCosmosMemoryClient from azure.cosmos.agent_memory.aio.processors import AsyncInProcessProcessor @@ -101,72 +100,17 @@ async def fake_upsert(body): class TestAsyncExtractRecentK: @pytest.mark.asyncio - @pytest.mark.parametrize( - ("n_facts", "batch_count", "counter_result", "expected_recent_k"), - [ - (1, 1, (0, 1), 1), - (1, 3, (0, 3), 3), - (5, 1, (4, 5), 5), - ], - ) - async def test_extract_recent_k_uses_max_threshold_and_batch_count( - self, - monkeypatch, - n_facts, - batch_count, - counter_result, - expected_recent_k, - ): - monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", str(n_facts)) - monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") - monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") - - processor = AsyncInProcessProcessor(pipeline=MagicMock()) - processor.process_extract_memories = MagicMock(return_value={}) - - client = AsyncCosmosMemoryClient(use_default_credential=False, processor=processor) - - async def fake_upsert(body): - return body - - client._memories_container_client = MagicMock() - client._memories_container_client.upsert_item = MagicMock(side_effect=fake_upsert) - client._turns_container_client = client._memories_container_client - client._summaries_container_client = client._memories_container_client - client._counter_container_client = MagicMock() - - with patch( - "azure.cosmos.agent_memory._counters.increment_counter_async", - return_value=counter_result, - ): - for i in range(batch_count): - client.add_local(user_id="u1", role="user", thread_id="t1", content=f"hi {i}") - await client.push_to_cosmos() - await asyncio.gather(*list(client._background_tasks), return_exceptions=True) - - processor.process_extract_memories.assert_called_once_with( - user_id="u1", - thread_id="t1", - recent_k=expected_recent_k, - ) - @pytest.mark.asyncio - @pytest.mark.parametrize( - ("counter_result", "watermark", "expected_recent_k"), - [ - ((5, 10), 5, 5), # backlog = new - watermark - ((0, 1), 1, 1), # new == watermark -> floored to 1 - ((98, 100), 0, 100), # large backlog is NOT capped - ((20, 30), None, 30), # BOOTSTRAP: no watermark -> base=0 -> recent_k = new_count - ], - ) - async def test_extract_recent_k_uses_watermark(self, monkeypatch, counter_result, watermark, expected_recent_k): + async def test_extract_fires_without_recent_k_or_watermark(self, monkeypatch): + """Async: extraction covers all un-extracted turns (extracted_at gated) and + batches internally, so it fires with NO recent_k and NO success watermark.""" monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") processor = AsyncInProcessProcessor(pipeline=MagicMock()) processor.process_extract_memories = MagicMock(return_value={}) + client = AsyncCosmosMemoryClient(use_default_credential=False, processor=processor) client._memories_container_client = MagicMock() client._memories_container_client.upsert_item = AsyncMock(side_effect=lambda body: body) @@ -174,37 +118,26 @@ async def test_extract_recent_k_uses_watermark(self, monkeypatch, counter_result client._summaries_container_client = client._memories_container_client client._counter_container_client = MagicMock() - with ( - patch( - "azure.cosmos.agent_memory._counters.increment_counter_async", - return_value=counter_result, - ), - patch( - "azure.cosmos.agent_memory._counters.read_extract_watermark_async", - new=AsyncMock(return_value=watermark), - ), - patch( - "azure.cosmos.agent_memory._counters.advance_extract_watermark_async", - new=AsyncMock(), - ) as advance, + with patch( + "azure.cosmos.agent_memory._counters.increment_counter_async", + return_value=(0, 1), ): client.add_local(user_id="u1", role="user", thread_id="t1", content="hi") await client.push_to_cosmos() await asyncio.gather(*list(client._background_tasks), return_exceptions=True) - processor.process_extract_memories.assert_called_once_with( - user_id="u1", thread_id="t1", recent_k=expected_recent_k - ) - advance.assert_awaited_once() + processor.process_extract_memories.assert_called_once_with(user_id="u1", thread_id="t1") @pytest.mark.asyncio - async def test_watermark_not_advanced_when_extract_fails(self, monkeypatch): + async def test_extract_failure_stamps_failure(self, monkeypatch): + """A total async extract failure is recorded via stamp_failure_async.""" monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") processor = AsyncInProcessProcessor(pipeline=MagicMock()) processor.process_extract_memories = MagicMock(side_effect=RuntimeError("llm down")) + client = AsyncCosmosMemoryClient(use_default_credential=False, processor=processor) client._memories_container_client = MagicMock() client._memories_container_client.upsert_item = AsyncMock(side_effect=lambda body: body) @@ -217,14 +150,6 @@ async def test_watermark_not_advanced_when_extract_fails(self, monkeypatch): "azure.cosmos.agent_memory._counters.increment_counter_async", return_value=(0, 1), ), - patch( - "azure.cosmos.agent_memory._counters.read_extract_watermark_async", - new=AsyncMock(return_value=None), - ), - patch( - "azure.cosmos.agent_memory._counters.advance_extract_watermark_async", - new=AsyncMock(), - ) as advance, patch( "azure.cosmos.agent_memory._counters.stamp_failure_async", new=AsyncMock(), @@ -234,96 +159,7 @@ async def test_watermark_not_advanced_when_extract_fails(self, monkeypatch): await client.push_to_cosmos() await asyncio.gather(*list(client._background_tasks), return_exceptions=True) - advance.assert_not_awaited() - stamp.assert_awaited_once() - - @pytest.mark.asyncio - async def test_watermark_round_trip_fail_then_succeed_no_strand(self, monkeypatch): - """Async stateful round-trip against a REAL in-memory counter: first - extract fails, second succeeds and must cover EVERY turn so far (20), - not just its own batch (10) — the bootstrap strand regression.""" - monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") - monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") - monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") - - counter = _AsyncFakeCounterContainer() - recorded: list[int] = [] - - def extract(*, user_id, thread_id, recent_k): - recorded.append(recent_k) - if len(recorded) == 1: - raise RuntimeError("transient LLM outage") - return {} - - processor = AsyncInProcessProcessor(pipeline=MagicMock()) - processor.process_extract_memories = MagicMock(side_effect=extract) - client = AsyncCosmosMemoryClient(use_default_credential=False, processor=processor) - client._memories_container_client = MagicMock() - client._memories_container_client.upsert_item = AsyncMock(side_effect=lambda body: body) - client._turns_container_client = client._memories_container_client - client._summaries_container_client = client._memories_container_client - client._counter_container_client = counter - - for i in range(10): - client.add_local(user_id="u1", role="user", thread_id="t1", content=f"a{i}") - await client.push_to_cosmos() - await asyncio.gather(*list(client._background_tasks), return_exceptions=True) - - for i in range(10): - client.add_local(user_id="u1", role="user", thread_id="t1", content=f"b{i}") - await client.push_to_cosmos() - await asyncio.gather(*list(client._background_tasks), return_exceptions=True) - - assert recorded == [10, 20] - cid = _counters.thread_counter_id("u1", "t1") - assert await _counters.read_extract_watermark_async(counter, cid, "u1", "t1") == 20 - - @pytest.mark.asyncio - async def test_reconcile_full_rebuild_on_persisted_counter_cadence(self, monkeypatch): - """Async symmetry: in-process auto-trigger requests full_rebuild on the - persisted-counter cadence (every 2 turns here), like the durable backend.""" - monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") - monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") - monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") - monkeypatch.setenv("DEDUP_EVERY_N", "1") - monkeypatch.setattr("azure.cosmos.agent_memory.thresholds.get_dedup_full_recluster_every_n", lambda: 2) - - rebuilds: list[bool] = [] - processor = AsyncInProcessProcessor(pipeline=MagicMock()) - processor.process_extract_memories = MagicMock(return_value={}) - processor.synthesize_procedural = MagicMock(return_value=None) - processor.process_reconcile = MagicMock( - side_effect=lambda *, user_id, full_rebuild=False: rebuilds.append(full_rebuild) - ) - client = AsyncCosmosMemoryClient(use_default_credential=False, processor=processor) - client._memories_container_client = MagicMock() - client._memories_container_client.upsert_item = AsyncMock(side_effect=lambda body: body) - client._turns_container_client = client._memories_container_client - client._summaries_container_client = client._memories_container_client - client._counter_container_client = MagicMock() - - with ( - patch( - "azure.cosmos.agent_memory._counters.increment_counter_async", - new=AsyncMock(side_effect=[(0, 1), (1, 2)]), - ), - patch( - "azure.cosmos.agent_memory._counters.read_extract_watermark_async", - new=AsyncMock(return_value=None), - ), - patch( - "azure.cosmos.agent_memory._counters.advance_extract_watermark_async", - new=AsyncMock(), - ), - ): - client.add_local(user_id="u1", role="user", thread_id="t1", content="a") - await client.push_to_cosmos() - await asyncio.gather(*list(client._background_tasks), return_exceptions=True) - client.add_local(user_id="u1", role="user", thread_id="t1", content="b") - await client.push_to_cosmos() - await asyncio.gather(*list(client._background_tasks), return_exceptions=True) - - assert rebuilds == [False, True] + stamp.assert_called_once() class TestPushToCosmosUnflushedDelta: diff --git a/tests/unit/aio/test_cosmos_memory_client.py b/tests/unit/aio/test_cosmos_memory_client.py index 4e728d7..e3dfb3e 100644 --- a/tests/unit/aio/test_cosmos_memory_client.py +++ b/tests/unit/aio/test_cosmos_memory_client.py @@ -218,7 +218,7 @@ def test_add_local_turn_requires_thread_id(self): mem = _make_client() with pytest.raises(ValidationError, match="thread_id is required"): mem.add_local(user_id="u1", role="user", content="hi") - # Validation must run BEFORE append — otherwise an orphan turn + # Validation must run BEFORE append - otherwise an orphan turn # with thread_id=None would persist and pollute pk on push. assert mem.local_memory == [] assert mem._unflushed_turn_counts == {} diff --git a/tests/unit/aio/test_embeddings.py b/tests/unit/aio/test_embeddings.py index 165e6f6..75e8e0f 100644 --- a/tests/unit/aio/test_embeddings.py +++ b/tests/unit/aio/test_embeddings.py @@ -35,7 +35,7 @@ def client(): # =================================================================== -# generate() — success +# generate() - success # =================================================================== @@ -51,7 +51,7 @@ async def test_generate_success(client): # =================================================================== -# generate() — lazy init +# generate() - lazy init # =================================================================== @@ -77,7 +77,7 @@ async def test_generate_lazy_init(): # =================================================================== -# generate() — api_key vs credential auth +# generate() - api_key vs credential auth # =================================================================== @@ -124,7 +124,7 @@ async def test_generate_credential_auth(): # =================================================================== -# generate() — missing endpoint +# generate() - missing endpoint # =================================================================== @@ -145,7 +145,7 @@ async def test_generate_missing_credential_and_api_key(): # =================================================================== -# generate() — API failure +# generate() - API failure # =================================================================== @@ -200,7 +200,7 @@ async def test_generate_batch_api_failure(client): # =================================================================== -# generate_batch() — N=16 chunk guard +# generate_batch() - N=16 chunk guard # =================================================================== diff --git a/tests/unit/aio/test_procedural_synthesis.py b/tests/unit/aio/test_procedural_synthesis.py index af2307c..513208f 100644 --- a/tests/unit/aio/test_procedural_synthesis.py +++ b/tests/unit/aio/test_procedural_synthesis.py @@ -3,7 +3,7 @@ The procedural-synthesis business logic is covered exhaustively by sync tests in ``tests/unit/test_procedural_synthesis.py`` against ``PipelineService``; ``AsyncPipelineService`` is a 1:1 async mirror. -These tests verify async wiring — that the client awaits the pipeline +These tests verify async wiring - that the client awaits the pipeline correctly, that the durable-processor branch short-circuits, and that the store-backed procedural reads work over async iterators. """ diff --git a/tests/unit/aio/test_process_now.py b/tests/unit/aio/test_process_now.py index d0aa22e..83d8be5 100644 --- a/tests/unit/aio/test_process_now.py +++ b/tests/unit/aio/test_process_now.py @@ -53,8 +53,8 @@ async def test_process_now_with_inprocess_invokes_full_pipeline(): assert pipeline.reconcile_memories.await_count == 2 pipeline.reconcile_memories.assert_has_awaits( [ - call("u", n=50, memory_type="fact", full_rebuild=False), - call("u", n=50, memory_type="episodic", full_rebuild=False), + call("u", n=50, memory_type="fact"), + call("u", n=50, memory_type="episodic"), ] ) pipeline.synthesize_procedural.assert_awaited_once_with("u", force=False) @@ -176,7 +176,7 @@ async def test_process_now_propagates_permanent_user_summary_failure(): @pytest.mark.asyncio async def test_process_now_with_durable_skips_tail_steps(): - """Durable mode must NOT call synthesize_procedural or generate_user_summary — + """Durable mode must NOT call synthesize_procedural or generate_user_summary - those are driven by the change-feed-fed sibling Function app.""" client = _connected(processor=AsyncDurableFunctionProcessor()) pipeline = AsyncMock() diff --git a/tests/unit/aio/test_reconcile_telemetry.py b/tests/unit/aio/test_reconcile_telemetry.py index cea2326..ef2fdbd 100644 --- a/tests/unit/aio/test_reconcile_telemetry.py +++ b/tests/unit/aio/test_reconcile_telemetry.py @@ -20,8 +20,8 @@ @pytest.fixture(autouse=True) def _pin_async_legacy_reconcile(monkeypatch): monkeypatch.setattr( - "azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_reconcile_mode", - lambda: "full_pool", + "azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_vector_enabled", + lambda: False, ) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 7143ebe..6123d73 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -1,7 +1,7 @@ """Unit-test conftest: install lightweight stubs for Azure Functions packages. ``azure-functions`` and ``azure-durable-functions`` are Azure Functions -runtime packages that are not installed in the library's dev dependencies — +runtime packages that are not installed in the library's dev dependencies - they're only needed when *running* the function app, not when testing the helper logic inside it. This conftest inserts minimal stubs into ``sys.modules`` before any test module is imported so that @@ -15,7 +15,7 @@ def _install_azure_functions_stubs() -> None: # Inject stub sub-modules directly into sys.modules. Do NOT create or - # overwrite the top-level ``azure`` entry — that is a real namespace + # overwrite the top-level ``azure`` entry - that is a real namespace # package provided by azure-cosmos (a core dependency), and replacing it # with a plain module breaks its sub-package imports. diff --git a/tests/unit/function_app/conftest.py b/tests/unit/function_app/conftest.py index fb62332..e12faa5 100644 --- a/tests/unit/function_app/conftest.py +++ b/tests/unit/function_app/conftest.py @@ -1,7 +1,7 @@ """Path bootstrap for function-app unit tests. The function-app source lives in ``function_app/`` (a sibling of the SDK -``azure/cosmos/agent_memory/``) and is *not* a package — Azure Functions discovers +``azure/cosmos/agent_memory/``) and is *not* a package - Azure Functions discovers modules by file name. We add the directory to ``sys.path`` here so tests can ``import shared.counters``, ``import triggers.change_feed`` etc. without each test file having to repeat the ``sys.path`` dance. diff --git a/tests/unit/function_app/test_change_feed.py b/tests/unit/function_app/test_change_feed.py index a70bcb3..18b6dbb 100644 --- a/tests/unit/function_app/test_change_feed.py +++ b/tests/unit/function_app/test_change_feed.py @@ -416,7 +416,7 @@ def test_per_thread_grouping_is_correct(): _turn(user_id="u1", thread_id="t1"), _turn(user_id="u1", thread_id="t1"), # t1 crosses 4 _turn(user_id="u1", thread_id="t2"), - _turn(user_id="u1", thread_id="t2"), # t2 only at 2 — no cross + _turn(user_id="u1", thread_id="t2"), # t2 only at 2 - no cross ] asyncio.run(process_changefeed_batch(docs, starter, counter_container=container)) @@ -425,7 +425,7 @@ def test_per_thread_grouping_is_correct(): # t1 crossed → summary + extract started for t1 only. assert ("ThreadSummaryOrchestrator", "thread_summary:u1:t1:4") in started assert ("ExtractMemoriesOrchestrator", "extract:u1:t1:4") in started - # t2 below threshold — no orchestrators for t2. + # t2 below threshold - no orchestrators for t2. assert not any("t2" in iid for _, iid in started) @@ -485,16 +485,16 @@ def test_lsn_replay_does_not_double_increment(): # Layer 2 (deterministic instance ID): on replay the counter helper still # reports the same ``(old, new)`` it returned the first time, so the - # threshold is "crossed" again — but with the IDENTICAL deterministic + # threshold is "crossed" again - but with the IDENTICAL deterministic # instance id. Azure Durable Functions then dedups the duplicate # ``start_new`` server-side. We assert the determinism here. summary_starts = [c for c in starter.start_new.await_args_list if c.args[0] == "ThreadSummaryOrchestrator"] - assert len(summary_starts) == 2 # same id sent twice — durable dedups + assert len(summary_starts) == 2 # same id sent twice - durable dedups assert all(c.kwargs["instance_id"] == "thread_summary:u1:t1:4" for c in summary_starts) # --------------------------------------------------------------------------- -# MEMORY_PROCESSOR_OWNER exclusivity — the change-feed trigger must +# MEMORY_PROCESSOR_OWNER exclusivity - the change-feed trigger must # respect the owner env var the same way the SDK auto-trigger does, so a # shared Cosmos container is processed by exactly one backend. # --------------------------------------------------------------------------- @@ -553,7 +553,7 @@ def test_runs_normally_when_owner_durable(): # Counter was written. assert container._state["thread:u1:t1"]["count"] == 2 - # Threshold (2) crossed — orchestrator started. + # Threshold (2) crossed - orchestrator started. summary_starts = [c for c in starter.start_new.await_args_list if c.args[0] == "ThreadSummaryOrchestrator"] assert len(summary_starts) == 1 @@ -572,7 +572,7 @@ def test_skips_when_owner_unset(monkeypatch): This is the protection against the day-one footgun where a customer deploys the Function App next to an existing SDK install without - configuring the env var — without this guard both backends would + configuring the env var - without this guard both backends would race on the same writes. """ monkeypatch.delenv("MEMORY_PROCESSOR_OWNER", raising=False) @@ -616,8 +616,8 @@ def _extract_payload(call): def test_reconcile_flag_set_only_when_n_facts_times_n_dedup_threshold_crosses(): """Reconcile threshold = FACT_EXTRACTION_EVERY_N * DEDUP_EVERY_N (here 1 * 5 = 5). The change-feed signals reconcile via the - ``reconcile`` flag on the orchestrator payload — never as a separate - dispatch — so DEDUP_EVERY_N is honored on the FA path.""" + ``reconcile`` flag on the orchestrator payload - never as a separate + dispatch - so DEDUP_EVERY_N is honored on the FA path.""" starter = _make_starter() container = _make_counter_container_starting_at() @@ -648,46 +648,6 @@ def test_reconcile_flag_set_only_when_n_facts_times_n_dedup_threshold_crosses(): assert payload.get("recent_k") == 1 -@patch.dict( - os.environ, - { - "THREAD_SUMMARY_EVERY_N": "0", - "FACT_EXTRACTION_EVERY_N": "1", - "USER_SUMMARY_EVERY_N": "0", - "DEDUP_EVERY_N": "1", - }, - clear=False, -) -def test_full_rebuild_flag_set_on_persisted_counter_cadence(): - """The full-pool backstop is driven by the PERSISTED counter (durable-safe), - not the in-memory per-worker sweep counter: full_rebuild=True every - (n_facts * n_dedup * DEDUP_FULL_RECLUSTER_EVERY_N) turns. Here that's - 1 * 1 * 2 = every 2 turns.""" - with patch( - "azure.cosmos.agent_memory.thresholds.get_dedup_full_recluster_every_n", - return_value=2, - ): - starter = _make_starter() - container = _make_counter_container_starting_at() - - # Turn 1: counter 0->1. Reconcile crosses (n=1), full does NOT (n=2). - asyncio.run(process_changefeed_batch([_turn()], starter, counter_container=container)) - p1 = _extract_payload( - next(c for c in starter.start_new.await_args_list if c.args[0] == "ExtractMemoriesOrchestrator") - ) - assert p1.get("reconcile") is True - assert p1.get("full_rebuild") is False - - # Turn 2: counter 1->2. Full backstop threshold (2) crossed. - starter.start_new.reset_mock() - asyncio.run(process_changefeed_batch([_turn()], starter, counter_container=container)) - p2 = _extract_payload( - next(c for c in starter.start_new.await_args_list if c.args[0] == "ExtractMemoriesOrchestrator") - ) - assert p2.get("reconcile") is True - assert p2.get("full_rebuild") is True - - @patch.dict( os.environ, { diff --git a/tests/unit/function_app/test_counters.py b/tests/unit/function_app/test_counters.py index 6e48fbd..1efc96e 100644 --- a/tests/unit/function_app/test_counters.py +++ b/tests/unit/function_app/test_counters.py @@ -5,7 +5,7 @@ * ``increment_counter_by`` first-write path, happy path, ETag retry path, ETag exhaustion (skip), and LSN-based replay detection. -The Cosmos container is mocked — no Azure dependency is required at test time. +The Cosmos container is mocked - no Azure dependency is required at test time. """ from __future__ import annotations @@ -80,7 +80,7 @@ def test_counter_id_helpers(): # --------------------------------------------------------------------------- -# increment_counter_by — fixtures +# increment_counter_by - fixtures # --------------------------------------------------------------------------- @@ -100,7 +100,7 @@ def _http_error(status_code: int) -> CosmosHttpResponseError: # --------------------------------------------------------------------------- -# increment_counter_by — first-write path +# increment_counter_by - first-write path # --------------------------------------------------------------------------- @@ -138,7 +138,7 @@ def test_create_conflict_then_succeeds(self): # --------------------------------------------------------------------------- -# increment_counter_by — happy / etag-retry / replay paths +# increment_counter_by - happy / etag-retry / replay paths # --------------------------------------------------------------------------- @@ -180,7 +180,7 @@ def test_etag_conflict_retries_then_succeeds(self): old, new = asyncio.run(increment_counter_by(container, "thread:u1:t1", "u1", "t1", 2)) - # We restart from the *latest* read — count went 7 → 9. + # We restart from the *latest* read - count went 7 → 9. assert (old, new) == (7, 9) assert container.read_item.await_count == 2 assert container.upsert_item.await_count == 2 @@ -190,7 +190,7 @@ def test_etag_conflict_exhausted_raises(self): batch retries (at-least-once redelivery + LSN replay protection make this safe). Silently returning ``(old, old)`` would advance the lease without ever firing the orchestrator the increment was supposed to - trigger — a permanent threshold-miss bug.""" + trigger - a permanent threshold-miss bug.""" from azure.cosmos.exceptions import CosmosHttpResponseError container = _make_container() @@ -263,7 +263,7 @@ def test_different_lsn_does_not_trigger_replay(self): def test_out_of_order_replay_is_noop(self): """When the redelivered batch's LSN is *less than* the stored LSN (lease re-balance / host crash redelivering an old batch after a - newer one already landed), the increment is a no-op — return + newer one already landed), the increment is a no-op - return (current, current) so threshold-crossing logic doesn't fire a spurious extract/dedup. We use ``>=`` not ``==`` for replay detection.""" diff --git a/tests/unit/function_app/test_orchestrators.py b/tests/unit/function_app/test_orchestrators.py index 808e9e1..634514b 100644 --- a/tests/unit/function_app/test_orchestrators.py +++ b/tests/unit/function_app/test_orchestrators.py @@ -235,7 +235,7 @@ def test_chains_reconcile_when_flag_true(self, _retry): names = [c[0] for c in ctx._yielded_calls] assert names == ["em_Extract", "em_Dedup", "em_Persist", "em_ReconcileMemories"] - assert ctx._yielded_calls[3][2] == {"user_id": "u1", "full_rebuild": False} + assert ctx._yielded_calls[3][2] == {"user_id": "u1"} assert [s[0] for s in ctx._yielded_sub_orchestrators] == [ "SynthesizeProceduralOrchestrator", ] @@ -267,7 +267,7 @@ def boom_after_sub(name, retry, sub_payload, *args, **kwargs): gen.send({"facts": [{"id": "f1"}], "episodic": [], "updates": []}) # Yield 4: em_ReconcileMemories gen.send({"fact_count": 2, "episodic_count": 0, "updated_count": 0}) - # Yield 5: SynthesizeProceduralOrchestrator — throw an exception + # Yield 5: SynthesizeProceduralOrchestrator - throw an exception gen.send( { "fact": {"kept": 0, "merged": 1, "contradicted": 0}, @@ -428,30 +428,15 @@ def test_em_reconcile_memories_reconciles_fact_and_episodic(self): ): result = em_mod.em_ReconcileMemories({"user_id": "u1"}) - # full_rebuild defaults False when the change-feed didn't request a backstop. assert pipeline.reconcile_memories.call_args_list == [ - call(user_id="u1", n=17, memory_type="fact", full_rebuild=False), - call(user_id="u1", n=17, memory_type="episodic", full_rebuild=False), + call(user_id="u1", n=17, memory_type="fact"), + call(user_id="u1", n=17, memory_type="episodic"), ] assert result == { "fact": {"kept": 2, "merged": 1, "contradicted": 0}, "episodic": {"kept": 1, "merged": 0, "contradicted": 0}, } - def test_em_reconcile_memories_forwards_full_rebuild(self): - pipeline = MagicMock() - pipeline.reconcile_memories.side_effect = [{"kept": 0}, {"kept": 0}] - with ( - patch.object(em_mod, "get_pipeline", return_value=pipeline), - patch("azure.cosmos.agent_memory.thresholds.get_dedup_pool_size", return_value=17), - ): - em_mod.em_ReconcileMemories({"user_id": "u1", "full_rebuild": True}) - - assert pipeline.reconcile_memories.call_args_list == [ - call(user_id="u1", n=17, memory_type="fact", full_rebuild=True), - call(user_id="u1", n=17, memory_type="episodic", full_rebuild=True), - ] - def test_em_advance_extract_watermark_stamps_counter(self): import asyncio diff --git a/tests/unit/function_app/test_retry.py b/tests/unit/function_app/test_retry.py index 3100b55..1fb6ecc 100644 --- a/tests/unit/function_app/test_retry.py +++ b/tests/unit/function_app/test_retry.py @@ -49,7 +49,7 @@ def test_default_retry_options_uses_two_second_first_interval(): def test_default_retry_options_only_passes_supported_kwargs(): - """The real SDK constructor accepts only two kwargs — pinning that here + """The real SDK constructor accepts only two kwargs - pinning that here so a future change cannot silently re-introduce the original bug.""" with _patch_retry_ctor() as ctor: default_retry_options() @@ -61,7 +61,7 @@ def test_default_retry_options_only_passes_supported_kwargs(): def test_default_retry_options_is_called_freshly_each_time(): - """No caching — each call constructs a new RetryOptions.""" + """No caching - each call constructs a new RetryOptions.""" with _patch_retry_ctor(side_effect=lambda **kw: MagicMock()) as ctor: a = default_retry_options() b = default_retry_options() diff --git a/tests/unit/processors/test_durable.py b/tests/unit/processors/test_durable.py index 6625b57..cf65591 100644 --- a/tests/unit/processors/test_durable.py +++ b/tests/unit/processors/test_durable.py @@ -1,4 +1,4 @@ -"""Tests for DurableFunctionProcessor — verifies all methods are no-ops.""" +"""Tests for DurableFunctionProcessor - verifies all methods are no-ops.""" from __future__ import annotations diff --git a/tests/unit/processors/test_inprocess.py b/tests/unit/processors/test_inprocess.py index 28ddc2f..d6f60a0 100644 --- a/tests/unit/processors/test_inprocess.py +++ b/tests/unit/processors/test_inprocess.py @@ -1,4 +1,4 @@ -"""Tests for InProcessProcessor — verifies pipeline delegation order.""" +"""Tests for InProcessProcessor - verifies pipeline delegation order.""" from __future__ import annotations diff --git a/tests/unit/prompts/test_schema_prompty_conformance.py b/tests/unit/prompts/test_schema_prompty_conformance.py index 6d8093c..a66ad02 100644 --- a/tests/unit/prompts/test_schema_prompty_conformance.py +++ b/tests/unit/prompts/test_schema_prompty_conformance.py @@ -72,7 +72,7 @@ def _walk_required(data: object, schema: object, path: str) -> list[str]: Walks into object ``properties`` and array ``items`` (validating only the first item, since examples are illustrative not exhaustive). Returns a - flat list of human-readable error strings — empty list ⇒ conformant. + flat list of human-readable error strings - empty list ⇒ conformant. """ if not isinstance(schema, dict): return [] @@ -167,7 +167,7 @@ def test_prompty_nested_required_fields_present(filename: str, schema_entry: tup pytest.fail(f"\n{filename}: nested required keys missing in example JSON:\n " + "\n ".join(errors)) return - pytest.skip(f"{filename}: no top-level matching block — covered by shape test") + pytest.skip(f"{filename}: no top-level matching block - covered by shape test") def test_every_registered_prompty_file_exists() -> None: @@ -181,6 +181,6 @@ def test_every_prompty_file_has_a_registered_schema() -> None: prompty_files = {path.name for path in _PROMPTS_DIR.glob("*.prompty") if not path.name.startswith("_")} unregistered = prompty_files - set(PROMPTY_SCHEMAS) assert not unregistered, ( - f"These prompty files have no entry in PROMPTY_SCHEMAS — " + f"These prompty files have no entry in PROMPTY_SCHEMAS - " f"strict response_format will not be applied: {sorted(unregistered)}" ) diff --git a/tests/unit/services/test_chaos_extract_persist.py b/tests/unit/services/test_chaos_extract_persist.py index 82c61d0..33afa5e 100644 --- a/tests/unit/services/test_chaos_extract_persist.py +++ b/tests/unit/services/test_chaos_extract_persist.py @@ -17,18 +17,10 @@ def _pin_legacy_extract_dedup(monkeypatch): "azure.cosmos.agent_memory.thresholds.get_dedup_vector_enabled", lambda: False, ) - monkeypatch.setattr( - "azure.cosmos.agent_memory.thresholds.get_dedup_context_vector_enabled", - lambda: False, - ) monkeypatch.setattr( "azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_vector_enabled", lambda: False, ) - monkeypatch.setattr( - "azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_context_vector_enabled", - lambda: False, - ) class _FlakyContainer: diff --git a/tests/unit/services/test_dedup_vector.py b/tests/unit/services/test_dedup_vector.py index 6b4c4bd..80fe012 100644 --- a/tests/unit/services/test_dedup_vector.py +++ b/tests/unit/services/test_dedup_vector.py @@ -121,23 +121,24 @@ def query_items(*, query: str, parameters, **kwargs): assert [c["id"] for c in out] == ["far", "near"] -def test_dedup_extracted_vector_ladder_and_intra_batch() -> None: +def test_dedup_extracted_folds_near_dup_in_place_and_keeps_novel() -> None: + # Write-time in-place dedup: a new fact whose nearest active neighbor is at/above + # DEDUP_SIM_HIGH is folded into that neighbor in place (dropped from the ADD set); + # a fact with no close neighbor is novel and passes through to persist. p = _make_pipeline() - p._embed_batch.return_value = [[1.0, 0.0], [0.0, 1.0], [1.0, 0.0], [0.99, 0.01]] - p._vector_candidates = MagicMock( + p._vector_distance_function = MagicMock(return_value="cosine") + p._embed_batch.return_value = [[1.0, 0.0], [0.0, 1.0]] + p._nearest_active_full = MagicMock( side_effect=[ - [{"id": "existing-high", "content": "same", "type": "fact", "score": 0.99}], - [{"id": "existing-mid", "content": "near", "type": "fact", "score": 0.85}], - [{"id": "existing-low", "content": "far", "type": "fact", "score": 0.20}], - [{"id": "existing-low-2", "content": "far", "type": "fact", "score": 0.10}], + ({"id": "existing-1", "content": "same", "type": "fact"}, 0.99), + (None, 0.0), ] ) + p._apply_inplace_update = MagicMock(return_value=True) extracted = { "facts": [ - _doc("f-high", "drop against existing"), - _doc("f-mid", "tag against existing"), - _doc("f-clean", "keep clean"), - _doc("f-intra", "near-dup in batch is now kept (deferred to reconcile)"), + _doc("f-dup", "restatement of existing"), + _doc("f-novel", "brand new fact"), ], "episodic": [], "updates": [], @@ -145,173 +146,139 @@ def test_dedup_extracted_vector_ladder_and_intra_batch() -> None: out = p.dedup_extracted_memories("u1", extracted) - # Intra-batch new-vs-new dedup was removed: f-intra is compared only against - # persisted memories (Cosmos) -> novel here -> kept; reconcile catches any - # same-batch near-dups later. - assert [doc["id"] for doc in out["facts"]] == ["f-mid", "f-clean", "f-intra"] - assert out["facts"][0]["tags"][-1] == "sys:dup-candidate" - assert out["facts"][0]["metadata"]["dup_of"] == "existing-mid" - assert out["facts"][0]["metadata"]["dup_score"] == 0.85 - assert "sys:dup-candidate" not in out["facts"][1]["tags"] - assert all("embedding" in doc for doc in out["facts"]) - assert out["updates"][-1]["vector_dedup_skipped"] == 1 - assert out["updates"][-1]["dup_candidates_tagged"] == 1 - - -def test_candidate_mode_clears_tags_only_for_survivors() -> None: - # Latent-bug regression: a source consumed (superseded) by a merge must NOT be - # re-upserted by tag-clearing, which would resurrect it without superseded_by. - p = _make_pipeline() - f1 = _doc("f1", "a", tags=["sys:fact", "sys:dup-candidate"]) - f2 = _doc("f2", "b", tags=["sys:fact", "sys:dup-candidate"]) - f3 = _doc("f3", "c", tags=["sys:fact", "sys:dup-candidate"]) - p._build_candidate_clusters = MagicMock(return_value=([[f1, f2, f3]], 3, [f1, f2, f3])) - p._reconcile_pool = MagicMock(return_value=({"kept": 1, "merged": 2, "contradicted": 0}, {"f1", "f2"})) - cleared: list[str] = [] - p._clear_dup_candidate_tags = MagicMock(side_effect=lambda docs: cleared.extend(d["id"] for d in docs)) - p._emit_reconcile_outcome = MagicMock() - - result = p._reconcile_candidate_mode("u1", n=50, memory_type="fact", started_at=0.0) - - assert cleared == ["f3"] - assert result == { - "kept": 1, - "merged": 2, - "contradicted": 0, - "reconcile_clusters_sent": 1, - "reconcile_llm_calls_saved": 2, - } + # f-dup folded in place (removed from ADD set); f-novel kept. + assert [doc["id"] for doc in out["facts"]] == ["f-novel"] + p._apply_inplace_update.assert_called_once() + target, new_doc = p._apply_inplace_update.call_args.args + assert target["id"] == "existing-1" + assert new_doc["id"] == "f-dup" + assert out["updates"][-1]["inplace_updated"] == 1 -def test_candidate_mode_clears_tags_on_orphan_seeds() -> None: - # Seeds tagged sys:dup-candidate that never join a cluster (no near-duplicate) - # must have their stale tag cleared so future sweeps don't re-scan them forever. +def test_dedup_extracted_failed_inplace_update_keeps_new_doc() -> None: + # If the in-place upsert fails, the new doc must NOT be lost - it stays in the + # result so persist ADDs it as a novel record. p = _make_pipeline() - orphan = _doc("orphan", "lonely", tags=["sys:fact", "sys:dup-candidate"]) - c1 = _doc("c1", "a", tags=["sys:fact", "sys:dup-candidate"]) - c2 = _doc("c2", "b", tags=["sys:fact", "sys:dup-candidate"]) - p._build_candidate_clusters = MagicMock(return_value=([[c1, c2]], 3, [orphan, c1, c2])) - p._reconcile_pool = MagicMock(return_value=({"kept": 2, "merged": 0, "contradicted": 0}, set())) - cleared: list[str] = [] - p._clear_dup_candidate_tags = MagicMock(side_effect=lambda docs: cleared.extend(d["id"] for d in docs)) - p._emit_reconcile_outcome = MagicMock() + p._vector_distance_function = MagicMock(return_value="cosine") + p._embed_batch.return_value = [[1.0, 0.0]] + p._nearest_active_full = MagicMock(return_value=({"id": "existing-1", "content": "same", "type": "fact"}, 0.99)) + p._apply_inplace_update = MagicMock(return_value=False) + extracted = {"facts": [_doc("f-dup", "restatement")], "episodic": [], "updates": []} - p._reconcile_candidate_mode("u1", n=50, memory_type="fact", started_at=0.0) + out = p.dedup_extracted_memories("u1", extracted) - # Cluster survivors (c1, c2) plus the orphan all get cleared. - assert set(cleared) == {"c1", "c2", "orphan"} + assert [doc["id"] for doc in out["facts"]] == ["f-dup"] + assert all(op.get("op") != "stats" or "inplace_updated" not in op for op in out["updates"]) -def test_euclidean_disables_near_exact_autodrop() -> None: - # On euclidean distance the cosine-calibrated DEDUP_SIM_HIGH auto-drop is - # disabled: a near-identical existing memory must NOT silently drop the new - # one. It falls through to borderline tagging (LLM reconcile adjudicates). +def test_dedup_extracted_below_threshold_is_novel() -> None: + # A neighbor below DEDUP_SIM_HIGH is not a near-duplicate: the new fact is novel + # and no in-place update happens. p = _make_pipeline() - p._vector_distance_function = MagicMock(return_value="euclidean") + p._vector_distance_function = MagicMock(return_value="cosine") p._embed_batch.return_value = [[1.0, 0.0]] - # euclidean score = distance; 0.05 is "near-exact" (would drop under cosine rules). - p._vector_candidates = MagicMock( - return_value=[{"id": "existing", "content": "same", "type": "fact", "score": 0.05}] - ) - extracted = {"facts": [_doc("f-new", "near identical")], "episodic": [], "updates": []} + p._nearest_active_full = MagicMock(return_value=({"id": "existing-1", "content": "near", "type": "fact"}, 0.85)) + p._apply_inplace_update = MagicMock(return_value=True) + extracted = {"facts": [_doc("f-new", "somewhat similar")], "episodic": [], "updates": []} out = p.dedup_extracted_memories("u1", extracted) - # Not dropped — kept and tagged for LLM reconcile instead. assert [doc["id"] for doc in out["facts"]] == ["f-new"] - assert out["facts"][0]["tags"][-1] == "sys:dup-candidate" - assert out["updates"][-1]["vector_dedup_skipped"] == 0 - assert out["updates"][-1]["dup_candidates_tagged"] == 1 + p._apply_inplace_update.assert_not_called() -def test_candidate_mode_has_no_inmemory_backstop() -> None: - # The periodic full-pool backstop is no longer driven by an in-memory sweep - # counter (unreliable on the FA per-worker singleton). Candidate mode does - # ONLY clustering now; the full-pool pass is requested by the caller via - # full_rebuild on a persisted-counter cadence. - p = _make_pipeline() - assert not hasattr(p, "_next_reconcile_sweep") - p._build_candidate_clusters = MagicMock(return_value=([], 0, [])) - p._active_memories_for_reconcile = MagicMock() - p._emit_reconcile_outcome = MagicMock() - - # Many sweeps in a row never escalate to a full-pool pass on their own. - for _ in range(30): - p._reconcile_candidate_mode("u1", n=50, memory_type="fact", started_at=0.0) - p._build_candidate_clusters.assert_called() - p._active_memories_for_reconcile.assert_not_called() - - -def test_full_rebuild_bypasses_candidate_mode(monkeypatch) -> None: - # Public reconcile(full_rebuild=True) must take the full-pool single-LLM-pass - # path even under candidate mode, so it catches dissimilar contradictions. - monkeypatch.setattr("azure.cosmos.agent_memory.thresholds.get_dedup_reconcile_mode", lambda: "candidate") +def test_dedup_second_batch_dup_of_same_target_is_dropped_once() -> None: + # Two new facts that both fold into the SAME existing neighbor: only the first + # refreshes it; the second is dropped without re-writing the target. p = _make_pipeline() - pool = [_doc("f1", "User is vegetarian"), _doc("f2", "User loves steak")] - p._active_memories_for_reconcile = MagicMock(return_value=pool) - p._reconcile_pool = MagicMock(return_value=({"kept": 1, "merged": 0, "contradicted": 1}, set())) - p._reconcile_candidate_mode = MagicMock() - p._emit_reconcile_outcome = MagicMock() + p._vector_distance_function = MagicMock(return_value="cosine") + p._embed_batch.return_value = [[1.0, 0.0], [1.0, 0.0]] + p._nearest_active_full = MagicMock( + side_effect=[ + ({"id": "existing-1", "content": "same", "type": "fact"}, 0.99), + ({"id": "existing-1", "content": "same", "type": "fact"}, 0.98), + ] + ) + p._apply_inplace_update = MagicMock(return_value=True) + extracted = { + "facts": [_doc("f-a", "restate one"), _doc("f-b", "restate two")], + "episodic": [], + "updates": [], + } - result = p.reconcile_memories("u1", n=50, memory_type="fact", full_rebuild=True) + out = p.dedup_extracted_memories("u1", extracted) - p._reconcile_candidate_mode.assert_not_called() - p._reconcile_pool.assert_called_once_with("u1", "fact", pool) - assert result["contradicted"] == 1 + assert out["facts"] == [] + assert p._apply_inplace_update.call_count == 1 + assert out["updates"][-1]["inplace_updated"] == 1 -def test_full_rebuild_clears_survivor_tags(monkeypatch) -> None: - # full_rebuild full-pool path must clear sys:dup-candidate on survivors so it - # doesn't leave stale tags/metadata on user-visible memories. - monkeypatch.setattr("azure.cosmos.agent_memory.thresholds.get_dedup_reconcile_mode", lambda: "candidate") +def test_euclidean_disables_inplace_folding() -> None: + # On euclidean distance the cosine-calibrated DEDUP_SIM_HIGH is not comparable, + # so in-place folding is disabled and every extracted doc passes through as-is. p = _make_pipeline() - pool = [ - _doc("f1", "a", tags=["sys:fact", "sys:dup-candidate"]), - _doc("f2", "b", tags=["sys:fact", "sys:dup-candidate"]), - ] - p._active_memories_for_reconcile = MagicMock(return_value=pool) - # f1 consumed (superseded), f2 survives. - p._reconcile_pool = MagicMock(return_value=({"kept": 1, "merged": 1, "contradicted": 0}, {"f1"})) - cleared: list[str] = [] - p._clear_dup_candidate_tags = MagicMock(side_effect=lambda docs: cleared.extend(d["id"] for d in docs)) - p._emit_reconcile_outcome = MagicMock() + p._vector_distance_function = MagicMock(return_value="euclidean") + p._embed_batch.return_value = [[1.0, 0.0]] + p._nearest_active_full = MagicMock() + p._apply_inplace_update = MagicMock() + extracted = {"facts": [_doc("f-new", "near identical")], "episodic": [], "updates": []} - p.reconcile_memories("u1", n=50, memory_type="fact", full_rebuild=True) + out = p.dedup_extracted_memories("u1", extracted) - # Survivor f2 cleared; consumed f1 not re-upserted (would resurrect it). - assert cleared == ["f2"] + assert [doc["id"] for doc in out["facts"]] == ["f-new"] + p._nearest_active_full.assert_not_called() + p._apply_inplace_update.assert_not_called() -def test_sweep_survives_one_cluster_failure() -> None: - # A truncated/malformed LLM response on one cluster must not abort the sweep: - # remaining clusters still reconcile, orphan clearing still runs, and the failed - # cluster's tags are RETAINED (not cleared) so it retries next sweep. +def test_apply_inplace_update_recency_wins_and_unions() -> None: + # The refreshed doc keeps the neighbor's id but takes the new content/embedding, + # max salience/confidence, unioned tags (minus sys:dup-candidate), bumped updated_at. p = _make_pipeline() - c1 = [ - _doc("a1", "x", tags=["sys:fact", "sys:dup-candidate"]), - _doc("a2", "y", tags=["sys:fact", "sys:dup-candidate"]), - ] - c2 = [ - _doc("b1", "p", tags=["sys:fact", "sys:dup-candidate"]), - _doc("b2", "q", tags=["sys:fact", "sys:dup-candidate"]), - ] - orphan = _doc("o1", "lonely", tags=["sys:fact", "sys:dup-candidate"]) - seeds = [c1[0], c1[1], c2[0], c2[1], orphan] - p._build_candidate_clusters = MagicMock(return_value=([c1, c2], 5, seeds)) - p._reconcile_pool = MagicMock( - side_effect=[RuntimeError("truncated LLM response"), ({"kept": 2, "merged": 0, "contradicted": 0}, set())] + neighbor = _doc("existing-1", "old content", confidence=0.6, salience=0.5, tags=["sys:fact", "topic:a"]) + neighbor["_etag"] = "etag-xyz" + new_doc = _doc( + "f-new", + "new richer content", + confidence=0.9, + salience=0.8, + tags=["sys:fact", "topic:b", "sys:dup-candidate"], + embedding=[0.5, 0.5], ) - cleared: list[str] = [] - p._clear_dup_candidate_tags = MagicMock(side_effect=lambda docs: cleared.extend(d["id"] for d in docs)) - p._emit_reconcile_outcome = MagicMock() - result = p._reconcile_candidate_mode("u1", n=50, memory_type="fact", started_at=0.0) + ok = p._apply_inplace_update(neighbor, new_doc) - # c1 failed -> its tags retained (not cleared) for retry; c2 survivors + orphan cleared. - assert "a1" not in cleared and "a2" not in cleared - assert {"b1", "b2", "o1"} <= set(cleared) - assert result["reconcile_clusters_sent"] == 2 - assert result["kept"] == 2 + assert ok is True + written = p._memories_container.upsert_item.call_args.kwargs["body"] + assert written["id"] == "existing-1" + assert written["content"] == "new richer content" + assert written["embedding"] == [0.5, 0.5] + assert written["salience"] == 0.8 + assert written["confidence"] == 0.9 + assert "topic:a" in written["tags"] and "topic:b" in written["tags"] + assert "sys:dup-candidate" not in written["tags"] + assert "_etag" not in written + assert written["updated_at"] != neighbor["updated_at"] + + +def test_nearest_active_full_returns_full_doc_and_skips_excluded() -> None: + p = _make_pipeline() + doc_a = _doc("a", "first") + doc_b = _doc("b", "second") + + def query_items(*, query: str, parameters, **kwargs): + del query, parameters, kwargs + return iter( + [ + {"doc": doc_a, "score": 0.99}, + {"doc": doc_b, "score": 0.80}, + ] + ) + + p._memories_container.query_items.side_effect = query_items + # Exclude the closest -> falls through to the next candidate. + neighbor, score = p._nearest_active_full(user_id="u1", embedding=[1.0, 0.0], memory_type="fact", exclude_ids={"a"}) + assert neighbor["id"] == "b" + assert score == 0.80 def test_dedup_extracted_flag_off_is_noop(monkeypatch) -> None: @@ -325,60 +292,46 @@ def test_dedup_extracted_flag_off_is_noop(monkeypatch) -> None: p._embed_batch.assert_not_called() -def test_reconcile_memory_type_routing_episodic_and_procedural(monkeypatch) -> None: - monkeypatch.setattr("azure.cosmos.agent_memory.thresholds.get_dedup_reconcile_mode", lambda: "full_pool") +def test_reconcile_memory_type_routing_episodic_and_procedural() -> None: + # Episodic and procedural reconcile are no-ops (their near-dups fold at write + # time / have no contradiction semantics): no LLM call, zeroed counts. p = _make_pipeline() - episodes = [_doc("e1", "episode one", "episodic"), _doc("e2", "episode two", "episodic")] - p._memories_container.query_items.return_value = iter(episodes) - p._run_prompty.return_value = json.dumps({"duplicate_groups": [], "kept_ids": ["e1", "e2"]}) - result = p.reconcile_memories("u1", memory_type="episodic") + episodic_result = p.reconcile_memories("u1", memory_type="episodic") + assert episodic_result == {"kept": 0, "merged": 0, "contradicted": 0} - assert result["contradicted"] == 0 - assert p._run_prompty.call_args.args[0] == "dedup_episodic.prompty" - assert "episodics_text" in p._run_prompty.call_args.kwargs["inputs"] + procedural_result = p.reconcile_memories("u1", memory_type="procedural") + assert procedural_result == {"kept": 0, "merged": 0, "contradicted": 0} - p._run_prompty.reset_mock() - assert p.reconcile_memories("u1", memory_type="procedural")["reconcile_clusters_sent"] == 0 p._run_prompty.assert_not_called() -def test_candidate_mode_connected_components() -> None: +def test_reconcile_fact_contradiction_only() -> None: + # The fact reconcile path applies only contradicted_pairs; duplicate_groups in + # the LLM response are ignored (write-time in-place dedup owns paraphrases). p = _make_pipeline() - seed = _doc( - "f1", - "User likes coffee", - embedding=[1.0, 0.0], - tags=["sys:fact", "sys:dup-candidate"], - metadata={"category": "preference", "dup_of": "f2", "dup_score": 0.9}, - ) - neighbor = _doc("f2", "User loves coffee", embedding=[0.95, 0.05]) - - def query_items(*, query: str, parameters: list[dict[str, Any]], **kwargs: Any): - del kwargs - params = {p["name"]: p["value"] for p in parameters} - if "ARRAY_CONTAINS" in query: - return iter([seed]) - ids = {value for name, value in params.items() if name.startswith("@id")} - if ids: - return iter([doc for doc in (seed, neighbor) if doc["id"] in ids]) - return iter([seed, neighbor]) - - p._memories_container.query_items.side_effect = query_items - p._vector_candidates = MagicMock( - return_value=[{"id": "f2", "content": neighbor["content"], "type": "fact", "score": 0.9}] - ) - p._run_prompty.return_value = json.dumps( - { - "duplicate_groups": [{"merged_content": "User likes coffee.", "source_ids": ["f1", "f2"]}], - "contradicted_pairs": [], - "kept_ids": [], - } + facts = [ + _doc("f1", "User's deadline is March 1", created_at="2024-01-01T00:00:00+00:00"), + _doc("f2", "User's deadline is March 15", created_at="2024-02-01T00:00:00+00:00"), + ] + p._memories_container.query_items.return_value = iter(facts) + p._run_prompty = MagicMock( + return_value=json.dumps( + { + "duplicate_groups": [{"merged_content": "ignored", "source_ids": ["f1", "f2"]}], + "contradicted_pairs": [{"winner_id": "f2", "loser_id": "f1", "reason": "more recent"}], + "kept_ids": ["f2"], + } + ) ) result = p.reconcile_memories("u1", memory_type="fact") - assert result["reconcile_clusters_sent"] == 1 - assert result["merged"] == 2 - p._run_prompty.assert_called_once() + assert result == {"kept": 1, "merged": 0, "contradicted": 1} + # No merged doc upserted; only the loser superseded. + p._upsert_memory.assert_not_called() + assert p._mark_superseded.call_count == 1 + assert p._mark_superseded.call_args.args[0]["id"] == "f1" + assert p._mark_superseded.call_args.args[1] == "f2" + assert p._mark_superseded.call_args.kwargs["reason"] == "contradict" assert p._run_prompty.call_args.args[0] == "dedup.prompty" diff --git a/tests/unit/services/test_extract_dry.py b/tests/unit/services/test_extract_dry.py index 635bcf9..97bae44 100644 --- a/tests/unit/services/test_extract_dry.py +++ b/tests/unit/services/test_extract_dry.py @@ -246,102 +246,25 @@ def test_extract_memories_dry_is_byte_deterministic_for_same_llm_response() -> N ) -def test_extract_memories_dry_stage1_searches_user_turn_text_by_default() -> None: - chat = _SyncChat([_response()]) - memories_store = _Store([]) - memories_store.search_results = [ - { - "id": "memory-hybrid", - "content": "Existing hybrid memory from search.", - "type": "fact", - "salience": 0.7, - } - ] - turns = [_turn(1), _turn(2)] - turns_store = _Store(turns) - service = PipelineService( - memories_store, - chat, - _SyncEmbeddings(), - containers=_containers_for_store(memories_store, turns_store=turns_store), - ) - - service.extract_memories_dry("u1", "t1") - - assert memories_store.search_calls == [ - { - "search_terms": "\n".join(turn["content"] for turn in turns), - "user_id": "u1", - "memory_types": ["fact"], - "top_k": 10, - "include_superseded": False, - } - ] - assert "Existing hybrid memory from search." in json.dumps(chat.messages) - - -def test_extract_memories_dry_stage1_falls_back_to_transcript_without_user_turns() -> None: +def test_extract_memories_dry_does_not_call_store_search() -> None: + """Extraction is single-pass and existing-memory-free: it must not issue a + dedup-context vector search (dedup is handled by hash + reconciliation).""" memories_store = _Store([]) - turns = [ - { - **_turn(1), - "id": "assistant-turn-1", - "role": "assistant", - "content": "Assistant response with no user-role content.", - } - ] + memories_store.search_results = [{"id": "should-not-appear", "content": "x", "type": "fact"}] service = PipelineService( memories_store, _SyncChat([_response()]), _SyncEmbeddings(), - containers=_containers_for_store(memories_store, turns_store=_Store(turns)), - ) - - service.extract_memories_dry("u1", "t1") - - assert memories_store.search_calls == [ - { - "search_terms": service._build_transcript(turns), - "user_id": "u1", - "memory_types": ["fact"], - "top_k": 10, - "include_superseded": False, - } - ] - - -def test_extract_memories_dry_stage1_legacy_context_does_not_call_search(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr("azure.cosmos.agent_memory.thresholds.get_dedup_context_vector_enabled", lambda: False) - chat = _SyncChat([_response()]) - memories_store = _Store( - [ - { - "id": "legacy-memory", - "user_id": "u1", - "type": "fact", - "content": "Existing legacy memory from load.", - "salience": 0.6, - } - ] - ) - service = PipelineService( - memories_store, - chat, - _SyncEmbeddings(), containers=_containers_for_store(memories_store, turns_store=_Store([_turn(1)])), ) service.extract_memories_dry("u1", "t1") assert memories_store.search_calls == [] - assert "Existing legacy memory from load." in json.dumps(chat.messages) @pytest.mark.asyncio -async def test_async_extract_memories_dry_shape_is_small_and_has_no_embeddings(monkeypatch) -> None: - monkeypatch.setattr( - "azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_context_vector_enabled", lambda: False - ) +async def test_async_extract_memories_dry_shape_is_small_and_has_no_embeddings() -> None: chat = _AsyncChat([_response()]) embeddings = _AsyncEmbeddings() memories_store = _AsyncStore([]) @@ -362,10 +285,7 @@ async def test_async_extract_memories_dry_shape_is_small_and_has_no_embeddings(m @pytest.mark.asyncio -async def test_async_extract_memories_dry_is_byte_deterministic_for_same_llm_response(monkeypatch) -> None: - monkeypatch.setattr( - "azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_context_vector_enabled", lambda: False - ) +async def test_async_extract_memories_dry_is_byte_deterministic_for_same_llm_response() -> None: store = _AsyncStore([]) turns_store = _AsyncStore([_turn(1)]) service = AsyncPipelineService( @@ -384,172 +304,115 @@ async def test_async_extract_memories_dry_is_byte_deterministic_for_same_llm_res @pytest.mark.asyncio -async def test_async_extract_memories_dry_stage1_searches_user_turn_text_by_default() -> None: +async def test_async_extract_memories_dry_does_not_call_store_search() -> None: store = _AsyncStore([]) - store.search = AsyncMock( - return_value=[ - { - "id": "fact-1", - "content": "The user prefers dark mode.", - "type": "fact", - "salience": 0.7, - } - ] - ) - turns_store = _AsyncStore([_turn(1)]) - service = AsyncPipelineService( - store, - _AsyncChat([_response()]), - _AsyncEmbeddings(), - containers=_async_containers_for_store(store, turns_store=turns_store), - ) - service._vector_candidates = AsyncMock(return_value=[]) - turns = [_turn(1)] - - await service.extract_memories_dry("u1", "t1") - - store.search.assert_awaited_once_with( - search_terms="\n".join(turn["content"] for turn in turns), - user_id="u1", - memory_types=["fact"], - top_k=10, - ) - service._vector_candidates.assert_not_awaited() - -@pytest.mark.asyncio -async def test_async_extract_memories_dry_stage1_falls_back_to_transcript_without_user_turns() -> None: - store = _AsyncStore([]) store.search = AsyncMock(return_value=[]) - turns = [ - { - **_turn(1), - "id": "assistant-turn-1", - "role": "assistant", - "content": "Assistant response with no user-role content.", - } - ] - turns_store = _AsyncStore(turns) service = AsyncPipelineService( store, _AsyncChat([_response()]), _AsyncEmbeddings(), - containers=_async_containers_for_store(store, turns_store=turns_store), + containers=_async_containers_for_store(store, turns_store=_AsyncStore([_turn(1)])), ) - transcript = service._build_transcript(turns) await service.extract_memories_dry("u1", "t1") - store.search.assert_awaited_once_with( - search_terms=transcript, - user_id="u1", - memory_types=["fact"], - top_k=10, - ) + store.search.assert_not_awaited() -@pytest.mark.asyncio -async def test_async_extract_memories_dry_stage1_legacy_path_when_context_vector_unset(monkeypatch) -> None: - monkeypatch.setattr( - "azure.cosmos.agent_memory.aio.services.pipeline.get_dedup_context_vector_enabled", lambda: False - ) - store = _AsyncStore( - [ +class _BatchChat: + """Returns a unique fact per call; optionally raises on a chosen call index.""" + + def __init__(self, fail_on_call=None, error=None): + self.calls = 0 + self.fail_on_call = fail_on_call + self.error = error + + def generate(self, messages, **opts): + del messages, opts + self.calls += 1 + if self.fail_on_call and self.calls == self.fail_on_call: + raise self.error + return json.dumps( { - "id": "fact-1", - "user_id": "u1", - "type": "fact", - "content": "Existing fact.", - "content_hash": "hash-1", + "facts": [ + { + "text": f"Fact from call {self.calls}.", + "category": "other", + "confidence": 0.9, + "salience": 0.8, + "temporal_context": None, + "tags": [], + } + ], + "episodic": [], } - ] - ) - store.search = AsyncMock(return_value=[]) - turns_store = _AsyncStore([_turn(1)]) - service = AsyncPipelineService( - store, - _AsyncChat([_response()]), - _AsyncEmbeddings(), - containers=_async_containers_for_store(store, turns_store=turns_store), - ) + ) - await service.extract_memories_dry("u1", "t1") - store.search.assert_not_awaited() +def _one_turn_per_batch(monkeypatch): + # Force each small turn into its own extraction batch. + monkeypatch.setattr("azure.cosmos.agent_memory.thresholds.get_extraction_batch_max_tokens", lambda: 5) -def test_extract_memories_dry_stage1_search_failure_falls_back_to_hash_memories() -> None: - """A failing dedup-context vector search must not abort extraction; it - falls back to the hash-loaded existing memories (option 3 resilience).""" - chat = _SyncChat([_response()]) - memories_store = _Store( - [ - { - "id": "hash-memory", - "user_id": "u1", - "type": "fact", - "content": "Existing hash-based memory from load.", - "content_hash": "h1", - "salience": 0.6, - } - ] +def test_extract_batches_run_independently_one_call_per_batch(monkeypatch) -> None: + _one_turn_per_batch(monkeypatch) + chat = _BatchChat() + memories_store = _Store([]) + turns_store = _Store([_turn(i) for i in range(3)]) + service = PipelineService( + memories_store, + chat, + _SyncEmbeddings(), + containers=_containers_for_store(memories_store, turns_store=turns_store), ) - def _boom(**_kwargs): - raise RuntimeError("vector search down") + out = service.extract_memories_dry("u1", "t1") + + assert chat.calls == 3 # one LLM call per batch + assert len(out["facts"]) == 3 + assert len(out["processed_turn_docs"]) == 3 + - memories_store.search = _boom +def test_extract_quarantines_non_retryable_batch_but_keeps_others(monkeypatch) -> None: + _one_turn_per_batch(monkeypatch) + chat = _BatchChat(fail_on_call=2, error=Exception("Error code: 400 content_filter")) + memories_store = _Store([]) + turns_store = _Store([_turn(i) for i in range(3)]) service = PipelineService( memories_store, chat, _SyncEmbeddings(), - containers=_containers_for_store(memories_store, turns_store=_Store([_turn(1)])), + containers=_containers_for_store(memories_store, turns_store=turns_store), ) - # Must not raise despite the Stage-1 vector search failing. - output = service.extract_memories_dry("u1", "t1") + out = service.extract_memories_dry("u1", "t1") - assert output["facts"] # extraction proceeded - # Fallback surfaced the hash-loaded memory into the extraction prompt. - assert "Existing hash-based memory from load." in json.dumps(chat.messages) + # Batches 1 and 3 produced facts; batch 2 was quarantined (no fact) ... + assert len(out["facts"]) == 2 + # ... but ALL 3 turns are stamped (quarantined turn included) so it never re-poisons. + assert len(out["processed_turn_docs"]) == 3 + stats = [u for u in out["updates"] if u.get("op") == "stats" and "quarantined_turn_count" in u] + assert stats and stats[0]["quarantined_turn_count"] == 1 -@pytest.mark.asyncio -async def test_async_extract_memories_dry_stage1_search_failure_falls_back_to_hash_memories() -> None: - class _RecordingAsyncChat(_AsyncChat): - def __init__(self, responses): - super().__init__(responses) - self.messages: list = [] - - async def generate(self, messages, **opts): - self.messages.append(messages) - del opts - self.calls += 1 - return json.dumps(self.responses.pop(0)) - - chat = _RecordingAsyncChat([_response()]) - store = _AsyncStore( - [ - { - "id": "hash-memory", - "user_id": "u1", - "type": "fact", - "content": "Existing hash-based memory from load.", - "content_hash": "h1", - "salience": 0.6, - } - ] - ) - store.search = AsyncMock(side_effect=RuntimeError("vector search down")) - service = AsyncPipelineService( - store, +def test_extract_defers_retryable_batch_leaving_turns_unstamped(monkeypatch) -> None: + _one_turn_per_batch(monkeypatch) + chat = _BatchChat(fail_on_call=2, error=Exception("Error code: 429 rate limit")) + memories_store = _Store([]) + turns_store = _Store([_turn(i) for i in range(3)]) + service = PipelineService( + memories_store, chat, - _AsyncEmbeddings(), - containers=_async_containers_for_store(store, turns_store=_AsyncStore([_turn(1)])), + _SyncEmbeddings(), + containers=_containers_for_store(memories_store, turns_store=turns_store), ) - output = await service.extract_memories_dry("u1", "t1") + out = service.extract_memories_dry("u1", "t1") - assert output["facts"] # extraction proceeded - store.search.assert_awaited_once() - assert "Existing hash-based memory from load." in json.dumps(chat.messages) + # Batches 1 and 3 produced facts; batch 2 deferred (retryable) ... + assert len(out["facts"]) == 2 + # ... its turn is NOT in processed_turn_docs, so it stays un-extracted and retries. + assert len(out["processed_turn_docs"]) == 2 + stats = [u for u in out["updates"] if u.get("op") == "stats" and "deferred_turn_count" in u] + assert stats and stats[0]["deferred_turn_count"] == 1 diff --git a/tests/unit/services/test_extraction_batching.py b/tests/unit/services/test_extraction_batching.py new file mode 100644 index 0000000..5793b17 --- /dev/null +++ b/tests/unit/services/test_extraction_batching.py @@ -0,0 +1,58 @@ +"""Unit tests for Option-4 extraction helpers: failure classification + batching.""" + +from __future__ import annotations + +from azure.cosmos.agent_memory.services._pipeline_helpers import ( + batch_turns_by_tokens, + is_retryable_llm_error, +) + + +class TestIsRetryableLLMError: + def test_content_filter_is_non_retryable(self) -> None: + assert is_retryable_llm_error(Exception("Error code: 400 - content_filter triggered")) is False + + def test_content_management_policy_is_non_retryable(self) -> None: + assert is_retryable_llm_error(Exception("blocked by Azure content management policy")) is False + + def test_context_length_is_non_retryable(self) -> None: + assert is_retryable_llm_error(Exception("context_length_exceeded")) is False + assert is_retryable_llm_error(Exception("This model's maximum context length is 8192")) is False + + def test_rate_limit_is_retryable(self) -> None: + assert is_retryable_llm_error(Exception("Error code: 429 - rate limit exceeded")) is True + + def test_server_error_is_retryable(self) -> None: + assert is_retryable_llm_error(Exception("Error code: 503 - service unavailable")) is True + + def test_unknown_error_defaults_retryable(self) -> None: + # Conservative: never quarantine (drop) turns on an unclassified error. + assert is_retryable_llm_error(Exception("something weird happened")) is True + + +class TestBatchTurnsByTokens: + def _turns(self, n, content="word " * 10): + return [{"id": f"t{i}", "content": content} for i in range(n)] + + def test_empty_returns_empty(self) -> None: + assert batch_turns_by_tokens([], 1000) == [] + + def test_small_turns_single_batch_when_budget_large(self) -> None: + turns = self._turns(5) + batches = batch_turns_by_tokens(turns, 10_000) + assert len(batches) == 1 + assert len(batches[0]) == 5 + + def test_splits_when_budget_small(self) -> None: + turns = self._turns(6, content="word " * 20) # ~20 tokens each + batches = batch_turns_by_tokens(turns, 25) # ~1 turn per batch + assert len(batches) == 6 + assert all(len(b) == 1 for b in batches) + # order + coverage preserved + assert [t["id"] for b in batches for t in b] == [f"t{i}" for i in range(6)] + + def test_oversized_single_turn_is_own_batch(self) -> None: + turns = [{"id": "big", "content": "word " * 5000}, {"id": "small", "content": "hi"}] + batches = batch_turns_by_tokens(turns, 100) + assert [t["id"] for b in batches for t in b] == ["big", "small"] + assert len(batches) == 2 diff --git a/tests/unit/services/test_parse_llm_json.py b/tests/unit/services/test_parse_llm_json.py index 9352c75..2f7ddd2 100644 --- a/tests/unit/services/test_parse_llm_json.py +++ b/tests/unit/services/test_parse_llm_json.py @@ -1,4 +1,4 @@ -"""Unit tests for parse_llm_json — resilient LLM JSON decoding (Issue A).""" +"""Unit tests for parse_llm_json - resilient LLM JSON decoding (Issue A).""" from __future__ import annotations diff --git a/tests/unit/services/test_pipeline_service.py b/tests/unit/services/test_pipeline_service.py index 27518b5..0800ab1 100644 --- a/tests/unit/services/test_pipeline_service.py +++ b/tests/unit/services/test_pipeline_service.py @@ -12,18 +12,10 @@ @pytest.fixture(autouse=True) def _pin_legacy_dedup_paths(monkeypatch): - monkeypatch.setattr( - "azure.cosmos.agent_memory.thresholds.get_dedup_reconcile_mode", - lambda: "full_pool", - ) monkeypatch.setattr( "azure.cosmos.agent_memory.thresholds.get_dedup_vector_enabled", lambda: False, ) - monkeypatch.setattr( - "azure.cosmos.agent_memory.thresholds.get_dedup_context_vector_enabled", - lambda: False, - ) class FakeLLMService: @@ -247,7 +239,7 @@ def test_extract_memories_happy_path_writes_fact_and_episodic() -> None: def test_extract_memories_creates_new_fact_without_superseding() -> None: - # Extract no longer detects updates/contradictions — it just adds new facts. + # Extract no longer detects updates/contradictions - it just adds new facts. # Contradiction resolution happens later in reconcile, not at extract time. old = _fact("old_fact", "The user prefers light mode.") store = FakeStore([old]) @@ -304,7 +296,7 @@ def test_synthesize_procedural_produces_procedural_memory() -> None: assert store.upserts == [proc] -def test_reconcile_memories_returns_three_bucket_counts() -> None: +def test_reconcile_memories_returns_contradiction_counts() -> None: store = FakeStore( [ _fact("f1", "User likes coffee", ts=1), @@ -316,6 +308,8 @@ def test_reconcile_memories_returns_three_bucket_counts() -> None: llm = FakeLLMService( [ { + # duplicate_groups are ignored now (paraphrases fold at write time); + # only contradicted_pairs are applied. "duplicate_groups": [ {"source_ids": ["f1", "f2"], "merged_content": "The user prefers coffee.", "confidence": 0.9} ], @@ -327,7 +321,9 @@ def test_reconcile_memories_returns_three_bucket_counts() -> None: result = _pipeline(store, llm).reconcile_memories("u1", n=4) - assert result == {"kept": 1, "merged": 2, "contradicted": 1} + # f1/f2 are NOT merged (paraphrase folding moved to write time); only f4 loses a + # contradiction. Survivors: f1, f2, f3. + assert result == {"kept": 3, "merged": 0, "contradicted": 1} def test_reconcile_memories_tombstones_losers_with_reasons() -> None: @@ -352,8 +348,10 @@ def test_reconcile_memories_tombstones_losers_with_reasons() -> None: _pipeline(store, llm).reconcile_memories("u1", n=4) by_id = {doc["id"]: doc for doc in store.docs} - assert by_id["f1"]["supersede_reason"] == "duplicate" - assert by_id["f2"]["supersede_reason"] == "duplicate" + # Paraphrases f1/f2 are left untouched (no write-time merge in reconcile). + assert by_id["f1"].get("supersede_reason") is None + assert by_id["f2"].get("supersede_reason") is None + # Only the contradiction loser is tombstoned. assert by_id["f4"]["supersede_reason"] == "contradict" assert by_id["f4"]["superseded_by"] == "f3" diff --git a/tests/unit/services/test_transcript_metadata.py b/tests/unit/services/test_transcript_metadata.py index fa67bb8..e6b9f85 100644 --- a/tests/unit/services/test_transcript_metadata.py +++ b/tests/unit/services/test_transcript_metadata.py @@ -153,7 +153,7 @@ async def test_async_pipeline_propagates_allowlist(self) -> None: class TestStringInputRejected: - """A bare ``str`` is iterable char-by-char — passing it would silently + """A bare ``str`` is iterable char-by-char - passing it would silently produce a per-character allow-list. Reject with a clear ``TypeError``.""" def test_build_transcript_rejects_str(self) -> None: @@ -211,7 +211,7 @@ def test_client_ctor_rejects_str(self) -> None: class TestCompactJsonSerialization: - """Token-reduction is the whole point — emit compact JSON, not pretty.""" + """Token-reduction is the whole point - emit compact JSON, not pretty.""" def test_no_whitespace_between_separators(self) -> None: items = [_turn("user", "Hi", a=1, b=2)] @@ -231,7 +231,7 @@ def test_non_ascii_preserved(self) -> None: class TestNonSerializableMetadataCoerced: - """Real-world metadata carries datetimes, UUIDs, etc. — ``default=str`` + """Real-world metadata carries datetimes, UUIDs, etc. - ``default=str`` keeps a single bad value from blowing up the entire extraction.""" def test_datetime_value_does_not_crash(self) -> None: diff --git a/tests/unit/store/test_get_memory_history.py b/tests/unit/store/test_get_memory_history.py new file mode 100644 index 0000000..1ecbad3 --- /dev/null +++ b/tests/unit/store/test_get_memory_history.py @@ -0,0 +1,150 @@ +"""Unit tests for get_memory_history - supersession-chain walking. + +Covers the sync store logic (chain walking, ordering, partition scoping, +validation, cycle guard) plus the shared projection enrichment. The async +store shares the same helper module, so a single async smoke test guards +parity. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest + +from azure.cosmos.agent_memory._container_routing import ContainerKey +from azure.cosmos.agent_memory.exceptions import ValidationError +from azure.cosmos.agent_memory.store import MemoryStore +from azure.cosmos.agent_memory.store._search_helpers import MEMORY_PROJECTION + + +def _containers(*, memories=None): + return { + ContainerKey.TURNS: MagicMock(), + ContainerKey.MEMORIES: memories if memories is not None else MagicMock(), + ContainerKey.SUMMARIES: MagicMock(), + } + + +def _fact(fact_id, *, superseded_by=None, superseded_at=None, reason=None, content=""): + return { + "id": fact_id, + "user_id": "u1", + "thread_id": "t1", + "type": "fact", + "content": content or fact_id, + "metadata": {"category": "biographical"}, + "created_at": "2025-01-01T00:00:00+00:00", + "superseded_by": superseded_by, + "superseded_at": superseded_at, + "supersede_reason": reason, + } + + +# --------------------------------------------------------------------------- +# Projection enrichment (Change 1) +# --------------------------------------------------------------------------- + + +def test_projection_includes_supersession_audit_fields(): + assert "c.superseded_at" in MEMORY_PROJECTION + assert "c.supersede_reason" in MEMORY_PROJECTION + # still carries the fields callers already relied on + assert "c.superseded_by" in MEMORY_PROJECTION + assert "c.created_at" in MEMORY_PROJECTION + + +# --------------------------------------------------------------------------- +# get_memory_history +# --------------------------------------------------------------------------- + + +def test_single_level_history_returns_direct_predecessor(): + memories = MagicMock() + # First hop finds the doc superseded by 'current'; second hop finds nothing. + memories.query_items.side_effect = [ + [_fact("v1", superseded_at="2025-02-01T00:00:00+00:00", reason="update")], + [], + ] + store = MemoryStore(containers=_containers(memories=memories)) + + history = store.get_memory_history("current", user_id="u1", thread_id="t1") + + assert [d["id"] for d in history] == ["v1"] + assert history[0]["supersede_reason"] == "update" + + +def test_multi_level_chain_walks_transitively_newest_first(): + memories = MagicMock() + memories.query_items.side_effect = [ + [_fact("v2", superseded_at="2025-03-01T00:00:00+00:00", reason="contradict")], + [_fact("v1", superseded_at="2025-02-01T00:00:00+00:00", reason="update")], + [], + ] + store = MemoryStore(containers=_containers(memories=memories)) + + history = store.get_memory_history("current", user_id="u1", thread_id="t1") + + # Ordered most-recently-superseded first. + assert [d["id"] for d in history] == ["v2", "v1"] + + +def test_no_history_returns_empty_list(): + memories = MagicMock() + memories.query_items.return_value = [] + store = MemoryStore(containers=_containers(memories=memories)) + + assert store.get_memory_history("current", user_id="u1", thread_id="t1") == [] + + +def test_cycle_is_bounded_and_deduplicated(): + memories = MagicMock() + # Pathological: the predecessor points back at an already-seen id. + memories.query_items.side_effect = [ + [_fact("v1", superseded_at="2025-02-01T00:00:00+00:00")], + [_fact("current")], # 'current' already in seen -> skipped, loop stops + [], + ] + store = MemoryStore(containers=_containers(memories=memories)) + + history = store.get_memory_history("current", user_id="u1", thread_id="t1", max_depth=50) + + assert [d["id"] for d in history] == ["v1"] + + +def test_thread_scoped_query_uses_single_partition(): + memories = MagicMock() + memories.query_items.return_value = [] + store = MemoryStore(containers=_containers(memories=memories)) + + store.get_memory_history("current", user_id="u1", thread_id="t1") + + kwargs = memories.query_items.call_args.kwargs + assert kwargs["partition_key"] == ["u1", "t1"] + assert "enable_cross_partition_query" not in kwargs + assert "c.thread_id = @thread_id" in kwargs["query"] + + +def test_history_without_thread_id_fans_out_cross_partition(): + memories = MagicMock() + memories.query_items.return_value = [] + store = MemoryStore(containers=_containers(memories=memories)) + + store.get_memory_history("current", user_id="u1") + + kwargs = memories.query_items.call_args.kwargs + assert kwargs.get("enable_cross_partition_query") is True + assert "c.thread_id = @thread_id" not in kwargs["query"] + assert "c.superseded_by IN (@sid0)" in kwargs["query"] + + +def test_missing_memory_id_raises(): + store = MemoryStore(containers=_containers()) + with pytest.raises(ValidationError): + store.get_memory_history("", user_id="u1") + + +def test_missing_user_id_raises(): + store = MemoryStore(containers=_containers()) + with pytest.raises(ValidationError): + store.get_memory_history("current", user_id="") diff --git a/tests/unit/store/test_memory_store.py b/tests/unit/store/test_memory_store.py index e4371e1..80a61c9 100644 --- a/tests/unit/store/test_memory_store.py +++ b/tests/unit/store/test_memory_store.py @@ -489,7 +489,7 @@ def test_search_with_episodic_and_thread_id_forces_cross_partition(): call_kwargs = memories.query_items.call_args.kwargs # When user-scoped types are in scope, search must fan out across # partitions instead of confining to [u1, t1] (where no episodic - # ever lives — they all use the "__episodic__" sentinel partition). + # ever lives - they all use the "__episodic__" sentinel partition). assert call_kwargs.get("enable_cross_partition_query") is True assert "partition_key" not in call_kwargs assert "(c.thread_id = @thread_id OR c.type IN (@user_scoped_type_0))" in call_kwargs["query"] diff --git a/tests/unit/test_async_hygiene.py b/tests/unit/test_async_hygiene.py index 7debcf3..5f1f59f 100644 --- a/tests/unit/test_async_hygiene.py +++ b/tests/unit/test_async_hygiene.py @@ -21,7 +21,7 @@ # Each entry is ``(relative_path_under_aio, name_of_an_enclosing_function)``. # A call site matches if *any* function in its enclosing chain has the -# allowed name — this lets us tolerate nested closures (e.g. the inner +# allowed name - this lets us tolerate nested closures (e.g. the inner # ``_provider`` inside ``_make_sync_token_provider_for_async``) without # pinning to the inner name. ALLOWED_CALL_SITES: set[tuple[str, str]] = { @@ -93,8 +93,8 @@ def test_asyncio_to_thread_only_in_allowed_call_sites() -> None: """Every ``asyncio.to_thread`` call in ``aio/`` must be in the allowed set. Failure here means a new sync-to-async bridge was introduced. Either - rewrite the call to use a native async API, or — if the bridge is - truly necessary — add ``(file, function)`` to ``ALLOWED_CALL_SITES`` + rewrite the call to use a native async API, or - if the bridge is + truly necessary - add ``(file, function)`` to ``ALLOWED_CALL_SITES`` in this test along with a comment explaining why. """ hits = _collect_to_thread_calls() diff --git a/tests/unit/test_auto_trigger.py b/tests/unit/test_auto_trigger.py index 2f23420..a2891eb 100644 --- a/tests/unit/test_auto_trigger.py +++ b/tests/unit/test_auto_trigger.py @@ -10,10 +10,8 @@ from unittest.mock import MagicMock, patch -import pytest from azure.cosmos.exceptions import CosmosResourceNotFoundError -from azure.cosmos.agent_memory import _counters from azure.cosmos.agent_memory.cosmos_memory_client import CosmosMemoryClient from azure.cosmos.agent_memory.processors import DurableFunctionProcessor, InProcessProcessor @@ -85,7 +83,7 @@ def test_push_to_cosmos_fires_inprocess_trigger_per_turn(monkeypatch): client.add_local(user_id="u1", role="user", thread_id="t1", content="hi") client.push_to_cosmos() - pipeline.extract_memories.assert_called_once_with("u1", "t1", recent_k=1) + pipeline.extract_memories.assert_called_once_with("u1", "t1", recent_k=None) def test_push_to_cosmos_durable_does_not_fire_trigger(monkeypatch): @@ -156,7 +154,7 @@ def test_push_to_cosmos_skips_when_counter_container_unavailable(monkeypatch): # --------------------------------------------------------------------------- -# Per-step trigger gating — each *_EVERY_N fires its own pipeline step +# Per-step trigger gating - each *_EVERY_N fires its own pipeline step # independently. The InProcess backend mirrors the function-app # split-orchestrator behavior so the two backends produce the same memory # contents for the same chat history. @@ -185,38 +183,15 @@ def test_extract_fires_independently_of_summary(self, monkeypatch): client.add_local(user_id="u1", role="user", thread_id="t1", content="hi") client.push_to_cosmos() - processor.process_extract_memories.assert_called_once_with(user_id="u1", thread_id="t1", recent_k=1) + processor.process_extract_memories.assert_called_once_with(user_id="u1", thread_id="t1") processor.process_thread_summary.assert_not_called() processor.process_user_summary.assert_not_called() - @pytest.mark.parametrize( - ("n_facts", "batch_count", "counter_result", "watermark", "expected_recent_k"), - [ - (1, 1, (0, 1), None, 1), - (1, 3, (0, 3), None, 3), - (5, 1, (4, 5), None, 5), - (1, 1, (5, 10), 5, 5), - # Large backlog is NOT capped: recent_k spans every turn since the - # watermark (newest-recent_k slice covers exactly those), so the - # watermark can advance to new_count with no stranded turns. - (1, 1, (98, 100), 0, 100), - # BOOTSTRAP regression: no watermark yet but the counter is already - # ahead of this batch (earlier extracts failed). base=0 so recent_k = - # new_count (30) covers ALL turns — the old fallback max(n_facts, - # batch_count) would return 2 and strand turns 1-28 forever. - (1, 2, (20, 30), None, 30), - ], - ) - def test_extract_recent_k_uses_watermark_then_falls_back( - self, - monkeypatch, - n_facts, - batch_count, - counter_result, - watermark, - expected_recent_k, - ): - monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", str(n_facts)) + def test_extract_fires_without_recent_k_or_watermark(self, monkeypatch): + """Extraction now covers all un-extracted turns (gated per-turn by + extracted_at) and batches internally, so the auto-trigger fires it with + NO recent_k and tracks NO success-gated watermark.""" + monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") @@ -226,74 +201,19 @@ def test_extract_recent_k_uses_watermark_then_falls_back( client = _connected(processor=processor) client._counter_container_client = MagicMock() - with ( - patch( - "azure.cosmos.agent_memory._counters.increment_counter_sync", - return_value=counter_result, - ), - patch( - "azure.cosmos.agent_memory._counters.read_extract_watermark_sync", - return_value=watermark, - ), - patch( - "azure.cosmos.agent_memory._counters.advance_extract_watermark_sync", - ) as advance, + with patch( + "azure.cosmos.agent_memory._counters.increment_counter_sync", + return_value=(0, 1), ): - for i in range(batch_count): - client.add_local(user_id="u1", role="user", thread_id="t1", content=f"hi {i}") + client.add_local(user_id="u1", role="user", thread_id="t1", content="hi") client.push_to_cosmos() - processor.process_extract_memories.assert_called_once_with( - user_id="u1", - thread_id="t1", - recent_k=expected_recent_k, - ) - advance.assert_called_once() - - def test_watermark_round_trip_fail_then_succeed_no_strand(self, monkeypatch): - """End-to-end round-trip against a REAL in-memory counter (no constant - mocks): a thread's first extract fails, the second succeeds, and the - second must cover EVERY turn so far — not just its own batch — so turns - from the failed batch are never stranded. This is the bootstrap case the - constant-mock tests could not catch.""" - monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") - monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") - monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") - - counter = _FakeCounterContainer() - recorded_recent_k: list[int] = [] - - def extract(*, user_id, thread_id, recent_k): - recorded_recent_k.append(recent_k) - if len(recorded_recent_k) == 1: - raise RuntimeError("transient LLM outage") # first extract fails - return {} - - processor = InProcessProcessor(pipeline=MagicMock()) - processor.process_extract_memories = MagicMock(side_effect=extract) - client = _connected(processor=processor) - client._counter_container_client = counter - - # Batch 1: 10 turns -> counter 0->10, extract FAILS (watermark not advanced). - for i in range(10): - client.add_local(user_id="u1", role="user", thread_id="t1", content=f"a{i}") - client.push_to_cosmos() - - # Batch 2: 10 turns -> counter 10->20, extract SUCCEEDS. - for i in range(10): - client.add_local(user_id="u1", role="user", thread_id="t1", content=f"b{i}") - client.push_to_cosmos() - - # First fired with 10 (all turns so far); second with 20 (ALL turns, since - # the failed first extract left the watermark unset) — NOT 10. - assert recorded_recent_k == [10, 20] - # Watermark now seeded at the full count after the successful extract. - cid = _counters.thread_counter_id("u1", "t1") - assert _counters.read_extract_watermark_sync(counter, cid, "u1", "t1") == 20 + processor.process_extract_memories.assert_called_once_with(user_id="u1", thread_id="t1") - def test_watermark_not_advanced_when_extract_fails(self, monkeypatch): - """advance-on-success: a failing extract must NOT move the watermark, so - the skipped turns are retried next sweep; failure is stamped instead.""" + def test_extract_failure_stamps_failure(self, monkeypatch): + """A total extract failure is recorded via stamp_failure (telemetry); there + is no watermark to (not) advance - per-batch failures are handled inside + the pipeline, so this outer path only sees unexpected total failures.""" monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") @@ -309,13 +229,6 @@ def test_watermark_not_advanced_when_extract_fails(self, monkeypatch): "azure.cosmos.agent_memory._counters.increment_counter_sync", return_value=(0, 1), ), - patch( - "azure.cosmos.agent_memory._counters.read_extract_watermark_sync", - return_value=None, - ), - patch( - "azure.cosmos.agent_memory._counters.advance_extract_watermark_sync", - ) as advance, patch( "azure.cosmos.agent_memory._counters.stamp_failure_sync", ) as stamp, @@ -323,51 +236,8 @@ def test_watermark_not_advanced_when_extract_fails(self, monkeypatch): client.add_local(user_id="u1", role="user", thread_id="t1", content="hi") client.push_to_cosmos() - advance.assert_not_called() stamp.assert_called_once() - def test_reconcile_full_rebuild_on_persisted_counter_cadence(self, monkeypatch): - """Symmetry with the durable backend: the in-process auto-trigger requests - a full-pool reconcile (full_rebuild=True) on a PERSISTED-counter cadence — - every DEDUP_FULL_RECLUSTER_EVERY_N-th reconcile — not via an in-memory - per-instance sweep counter. Here that's every 2 turns.""" - monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") - monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") - monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") - monkeypatch.setenv("DEDUP_EVERY_N", "1") - monkeypatch.setattr("azure.cosmos.agent_memory.thresholds.get_dedup_full_recluster_every_n", lambda: 2) - - rebuilds: list[bool] = [] - processor = InProcessProcessor(pipeline=MagicMock()) - processor.process_extract_memories = MagicMock(return_value={}) - processor.synthesize_procedural = MagicMock(return_value=None) - processor.process_reconcile = MagicMock( - side_effect=lambda *, user_id, full_rebuild=False: rebuilds.append(full_rebuild) - ) - - client = _connected(processor=processor) - client._counter_container_client = MagicMock() - - with ( - patch( - "azure.cosmos.agent_memory._counters.increment_counter_sync", - side_effect=[(0, 1), (1, 2)], - ), - patch( - "azure.cosmos.agent_memory._counters.read_extract_watermark_sync", - return_value=None, - ), - patch( - "azure.cosmos.agent_memory._counters.advance_extract_watermark_sync", - ), - ): - client.add_local(user_id="u1", role="user", thread_id="t1", content="a") - client.push_to_cosmos() # counter 0->1: reconcile, full crosses 2? no - client.add_local(user_id="u1", role="user", thread_id="t1", content="b") - client.push_to_cosmos() # counter 1->2: full backstop threshold (2) crossed - - assert rebuilds == [False, True] - def test_summary_fires_independently_when_threshold_crossed(self, monkeypatch): """N_summary=10 boundary fires summary; N_facts=0 prevents extract.""" monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "0") @@ -416,7 +286,7 @@ def test_user_summary_fires_at_user_threshold(self, monkeypatch): # --------------------------------------------------------------------------- -# Owner exclusivity — MEMORY_PROCESSOR_OWNER ensures only one of +# Owner exclusivity - MEMORY_PROCESSOR_OWNER ensures only one of # {SDK auto-trigger, FA change-feed processor} runs against a shared # container, preventing double-extraction / double-dedup. # --------------------------------------------------------------------------- diff --git a/tests/unit/test_cosmos_memory_client.py b/tests/unit/test_cosmos_memory_client.py index 548730d..240ed58 100644 --- a/tests/unit/test_cosmos_memory_client.py +++ b/tests/unit/test_cosmos_memory_client.py @@ -75,8 +75,8 @@ def test_default_credential_created_when_flag_true(self): with patch.dict("sys.modules", {"azure.identity": mock_module}): mem = CosmosMemoryClient(use_default_credential=True) - # Two independent instances — one per consumer (cosmos + AI Foundry) - # — so close() can tear each down without affecting the other. + # Two independent instances - one per consumer (cosmos + AI Foundry) + # - so close() can tear each down without affecting the other. assert mock_module.DefaultAzureCredential.call_count == 2 assert mem._cosmos_credential is sentinel assert mem._ai_foundry_credential is sentinel @@ -205,7 +205,7 @@ def test_add_local_turn_requires_thread_id(self): mem = _make_client() with pytest.raises(ValidationError, match="thread_id is required"): mem.add_local(user_id="u1", role="user", content="hi") - # Validation must run BEFORE append — otherwise an orphan turn + # Validation must run BEFORE append - otherwise an orphan turn # with thread_id=None would persist and pollute pk on push. assert mem.local_memory == [] assert mem._unflushed_turn_counts == {} @@ -558,7 +558,7 @@ def test_constructor_rejects_invalid_throughput_mode(self): class TestAddCosmos: def test_add_cosmos(self): mem, container = _connected_client() - # Suppress cadence work — the trigger path is exercised in + # Suppress cadence work - the trigger path is exercised in # tests/unit/test_auto_trigger.py; this test just asserts the CRUD write. mem._maybe_auto_trigger = MagicMock() mem.add_cosmos(user_id="u1", role="user", content="hello", thread_id="t1") @@ -604,12 +604,12 @@ def test_add_cosmos_turn_triggers_cadence(self): trigger.assert_called_once_with({("u1", "t1"): 1}) def test_add_cosmos_swallows_cadence_failure(self): - """If the cadence trigger raises, the add_cosmos call must still succeed — + """If the cadence trigger raises, the add_cosmos call must still succeed - the user's turn was written; cadence is best-effort telemetry.""" mem, _ = _connected_client() mem._maybe_auto_trigger = MagicMock(side_effect=RuntimeError("boom")) - # Should NOT raise — the write succeeded. + # Should NOT raise - the write succeeded. result_id = mem.add_cosmos(user_id="u1", role="user", content="hi", thread_id="t1") assert isinstance(result_id, str) @@ -685,7 +685,7 @@ def _generate_batch(texts: list[str]) -> list[list[float]]: mem.push_to_cosmos() mem.push_to_cosmos() - # Second push should not re-embed — embedding is cached on local_memory. + # Second push should not re-embed - embedding is cached on local_memory. assert embed_calls == [["fact one"]] assert mem.local_memory[0]["embedding"] == [0.5, 0.6, 0.7] diff --git a/tests/unit/test_counters.py b/tests/unit/test_counters.py index 5ebd996..f25a6ab 100644 --- a/tests/unit/test_counters.py +++ b/tests/unit/test_counters.py @@ -1,7 +1,7 @@ """Tests for ``azure.cosmos.agent_memory._counters``. Covers counter-doc construction (LSN preservation, failure breadcrumbs) -and ``stamp_failure_sync`` — the helpers that let SDK and FA share a +and ``stamp_failure_sync`` - the helpers that let SDK and FA share a counter container without trampling each other. """ @@ -114,11 +114,11 @@ def test_writes_failure_fields_via_patch_item(self): assert "/last_failure_at" in op_paths def test_swallows_patch_errors(self): - """Failure stamping is best-effort breadcrumbing — must never raise.""" + """Failure stamping is best-effort breadcrumbing - must never raise.""" container = MagicMock() container.patch_item.side_effect = RuntimeError("cosmos down") stamp_failure_sync(container, "thread:u:t", "u", "t", "boom") - # Did not propagate — that's the contract. + # Did not propagate - that's the contract. def test_truncates_long_reason(self): container = MagicMock() diff --git a/tests/unit/test_embeddings.py b/tests/unit/test_embeddings.py index 9616138..9e6e1ee 100644 --- a/tests/unit/test_embeddings.py +++ b/tests/unit/test_embeddings.py @@ -193,7 +193,7 @@ def test_batch_api_failure(self, MockAOAI): # --------------------------------------------------------------------------- -# generate_batch() — N=16 chunk guard +# generate_batch() - N=16 chunk guard # --------------------------------------------------------------------------- @@ -250,7 +250,7 @@ def test_empty_list_no_api_call(self, MockAOAI): assert result == [] assert mock_client.embeddings.create.call_count == 0 - # Lazy-init must not even fire — empty short-circuit comes first. + # Lazy-init must not even fire - empty short-circuit comes first. MockAOAI.assert_not_called() @patch("openai.AzureOpenAI") diff --git a/tests/unit/test_memory_type_multi.py b/tests/unit/test_memory_type_multi.py index d00d7e2..f76f357 100644 --- a/tests/unit/test_memory_type_multi.py +++ b/tests/unit/test_memory_type_multi.py @@ -16,7 +16,7 @@ from azure.cosmos.agent_memory.cosmos_memory_client import CosmosMemoryClient # --------------------------------------------------------------------------- -# Helpers — a small replica of the patterns used in test_cosmos_memory_client. +# Helpers - a small replica of the patterns used in test_cosmos_memory_client. # --------------------------------------------------------------------------- @@ -162,7 +162,7 @@ def test_no_thread_id_no_or_clause(): # --------------------------------------------------------------------------- -# Sync client surface — verifies the list reaches the generated SQL. +# Sync client surface - verifies the list reaches the generated SQL. # --------------------------------------------------------------------------- diff --git a/tests/unit/test_models.py b/tests/unit/test_models.py index f89c9d5..d3a91b9 100644 --- a/tests/unit/test_models.py +++ b/tests/unit/test_models.py @@ -32,7 +32,7 @@ ) # --------------------------------------------------------------------------- -# Helpers — minimal valid kwargs per typed subclass +# Helpers - minimal valid kwargs per typed subclass # --------------------------------------------------------------------------- diff --git a/tests/unit/test_pipeline_confidence.py b/tests/unit/test_pipeline_confidence.py index 3dcf20e..d31c3a4 100644 --- a/tests/unit/test_pipeline_confidence.py +++ b/tests/unit/test_pipeline_confidence.py @@ -19,10 +19,6 @@ def _pin_legacy_extract_dedup(monkeypatch): "azure.cosmos.agent_memory.thresholds.get_dedup_vector_enabled", lambda: False, ) - monkeypatch.setattr( - "azure.cosmos.agent_memory.thresholds.get_dedup_context_vector_enabled", - lambda: False, - ) def _make_pipeline(llm_response: dict): @@ -115,35 +111,6 @@ def test_extract_defaults_confidence_to_half_when_missing(): assert doc["confidence"] == 0.5, f"missing default for {doc['type']} {doc['id']}" -def test_extract_routes_unclassified_to_fact_with_tag(): - pipeline, upserted = _make_pipeline( - { - "unclassified": [ - { - "text": "Weird ambiguous thing about the user", - "confidence": 0.45, - "salience": 0.4, - "tags": ["ambig"], - "reason": "could be fact or episodic", - } - ] - } - ) - - result = pipeline.extract_memories("u1", "t1") - - assert len(upserted) == 1 - doc = upserted[0] - assert doc["type"] == "fact" - assert "sys:unclassified" in doc["tags"] - assert "sys:fact" in doc["tags"] - assert "topic:ambig" in doc["tags"] - assert doc["confidence"] == pytest.approx(0.45) - assert doc["metadata"]["unclassified_reason"] == "could be fact or episodic" - assert result["unclassified_count"] == 1 - assert result["fact_count"] == 0 - - def test_extract_episodic_carries_confidence(): pipeline, upserted = _make_pipeline( { diff --git a/tests/unit/test_procedural_synthesis.py b/tests/unit/test_procedural_synthesis.py index 5930104..533f090 100644 --- a/tests/unit/test_procedural_synthesis.py +++ b/tests/unit/test_procedural_synthesis.py @@ -21,10 +21,6 @@ def _pin_legacy_extract_dedup(monkeypatch): "azure.cosmos.agent_memory.thresholds.get_dedup_vector_enabled", lambda: False, ) - monkeypatch.setattr( - "azure.cosmos.agent_memory.thresholds.get_dedup_context_vector_enabled", - lambda: False, - ) def _assert_iso8601(text: str) -> None: @@ -37,7 +33,7 @@ def _capture_upserts(): The pipeline writes new facts/episodics via ``create_item`` (for 409 idempotency) and writes new procedural versions via ``create_item + - bump-seq-retry``. Either way, tests want the persisted body — so the + bump-seq-retry``. Either way, tests want the persisted body - so the helper now wires the same capture to both side_effects. """ upserted: list[dict] = [] @@ -318,7 +314,6 @@ def test_synthesize_procedural_only_touches_memories_container(): [], [_fact_doc("f1", "Always use bullet points.", category="preference", salience=0.95)], [_episodic_doc("e1", lesson="Keep examples small.")], - [], ] memories_container.create_item.side_effect = lambda body: body containers = { @@ -334,7 +329,7 @@ def test_synthesize_procedural_only_touches_memories_container(): result = pipeline.synthesize_procedural("u1") assert result["status"] == "synthesized" - assert memories_container.query_items.call_count == 4 + assert memories_container.query_items.call_count == 3 memories_container.create_item.assert_called_once() turns_container.method_calls == [] summaries_container.method_calls == [] @@ -462,7 +457,7 @@ def test_synthesize_procedural_short_circuits_on_second_call_with_tied_salience( Before the SQL composite ORDER BY (salience DESC, created_at ASC, id ASC), Cosmos returned an arbitrary 50-of-60 per query; the prior/current source sets never matched and the LLM fired on every reconcile. The Python re-sort - that previously masked this in unit tests has been removed — the SQL itself + that previously masked this in unit tests has been removed - the SQL itself must be deterministic. """ fact_docs = [ @@ -647,7 +642,6 @@ def test_synthesize_procedural_retries_with_fresh_llm_call_when_winner_has_parti [prior_v1], fact_docs, [], - [], [ prior_v1, _procedural_doc( @@ -705,7 +699,7 @@ def _create(*, body): def test_synthesize_procedural_short_circuits_when_race_winner_covers_loser_sources(): """Common race case: both writers process the same source set. After 409, loser re-reads, finds winner's source_ids ⊇ loser's, and returns the - winner's doc as ``unchanged`` — no second LLM call, no wasted write. + winner's doc as ``unchanged`` - no second LLM call, no wasted write. """ from azure.cosmos.exceptions import CosmosResourceExistsError @@ -736,7 +730,6 @@ def test_synthesize_procedural_short_circuits_when_race_winner_covers_loser_sour [prior_v1], fact_docs, [], - [], [prior_v1, winner_v2], ] upserted, capture = _capture_upserts() diff --git a/tests/unit/test_process_now.py b/tests/unit/test_process_now.py index 3d71c56..11132bd 100644 --- a/tests/unit/test_process_now.py +++ b/tests/unit/test_process_now.py @@ -53,8 +53,8 @@ def test_process_now_with_inprocess_invokes_full_pipeline(): assert pipeline.reconcile_memories.call_count == 2 pipeline.reconcile_memories.assert_has_calls( [ - call("u1", n=50, memory_type="fact", full_rebuild=False), - call("u1", n=50, memory_type="episodic", full_rebuild=False), + call("u1", n=50, memory_type="fact"), + call("u1", n=50, memory_type="episodic"), ] ) pipeline.synthesize_procedural.assert_called_once_with(user_id="u1", force=False) @@ -109,7 +109,7 @@ def test_process_now_swallows_user_summary_failure(): def test_process_now_swallows_transient_http_error_by_status_code(): """Cosmos / HTTP exceptions with transient status codes (429, 503) must be - swallowed — they're infrastructure hiccups, not bugs.""" + swallowed - they're infrastructure hiccups, not bugs.""" class _FakeHttpExc(Exception): def __init__(self, status_code): @@ -136,7 +136,7 @@ def __init__(self, status_code): def test_process_now_propagates_permanent_procedural_failure(): """A non-transient failure (e.g. ``KeyError`` from a schema bug) must NOT - be silently swallowed — it should surface to the caller so config / + be silently swallowed - it should surface to the caller so config / programmer bugs do not turn into invisible ``WARNING`` lines.""" client = _connected() pipeline = MagicMock() @@ -157,7 +157,7 @@ def test_process_now_propagates_permanent_procedural_failure(): def test_process_now_propagates_permanent_user_summary_failure(): """ValidationError from generate_user_summary (e.g. schema bug) must - surface to the caller — silencing config bugs is a bug.""" + surface to the caller - silencing config bugs is a bug.""" client = _connected() pipeline = MagicMock() pipeline.generate_thread_summary.return_value = {"id": "s"} @@ -175,7 +175,7 @@ def test_process_now_propagates_permanent_user_summary_failure(): def test_process_now_with_durable_skips_tail_steps(): - """Durable mode must NOT call synthesize_procedural or generate_user_summary — + """Durable mode must NOT call synthesize_procedural or generate_user_summary - those are driven by the change-feed-fed sibling Function app.""" client = _connected(processor=DurableFunctionProcessor()) pipeline = MagicMock() @@ -256,7 +256,7 @@ def test_process_now_and_wait_durable_swallows_search_errors_until_timeout(): def test_process_now_and_wait_durable_propagates_non_cosmos_errors(): - """Non-Cosmos errors must NOT be silently swallowed in the polling loop — + """Non-Cosmos errors must NOT be silently swallowed in the polling loop - operators would otherwise wait the full timeout with no signal.""" client = _connected(processor=DurableFunctionProcessor()) _patch_get_thread(client, []) diff --git a/tests/unit/test_reconcile.py b/tests/unit/test_reconcile.py index f0a1c32..9a5351b 100644 --- a/tests/unit/test_reconcile.py +++ b/tests/unit/test_reconcile.py @@ -31,19 +31,12 @@ @pytest.fixture(autouse=True) def _pin_legacy_dedup_paths(monkeypatch): - """These tests cover the pre-candidate reconcile/extract code paths.""" - monkeypatch.setattr( - "azure.cosmos.agent_memory.thresholds.get_dedup_reconcile_mode", - lambda: "full_pool", - ) + """Disable write-time in-place folding so the extract-path tests here + exercise the plain ADD path deterministically.""" monkeypatch.setattr( "azure.cosmos.agent_memory.thresholds.get_dedup_vector_enabled", lambda: False, ) - monkeypatch.setattr( - "azure.cosmos.agent_memory.thresholds.get_dedup_context_vector_enabled", - lambda: False, - ) def _make_pipeline() -> PipelineService: @@ -157,143 +150,6 @@ def test_validates_n(self): with pytest.raises(ValidationError, match="<= 500"): p.reconcile_memories("u1", n=501) - def test_merge_falls_back_to_max_source_confidence_salience_when_llm_omits(self): - """If the LLM omits or zero-fills confidence/salience on a duplicate - group, the merged record must inherit max(source_*) so it doesn't - silently drop below ``min_confidence`` / ``min_salience`` filters.""" - p = _make_pipeline() - facts = [ - _fact("f1", "User likes aisle seats", confidence=0.92, salience=0.7), - _fact("f2", "User prefers aisle seats on flights", confidence=0.85, salience=0.65), - ] - p._container.query_items.return_value = iter(facts) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "duplicate_groups": [ - { - "merged_content": "User prefers aisle seats", - "source_ids": ["f1", "f2"], - # confidence/salience deliberately omitted - } - ], - "contradicted_pairs": [], - "kept_ids": [], - } - ) - ) - p._upsert_memory = MagicMock() - p._mark_superseded = MagicMock(return_value=True) - p._embeddings = MagicMock() - p._embeddings.generate.return_value = [0.0] - p.reconcile_memories("u1") - merged_doc = p._upsert_memory.call_args.args[0] - assert merged_doc["confidence"] == 0.92 # max of 0.92, 0.85 - assert merged_doc["salience"] == 0.7 # max of 0.7, 0.65 - - def test_merge_treats_zero_confidence_as_omitted(self): - """gpt-4o-mini sometimes echoes the literal placeholder. A zero - confidence/salience on the LLM output must trigger the same - max(source_*) fallback as omission.""" - p = _make_pipeline() - facts = [ - _fact("f1", "User likes coffee", confidence=0.9, salience=0.8), - _fact("f2", "User loves coffee in the mornings", confidence=0.95, salience=0.85), - ] - p._container.query_items.return_value = iter(facts) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "duplicate_groups": [ - { - "merged_content": "User loves coffee in the mornings", - "source_ids": ["f1", "f2"], - "confidence": 0.0, - "salience": 0.0, - } - ], - "contradicted_pairs": [], - "kept_ids": [], - } - ) - ) - p._upsert_memory = MagicMock() - p._mark_superseded = MagicMock(return_value=True) - p._embeddings = MagicMock() - p._embeddings.generate.return_value = [0.0] - p.reconcile_memories("u1") - merged_doc = p._upsert_memory.call_args.args[0] - assert merged_doc["confidence"] == 0.95 - assert merged_doc["salience"] == 0.85 - - def test_merged_doc_thread_id_picked_from_newest_source_by_ts(self): - """Merged record's thread_id (partition key) must come from the - source with the highest Cosmos ``_ts``, independent of the order - the LLM lists ``source_ids`` in.""" - p = _make_pipeline() - # f-old has lower _ts; f-new has higher. LLM lists them in the - # "wrong" order (old first) — pipeline must still pick f-new's - # thread_id for the merged doc. - f_old = _fact("f-old", "User likes coffee", thread_id="thread-old") - f_old["_ts"] = 100 - f_new = _fact("f-new", "User loves coffee", thread_id="thread-new") - f_new["_ts"] = 999 - p._container.query_items.return_value = iter([f_new, f_old]) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "duplicate_groups": [ - { - "merged_content": "User loves coffee", - "source_ids": ["f-old", "f-new"], - "confidence": 0.9, - "salience": 0.7, - } - ], - "contradicted_pairs": [], - "kept_ids": [], - } - ) - ) - p._upsert_memory = MagicMock() - p._mark_superseded = MagicMock(return_value=True) - p._embeddings = MagicMock() - p._embeddings.generate.return_value = [0.0] - p.reconcile_memories("u1") - merged_doc = p._upsert_memory.call_args.args[0] - assert merged_doc["thread_id"] == "thread-new" - - def test_synthetic_partition_is_user_scoped_when_no_thread_id(self): - """If every source somehow lacks a thread_id, the synthetic - fallback must be scoped per-user to avoid cross-tenant collisions.""" - p = _make_pipeline() - f1 = _fact("f1", "alpha", thread_id="") - f2 = _fact("f2", "alpha-restated", thread_id="") - p._container.query_items.return_value = iter([f1, f2]) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "duplicate_groups": [ - { - "merged_content": "alpha", - "source_ids": ["f1", "f2"], - "confidence": 0.9, - "salience": 0.7, - } - ], - "contradicted_pairs": [], - "kept_ids": [], - } - ) - ) - p._upsert_memory = MagicMock() - p._mark_superseded = MagicMock(return_value=True) - p._embeddings = MagicMock() - p._embeddings.generate.return_value = [0.0] - p.reconcile_memories("user-abc") - merged_doc = p._upsert_memory.call_args.args[0] - assert merged_doc["thread_id"] == "__reconciled__:user-abc" - def test_empty_pool(self): p = _make_pipeline() p._container.query_items.return_value = iter([]) @@ -310,43 +166,6 @@ def test_single_fact_no_op(self): assert result == {"kept": 1, "merged": 0, "contradicted": 0} p._run_prompty.assert_not_called() - def test_only_duplicates(self): - p = _make_pipeline() - facts = [ - _fact("f1", "User prefers aisle seats on flights"), - _fact("f2", "User likes aisle seats when flying"), - ] - p._container.query_items.return_value = iter(facts) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "duplicate_groups": [ - { - "merged_content": "User prefers aisle seats on flights", - "source_ids": ["f1", "f2"], - "confidence": 0.95, - "salience": 0.7, - } - ], - "contradicted_pairs": [], - "kept_ids": [], - } - ) - ) - - result = p.reconcile_memories("u1") - - assert result == {"kept": 0, "merged": 2, "contradicted": 0} - # merged doc upserted - assert p._upsert_memory.call_count == 1 - merged_doc = p._upsert_memory.call_args.args[0] - assert merged_doc["content"] == "User prefers aisle seats on flights" - assert "f1" in merged_doc["supersedes_ids"] and "f2" in merged_doc["supersedes_ids"] - # both sources marked superseded with reason=duplicate - assert p._mark_superseded.call_count == 2 - for call in p._mark_superseded.call_args_list: - assert call.kwargs["reason"] == "duplicate" - def test_only_contradictions(self): p = _make_pipeline() facts = [ @@ -375,104 +194,6 @@ def test_only_contradictions(self): assert call.args[1] == "f2" assert call.kwargs["reason"] == "contradict" - def test_mixed_pool_with_dangling_resolution(self): - """Loser of a contradiction is also a duplicate source. - - Pipeline must redirect the contradiction's ``loser_id`` through - ``source_to_merged_id`` and supersede the *merged* doc, not the - original (already-merged) source. - """ - p = _make_pipeline() - facts = [ - _fact("f1", "User prefers aisle seats on flights"), - _fact("f2", "User likes aisle seats when flying"), - _fact("f3", "User loves the window seat"), - ] - p._container.query_items.return_value = iter(facts) - - # The dangling-loser redirect resolves through the in-memory - # ``merged_docs_by_id`` cache populated when the merged doc is - # upserted — no second Cosmos query is issued. The upsert - # response carries the ``_etag`` that flows through the cache. - def upsert(doc): - snap = dict(doc) - snap["_etag"] = "merged-etag" - return snap - - p._upsert_memory.side_effect = upsert - - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "duplicate_groups": [ - { - "merged_content": "User prefers aisle seats on flights", - "source_ids": ["f1", "f2"], - "confidence": 0.95, - "salience": 0.7, - } - ], - "contradicted_pairs": [ - # Loser f2 was just merged into the merged doc; - # winner f3 must contradict the *merged* doc. - {"winner_id": "f3", "loser_id": "f2", "reason": "contradicts merged"} - ], - "kept_ids": ["f3"], - } - ) - ) - - result = p.reconcile_memories("u1") - - assert result["merged"] == 2 # f1, f2 marked dup - assert result["contradicted"] == 1 # merged doc marked contradiction - # mark_superseded calls: f1 (dup), f2 (dup), then merged doc (contradiction) - assert p._mark_superseded.call_count == 3 - last = p._mark_superseded.call_args_list[-1] - # The third call should target the merged doc (fetched via resolver) - # and use the merged record's id as the new winner id only if winner - # also collapsed; here winner=f3 stays as-is. - assert last.kwargs["reason"] == "contradict" - # winner remains f3 (was not in any dup group) - assert last.args[1] == "f3" - - def test_dangling_collapses_to_no_op(self): - """Both winner and loser absorbed into the same merged group → skip.""" - p = _make_pipeline() - facts = [ - _fact("f1", "User likes coffee"), - _fact("f2", "User likes coffee in the morning"), - ] - p._container.query_items.return_value = iter(facts) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "duplicate_groups": [ - { - "merged_content": "User likes coffee in the morning", - "source_ids": ["f1", "f2"], - "confidence": 0.9, - "salience": 0.6, - } - ], - "contradicted_pairs": [ - # Both f1 and f2 collapse to the same merged id → skip - {"winner_id": "f1", "loser_id": "f2", "reason": "irrelevant"} - ], - "kept_ids": [], - } - ) - ) - - result = p.reconcile_memories("u1") - - assert result["merged"] == 2 - assert result["contradicted"] == 0 - # No contradiction supersede call beyond the two duplicate marks - assert p._mark_superseded.call_count == 2 - for call in p._mark_superseded.call_args_list: - assert call.kwargs["reason"] == "duplicate" - def test_n_cap_honored(self): """Custom ``n`` is interpolated into the SQL query's TOP clause.""" p = _make_pipeline() @@ -618,7 +339,7 @@ def _build(self) -> PipelineService: def test_fact_not_dropped_when_only_procedural_has_same_hash(self): p = self._build() text = "Always reply in Spanish" - # Existing PROCEDURAL with that text — must NOT poison the FACT bucket. + # Existing PROCEDURAL with that text - must NOT poison the FACT bucket. existing = [ { "id": "proc_existing", @@ -690,52 +411,9 @@ def test_empty_thread_returns_full_dict_shape(self): assert out[key] == 0 -class TestReconcileEmbeddingFailureAborts: - """If embedding generation fails for the merged content, the duplicate - group must be aborted entirely — no upsert, no supersede — so we don't - create a search-index hole.""" - - def test_embedding_failure_skips_upsert_and_supersede(self): - p = PipelineService.__new__(PipelineService) - p._container = MagicMock() - p._memories_container = p._container - p._turns_container = p._container - p._summaries_container = p._container - facts = [ - _fact("f1", "alpha"), - _fact("f2", "alpha-restated"), - ] - p._container.query_items.return_value = iter(facts) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "duplicate_groups": [ - { - "merged_content": "alpha (consolidated)", - "source_ids": ["f1", "f2"], - "confidence": 0.9, - "salience": 0.7, - } - ], - "contradicted_pairs": [], - "kept_ids": [], - } - ) - ) - p._upsert_memory = MagicMock() - p._mark_superseded = MagicMock(return_value=True) - p._embeddings = MagicMock() - p._embeddings.generate.side_effect = RuntimeError("rate limit") - result = p.reconcile_memories("u1") - # Embedding failed → abort: nothing upserted, nothing superseded. - p._upsert_memory.assert_not_called() - p._mark_superseded.assert_not_called() - assert result == {"kept": 2, "merged": 0, "contradicted": 0} - - class TestReconcileSupersedeRaceCounting: """When ``_mark_superseded`` returns False (lost ETag race), the source - must NOT be added to ``source_to_merged_id`` or counted as consumed — + must NOT be added to ``source_to_merged_id`` or counted as consumed - otherwise contradictions get redirected to a doc that doesn't claim the source, and ``kept`` undercounts.""" @@ -778,7 +456,7 @@ def test_failed_supersede_does_not_consume_source(self): class TestReconcileWinnerValidation: - """Hallucinated ``winner_id`` must be refused — never write a dangling + """Hallucinated ``winner_id`` must be refused - never write a dangling ``superseded_by`` that breaks the audit trail.""" def test_hallucinated_winner_id_skipped(self): @@ -815,89 +493,6 @@ def test_hallucinated_winner_id_skipped(self): p._mark_superseded.assert_not_called() assert result == {"kept": 2, "merged": 0, "contradicted": 0} - def test_resolved_winner_via_merge_redirect_is_accepted(self): - """If winner_id refers to a fact that was just absorbed into a - duplicate group, the merged_id must satisfy the validation.""" - p = PipelineService.__new__(PipelineService) - p._container = MagicMock() - p._memories_container = p._container - p._turns_container = p._container - p._summaries_container = p._container - facts = [ - _fact("f1", "alpha"), - _fact("f2", "alpha-paraphrased"), - _fact("f3", "contradicts alpha"), - ] - p._container.query_items.return_value = iter(facts) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "duplicate_groups": [ - { - "merged_content": "alpha (consolidated)", - "source_ids": ["f1", "f2"], - "confidence": 0.9, - "salience": 0.7, - } - ], - "contradicted_pairs": [ - # winner_id=f1 was absorbed into the merged group; - # redirect must resolve and validate cleanly. - {"winner_id": "f1", "loser_id": "f3", "reason": "x"}, - ], - "kept_ids": [], - } - ) - ) - p._upsert_memory = MagicMock() - p._mark_superseded = MagicMock(return_value=True) - p._embeddings = MagicMock() - p._embeddings.generate.return_value = [0.0] - result = p.reconcile_memories("u1") - assert result["contradicted"] == 1 - - -class TestReconcileBoolNotNumeric: - """``True`` and ``False`` are instances of ``int`` in Python — they must - NOT be treated as numeric LLM-supplied confidence/salience.""" - - def test_bool_confidence_falls_back_to_max_source(self): - p = PipelineService.__new__(PipelineService) - p._container = MagicMock() - p._memories_container = p._container - p._turns_container = p._container - p._summaries_container = p._container - facts = [ - _fact("f1", "alpha", confidence=0.7, salience=0.5), - _fact("f2", "alpha-restated", confidence=0.85, salience=0.6), - ] - p._container.query_items.return_value = iter(facts) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "duplicate_groups": [ - { - "merged_content": "alpha", - "source_ids": ["f1", "f2"], - "confidence": True, # JSON boolean — must be ignored - "salience": False, - } - ], - "contradicted_pairs": [], - "kept_ids": [], - } - ) - ) - p._upsert_memory = MagicMock() - p._mark_superseded = MagicMock(return_value=True) - p._embeddings = MagicMock() - p._embeddings.generate.return_value = [0.0] - p.reconcile_memories("u1") - merged_doc = p._upsert_memory.call_args.args[0] - # NOT 1.0 (True coerced) and NOT 0.0 (False coerced); fallback to max source. - assert merged_doc["confidence"] == 0.85 - assert merged_doc["salience"] == 0.6 - class TestReconcileFactsTextEscapesContent: """Content with ``"`` or ``|`` must not break the prompt grammar.""" @@ -924,7 +519,7 @@ def _capture(name, inputs): p._mark_superseded = MagicMock(return_value=True) p._embeddings = MagicMock() p.reconcile_memories("u1") - # The embedded `"` must be escaped (\\") — not raw — and the + # The embedded `"` must be escaped (\\") - not raw - and the # Content: field must remain JSON-quoted so the LLM can parse it # as a single string even though the original text contained ``|``. text = captured["facts_text"] @@ -964,183 +559,8 @@ def test_pool_size_zero_falls_back_to_default(self, monkeypatch): assert get_dedup_pool_size() == DEFAULT_DEDUP_POOL_SIZE -class TestReconcileMergedIdDeterministic: - """RD#1: merged_id is deterministic on (user, content_hash) so cycles are idempotent.""" - - def test_same_merged_content_yields_same_id_across_runs(self): - import hashlib - - p = _make_pipeline() - upserts: list[dict] = [] - p._upsert_memory = MagicMock(side_effect=lambda doc: upserts.append(doc) or doc) - p._mark_superseded = MagicMock(return_value=True) - p._container.query_items.return_value = iter( - [_fact("a1", "User likes coffee"), _fact("a2", "User enjoys coffee")] - ) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "duplicate_groups": [ - { - "merged_content": "User likes coffee", - "source_ids": ["a1", "a2"], - } - ], - "contradicted_pairs": [], - "kept_ids": [], - } - ) - ) - p.reconcile_memories("u1") - first_id = upserts[0]["id"] - # Predict id from public formula: - ch = compute_content_hash("User likes coffee") - from azure.cosmos.agent_memory.services.pipeline import _ID_SEED_SEP - - seed = _ID_SEED_SEP.join(("u1", "merged", ch)) - expected = "fact_" + hashlib.sha256(seed.encode()).hexdigest()[:32] - assert first_id == expected - # Second run: different source ids, identical canonical merged - # content → identical merged id (idempotent upsert). - upserts.clear() - p._container.query_items.return_value = iter( - [_fact("b1", "User likes coffee"), _fact("b2", "user LIKES coffee")] - ) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "duplicate_groups": [ - { - "merged_content": "User likes coffee", - "source_ids": ["b1", "b2"], - } - ], - "contradicted_pairs": [], - "kept_ids": [], - } - ) - ) - p.reconcile_memories("u1") - assert upserts[0]["id"] == first_id - - -class TestReconcileSupersedesIdsFiltered: - """RD#4: hallucinated source_ids are scrubbed from supersedes_ids.""" - - def test_hallucinated_source_ids_filtered(self): - p = _make_pipeline() - upserts: list[dict] = [] - p._upsert_memory = MagicMock(side_effect=lambda doc: upserts.append(doc) or doc) - p._mark_superseded = MagicMock(return_value=True) - p._container.query_items.return_value = iter([_fact("real1", "X"), _fact("real2", "Y")]) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "duplicate_groups": [ - { - "merged_content": "X and Y", - "source_ids": ["real1", "real2", "ghost_id_404"], - } - ], - "contradicted_pairs": [], - "kept_ids": [], - } - ) - ) - p.reconcile_memories("u1") - assert upserts, "merged doc should have been upserted" - assert "ghost_id_404" not in upserts[0]["supersedes_ids"] - assert set(upserts[0]["supersedes_ids"]) == {"real1", "real2"} - - -class TestReconcileTransitiveSupersedes: - """RD#1 follow-on: prior chain hops survive into the new merged record.""" - - def test_prior_supersedes_ids_preserved_in_chain(self): - p = _make_pipeline() - upserts: list[dict] = [] - p._upsert_memory = MagicMock(side_effect=lambda doc: upserts.append(doc) or doc) - p._mark_superseded = MagicMock(return_value=True) - # f1 was itself a previously-merged record carrying its own provenance. - f1 = _fact("f1", "X v1") - f1["supersedes_ids"] = ["older_a", "older_b"] - f2 = _fact("f2", "X v2") - p._container.query_items.return_value = iter([f1, f2]) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "duplicate_groups": [{"merged_content": "X canonical", "source_ids": ["f1", "f2"]}], - "contradicted_pairs": [], - "kept_ids": [], - } - ) - ) - p.reconcile_memories("u1") - sup = upserts[0]["supersedes_ids"] - assert "f1" in sup and "f2" in sup - assert "older_a" in sup and "older_b" in sup - - -class TestReconcileMergedMetadata: - """RD#3: merged docs carry a positive merged_via signal.""" - - def test_merged_doc_has_metadata(self): - p = _make_pipeline() - upserts: list[dict] = [] - p._upsert_memory = MagicMock(side_effect=lambda doc: upserts.append(doc) or doc) - p._mark_superseded = MagicMock(return_value=True) - p._container.query_items.return_value = iter([_fact("x1", "A"), _fact("x2", "B")]) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "duplicate_groups": [{"merged_content": "A and B", "source_ids": ["x1", "x2"]}], - "contradicted_pairs": [], - "kept_ids": [], - } - ) - ) - p.reconcile_memories("u1") - meta = upserts[0].get("metadata") or {} - assert meta.get("merged_via") == "reconcile" - assert meta.get("merged_from_count") == 2 - - -class TestReconcileNoOrphanDeleteOnRace: - """When ALL supersede attempts lose the ETag race, the merged doc must - NOT be deleted: deleting it would orphan any sources whose - ``superseded_by`` was already pointed at this deterministic merged id - by the concurrent-reconcile winner — those sources would become - invisible to default reads (filter ``superseded_by IS NULL``) and to - the reconcile pool, causing permanent data loss. The merged doc is - idempotent (deterministic id) so leaving it in place is consistent.""" - - def test_orphan_merged_doc_is_not_deleted_when_no_supersede_succeeds(self): - p = _make_pipeline() - p._upsert_memory = MagicMock(side_effect=lambda doc: doc) - # Every supersede attempt loses the race → without the fix, the - # merged doc would be hard-deleted. With the fix, the merged doc - # stays as-is and the loss path is logged at INFO. - p._mark_superseded = MagicMock(return_value=False) - p._container.query_items.return_value = iter([_fact("o1", "A"), _fact("o2", "B")]) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "duplicate_groups": [{"merged_content": "A&B", "source_ids": ["o1", "o2"]}], - "contradicted_pairs": [], - "kept_ids": [], - } - ) - ) - result = p.reconcile_memories("u1") - # delete_item must NOT be called — orphan-delete path was the - # data-loss bug fixed in this round. - assert not p._container.delete_item.called - # No facts merged (no supersede succeeded). - assert result["merged"] == 0 - - class TestReconcileContradictionWinnerNotInKeptIds: - """RD#5+#13: contradiction winners are absent from kept_ids — must NOT trigger warning.""" + """RD#5+#13: contradiction winners are absent from kept_ids - must NOT trigger warning.""" def test_clean_contradiction_does_not_warn_about_kept_mismatch(self, caplog): import logging @@ -1180,132 +600,6 @@ def test_query_uses_is_null(self): assert "c.superseded_by = null" not in sql -class TestReconcileClampsConfidenceAndSalience: - """LLM emitting values outside (0, 1] (e.g. 1.05 from a model that - confused percent with [0,1]) must NOT propagate to MemoryRecord — the - Pydantic validator would reject and the blanket except in reconcile - would silently drop the entire merge group. Out-of-range values must - fall back to ``max(source.*)``.""" - - def test_out_of_range_confidence_falls_back_to_source_max(self): - p = _make_pipeline() - upserts: list[dict] = [] - p._upsert_memory = MagicMock(side_effect=lambda doc: upserts.append(doc) or doc) - p._container.query_items.return_value = iter( - [ - _fact("f1", "User likes coffee", confidence=0.7, salience=0.6), - _fact("f2", "User enjoys coffee", confidence=0.85, salience=0.8), - ] - ) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "duplicate_groups": [ - { - "merged_content": "User likes coffee", - "source_ids": ["f1", "f2"], - "confidence": 1.05, - "salience": 1.5, - } - ], - "contradicted_pairs": [], - "kept_ids": [], - } - ) - ) - result = p.reconcile_memories("u1") - assert result["merged"] == 2 - assert len(upserts) == 1 - assert upserts[0]["confidence"] == pytest.approx(0.85) - assert upserts[0]["salience"] == pytest.approx(0.8) - - -class TestReconcileEtagFlowsThroughMergedDocCache: - """``_upsert_memory`` must return the response (which carries the - fresh ``_etag``) so the in-memory ``merged_docs_by_id`` cache can - feed it into a downstream supersede on the contradiction-redirect - path. Without it, ``_mark_superseded`` falls through to - ``upsert_item`` with no concurrency protection. - - This test exercises the full path: a duplicate group folds f1 + f2 - into a fresh merged doc M, then a contradiction names f2 as the - loser. The pipeline must redirect the contradiction to M and call - ``_mark_superseded(M, ...)`` with M carrying the etag returned from - the upsert. - """ - - def test_supersede_on_merged_doc_receives_doc_with_etag(self): - p = _make_pipeline() - - def upsert_response(doc): - response = dict(doc) - response["_etag"] = "etag-from-cosmos" - return response - - p._upsert_memory = MagicMock(side_effect=upsert_response) - p._mark_superseded = MagicMock(return_value=True) - p._container.query_items.return_value = iter( - [ - _fact("f1", "User likes tea"), - _fact("f2", "User enjoys tea"), - _fact("f3", "User hates all hot drinks"), - ] - ) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "duplicate_groups": [{"merged_content": "User likes tea", "source_ids": ["f1", "f2"]}], - "contradicted_pairs": [{"winner_id": "f3", "loser_id": "f2", "reason": "contradicts merged"}], - "kept_ids": ["f3"], - } - ) - ) - p.reconcile_memories("u1") - - # The third call to _mark_superseded targets the merged doc - # (loser f2 was redirected through merged_docs_by_id). That - # doc must carry the _etag returned by _upsert_memory — proving - # the response actually flowed through the cache and not just - # the locally-built dict. - assert p._mark_superseded.call_count == 3 - contradiction_call = p._mark_superseded.call_args_list[-1] - merged_doc_passed = contradiction_call.args[0] - assert merged_doc_passed.get("_etag") == "etag-from-cosmos" - - -class TestReconcileSkipsSingleSourceDuplicateGroup: - """A `duplicate_group` with only one valid `source_id` is a no-op - masquerading as a merge — it would supersede a single fact with a - near-identical clone (extra row, no signal) and could redirect a - later contradiction's loser_id onto a merged doc that represents - nothing real. Skip such groups.""" - - def test_single_source_duplicate_group_does_not_create_merged_doc(self): - p = _make_pipeline() - p._upsert_memory = MagicMock(side_effect=lambda d: dict(d, _etag="e")) - p._mark_superseded = MagicMock(return_value=True) - p._container.query_items.return_value = iter([_fact("f1", "User likes tea"), _fact("f2", "User likes coffee")]) - p._run_prompty = MagicMock( - return_value=json.dumps( - { - "duplicate_groups": [ - {"merged_content": "User likes tea", "source_ids": ["f1"]}, - ], - "contradicted_pairs": [], - "kept_ids": ["f1", "f2"], - } - ) - ) - - out = p.reconcile_memories("u1") - - # No merged record created; no source superseded as a duplicate. - merged_upserts = [c for c in p._upsert_memory.call_args_list if c.args[0].get("type") == "fact"] - assert merged_upserts == [] - assert p._mark_superseded.call_count == 0 - assert out["merged"] == 0 - - class TestFactsTextHandlesNullConfidence: """Pool facts with ``confidence=None`` / ``salience=None`` (legacy docs from before these fields existed) must render as ``N/A`` in the @@ -1342,7 +636,7 @@ def capture_prompty(name, inputs): class TestExtractUpdateSelfCollapseGuard: """Procedural synthesis self-collapse guard: when synthesis would emit a proc id identical to the existing one, treat as a no-op. (Fact extract-time - UPDATE was removed — facts/contradictions are reconciled, not extract-tagged.)""" + UPDATE was removed - facts/contradictions are reconciled, not extract-tagged.)""" def _build(self) -> PipelineService: p = PipelineService.__new__(PipelineService) diff --git a/tests/unit/test_thresholds.py b/tests/unit/test_thresholds.py index d7f9d3d..d867d71 100644 --- a/tests/unit/test_thresholds.py +++ b/tests/unit/test_thresholds.py @@ -178,22 +178,10 @@ def test_processor_owner_invalid_uses_default(monkeypatch: pytest.MonkeyPatch) - def test_internalized_getters_return_fixed_constants_and_ignore_env(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("DEDUP_CONTEXT_VECTOR_ENABLED", "false") - monkeypatch.setenv("DEDUP_CONTEXT_TOPK", "7") + monkeypatch.setenv("EXTRACTION_BATCH_MAX_TOKENS", "999") monkeypatch.setenv("DEDUP_VECTOR_ENABLED", "false") monkeypatch.setenv("DEDUP_SIM_HIGH", "0.50") - monkeypatch.setenv("DEDUP_SIM_LOW", "0.40") - monkeypatch.setenv("DEDUP_CANDIDATE_TOPK", "8") - monkeypatch.setenv("DEDUP_RECONCILE_MODE", "full_pool") - monkeypatch.setenv("DEDUP_CLUSTER_SIM", "0.10") - monkeypatch.setenv("DEDUP_FULL_RECLUSTER_EVERY_N", "4") - - assert thresholds.get_dedup_context_vector_enabled() is True - assert thresholds.get_dedup_context_topk() == 10 + + assert thresholds.get_extraction_batch_max_tokens() == 7000 assert thresholds.get_dedup_vector_enabled() is True assert thresholds.get_dedup_sim_high() == 0.97 - assert thresholds.get_dedup_sim_low() == 0.80 - assert thresholds.get_dedup_candidate_topk() == 10 - assert thresholds.get_dedup_reconcile_mode() == "candidate" - assert thresholds.get_dedup_cluster_sim() == 0.60 - assert thresholds.get_dedup_full_recluster_every_n() == 12 From 4e6a491a03c3306e96fcd84b5f94f49f1325dece Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Fri, 10 Jul 2026 23:06:54 -0700 Subject: [PATCH 6/8] Resolving comments --- CHANGELOG.md | 21 +-- .../agent_memory/aio/services/pipeline.py | 52 +++++- azure/cosmos/agent_memory/prompts/_schemas.py | 29 ---- .../prompts/dedup_episodic.prompty | 160 ------------------ .../cosmos/agent_memory/services/pipeline.py | 75 +++++++- .../aio/services/test_dedup_vector_async.py | 21 ++- tests/unit/services/test_dedup_vector.py | 45 ++++- 7 files changed, 187 insertions(+), 216 deletions(-) delete mode 100644 azure/cosmos/agent_memory/prompts/dedup_episodic.prompty diff --git a/CHANGELOG.md b/CHANGELOG.md index 1fdb89e..35cc968 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,22 +3,15 @@ ## [0.3.0b1] (Unreleased) #### Features Added -* Contradiction-aware vector dedup and reconciliation. Extraction now runs a - vector-similarity dedup ladder: near-exact duplicates are auto-dropped and - borderline matches are tagged for review before they are written. A periodic - LLM reconcile pass (driven by `dedup.prompty` for facts and - `dedup_episodic.prompty` for episodic memories) then merges duplicate groups - and resolves contradictions, soft-deleting the losers with a supersede reason. - Reconciliation is distance-function aware (the destructive near-exact - auto-drop is disabled for `euclidean` containers). In-process reconcile now - covers both facts and episodic memories, matching the Durable backend, and its - cadence is derived from the persisted message counter. See [PR:#26](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/26) +* Write-time in-place deduplication: near-duplicate memories fold into the existing record (same id, newer content) instead of creating a new doc. See [PR:#26](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/26) +* Reconciliation now resolves contradictions only, soft-deleting the loser with `superseded_by`; `get_memory_history()` walks that chain. See [PR:#26](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/26) + +#### Bugs Fixed +* Fixed a re-extraction loop that re-extracted the whole conversation every cycle (turns were never stamped `extracted_at` when vector dedup was on). See [PR:#26](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/26) #### Other Changes -* `search_cosmos` and `search_turns` now always fuse vector similarity with - BM25 (full-text) ranking, falling back to vector-only for all-stopword - queries. The `hybrid_search` flag has been removed — hybrid ranking is the - default and requires no opt-in. See [PR:#26](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/26) +* Reworked the extraction prompt (anti-inference, preserve specifics, topic-grouped memories) and simplified the schema to `fact`/`episodic` with fixed fact categories. See [PR:#26](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/26) +* Token-bounded extraction batches with per-batch failure isolation; embedding inputs truncated to the model token budget. See [PR:#26](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/26) ## [0.2.0b3] (2026-07-08) diff --git a/azure/cosmos/agent_memory/aio/services/pipeline.py b/azure/cosmos/agent_memory/aio/services/pipeline.py index 124c8cb..fff87f5 100644 --- a/azure/cosmos/agent_memory/aio/services/pipeline.py +++ b/azure/cosmos/agent_memory/aio/services/pipeline.py @@ -273,6 +273,10 @@ async def _vector_distance_function(self) -> str: 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, @@ -280,8 +284,20 @@ async def _vector_distance_function(self) -> str: 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 in-place " + "near-duplicate folding this run to avoid mis-calibrated folds. Memories are " + "written as-is and deduped 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. @@ -738,8 +754,11 @@ async def dedup_extracted_memories(self, user_id: str, extracted: dict) -> dict: high = get_dedup_sim_high() distance_function = await self._vector_distance_function() - similarity_ok = vector_autodrop_supported(distance_function) - if not similarity_ok: + read_failed = getattr(self, "_distance_function_read_failed", False) + similarity_ok = (not read_failed) and vector_autodrop_supported(distance_function) + if read_failed: + self._warn_distance_policy_unavailable_once() + elif not similarity_ok: self._warn_euclidean_autodrop_once(distance_function) result = { @@ -843,7 +862,11 @@ async def _nearest_active_full( async def _apply_inplace_update(self, neighbor: dict[str, Any], new_doc: dict[str, Any]) -> bool: """Async mirror of the sync in-place refresh (recency-wins content+embedding).""" + from azure.core import MatchConditions + from azure.cosmos.exceptions import CosmosAccessConditionFailedError + try: + old_etag = neighbor.get("_etag") updated = dict(neighbor) for sys_prop in ("_rid", "_self", "_etag", "_attachments", "_ts"): updated.pop(sys_prop, None) @@ -868,8 +891,23 @@ async def _apply_inplace_update(self, neighbor: dict[str, Any], new_doc: dict[st if merged_tags: updated["tags"] = merged_tags - await self._upsert_item(self._memories_container, body=updated) + if old_etag and hasattr(self._memories_container, "replace_item"): + await self._replace_item( + self._memories_container, + item=updated["id"], + body=updated, + match_condition=MatchConditions.IfNotModified, + etag=old_etag, + ) + else: + await self._upsert_item(self._memories_container, body=updated) return True + except CosmosAccessConditionFailedError: + logger.info( + "in-place dedup update skipped (concurrent writer won) target_id=%s; keeping new doc as novel", + neighbor.get("id"), + ) + return False except Exception as exc: # noqa: BLE001 logger.warning( "in-place dedup update failed target_id=%s err=%s (keeping new doc as novel)", @@ -1648,6 +1686,14 @@ async def _reconcile_contradictions( pair, ) continue + # Guard chained contradictions (A>B then B>C) and re-supersession. + if winner_id in consumed_loser_ids or loser_id in consumed_loser_ids: + logger.info( + "reconcile_memories: skipping chained/duplicate contradiction pair %r " + "(winner or loser already superseded this pass)", + pair, + ) + continue loser_doc = facts_by_id.get(loser_id) if loser_doc is None: continue diff --git a/azure/cosmos/agent_memory/prompts/_schemas.py b/azure/cosmos/agent_memory/prompts/_schemas.py index 1d98908..2f43a12 100644 --- a/azure/cosmos/agent_memory/prompts/_schemas.py +++ b/azure/cosmos/agent_memory/prompts/_schemas.py @@ -51,34 +51,6 @@ } -# --------------------------------------------------------------------------- -# dedup_episodic.prompty - reconcile a pool of active episodic memories -# (MERGE-ONLY: same-event duplicates collapse; no contradiction/deletion) -# --------------------------------------------------------------------------- -DEDUP_EPISODIC_SCHEMA: dict[str, Any] = { - "type": "object", - "properties": { - "duplicate_groups": { - "type": "array", - "items": { - "type": "object", - "properties": { - "merged_content": {"type": "string"}, - "source_ids": {"type": "array", "items": {"type": "string"}}, - "confidence": {"type": ["number", "null"]}, - "salience": {"type": ["number", "null"]}, - }, - "required": ["merged_content", "source_ids", "confidence", "salience"], - "additionalProperties": False, - }, - }, - "kept_ids": {"type": "array", "items": {"type": "string"}}, - }, - "required": ["duplicate_groups", "kept_ids"], - "additionalProperties": False, -} - - # --------------------------------------------------------------------------- # extract_memories.prompty - extract facts + episodic + unclassified # --------------------------------------------------------------------------- @@ -267,7 +239,6 @@ # --------------------------------------------------------------------------- PROMPTY_SCHEMAS: dict[str, tuple[str, dict[str, Any]]] = { "dedup.prompty": ("DedupOutput", DEDUP_SCHEMA), - "dedup_episodic.prompty": ("DedupEpisodicOutput", DEDUP_EPISODIC_SCHEMA), "extract_memories.prompty": ("ExtractMemoriesOutput", EXTRACT_MEMORIES_SCHEMA), "summarize.prompty": ("SummarizeOutput", SUMMARIZE_SCHEMA), "summarize_update.prompty": ("SummarizeUpdateOutput", SUMMARIZE_UPDATE_SCHEMA), diff --git a/azure/cosmos/agent_memory/prompts/dedup_episodic.prompty b/azure/cosmos/agent_memory/prompts/dedup_episodic.prompty deleted file mode 100644 index 6008dd2..0000000 --- a/azure/cosmos/agent_memory/prompts/dedup_episodic.prompty +++ /dev/null @@ -1,160 +0,0 @@ ---- -name: dedup_episodic -version: v1 -description: Reconcile a pool of active episodic memories — collapse only true same-event duplicates. -model: - apiType: chat - options: - seed: 43 - maxOutputTokens: 16384 - additionalProperties: - response_format: - type: json_object -inputs: - episodics_text: - type: string ---- - -system: -You are a precision episodic-memory reconciliation system. You receive a pool of active episodic memories (each with an ID, content, confidence, salience, and creation timestamp) and must collapse only true same-event duplicates while leaving distinct experiences untouched. - -## Your Goal -Produce a clean merge-only reconciliation that: -1. Collapses repeated extractions of the same past experience into a single merged episode. -2. Preserves distinct occurrences, even when they involve the same action, tool, project, or outcome. -3. Leaves everything else untouched. - -## What is an EPISODIC MEMORY - -An episodic memory is a **past experience**: a specific event or bounded experience in the shape `situation → action_taken → outcome`, with a scope and an `outcome_valence`. It records what happened, not an atomic claim about what is generally true. - -Examples: -- "During the Q3 launch, Redis-backed sessions timed out under load → the team switched to DynamoDB → sessions stabilized" → episodic. -- "On Monday's load test, increasing worker count to 20 still failed with timeouts" → episodic. - -## Two Orthogonal Outcomes (mutually exclusive, exhaustive) - -Every input episode must end up in exactly one of these two places: -- `duplicate_groups[*].source_ids` — the episode is a re-extraction of the same event as one or more other episodes. -- `kept_ids` — the episode is not a true same-event duplicate. - -No input episode may be omitted from the output. Do NOT emit a `contradicted_pairs` key or any deletion bucket. - -## What is a DUPLICATE - -Two or more episodic memories are duplicates only when they describe the **same event re-extracted**: the same situation, the same action taken, and the same outcome. - -The duplicate bar is HIGH. Same topic, same tool, same action pattern, or same outcome is not enough. - -Resolution: emit one entry in `duplicate_groups` with: -- `merged_content` — a clean, self-contained same-event restatement. Do NOT concatenate the originals; synthesize. -- `source_ids` — every original ID that participates in the duplicate group. It must contain at least 2 IDs. -- `confidence` — the **max** confidence among the source episodes. -- `salience` — the **max** salience among the source episodes. - -## What is KEPT - -Episodes that are not true same-event duplicates go in `kept_ids`. - -Distinct occurrences MUST NOT merge. If the user tried the same action on Monday and again on Tuesday with the same outcome, those are separate experiences. Frequency and repetition are evidence and must be preserved. - -## No Contradiction Handling - -Past events are append-only ground truth. There is no contradiction handling and no deletion of episodes in this prompt. - -- Do NOT pick winners or losers. -- Do NOT soft-delete an episode because another episode appears to conflict with it. -- Do NOT emit `contradicted_pairs` or any contradiction/deletion bucket. -- Conflicting lessons are resolved elsewhere, not here. - -## Decision Guidelines - -1. **Conservative bias.** If you cannot confidently prove two episodes are the same event, put them both in `kept_ids`. Over-merging distinct experiences is worse than retaining redundancy. - -2. **Same event test.** Merge only when the memories align on the concrete situation/scope, action taken, and outcome. If the date, run, incident, session, customer, environment, or other occurrence marker differs, keep them separate. - -3. **Don't drop information.** A `merged_content` must preserve every material detail from its sources. If you cannot preserve everything in one clean same-event restatement, the episodes probably are not duplicates. - -4. **Don't fabricate.** Never introduce dates, outcomes, causes, tools, or entities that are not present in at least one source. - -## Input Format - -You will receive a numbered list of episodic memories. Each line has the form: -``` -N. ID: | Content: "" | Confidence: 0.85 | Salience: 0.7 | Created: -``` - -## Missing-Field Handling - -Some episodes may show `Confidence: N/A`, `Salience: N/A`, or `Created: N/A`. - -- Treat `N/A` confidence/salience as **unknown** — set those fields to `null` - in the `duplicate_groups` output (do NOT omit the keys); the pipeline will - fall back to `max(source.confidence)` / `max(source.salience)`. -- Created timestamps are for traceability only. Do not use recency to delete or prefer episodes. - -## Output Format - -You must output ONLY valid JSON matching this exact schema. No preamble, no explanation, no markdown fences — just the JSON object. - -```json -{ - "duplicate_groups": [ - {"merged_content": "", "source_ids": ["", ""], "confidence": 0.0, "salience": 0.0} - ], - "kept_ids": ["", ""] -} -``` - - -The `confidence` and `salience` values **must** be real numbers between 0 and 1 in the actual output — the `0.0` above is a structural placeholder, not a literal value to echo. Compute them as the maximum across source episodes in the group. If you cannot determine confidence or salience, set the field to `null` (the pipeline will fall back to `max(source.*)` from the source records). **Do not omit the keys** — the response schema requires both fields to appear on every group, with either a number or `null` as the value. - -If a bucket is empty, emit it as an empty array (`[]`) rather than omitting the key. Never include `contradicted_pairs`. - -## Worked Examples - -### Example 1: True same-event duplicate - -**Input pool:** -``` -1. ID: E1 | Content: "During the Q3 checkout launch, Redis sessions timed out under load; the team switched session storage to DynamoDB and checkout stabilized." | Confidence: 0.9 | Salience: 0.8 | Created: 2024-01-06T00:00:00Z -2. ID: E2 | Content: "In the Q3 checkout launch incident, Redis-backed sessions failed under load, so the team moved sessions to DynamoDB, which stabilized checkout." | Confidence: 0.85 | Salience: 0.9 | Created: 2024-01-07T00:00:00Z -3. ID: E3 | Content: "On Tuesday's checkout load test, raising Redis connection limits still produced session timeouts." | Confidence: 0.8 | Salience: 0.6 | Created: 2024-01-08T00:00:00Z -``` - -**Expected output:** -```json -{ - "duplicate_groups": [ - {"merged_content": "During the Q3 checkout launch, Redis-backed sessions timed out under load; the team moved session storage to DynamoDB, and checkout stabilized.", "source_ids": ["E1", "E2"], "confidence": 0.9, "salience": 0.9} - ], - "kept_ids": ["E3"] -} -``` - -E1 and E2 describe the same situation, action, and outcome. E3 is related but a different occurrence, so it stays kept. - -### Example 2: Same action on different days is not a duplicate - -**Input pool:** -``` -1. ID: E4 | Content: "On Monday's ingestion test, the team increased workers to 20, but the pipeline still timed out." | Confidence: 0.9 | Salience: 0.7 | Created: 2024-02-01T00:00:00Z -2. ID: E5 | Content: "On Tuesday's ingestion test, the team increased workers to 20, but the pipeline still timed out." | Confidence: 0.9 | Salience: 0.7 | Created: 2024-02-02T00:00:00Z -``` - -**Expected output:** -```json -{ - "duplicate_groups": [], - "kept_ids": ["E4", "E5"] -} -``` - -These episodes share the same action and outcome, but they occurred on different days. Preserve both because repetition is evidence. - -user: -Reconcile the following pool of active episodic memories: - -{{episodics_text}} - -Return JSON exactly matching the schema above. diff --git a/azure/cosmos/agent_memory/services/pipeline.py b/azure/cosmos/agent_memory/services/pipeline.py index a5058e1..e274dd4 100644 --- a/azure/cosmos/agent_memory/services/pipeline.py +++ b/azure/cosmos/agent_memory/services/pipeline.py @@ -302,6 +302,10 @@ def _vector_distance_function(self) -> str: # 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, @@ -309,6 +313,7 @@ def _vector_distance_function(self) -> str: 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: @@ -329,6 +334,17 @@ def _warn_euclidean_autodrop_once(self, distance_function: str) -> None: 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 in-place " + "near-duplicate folding this run to avoid mis-calibrated folds. Memories are " + "written as-is and deduped on a later run once the policy is readable." + ) + def _vector_candidates( self, *, @@ -807,6 +823,13 @@ def dedup_extracted_memories( document rather than minting a new one that a later reconcile sweep must merge and supersede. There is no ``sys:dup-candidate`` tagging and no clustering — reconcile only resolves contradictions. + + In-place folds commit here, before ``persist_extracted_memories`` writes + the novel docs. A crash between the two leaves some memories folded and + some novel docs unwritten, but the source turns are not stamped + ``extracted_at`` until persist completes, so the whole turn set is simply + re-extracted on the next run (folds are idempotent and re-ADDs hit + exact-hash/vector dedup) — no data is lost, only repeated. """ if not threshold_config.get_dedup_vector_enabled(): return extracted @@ -817,8 +840,16 @@ def dedup_extracted_memories( high = threshold_config.get_dedup_sim_high() distance_function = self._vector_distance_function() - similarity_ok = vector_autodrop_supported(distance_function) - if not similarity_ok: + read_failed = getattr(self, "_distance_function_read_failed", False) + # Skip destructive in-place folding when the container's distance policy + # could not be read: a defaulted cosine on a euclidean container would + # apply cosine thresholds to unbounded euclidean distances and fold + # unrelated memories. Everything ADDs; the next run (policy readable) + # dedups normally. + similarity_ok = (not read_failed) and vector_autodrop_supported(distance_function) + if read_failed: + self._warn_distance_policy_unavailable_once() + elif not similarity_ok: self._warn_euclidean_autodrop_once(distance_function) result: dict[str, list[dict[str, Any]]] = { @@ -934,12 +965,22 @@ def _nearest_active_full( def _apply_inplace_update(self, neighbor: dict[str, Any], new_doc: dict[str, Any]) -> bool: """Refresh an existing memory in place with a near-duplicate's content. + Uses ETag optimistic concurrency: the update is applied with + ``IfNotModified`` against the neighbor's ``_etag`` so a concurrent + supersede/refresh cannot be clobbered (which would resurrect a + soft-deleted memory or lose an update). On an ETag conflict — or any + other failure — returns False and the caller keeps the new doc as a + novel ADD, so nothing is lost. + Recency wins: the neighbor keeps its id / created_at / partition but takes the new doc's content + embedding, unions tags, and takes the max - salience/confidence. Returns True on a successful upsert; False otherwise - (the caller then keeps the new doc as a novel ADD, so nothing is lost). + salience/confidence. """ + from azure.core import MatchConditions + from azure.cosmos.exceptions import CosmosAccessConditionFailedError + try: + old_etag = neighbor.get("_etag") updated = dict(neighbor) for sys_prop in ("_rid", "_self", "_etag", "_attachments", "_ts"): updated.pop(sys_prop, None) @@ -964,8 +1005,22 @@ def _apply_inplace_update(self, neighbor: dict[str, Any], new_doc: dict[str, Any if merged_tags: updated["tags"] = merged_tags - self._memories_container.upsert_item(body=updated) + if old_etag and hasattr(self._memories_container, "replace_item"): + self._memories_container.replace_item( + item=updated["id"], + body=updated, + match_condition=MatchConditions.IfNotModified, + etag=old_etag, + ) + else: + self._memories_container.upsert_item(body=updated) return True + except CosmosAccessConditionFailedError: + logger.info( + "in-place dedup update skipped (concurrent writer won) target_id=%s; keeping new doc as novel", + neighbor.get("id"), + ) + return False except Exception as exc: # noqa: BLE001 logger.warning( "in-place dedup update failed target_id=%s err=%s (keeping new doc as novel)", @@ -1727,6 +1782,16 @@ def _reconcile_contradictions(self, user_id: str, memory_type: str, facts: list[ pair, ) continue + # Guard chained contradictions: never point a loser at a winner that + # was itself already superseded earlier this pass (A>B then B>C would + # tombstone C in favor of a now-dead B), and never re-supersede a loser. + if winner_id in consumed_loser_ids or loser_id in consumed_loser_ids: + logger.info( + "reconcile_memories: skipping chained/duplicate contradiction pair %r " + "(winner or loser already superseded this pass)", + pair, + ) + continue loser_doc = facts_by_id.get(loser_id) if loser_doc is None: continue diff --git a/tests/unit/aio/services/test_dedup_vector_async.py b/tests/unit/aio/services/test_dedup_vector_async.py index a193b78..efdad1c 100644 --- a/tests/unit/aio/services/test_dedup_vector_async.py +++ b/tests/unit/aio/services/test_dedup_vector_async.py @@ -187,7 +187,7 @@ async def test_euclidean_disables_inplace_folding(): @pytest.mark.asyncio async def test_apply_inplace_update_recency_wins_and_unions(): p = _service() - p._upsert_item = AsyncMock() + p._replace_item = AsyncMock() neighbor = _fact("existing-1", "old content", tags=["sys:fact", "topic:a"]) neighbor["confidence"] = 0.6 neighbor["salience"] = 0.5 @@ -202,9 +202,11 @@ async def test_apply_inplace_update_recency_wins_and_unions(): ok = await p._apply_inplace_update(neighbor, new_doc) assert ok is True - written = p._upsert_item.call_args.kwargs["body"] + call = p._replace_item.call_args + assert call.kwargs["etag"] == "etag-xyz" + written = call.kwargs["body"] assert written["id"] == "existing-1" - assert written["content"] == "new richer content" + assert written["content"] == "new richer content" # recency wins assert written["embedding"] == [0.5, 0.5] assert written["salience"] == 0.8 assert written["confidence"] == 0.9 @@ -213,6 +215,19 @@ async def test_apply_inplace_update_recency_wins_and_unions(): assert "_etag" not in written +@pytest.mark.asyncio +async def test_apply_inplace_update_etag_conflict_returns_false(): + from azure.cosmos.exceptions import CosmosAccessConditionFailedError + + p = _service() + p._replace_item = AsyncMock(side_effect=CosmosAccessConditionFailedError(message="etag")) + neighbor = _fact("existing-1", "old", tags=["sys:fact"]) + neighbor["_etag"] = "stale" + new_doc = _fact("f-new", "old restated", embedding=[0.5, 0.5], tags=["sys:fact"]) + + assert await p._apply_inplace_update(neighbor, new_doc) is False + + @pytest.mark.asyncio async def test_nearest_active_full_returns_full_doc_and_skips_excluded(): p = _service() diff --git a/tests/unit/services/test_dedup_vector.py b/tests/unit/services/test_dedup_vector.py index 80fe012..706595d 100644 --- a/tests/unit/services/test_dedup_vector.py +++ b/tests/unit/services/test_dedup_vector.py @@ -248,9 +248,12 @@ def test_apply_inplace_update_recency_wins_and_unions() -> None: ok = p._apply_inplace_update(neighbor, new_doc) assert ok is True - written = p._memories_container.upsert_item.call_args.kwargs["body"] + # ETag optimistic concurrency: goes through replace_item with IfNotModified. + call = p._memories_container.replace_item.call_args + assert call.kwargs["etag"] == "etag-xyz" + written = call.kwargs["body"] assert written["id"] == "existing-1" - assert written["content"] == "new richer content" + assert written["content"] == "new richer content" # recency wins assert written["embedding"] == [0.5, 0.5] assert written["salience"] == 0.8 assert written["confidence"] == 0.9 @@ -260,6 +263,19 @@ def test_apply_inplace_update_recency_wins_and_unions() -> None: assert written["updated_at"] != neighbor["updated_at"] +def test_apply_inplace_update_etag_conflict_returns_false() -> None: + # F1: a concurrent writer (ETag mismatch) must NOT clobber; caller ADDs novel. + from azure.cosmos.exceptions import CosmosAccessConditionFailedError + + p = _make_pipeline() + p._memories_container.replace_item.side_effect = CosmosAccessConditionFailedError(message="etag") + neighbor = _doc("existing-1", "old", tags=["sys:fact"]) + neighbor["_etag"] = "stale" + new_doc = _doc("f-new", "old restated", embedding=[0.5, 0.5], tags=["sys:fact"]) + + assert p._apply_inplace_update(neighbor, new_doc) is False + + def test_nearest_active_full_returns_full_doc_and_skips_excluded() -> None: p = _make_pipeline() doc_a = _doc("a", "first") @@ -335,3 +351,28 @@ def test_reconcile_fact_contradiction_only() -> None: assert p._mark_superseded.call_args.args[1] == "f2" assert p._mark_superseded.call_args.kwargs["reason"] == "contradict" assert p._run_prompty.call_args.args[0] == "dedup.prompty" + + +def test_reconcile_skips_chained_contradiction() -> None: + # F5: (A>B) then (B>C) must not tombstone C in favor of an already-dead B. + p = _make_pipeline() + facts = [_doc("A", "a"), _doc("B", "b"), _doc("C", "c")] + p._memories_container.query_items.return_value = iter(facts) + p._run_prompty = MagicMock( + return_value=json.dumps( + { + "contradicted_pairs": [ + {"winner_id": "A", "loser_id": "B", "reason": "x"}, + {"winner_id": "B", "loser_id": "C", "reason": "y"}, + ], + "kept_ids": [], + } + ) + ) + + result = p.reconcile_memories("u1", memory_type="fact") + + # Only the first pair applies; the chained (B>C) is skipped since B is dead. + assert result["contradicted"] == 1 + assert p._mark_superseded.call_count == 1 + assert p._mark_superseded.call_args.args[0]["id"] == "B" From b328df0358d51feb89e999d5f11de99fe4ca1f30 Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Sat, 11 Jul 2026 12:30:46 -0700 Subject: [PATCH 7/8] Resolving comments --- .../agent_memory/aio/services/pipeline.py | 14 ++++---- .../cosmos/agent_memory/services/pipeline.py | 14 ++++---- .../aio/services/test_dedup_vector_async.py | 25 ++++++++++++++ tests/unit/services/test_dedup_vector.py | 34 +++++++++++++++++-- tests/unit/test_procedural_synthesis.py | 1 - tests/unit/test_reconcile.py | 1 - 6 files changed, 68 insertions(+), 21 deletions(-) diff --git a/azure/cosmos/agent_memory/aio/services/pipeline.py b/azure/cosmos/agent_memory/aio/services/pipeline.py index fff87f5..0dfd5ee 100644 --- a/azure/cosmos/agent_memory/aio/services/pipeline.py +++ b/azure/cosmos/agent_memory/aio/services/pipeline.py @@ -458,7 +458,6 @@ def _empty_extract_counts() -> dict[str, int]: return { "fact_count": 0, "episodic_count": 0, - "unclassified_count": 0, "updated_count": 0, "contradicted_count": 0, "exact_dedup_skipped": 0, @@ -871,10 +870,12 @@ async def _apply_inplace_update(self, neighbor: dict[str, Any], new_doc: dict[st for sys_prop in ("_rid", "_self", "_etag", "_attachments", "_ts"): updated.pop(sys_prop, None) new_content = str(new_doc.get("content") or "") - updated["content"] = new_content - updated["content_hash"] = compute_content_hash(new_content) - if new_doc.get("embedding"): - updated["embedding"] = new_doc["embedding"] + old_content = str(neighbor.get("content") or "") + if len(new_content) >= len(old_content): + updated["content"] = new_content + updated["content_hash"] = compute_content_hash(new_content) + if new_doc.get("embedding"): + updated["embedding"] = new_doc["embedding"] updated["updated_at"] = datetime.now(timezone.utc).isoformat() new_sal = _max_or_none([neighbor.get("salience"), new_doc.get("salience")]) @@ -951,11 +952,8 @@ async def persist_extracted_memories( logger.info("persist_extracted_memories skipped existing id=%s", validated.get("id")) continue - tags = validated.get("tags", []) if doc_type == "episodic": result["episodic_count"] += 1 - elif "sys:unclassified" in tags: - result["unclassified_count"] += 1 elif doc_type == "fact": result["fact_count"] += 1 diff --git a/azure/cosmos/agent_memory/services/pipeline.py b/azure/cosmos/agent_memory/services/pipeline.py index e274dd4..2acf6a9 100644 --- a/azure/cosmos/agent_memory/services/pipeline.py +++ b/azure/cosmos/agent_memory/services/pipeline.py @@ -480,7 +480,6 @@ def _empty_extract_counts() -> dict[str, int]: return { "fact_count": 0, "episodic_count": 0, - "unclassified_count": 0, "updated_count": 0, "contradicted_count": 0, "exact_dedup_skipped": 0, @@ -985,10 +984,12 @@ def _apply_inplace_update(self, neighbor: dict[str, Any], new_doc: dict[str, Any for sys_prop in ("_rid", "_self", "_etag", "_attachments", "_ts"): updated.pop(sys_prop, None) new_content = str(new_doc.get("content") or "") - updated["content"] = new_content - updated["content_hash"] = compute_content_hash(new_content) - if new_doc.get("embedding"): - updated["embedding"] = new_doc["embedding"] + old_content = str(neighbor.get("content") or "") + if len(new_content) >= len(old_content): + updated["content"] = new_content + updated["content_hash"] = compute_content_hash(new_content) + if new_doc.get("embedding"): + updated["embedding"] = new_doc["embedding"] updated["updated_at"] = datetime.now(timezone.utc).isoformat() new_sal = _max_or_none([neighbor.get("salience"), new_doc.get("salience")]) @@ -1069,11 +1070,8 @@ def persist_extracted_memories( logger.info("persist_extracted_memories skipped existing id=%s", validated.get("id")) continue - tags = validated.get("tags", []) if doc_type == "episodic": result["episodic_count"] += 1 - elif "sys:unclassified" in tags: - result["unclassified_count"] += 1 elif doc_type == "fact": result["fact_count"] += 1 diff --git a/tests/unit/aio/services/test_dedup_vector_async.py b/tests/unit/aio/services/test_dedup_vector_async.py index efdad1c..3667be0 100644 --- a/tests/unit/aio/services/test_dedup_vector_async.py +++ b/tests/unit/aio/services/test_dedup_vector_async.py @@ -215,6 +215,31 @@ async def test_apply_inplace_update_recency_wins_and_unions(): assert "_etag" not in written +@pytest.mark.asyncio +async def test_apply_inplace_update_shorter_restatement_keeps_richer_content(): + p = _service() + p._replace_item = AsyncMock() + neighbor = _fact( + "existing-1", "March 1, room 204, deluxe suite", embedding=[0.1, 0.2], tags=["sys:fact", "topic:a"] + ) + neighbor["confidence"] = 0.6 + neighbor["salience"] = 0.5 + neighbor["_etag"] = "etag-xyz" + new_doc = _fact("f-new", "March 1", embedding=[0.5, 0.5], tags=["sys:fact", "topic:b"]) + new_doc["confidence"] = 0.9 + new_doc["salience"] = 0.8 + + ok = await p._apply_inplace_update(neighbor, new_doc) + + assert ok is True + written = p._replace_item.call_args.kwargs["body"] + assert written["content"] == "March 1, room 204, deluxe suite" # richer content kept + assert written["embedding"] == [0.1, 0.2] # matching embedding kept + assert written["salience"] == 0.8 # metadata still recency-wins + assert written["confidence"] == 0.9 + assert "topic:a" in written["tags"] and "topic:b" in written["tags"] + + @pytest.mark.asyncio async def test_apply_inplace_update_etag_conflict_returns_false(): from azure.cosmos.exceptions import CosmosAccessConditionFailedError diff --git a/tests/unit/services/test_dedup_vector.py b/tests/unit/services/test_dedup_vector.py index 706595d..9cabf3f 100644 --- a/tests/unit/services/test_dedup_vector.py +++ b/tests/unit/services/test_dedup_vector.py @@ -263,8 +263,36 @@ def test_apply_inplace_update_recency_wins_and_unions() -> None: assert written["updated_at"] != neighbor["updated_at"] -def test_apply_inplace_update_etag_conflict_returns_false() -> None: - # F1: a concurrent writer (ETag mismatch) must NOT clobber; caller ADDs novel. +def test_apply_inplace_update_shorter_restatement_keeps_richer_content() -> None: + p = _make_pipeline() + neighbor = _doc( + "existing-1", + "March 1, room 204, deluxe suite", + confidence=0.6, + salience=0.5, + tags=["sys:fact", "topic:a"], + embedding=[0.1, 0.2], + ) + neighbor["_etag"] = "etag-xyz" + new_doc = _doc( + "f-new", + "March 1", + confidence=0.9, + salience=0.8, + tags=["sys:fact", "topic:b"], + embedding=[0.5, 0.5], + ) + + ok = p._apply_inplace_update(neighbor, new_doc) + + assert ok is True + written = p._memories_container.replace_item.call_args.kwargs["body"] + assert written["content"] == "March 1, room 204, deluxe suite" # richer content kept + assert written["embedding"] == [0.1, 0.2] # matching embedding kept + assert written["salience"] == 0.8 # metadata still recency-wins + assert written["confidence"] == 0.9 + assert "topic:a" in written["tags"] and "topic:b" in written["tags"] + # A concurrent writer (ETag mismatch) must NOT clobber; caller ADDs novel. from azure.cosmos.exceptions import CosmosAccessConditionFailedError p = _make_pipeline() @@ -354,7 +382,7 @@ def test_reconcile_fact_contradiction_only() -> None: def test_reconcile_skips_chained_contradiction() -> None: - # F5: (A>B) then (B>C) must not tombstone C in favor of an already-dead B. + # (A>B) then (B>C) must not tombstone C in favor of an already-dead B. p = _make_pipeline() facts = [_doc("A", "a"), _doc("B", "b"), _doc("C", "c")] p._memories_container.query_items.return_value = iter(facts) diff --git a/tests/unit/test_procedural_synthesis.py b/tests/unit/test_procedural_synthesis.py index 533f090..e96ace4 100644 --- a/tests/unit/test_procedural_synthesis.py +++ b/tests/unit/test_procedural_synthesis.py @@ -242,7 +242,6 @@ def test_extract_memories_without_procedural_bucket_returns_new_count_shape(): assert result["fact_count"] == 1 assert result["episodic_count"] == 1 - assert result["unclassified_count"] == 0 assert legacy_fact_count_key not in result assert legacy_proc_key not in result assert all(doc["type"] != "procedural" for doc in upserted) diff --git a/tests/unit/test_reconcile.py b/tests/unit/test_reconcile.py index 9a5351b..81c563c 100644 --- a/tests/unit/test_reconcile.py +++ b/tests/unit/test_reconcile.py @@ -403,7 +403,6 @@ def test_empty_thread_returns_full_dict_shape(self): for key in ( "fact_count", "episodic_count", - "unclassified_count", "updated_count", "exact_dedup_skipped", ): From 8314298262315ba1d63d849aca78546d0689bf66 Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Sat, 11 Jul 2026 12:55:54 -0700 Subject: [PATCH 8/8] Resolving comments --- Docs/azure_testing.md | 34 ++--- Docs/concepts.md | 135 +++++++++--------- Docs/public_api.md | 4 +- README.md | 122 ++++++++-------- .../agent_memory/aio/services/pipeline.py | 11 +- .../cosmos/agent_memory/services/pipeline.py | 31 ++-- tests/unit/aio/test_auto_trigger.py | 1 - 7 files changed, 170 insertions(+), 168 deletions(-) diff --git a/Docs/azure_testing.md b/Docs/azure_testing.md index 590dc4b..6b1523b 100644 --- a/Docs/azure_testing.md +++ b/Docs/azure_testing.md @@ -6,13 +6,13 @@ This guide covers the minimum Azure resources, deployment steps, throughput sett ## Required Azure Services -| Service | Purpose | -|---------|---------| -| **Azure Cosmos DB for NoSQL** | Persistent memory store with vector and full-text indexes | -| **Azure OpenAI / AI Services** | Embeddings and chat generation | -| **Azure Functions** | Durable Functions orchestrator and activities | -| **Azure Storage Account** | Required by Azure Functions | -| **Application Insights** | Recommended for monitoring | +| Service | Purpose | +|--------------------------------|-----------------------------------------------------------| +| **Azure Cosmos DB for NoSQL** | Persistent memory store with vector and full-text indexes | +| **Azure OpenAI / AI Services** | Embeddings and chat generation | +| **Azure Functions** | Durable Functions orchestrator and activities | +| **Azure Storage Account** | Required by Azure Functions | +| **Application Insights** | Recommended for monitoring | --- @@ -277,7 +277,7 @@ Bring the environment up in this order: 9. write a few turns and verify a thread `summary` memory appears 10. write more turns and verify `fact`, `procedural`, and `episodic` memories appear 11. verify a per-user `user_summary` memory appears once `USER_SUMMARY_EVERY_N` turns have accumulated for that user -12. test deduplication by writing two near-duplicate facts and confirming the dedup orchestrator merges them +12. test deduplication by writing two near-duplicate facts and confirming the second folds into the first in place (one active record remains, its `updated_at` bumped) rather than creating a duplicate This keeps failures isolated and easier to diagnose. @@ -367,15 +367,15 @@ az functionapp log tail \ Common issues: -| Symptom | Likely Cause | -|---------|--------------| -| 401 / 403 from Cosmos DB | Missing Cosmos DB RBAC | -| 401 / 403 from Azure OpenAI | Missing OpenAI RBAC | -| Durable Function starts but fails | Missing app settings or downstream RBAC | -| `No memories found` | No turn memories exist, or all candidate turns predate the existing summary | -| Search is slow | Embedding latency, index choice, or region mismatch | -| Change feed trigger not firing | Verify `COSMOS_DB__accountEndpoint` is set and the function can write to the configured `COSMOS_DB_COUNTERS_CONTAINER` container | -| Auto-processing not starting | Check threshold settings are > 0 in Function App configuration | +| Symptom | Likely Cause | +|-----------------------------------|----------------------------------------------------------------------------------------------------------------------------------| +| 401 / 403 from Cosmos DB | Missing Cosmos DB RBAC | +| 401 / 403 from Azure OpenAI | Missing OpenAI RBAC | +| Durable Function starts but fails | Missing app settings or downstream RBAC | +| `No memories found` | No turn memories exist, or all candidate turns predate the existing summary | +| Search is slow | Embedding latency, index choice, or region mismatch | +| Change feed trigger not firing | Verify `COSMOS_DB__accountEndpoint` is set and the function can write to the configured `COSMOS_DB_COUNTERS_CONTAINER` container | +| Auto-processing not starting | Check threshold settings are > 0 in Function App configuration | Recommended checks: diff --git a/Docs/concepts.md b/Docs/concepts.md index b72807c..5ac1942 100644 --- a/Docs/concepts.md +++ b/Docs/concepts.md @@ -8,17 +8,17 @@ Agent Memory Toolkit stores conversation data as memory documents and builds hig Every memory uses the same base shape: -| Field | Description | -|-------|-------------| -| `id` | Unique identifier | -| `user_id` | User the memory belongs to | -| `thread_id` | Conversation thread | -| `role` | `user`, `agent`, `tool`, or `system` | -| `type` | `turn`, `summary`, `fact`, or `user_summary` | -| `content` | Main text payload | -| `embedding` | Vector used for semantic search | -| `metadata` | Extra context (e.g. `tool_name`, `tool_call_id`, edit reasons) | -| `created_at` | ISO 8601 timestamp | +| Field | Description | +|--------------|----------------------------------------------------------------| +| `id` | Unique identifier | +| `user_id` | User the memory belongs to | +| `thread_id` | Conversation thread | +| `role` | `user`, `agent`, `tool`, or `system` | +| `type` | `turn`, `summary`, `fact`, or `user_summary` | +| `content` | Main text payload | +| `embedding` | Vector used for semantic search | +| `metadata` | Extra context (e.g. `tool_name`, `tool_call_id`, edit reasons) | +| `created_at` | ISO 8601 timestamp | --- @@ -64,12 +64,12 @@ Like thread summaries, user summaries update incrementally by merging the existi ## Short-Term vs. Long-Term Memory -| | Short-Term | Long-Term | -|---|---|---| -| **What** | Turn messages | Summaries, facts, user summaries | -| **Granularity** | Per message | Per thread, per fact, or per user | -| **Created by** | `add_local()` / `add_cosmos()` / `push_to_cosmos()` | `generate_thread_summary()` / `extract_facts()` / `generate_user_summary()` | -| **Purpose** | Replay recent context | Compact recall and semantic retrieval | +| | Short-Term | Long-Term | +|-----------------|-----------------------------------------------------|-----------------------------------------------------------------------------| +| **What** | Turn messages | Summaries, facts, user summaries | +| **Granularity** | Per message | Per thread, per fact, or per user | +| **Created by** | `add_local()` / `add_cosmos()` / `push_to_cosmos()` | `generate_thread_summary()` / `extract_facts()` / `generate_user_summary()` | +| **Purpose** | Replay recent context | Compact recall and semantic retrieval | Common pattern: keep turns during an active conversation, then generate summaries or facts when the thread gets long or is complete. @@ -79,12 +79,12 @@ Common pattern: keep turns during an active conversation, then generate summarie A thread is the unit of conversation. `get_thread()` returns the memories for one `thread_id`, optionally limited to the most recent `k` entries. `get_memories()` also supports a `thread_id` filter to retrieve memories from a specific thread. -| Role | Meaning | -|------|---------| -| `user` | Human message | -| `agent` | Assistant message | -| `tool` | Tool output (metadata can include `tool_name` and `tool_call_id`) | -| `system` | Generated artifacts such as summaries and facts | +| Role | Meaning | +|----------|-------------------------------------------------------------------| +| `user` | Human message | +| `agent` | Assistant message | +| `tool` | Tool output (metadata can include `tool_name` and `tool_call_id`) | +| `system` | Generated artifacts such as summaries and facts | --- @@ -118,34 +118,27 @@ Prompts for summarization and fact extraction live in `azure_functions/prompts/` ## Memory Reconciliation -Reconciliation runs in **two complementary tiers**: a cheap, LLM-free **vector-floor dedup ladder** applied to freshly-extracted memories before they persist, and a periodic **LLM reconcile** that runs in a **dual mode** (cheap candidate clusters most sweeps, a full-pool backstop occasionally). +Two independent mechanisms keep the fact pool clean and convergent: a cheap, LLM-free **write-time in-place dedup** that folds near-duplicate restatements into their existing record before they persist, and a periodic **LLM contradiction pass** that resolves opposing claims. Paraphrases are handled entirely at write time, so the LLM pass never merges duplicates — it only adjudicates contradictions. -### Vector-floor dedup ladder (write path, LLM-free) +### Write-time in-place dedup (LLM-free) -Between extraction and persist, `dedup_extracted_memories` compares each new fact/episodic memory against the user's existing active memories of the same type using Cosmos `VectorDistance` (pure vector, no hybrid). Each new memory takes one rung of a similarity ladder: +Between extraction and persist, `dedup_extracted_memories` compares each newly extracted fact/episodic doc against its single nearest active same-type memory using Cosmos `VectorDistance` (pure vector, no hybrid): -| band | condition (cosine) | action | -|------|--------------------|--------| -| exact | `content_hash` hit | skip (Stage 0, free) | -| near-exact | `s ≥ DEDUP_SIM_HIGH` (0.97) | **auto-skip** the new memory (no LLM); logged for audit | -| borderline | `DEDUP_SIM_LOW ≤ s < DEDUP_SIM_HIGH` (0.80–0.97) | persist, tag `sys:dup-candidate` + stash `dup_of`/`dup_score` for the LLM reconcile | -| novel | `s < DEDUP_SIM_LOW` | persist clean | +| condition (cosine) | action | +|-----------------------------|-----------------------------------------------------------------------------------| +| `content_hash` hit | skip — identical re-extraction, no vector query, no write (`exact_dedup_skipped`) | +| `s ≥ DEDUP_SIM_HIGH` (0.97) | **fold in place** into the existing neighbor (no LLM) | +| `s < DEDUP_SIM_HIGH` | persist as a novel record | -The thresholds are calibrated for **cosine/dotproduct** on normalized embeddings. On a container whose `distanceFunction` is **euclidean**, the destructive near-exact auto-skip is **disabled** (one-shot warning) and those memories fall through to borderline tagging so the LLM adjudicates — euclidean distances aren't a bounded [0,1] similarity and would mis-fire the cosine-tuned drop. +A **fold** refreshes the existing neighbor rather than minting a new doc: it keeps the neighbor's `id` / `created_at` / partition, unions tags, takes the max salience/confidence, and bumps `updated_at`. Content and embedding are recency-wins **except** that a shorter restatement never overwrites longer content, so specifics captured by the richer record aren't dropped. The write is applied with ETag optimistic concurrency (`IfNotModified`); on an ETag conflict the fold is abandoned and the new doc is added as novel, so a concurrent supersede/refresh is never clobbered. This makes the write path convergent — a restatement updates one document instead of creating a duplicate that a later sweep must merge and supersede. -### Dual-mode LLM reconcile +The threshold is calibrated for **cosine/dotproduct** on normalized embeddings. On a container whose `distanceFunction` is **euclidean** — or when the distance policy can't be read — the destructive in-place fold is **disabled** (one-shot warning) and every extracted doc persists as novel, because cosine thresholds don't translate to unbounded euclidean distances. -`reconcile_memories(user_id, n=50, *, memory_type="fact", full_rebuild=False)` identifies two orthogonal outcomes: +### LLM contradiction reconcile -- **Duplicates** — facts restating the same claim. Resolution: collapse into one merged fact; originals soft-deleted with `supersede_reason="duplicate"` and `superseded_by` set to the merged fact. -- **Contradictions** — facts asserting opposing claims about the same subject. Resolution: keep the winner (more recent first, higher confidence as tiebreaker), soft-delete the loser with `supersede_reason="contradict"`. +`reconcile_memories(user_id, n=50, *, memory_type="fact")` loads up to `n` active (non-superseded) facts, most recent first, and asks the dedup prompt to identify **contradicted pairs** — opposing claims about the same subject (e.g. "deadline March 1" vs "March 15"). Each loser is soft-deleted with `supersede_reason="contradict"` and `superseded_by` set to the winner (more recent wins, higher confidence as tiebreaker). Chained contradictions are guarded, so a fact already superseded in this pass can't be used to tombstone a third. -It runs in one of two modes: - -- **Candidate mode** (default auto sweeps) — builds connected-component clusters from the `sys:dup-candidate` seeds + their vector neighbors (edge threshold `DEDUP_CLUSTER_SIM`, 0.60) and sends **only those clusters** to the LLM. Cheap, but keyed on near-duplicate similarity. Tagged seeds that never join a cluster have their stale tag cleared so they aren't re-scanned forever. -- **Full-pool backstop** — every `DEDUP_FULL_RECLUSTER_EVERY_N`-th sweep (default 12), and on any explicit `reconcile(full_rebuild=True)`, the **entire** active pool goes into one LLM pass. This is the only path that catches **dissimilar contradictions** — paraphrased ("prefers aisle seats") and contradictory ("vegetarian" vs "loves steak") facts have very different embedding vectors and would never co-occur in a cosine cluster, so candidate mode alone can't link them. - -Both modes return `{"kept": int, "merged": int, "contradicted": int}`. In-process and durable backends reconcile **both** facts and episodic memories so episodic duplicates don't accrue forever. +Near-duplicate **paraphrases are not merged here** — write-time in-place dedup already folds them before they land — so reconcile is a single, bounded, convergent pass: no clustering, no candidate/full-pool modes, no synthesized merged documents, and no re-merge churn. `episodic` and `procedural` types are no-ops (episodic has no contradiction semantics; its near-dups fold at write time). The pass runs over one flat pool of the most-recent active facts and returns `{"kept": int, "merged": int, "contradicted": int}` — `merged` is always `0`, retained for backward-compatible callers. ### Loser preservation @@ -155,15 +148,23 @@ Soft-deleted facts stay in the container with their `supersede_reason`, `superse Each fact written by `extract_memories` carries a `content_hash` (SHA-256 of normalized content, truncated to 32 hex chars; lowercase, whitespace-collapsed). Before upserting a freshly-extracted fact, the pipeline checks the hash against existing active facts and short-circuits if a match exists, incrementing the `exact_dedup_skipped` metric. This catches identical re-extractions cheaply without an LLM call. -### Extraction watermark (`recent_k`) +### Extraction gating (`extracted_at` + `recent_k`) + +Extraction only ever reads turns not yet stamped `extracted_at`; `persist` stamps every turn it consumed once the extracted memories are durably written. This `extracted_at` gate is the authoritative "what's left to extract" signal and is shared by both backends — a failed or lagging extract leaves its turns unstamped, so the full backlog is retried on the next run and no turns are skipped. -The auto-trigger paths size `recent_k` (how many recent turns extraction reads) from a per-thread **watermark** (`last_extract_count` on the counter doc): `recent_k = current_count − last_extract_count` (with `last_extract_count` treated as `0` before the first successful extract). The newest-`recent_k` turns are exactly the turns added since the last successful extract, and the watermark advances **only after a successful extract** — so under normal operation no turns are skipped when extraction lags or transiently fails: a failed run leaves the watermark put and the full backlog is retried next sweep. The window is deliberately **not** capped by `DEDUP_POOL_SIZE` (that knob governs the reconcile prompt, not the extraction window) — capping would extract only the newest N and silently strand the oldest backlog turns. +- **In-process SDK auto-trigger** — calls extraction with `recent_k=None`, so it drains the entire unextracted backlog each run and relies purely on the `extracted_at` gate. +- **Change-feed / Durable Function App** — additionally **sizes** `recent_k` from a per-thread **watermark** (`last_extract_count` on the counter doc): `recent_k = current_count − last_extract_count` (with `last_extract_count` treated as `0` before the first successful extract), then still applies the `extracted_at` filter. The watermark advances **only after a successful extract**. The window is deliberately **not** capped by `DEDUP_POOL_SIZE` (that knob governs the reconcile pool, not the extraction window) — capping would extract only the newest N and silently strand the oldest backlog turns. -> **Caveat (rare):** the SDK's inline counter increment is best-effort — under sustained optimistic-concurrency contention it can drop an increment rather than block the user's write path (see `increment_counter_sync`). A dropped increment leaves `current_count` lagging the true turn count, which can in turn under-cover a later extraction window. This is the one case where the "no turns skipped" property does not hold; the Function App backend avoids it by raising to force change-feed redelivery. +> **Caveat (rare, change-feed path):** the counter increment is best-effort — under sustained optimistic-concurrency contention it can drop an increment rather than block the user's write path (see `increment_counter_sync`). A dropped increment leaves `current_count` lagging the true turn count, which can under-size the watermark-derived `recent_k`. The Function App backend avoids stranding turns by raising to force change-feed redelivery; the in-process path is unaffected because it uses `recent_k=None` and the `extracted_at` gate, not the counter, to decide what to extract. ### Tunable -`DEDUP_EVERY_N` (default 5) controls how often reconcile runs in the auto-trigger path. Set to `0` to disable. The candidate cap `n` (default `DEDUP_POOL_SIZE`, 50) is tunable per call; larger values give the LLM a wider view at higher token cost. `DEDUP_FULL_RECLUSTER_EVERY_N` (default 12) sets how often the full-pool backstop fires. +Only two reconcile knobs are operator-configurable: + +- `DEDUP_EVERY_N` (default `5`) — how often reconcile runs in the auto-trigger path (every Nth **extract**, not every Nth turn). Set to `0` to disable. +- `DEDUP_POOL_SIZE` (default `50`, hard cap `500`) — the pool size `n` passed to `reconcile_memories`; also overridable per call. Larger values give the LLM a wider view at higher token cost. + +The similarity threshold (`DEDUP_SIM_HIGH`, `0.97`) and the vector-fold on/off switch (`DEDUP_VECTOR_ENABLED`) ship as **fixed internal constants** in `azure.cosmos.agent_memory.thresholds` — they have no env plumbing and ignore any environment variable, so the write-time fold behavior is not operator-tunable today. > **Indexing note.** The reconcile pool query orders by `created_at` (matching the prompt's "more recent first" tiebreaker). Cosmos's default indexing policy includes every property, so this works out of the box. If you customize the indexing policy to reduce write RU, ensure `/created_at/?` remains indexed or the query will fail with a 400 (`Order-by over a non-indexed path`). @@ -195,23 +196,23 @@ on_memory_change trigger ### Threshold settings -| Setting | Scope | Default | -|---------|-------|---------| -| `THREAD_SUMMARY_EVERY_N` | Per `(user_id, thread_id)` | `0` (disabled) | -| `FACT_EXTRACTION_EVERY_N` | Per `(user_id, thread_id)` | `0` (disabled) | -| `USER_SUMMARY_EVERY_N` | Per `user_id` (across all threads) | `0` (disabled) | +| Setting | Scope | Default | +|---------------------------|------------------------------------|----------------| +| `THREAD_SUMMARY_EVERY_N` | Per `(user_id, thread_id)` | `0` (disabled) | +| `FACT_EXTRACTION_EVERY_N` | Per `(user_id, thread_id)` | `0` (disabled) | +| `USER_SUMMARY_EVERY_N` | Per `user_id` (across all threads) | `0` (disabled) | Set any value to `0` to disable that processing type. For example, setting `THREAD_SUMMARY_EVERY_N=5` generates a thread summary every 5 new turns in each thread. ### Required containers -| Container | Partition Key | Purpose | -|-----------|---------------|---------| -| `memories` | `/user_id`, `/thread_id` (hierarchical) | Durable derived memories (`fact`, `episodic`, `procedural`) | -| `memories_turns` | `/user_id`, `/thread_id` (hierarchical) | Raw conversation turns (`turn`) — append-only, TTL-pruned | -| `memories_summaries` | `/user_id`, `/thread_id` (hierarchical) | Thread + user summaries (`thread_summary`, `user_summary`) | -| `counter` | `/user_id`, `/thread_id` (hierarchical) | Message count tracking for automatic processing | -| `leases` | `/id` | Change feed checkpointing container created by `create_memory_store()` | +| Container | Partition Key | Purpose | +|----------------------|-----------------------------------------|------------------------------------------------------------------------| +| `memories` | `/user_id`, `/thread_id` (hierarchical) | Durable derived memories (`fact`, `episodic`, `procedural`) | +| `memories_turns` | `/user_id`, `/thread_id` (hierarchical) | Raw conversation turns (`turn`) — append-only, TTL-pruned | +| `memories_summaries` | `/user_id`, `/thread_id` (hierarchical) | Thread + user summaries (`thread_summary`, `user_summary`) | +| `counter` | `/user_id`, `/thread_id` (hierarchical) | Message count tracking for automatic processing | +| `leases` | `/id` | Change feed checkpointing container created by `create_memory_store()` | ### Throughput configuration @@ -224,10 +225,10 @@ This keeps the change feed dependencies aligned with the main memory store inste ### Push vs. pull -| Mode | Trigger | Use case | -|------|---------|----------| -| **On-demand (pull)** | SDK call (`generate_thread_summary()`, etc.) | Explicit control over when processing happens | -| **Automatic (push)** | Change feed trigger | Fire-and-forget — processing happens in the background as turns are written | +| Mode | Trigger | Use case | +|----------------------|----------------------------------------------|-----------------------------------------------------------------------------| +| **On-demand (pull)** | SDK call (`generate_thread_summary()`, etc.) | Explicit control over when processing happens | +| **Automatic (push)** | Change feed trigger | Fire-and-forget — processing happens in the background as turns are written | Both modes use the same Durable Functions orchestrator and activities, so prompts, incremental update logic, and stored outputs are identical. @@ -235,10 +236,10 @@ Both modes use the same Durable Functions orchestrator and activities, so prompt ## Local vs. Cloud Storage -| Backend | Use Case | Persistence | -|---------|----------|-------------| -| **Local (in-memory)** | Development and quick testing | Lost on process exit | -| **Cosmos DB** | Production, shared access, semantic search | Durable | +| Backend | Use Case | Persistence | +|-----------------------|--------------------------------------------|----------------------| +| **Local (in-memory)** | Development and quick testing | Lost on process exit | +| **Cosmos DB** | Production, shared access, semantic search | Durable | Local storage is enough for CRUD testing. Cosmos DB is required for persistence, vector search, and the processing pipeline. diff --git a/Docs/public_api.md b/Docs/public_api.md index bec37e4..96337d3 100644 --- a/Docs/public_api.md +++ b/Docs/public_api.md @@ -52,7 +52,7 @@ - `synthesize_procedural(user_id, *, force=False) -> dict` — synthesize the procedural prompt. - `generate_thread_summary(user_id, thread_id, recent_k=None, **kwargs) -> dict` — generate and persist a thread summary. - `generate_user_summary(user_id, thread_ids=None, recent_k=None, **kwargs) -> dict` — generate and persist a user summary. -- `reconcile(user_id, n=None) -> dict[str, int]` — reconcile duplicate or contradictory facts. +- `reconcile(user_id, n=None) -> dict[str, int]` — resolve contradictory facts (paraphrases fold at write time). - `process_now(*, user_id, thread_id) -> ProcessThreadResult` — run the configured processor immediately. - `process_now_and_wait(*, user_id, thread_id, timeout=30.0) -> bool` — process and wait for a summary. @@ -106,7 +106,7 @@ Local-buffer methods remain synchronous in-memory operations; Cosmos, retrieval, - `async synthesize_procedural(user_id, *, force=False) -> dict` — synthesize the procedural prompt. - `async generate_thread_summary(user_id, thread_id, recent_k=None, **kwargs) -> dict` — generate and persist a thread summary. - `async generate_user_summary(user_id, thread_ids=None, recent_k=None, **kwargs) -> dict` — generate and persist a user summary. -- `async reconcile(user_id, n=None) -> dict[str, int]` — reconcile duplicate or contradictory facts. +- `async reconcile(user_id, n=None) -> dict[str, int]` — resolve contradictory facts (paraphrases fold at write time). - `async process_now(*, user_id, thread_id) -> ProcessThreadResult` — run the configured processor immediately. - `async process_now_and_wait(*, user_id, thread_id, timeout=30.0) -> bool` — process and wait for a summary. diff --git a/README.md b/README.md index 971e62e..1927225 100644 --- a/README.md +++ b/README.md @@ -126,44 +126,44 @@ See [`Samples/`](Samples/) for end-to-end scenarios (chat memory, RAG, multi-age ## Concepts in 60 seconds -| Concept | What it is | API | -|---|---|---| -| **Turn** | One message (user or assistant) — the raw conversation atom | `add_cosmos(...)`, `add_local(...)` | -| **Thread summary** | LLM-generated, incrementally updated rollup of a single thread | `generate_thread_summary(...)` | -| **Fact** | Discrete, independently searchable assertion extracted from turns | `extract_memories(...)` | -| **Procedural** | Behavioral rule / instruction the user wants followed | `extract_memories(...)` | -| **Episodic** | Past situation → action → outcome experience (90-day TTL) | `extract_memories(...)` | -| **User summary** | Cross-thread profile of what's known about a user | `generate_user_summary(...)`, `get_user_summary(...)` | -| **Search** | Vector + full-text + filter over `fact` / `episodic` / `procedural` | `search_cosmos(...)` | -| **Process now** | Run the full pipeline (summary → facts → user profile) for recent turns | `process_now(...)`, `process_now_and_wait(...)` | +| Concept | What it is | API | +|--------------------|-------------------------------------------------------------------------|-------------------------------------------------------| +| **Turn** | One message (user or assistant) — the raw conversation atom | `add_cosmos(...)`, `add_local(...)` | +| **Thread summary** | LLM-generated, incrementally updated rollup of a single thread | `generate_thread_summary(...)` | +| **Fact** | Discrete, independently searchable assertion extracted from turns | `extract_memories(...)` | +| **Procedural** | Behavioral rule / instruction the user wants followed | `extract_memories(...)` | +| **Episodic** | Past situation → action → outcome experience (90-day TTL) | `extract_memories(...)` | +| **User summary** | Cross-thread profile of what's known about a user | `generate_user_summary(...)`, `get_user_summary(...)` | +| **Search** | Vector + full-text + filter over `fact` / `episodic` / `procedural` | `search_cosmos(...)` | +| **Process now** | Run the full pipeline (summary → facts → user profile) for recent turns | `process_now(...)`, `process_now_and_wait(...)` | AgentMemoryToolkit uses 3-container Cosmos topology, all partitioned by hierarchical `(user_id, thread_id)` keys: -| Container | Holds | Notes | -|---|---|---| -| `memories_turns` | raw `turn` documents | append-only conversation timeline | -| `memories` | `fact`, `episodic`, `procedural` documents | vector + full-text retrieval path | -| `memories_summaries` | thread and user summaries | latest-summary point/read path | +| Container | Holds | Notes | +|----------------------|--------------------------------------------|-----------------------------------| +| `memories_turns` | raw `turn` documents | append-only conversation timeline | +| `memories` | `fact`, `episodic`, `procedural` documents | vector + full-text retrieval path | +| `memories_summaries` | thread and user summaries | latest-summary point/read path | ### Memory Type Taxonomy The `extract_memories` pipeline classifies each item it pulls from the conversation into one of four buckets. Every memory carries a top-level `confidence` (0.0–1.0) so retrieval can suppress weakly-grounded extractions. -| Bucket | Meaning | Storage type | TTL | -|---|---|---|---| -| Fact | Declarative knowledge ("user prefers dark mode") | `type="fact"` | none | -| Procedural | Behavioral rule ("always confirm before deleting") | `type="procedural"` | none | -| Episodic | Past experience: situation → action → outcome | `type="episodic"` | 90 days | -| Unclassified | Item worth keeping but the LLM couldn't confidently classify | `type="fact"` + tag `sys:unclassified` | none | +| Bucket | Meaning | Storage type | TTL | +|--------------|--------------------------------------------------------------|----------------------------------------|---------| +| Fact | Declarative knowledge ("user prefers dark mode") | `type="fact"` | none | +| Procedural | Behavioral rule ("always confirm before deleting") | `type="procedural"` | none | +| Episodic | Past experience: situation → action → outcome | `type="episodic"` | 90 days | +| Unclassified | Item worth keeping but the LLM couldn't confidently classify | `type="fact"` + tag `sys:unclassified` | none | #### Confidence Scale -| Range | Meaning | -|---|---| -| 0.9–1.0 | Directly stated and unambiguous | -| 0.7–0.9 | Clearly implied, no contradicting evidence | +| Range | Meaning | +|---------|----------------------------------------------------| +| 0.9–1.0 | Directly stated and unambiguous | +| 0.7–0.9 | Clearly implied, no contradicting evidence | | 0.5–0.7 | Inferred from context — plausible but not explicit | -| < 0.5 | Should be in `unclassified` instead | +| < 0.5 | Should be in `unclassified` instead | Filter at retrieval time: @@ -174,28 +174,28 @@ high_conf_facts = memory.get_memories(user_id="u1", memory_types=["fact"], min_c ### Memory Reconciliation -`reconcile(user_id, n=50)` (on the public client; underlying pipeline method is `ProcessingPipeline.reconcile_memories`) collapses paraphrased duplicates and resolves semantic contradictions in a single LLM pass over the N most-recent active facts. Both outcomes soft-delete the loser with a `supersede_reason` of `"duplicate"` or `"contradict"`. See [Docs/concepts.md](Docs/concepts.md#memory-reconciliation) for details. +`reconcile(user_id, n=50)` (on the public client; underlying pipeline method is `ProcessingPipeline.reconcile_memories`) resolves **semantic contradictions** in a single LLM pass over the N most-recent active facts, soft-deleting each loser with `supersede_reason="contradict"`. Paraphrased duplicates are *not* handled here — they are folded in place at write time by the LLM-free vector dedup (see below), so reconcile stays a bounded, convergent contradiction pass. See [Docs/concepts.md](Docs/concepts.md#memory-reconciliation) for details. > **Cost note.** Each reconciliation makes one LLM call covering up to `n` facts (default 50, hard cap 500). With auto-trigger, this fires every `FACT_EXTRACTION_EVERY_N × DEDUP_EVERY_N` turns per user, with `n` taken from `DEDUP_POOL_SIZE`. The previous cosine-cluster pre-filter was removed deliberately — it could not catch semantic contradictions like "vegetarian" vs "ribeye steak" — so the LLM is now invoked whenever there are ≥ 2 active facts. To bound LLM cost more tightly: raise `DEDUP_EVERY_N` (lower frequency — reconcile fires every Nth extraction, so a *higher* N means *less often*), lower `DEDUP_POOL_SIZE` (smaller per-call pool), or override `n` per call when invoking `reconcile()` directly. -| New `MemoryRecord` field | Meaning | -|---|---| -| `content_hash` | SHA-256 of normalized content; enables write-time exact-dedup short-circuit | -| `supersede_reason` | `"duplicate"` or `"contradict"` (None for live records) | -| `superseded_at` | ISO timestamp when the supersede happened (None for live records) | -| `superseded_by` | Id of the record that replaced this one (existing field) | +| New `MemoryRecord` field | Meaning | +|--------------------------|-----------------------------------------------------------------------------| +| `content_hash` | SHA-256 of normalized content; enables write-time exact-dedup short-circuit | +| `supersede_reason` | `"contradict"` (None for live records) | +| `superseded_at` | ISO timestamp when the supersede happened (None for live records) | +| `superseded_by` | Id of the record that replaced this one (existing field) | ### Auto-trigger (per-turn extraction) By default, the **InProcess processor** runs each pipeline step independently as its own threshold trips inside `push_to_cosmos()`: -| Env var | Default | Step that fires | Async behavior | -|---|---|---|---| -| `FACT_EXTRACTION_EVERY_N` | `1` (every turn) | `process_extract_memories` | scheduled via `asyncio.create_task` | -| `DEDUP_EVERY_N` | `5` | `process_reconcile` (fires every Nth extract → effectively every `FACT_EXTRACTION_EVERY_N × DEDUP_EVERY_N` turns) | scheduled via `asyncio.create_task` | -| `DEDUP_POOL_SIZE` | `50` | pool size (`n`) passed to `process_reconcile` from the auto-trigger; hard-capped at `500` | n/a (per-call) | -| `THREAD_SUMMARY_EVERY_N` | `10` | `process_thread_summary` | scheduled via `asyncio.create_task` | -| `USER_SUMMARY_EVERY_N` | `20` | `process_user_summary` | scheduled via `asyncio.create_task` | +| Env var | Default | Step that fires | Async behavior | +|---------------------------|------------------|-------------------------------------------------------------------------------------------------------------------|-------------------------------------| +| `FACT_EXTRACTION_EVERY_N` | `1` (every turn) | `process_extract_memories` | scheduled via `asyncio.create_task` | +| `DEDUP_EVERY_N` | `5` | `process_reconcile` (fires every Nth extract → effectively every `FACT_EXTRACTION_EVERY_N × DEDUP_EVERY_N` turns) | scheduled via `asyncio.create_task` | +| `DEDUP_POOL_SIZE` | `50` | pool size (`n`) passed to `process_reconcile` from the auto-trigger; hard-capped at `500` | n/a (per-call) | +| `THREAD_SUMMARY_EVERY_N` | `10` | `process_thread_summary` | scheduled via `asyncio.create_task` | +| `USER_SUMMARY_EVERY_N` | `20` | `process_user_summary` | scheduled via `asyncio.create_task` | Each `*_EVERY_N=0` disables only that step. Dedup is gated independently of extract because cross-thread dedup is dramatically more expensive than per-thread extract (it reads every active fact for the user) — running it on every extract slammed AI Foundry. The Durable backend uses the same defaults via the change-feed function app (the function-app `azd` deploy bumps `FACT_EXTRACTION_EVERY_N` to `5` since the FA path is intended for higher-volume workloads). Calling `process_now()` is normally redundant — it remains as an explicit "process now" hook for tests, manual workflows, and operators who set every threshold to `0`. @@ -207,11 +207,11 @@ Both the SDK auto-trigger and the function-app change-feed processor write into Set the env var on **both sides** to make ownership explicit: -| `MEMORY_PROCESSOR_OWNER` | SDK behavior | Function-app behavior | -|---|---|---| -| _unset_ (default) | runs auto-trigger | runs orchestrator (today's behavior) | -| `inprocess` | runs auto-trigger | change-feed trigger skips batch + logs | -| `durable` | auto-trigger logs warning + skips | runs orchestrator | +| `MEMORY_PROCESSOR_OWNER` | SDK behavior | Function-app behavior | +|--------------------------|-----------------------------------|----------------------------------------| +| _unset_ (default) | runs auto-trigger | runs orchestrator (today's behavior) | +| `inprocess` | runs auto-trigger | change-feed trigger skips batch + logs | +| `durable` | auto-trigger logs warning + skips | runs orchestrator | The default (unset) preserves backward compatibility. For any production deployment we recommend setting it on both sides so a misconfiguration produces a loud log line instead of silent double-work. @@ -223,12 +223,12 @@ The default (unset) preserves backward compatibility. For any production deploym Pick at construction time via the `processor=` kwarg. -| | `InProcessProcessor` (default) | `DurableFunctionProcessor` | -|---|---|---| -| Infra | None — just `pip install` | Sibling Azure Function app | -| Best for | Prototypes, low TPS, single-agent | Fleet / multi-agent / high TPS | -| `process_now()` | Synchronous, returns when done | No-op (work runs async on change feed) | -| `process_now_and_wait()` | Returns immediately after flush | Polls until summary visible (RU-costly; tests/demos) | +| | `InProcessProcessor` (default) | `DurableFunctionProcessor` | +|--------------------------|-----------------------------------|------------------------------------------------------| +| Infra | None — just `pip install` | Sibling Azure Function app | +| Best for | Prototypes, low TPS, single-agent | Fleet / multi-agent / high TPS | +| `process_now()` | Synchronous, returns when done | No-op (work runs async on change feed) | +| `process_now_and_wait()` | Returns immediately after flush | Polls until summary visible (RU-costly; tests/demos) | ```python from azure.cosmos.agent_memory import CosmosMemoryClient, DurableFunctionProcessor @@ -265,16 +265,16 @@ memory = CosmosMemoryClient(..., processor=DurableFunctionProcessor()) ## Public API reference -| Symbol | Module | Purpose | -|---|---|---| -| `CosmosMemoryClient` | `azure.cosmos.agent_memory` | Sync client — local CRUD, Cosmos DB I/O, processing | -| `AsyncCosmosMemoryClient` | `azure.cosmos.agent_memory.aio` | Async mirror | -| `MemoryProcessor` | `azure.cosmos.agent_memory` | Protocol that any processor backend implements | -| `InProcessProcessor` | `azure.cosmos.agent_memory` | Default backend — runs the pipeline in-process | -| `DurableFunctionProcessor` | `azure.cosmos.agent_memory` | Marker backend — work runs in sibling Function app via change feed | -| `client.process_now()` | — | Run the pipeline for recent turns (in-process) or no-op (remote) | -| `client.process_now_and_wait()` | — | Opt-in poll until processing completes; useful for tests/demos with the remote backend | -| `MemoryRecord`, `MemoryType`, `Role` | `azure.cosmos.agent_memory` | Pydantic models / enums | +| Symbol | Module | Purpose | +|--------------------------------------|---------------------------------|----------------------------------------------------------------------------------------| +| `CosmosMemoryClient` | `azure.cosmos.agent_memory` | Sync client — local CRUD, Cosmos DB I/O, processing | +| `AsyncCosmosMemoryClient` | `azure.cosmos.agent_memory.aio` | Async mirror | +| `MemoryProcessor` | `azure.cosmos.agent_memory` | Protocol that any processor backend implements | +| `InProcessProcessor` | `azure.cosmos.agent_memory` | Default backend — runs the pipeline in-process | +| `DurableFunctionProcessor` | `azure.cosmos.agent_memory` | Marker backend — work runs in sibling Function app via change feed | +| `client.process_now()` | — | Run the pipeline for recent turns (in-process) or no-op (remote) | +| `client.process_now_and_wait()` | — | Opt-in poll until processing completes; useful for tests/demos with the remote backend | +| `MemoryRecord`, `MemoryType`, `Role` | `azure.cosmos.agent_memory` | Pydantic models / enums | Async equivalents (`AsyncInProcessProcessor`, `AsyncDurableFunctionProcessor`) live in `azure.cosmos.agent_memory.aio`. diff --git a/azure/cosmos/agent_memory/aio/services/pipeline.py b/azure/cosmos/agent_memory/aio/services/pipeline.py index 0dfd5ee..f28ad5d 100644 --- a/azure/cosmos/agent_memory/aio/services/pipeline.py +++ b/azure/cosmos/agent_memory/aio/services/pipeline.py @@ -800,7 +800,7 @@ async def dedup_extracted_memories(self, user_id: str, extracted: dict) -> dict: exclude_ids={doc_id} | set(doc.get("supersedes_ids") or []), ) if not neighbor or not vector_similarity_at_least(score, high, distance_function): - continue # novel — leave in result for persist to ADD + continue # novel - leave in result for persist to ADD neighbor_id = str(neighbor.get("id") or "") if not neighbor_id: @@ -1608,7 +1608,7 @@ async def reconcile_memories(self, user_id: str, n: int = 50, *, memory_type: st Async mirror of the sync contradiction-only reconcile. Near-duplicate paraphrases are folded in place at write time (:meth:`dedup_extracted_memories`); this pass only supersedes the loser - of each ``contradicted_pairs`` entry — no clustering, no merged + of each ``contradicted_pairs`` entry - no clustering, no merged documents, no re-merge churn. Episodic and procedural types are no-ops. Returns ``{"kept", "merged", "contradicted"}`` with ``merged`` always 0. """ @@ -1643,9 +1643,10 @@ async def _reconcile_contradictions( ) -> dict[str, int]: """Async mirror: resolve only ``contradicted_pairs`` within the pool. - ``duplicate_groups`` are ignored (write-time in-place dedup handles - paraphrases), so no merged docs are minted and the pass is convergent. - Returns ``{"kept", "merged": 0, "contradicted"}``. + Paraphrases are not merged here; the dedup prompt/schema no longer emits + duplicate groups because write-time in-place dedup already folds them, so + no merged docs are minted and the pass is convergent. Returns + ``{"kept", "merged": 0, "contradicted"}``. """ if len(facts) <= 1: return {"kept": len(facts), "merged": 0, "contradicted": 0} diff --git a/azure/cosmos/agent_memory/services/pipeline.py b/azure/cosmos/agent_memory/services/pipeline.py index 2acf6a9..c99a2c2 100644 --- a/azure/cosmos/agent_memory/services/pipeline.py +++ b/azure/cosmos/agent_memory/services/pipeline.py @@ -812,23 +812,23 @@ def dedup_extracted_memories( For each newly extracted fact/episodic doc we find its single nearest active same-type neighbor. If similarity is at/above ``SIM_HIGH`` the new - doc is a near-duplicate: we refresh the existing neighbor in place — + doc is a near-duplicate: we refresh the existing neighbor in place - recency-wins content + embedding, unioned tags, max salience/confidence, - bumped ``updated_at`` — keeping its id, and drop the new doc so it is not + bumped ``updated_at`` - keeping its id, and drop the new doc so it is not written as a fresh record. Everything below the threshold is novel and flows through to ``persist_extracted_memories`` unchanged. This makes the write path convergent: a restatement updates one existing document rather than minting a new one that a later reconcile sweep must merge and supersede. There is no ``sys:dup-candidate`` tagging and no - clustering — reconcile only resolves contradictions. + clustering - reconcile only resolves contradictions. In-place folds commit here, before ``persist_extracted_memories`` writes the novel docs. A crash between the two leaves some memories folded and some novel docs unwritten, but the source turns are not stamped ``extracted_at`` until persist completes, so the whole turn set is simply re-extracted on the next run (folds are idempotent and re-ADDs hit - exact-hash/vector dedup) — no data is lost, only repeated. + exact-hash/vector dedup) - no data is lost, only repeated. """ if not threshold_config.get_dedup_vector_enabled(): return extracted @@ -891,7 +891,7 @@ def dedup_extracted_memories( exclude_ids={doc_id} | set(doc.get("supersedes_ids") or []), ) if not neighbor or not vector_similarity_at_least(score, high, distance_function): - continue # novel — leave in result for persist to ADD + continue # novel - leave in result for persist to ADD neighbor_id = str(neighbor.get("id") or "") if not neighbor_id: @@ -906,7 +906,7 @@ def dedup_extracted_memories( inplace_updated += 1 folded_ids.add(doc_id) # If the in-place update failed, leave the doc in result so persist - # ADDs it as a novel record — never silently lose an extraction. + # ADDs it as a novel record - never silently lose an extraction. if folded_ids: for bucket in ("facts", "episodic"): @@ -967,8 +967,8 @@ def _apply_inplace_update(self, neighbor: dict[str, Any], new_doc: dict[str, Any Uses ETag optimistic concurrency: the update is applied with ``IfNotModified`` against the neighbor's ``_etag`` so a concurrent supersede/refresh cannot be clobbered (which would resurrect a - soft-deleted memory or lose an update). On an ETag conflict — or any - other failure — returns False and the caller keeps the new doc as a + soft-deleted memory or lose an update). On an ETag conflict - or any + other failure - returns False and the caller keeps the new doc as a novel ADD, so nothing is lost. Recency wins: the neighbor keeps its id / created_at / partition but takes @@ -1692,7 +1692,7 @@ def reconcile_memories(self, user_id: str, n: int = 50, *, memory_type: str = "f """Resolve contradictions among a user's most-recent active memories. Loads up to ``n`` active (non-superseded) ``memory_type`` records and - asks the dedup prompt to identify ``contradicted_pairs`` — opposing + asks the dedup prompt to identify ``contradicted_pairs`` - opposing claims about the same subject (e.g. "deadline March 1" vs "March 15"). Each loser is soft-deleted with ``supersede_reason="contradict"`` and ``superseded_by`` set to the winner. @@ -1700,7 +1700,7 @@ def reconcile_memories(self, user_id: str, n: int = 50, *, memory_type: str = "f Near-duplicate *paraphrases* are no longer merged here: the write-time in-place dedup (:meth:`dedup_extracted_memories`) folds restatements into their canonical record before they land, so reconcile is a bounded, - convergent contradiction pass — no clustering, no synthesized merged + convergent contradiction pass - no clustering, no synthesized merged documents, no re-merge churn. Episodic and procedural types are no-ops (episodic has no contradiction semantics; its near-dups fold at write time). @@ -1736,11 +1736,12 @@ def reconcile_memories(self, user_id: str, n: int = 50, *, memory_type: str = "f def _reconcile_contradictions(self, user_id: str, memory_type: str, facts: list[dict[str, Any]]) -> dict[str, int]: """Resolve contradictions within an explicit pool of same-type memories. - Runs the dedup prompt over the pool and applies ONLY its - ``contradicted_pairs`` — the loser of each pair is soft-deleted with - ``superseded_by`` set to the winner. ``duplicate_groups`` in the response - are ignored (write-time in-place dedup already handles paraphrases), so - no merged documents are minted and the pass is convergent. Returns + Runs the dedup prompt over the pool and applies its + ``contradicted_pairs`` - the loser of each pair is soft-deleted with + ``superseded_by`` set to the winner. Paraphrases are not merged here; + the dedup prompt/schema no longer emits duplicate groups because + write-time in-place dedup already folds them, so no merged documents are + minted and the pass is convergent. Returns ``{"kept", "merged": 0, "contradicted"}``. """ if len(facts) <= 1: diff --git a/tests/unit/aio/test_auto_trigger.py b/tests/unit/aio/test_auto_trigger.py index 3ab0093..668b134 100644 --- a/tests/unit/aio/test_auto_trigger.py +++ b/tests/unit/aio/test_auto_trigger.py @@ -99,7 +99,6 @@ async def fake_upsert(body): class TestAsyncExtractRecentK: - @pytest.mark.asyncio @pytest.mark.asyncio async def test_extract_fires_without_recent_k_or_watermark(self, monkeypatch): """Async: extraction covers all un-extracted turns (extracted_at gated) and