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 32e7342..35cc968 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,17 @@ ## Release History -## [Unreleased] +## [0.3.0b1] (Unreleased) + +#### Features Added +* 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 +* 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) @@ -17,7 +28,6 @@ `DEDUP_EVERY_N`, `THREAD_SUMMARY_EVERY_N`, `USER_SUMMARY_EVERY_N`); any key not present falls back to the environment/defaults, and `None` preserves today's env-only behavior. See [PR:#29](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/29) - ## [0.2.0b2] (2026-07-01) #### Features Added 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 e3b646f..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,26 +118,53 @@ 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: +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. -- **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. +### Write-time in-place dedup (LLM-free) -### Why one pass +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): -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}`. +| 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 | + +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. + +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. + +### LLM contradiction reconcile + +`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. + +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 -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 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. + +- **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, 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_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. +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`). @@ -169,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 @@ -198,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. @@ -209,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/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..96337d3 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. @@ -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. @@ -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. @@ -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/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/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/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 2d9a518..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,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,23 +104,23 @@ 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) 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") @@ -133,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") @@ -148,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") @@ -163,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/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/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 ba154e8..5c474e0 100644 --- a/azure/cosmos/agent_memory/_base/base_client.py +++ b/azure/cosmos/agent_memory/_base/base_client.py @@ -279,7 +279,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`` @@ -291,7 +291,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 2223d06..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", + "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, ) @@ -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. @@ -303,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. diff --git a/azure/cosmos/agent_memory/_embedding_tokens.py b/azure/cosmos/agent_memory/_embedding_tokens.py new file mode 100644 index 0000000..1cc8e17 --- /dev/null +++ b/azure/cosmos/agent_memory/_embedding_tokens.py @@ -0,0 +1,119 @@ +"""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 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, + 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/_utils.py b/azure/cosmos/agent_memory/_utils.py index 3cae259..d9a97fb 100644 --- a/azure/cosmos/agent_memory/_utils.py +++ b/azure/cosmos/agent_memory/_utils.py @@ -113,7 +113,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()) @@ -260,11 +260,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=( @@ -277,11 +278,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=( @@ -294,17 +296,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=( @@ -316,21 +317,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( @@ -343,28 +402,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( @@ -526,10 +572,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..60df3ae 100644 --- a/azure/cosmos/agent_memory/aio/auto_trigger.py +++ b/azure/cosmos/agent_memory/aio/auto_trigger.py @@ -110,6 +110,7 @@ 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), @@ -125,6 +126,7 @@ async def _fire_thread_steps( user_id: str, thread_id: str, *, + new_count: int, fire_extract: bool, fire_summary: bool, fire_dedup: bool, @@ -138,14 +140,21 @@ async def _fire_thread_steps( default=True, ) ) + if fire_extract: + try: + 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( + 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}, ), - (fire_dedup, "process_reconcile", processor.process_reconcile, {"user_id": user_id}), ( fire_procedural, "synthesize_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 7b24172..d922221 100644 --- a/azure/cosmos/agent_memory/aio/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/aio/cosmos_memory_client.py @@ -590,7 +590,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. @@ -708,7 +708,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, @@ -726,7 +725,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, @@ -744,7 +742,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, @@ -766,7 +763,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, @@ -775,6 +771,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, @@ -881,12 +902,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]: @@ -939,7 +971,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 503de8d..1811826 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, @@ -144,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 cdab0c0..aacc427 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( diff --git a/azure/cosmos/agent_memory/aio/processors/durable.py b/azure/cosmos/agent_memory/aio/processors/durable.py index 677855d..ed38d8f 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", diff --git a/azure/cosmos/agent_memory/aio/processors/inprocess.py b/azure/cosmos/agent_memory/aio/processors/inprocess.py index 9e10071..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 @@ -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( @@ -116,15 +115,26 @@ async def process_user_summary( async def process_reconcile(self, *, user_id: str) -> 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()) + + 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 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) + total += self._extract_reconcile_count(reconciled) + return total @staticmethod 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 0db0475..f28ad5d 100644 --- a/azure/cosmos/agent_memory/aio/services/pipeline.py +++ b/azure/cosmos/agent_memory/aio/services/pipeline.py @@ -4,13 +4,12 @@ :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. """ 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, @@ -50,17 +55,16 @@ 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 ( - 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 +72,11 @@ 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_sim_high, + get_dedup_vector_enabled, + get_extraction_batch_max_tokens, +) logger = get_logger("azure.cosmos.agent_memory.pipeline.aio") @@ -104,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: @@ -252,11 +258,123 @@ async def _embed_one(self, text: str) -> list[float]: async def _embed_batch(self, texts: list[str]) -> list[list[float]]: return await self._embeddings.generate_batch(texts) + async def _vector_distance_function(self) -> str: + """Return the container's configured Cosmos ``distanceFunction`` (cached). + + Read from the container's vector embedding policy (``await container.read()``) + - the authoritative, immutable source set when the container was created. + Drives the ORDER BY direction and similarity-threshold comparisons so dedup + never silently assumes cosine. Falls back to cosine when the policy can't be + read (e.g. ``__new__``-built test instances with mocked containers). + """ + fn = getattr(self, "_distance_function_cache", None) + if fn is not None: + return fn + try: + props = await self._memories_container.read() + except Exception: + # See sync pipeline: don't cache a defaulted cosine, and flag the + # failure so the destructive in-place fold path skips this run rather + # than mis-applying cosine bands to (possibly euclidean) distances. + self._distance_function_read_failed = True + logger.debug( + "vector dedup: could not read container vector policy; defaulting to cosine (not cached)", + exc_info=True, + ) + return "cosine" + fn = distance_function_from_container_properties(props) + self._distance_function_cache = fn + self._distance_function_read_failed = False + return fn + + def _warn_distance_policy_unavailable_once(self) -> None: + """One-shot WARN that in-place folding was skipped (policy unreadable).""" + if getattr(self, "_warned_distance_policy_unavailable", False): + return + self._warned_distance_policy_unavailable = True + logger.warning( + "vector dedup: container vector policy could not be read; skipping 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. + + 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. 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) @@ -298,7 +416,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. """ @@ -340,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, @@ -354,58 +471,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,35 +529,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": []} - existing_facts, existing_episodics = await asyncio.gather( - self._load_existing_memories(user_id, ["fact"]), - self._load_existing_memories(user_id, ["episodic"]), - ) + 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: - existing_text = "\n".join( - f"- [ID: {mem['id']}] {mem.get('content', '')} (type=fact, salience={mem.get('salience', 'N/A')})" - for mem in existing_facts - ) - 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, - }, - ) - parsed = self._parse_llm_json(response_text) - facts = parsed.get("facts", []) - episodic = parsed.get("episodic", []) - unclassified = parsed.get("unclassified", []) + # 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: + 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 + 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]] = [] @@ -502,14 +588,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, @@ -534,10 +619,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, @@ -546,22 +628,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 +647,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 +673,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 +684,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, @@ -649,52 +702,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_facts, + processed_turns, + existing_for_hash, user_id=user_id, thread_id=thread_id, logger=logger, @@ -704,7 +728,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", @@ -716,6 +740,183 @@ async def extract_memories_dry( ) return result + async def dedup_extracted_memories(self, user_id: str, extracted: dict) -> dict: + """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: + raise ValidationError("user_id is required") + if not isinstance(extracted, dict): + raise ValidationError("extracted must be a dict") + + high = get_dedup_sim_high() + distance_function = await self._vector_distance_function() + 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 = { + "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 + + 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 []), + ) + 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).""" + 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) + new_content = str(new_doc.get("content") or "") + 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")]) + 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 + + 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)", + neighbor.get("id"), + exc, + ) + return False + async def persist_extracted_memories( self, user_id: str, @@ -751,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 @@ -763,45 +961,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) - 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 "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 @@ -834,7 +1003,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) - return await self.persist_extracted_memories(user_id, extracted) + # 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) + 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, @@ -943,38 +1138,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()] @@ -993,7 +1162,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): @@ -1031,12 +1200,12 @@ 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: 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, @@ -1392,24 +1561,56 @@ def _emit_reconcile_outcome( }, ) - async def reconcile_memories(self, user_id: str, n: int = 50) -> 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. + 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}, + ], + ) - Returns ``{"kept": int, "merged": int, "contradicted": int}`` where - ``merged`` and ``contradicted`` count the *losers* that were - soft-deleted (duplicates and contradictions respectively). + 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 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 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") @@ -1417,52 +1618,39 @@ 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 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", user_id, n) + logger.info("reconcile_memories started user_id=%s n=%d memory_type=%s", user_id, n, memory_type) - # ---- 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 = await self._reconcile_contradictions(user_id, memory_type, facts) + self._emit_reconcile_outcome( + started_at=started_at, + user_id=user_id, + candidates=len(facts), + result=result, ) + return result + + async def _reconcile_contradictions( + self, user_id: str, memory_type: str, facts: list[dict[str, Any]] + ) -> dict[str, int]: + """Async mirror: resolve only ``contradicted_pairs`` within the pool. + 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: - logger.info( - "reconcile_memories: %d facts, nothing to reconcile", - len(facts), - ) - 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 - - # ---- 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) @@ -1474,357 +1662,51 @@ async def reconcile_memories(self, user_id: str, n: int = 50) -> dict[str, int]: 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 = await self._run_prompty( - "dedup.prompty", - inputs={"facts_text": facts_text}, - ) + 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 = 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 []) 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 not in seen_tags: - seen_tags.add(t) - merged_tags.append(t) - if not merged_tags: - merged_tags = ["sys:fact"] - - # 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_id = "fact_" + hashlib.sha256(merged_id_seed.encode()).hexdigest()[:32] - - try: - merged_record = construct_internal( - FactRecord, - { - "id": merged_id, - "user_id": user_id, - "role": "system", - "type": "fact", - "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": { - "category": "preference", - "merged_via": "reconcile", - "merged_from_count": len(valid_source_ids), - }, - **self._prompt_lineage("dedup.prompty"), - "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, + # 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(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) - self._emit_reconcile_outcome( - started_at=started_at, - user_id=user_id, - candidates=len(facts), - result=result, + 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 diff --git a/azure/cosmos/agent_memory/aio/store/memory_store.py b/azure/cosmos/agent_memory/aio/store/memory_store.py index 4f725f9..4d5942b 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 ( @@ -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 @@ -804,7 +805,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 +823,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 +845,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): @@ -862,13 +867,71 @@ 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, 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 +941,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 +953,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 +970,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..ee41322 100644 --- a/azure/cosmos/agent_memory/auto_trigger.py +++ b/azure/cosmos/agent_memory/auto_trigger.py @@ -154,6 +154,7 @@ 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), @@ -169,6 +170,7 @@ def _fire_thread_steps( user_id: str, thread_id: str, *, + new_count: int, fire_extract: bool, fire_summary: bool, fire_dedup: bool, @@ -182,13 +184,20 @@ def _fire_thread_steps( default=True, ) ) + if fire_extract: + try: + 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( + 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), ), - (fire_dedup, "process_reconcile", lambda: processor.process_reconcile(user_id=user_id)), ( fire_procedural, "synthesize_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 84114cd..5fec042 100644 --- a/azure/cosmos/agent_memory/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/cosmos_memory_client.py @@ -543,7 +543,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. @@ -663,7 +663,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, @@ -686,7 +685,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, @@ -704,7 +702,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, @@ -726,7 +723,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, @@ -735,6 +731,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, @@ -843,15 +864,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, @@ -912,7 +948,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 493c25f..65f08ca 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, @@ -144,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 cd257b7..164c948 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( diff --git a/azure/cosmos/agent_memory/processors/durable.py b/azure/cosmos/agent_memory/processors/durable.py index d63b24f..3a2b580 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", diff --git a/azure/cosmos/agent_memory/processors/inprocess.py b/azure/cosmos/agent_memory/processors/inprocess.py index 9684ce0..734d6b0 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( @@ -127,22 +126,33 @@ def process_user_summary( return UserSummaryResult(summary=summary if isinstance(summary, dict) else None) def process_reconcile(self, *, user_id: str) -> int: - """Run reconciliation standalone. Returns count of facts merged + contradicted. + """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. """ 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()) + + 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 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) + total += self._extract_reconcile_count(reconciled) + return total @staticmethod 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 1b34973..2f43a12 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, } # --------------------------------------------------------------------------- -# extract_memories.prompty — extract facts + episodic + unclassified +# extract_memories.prompty - extract facts + episodic + unclassified # --------------------------------------------------------------------------- _FACT_ITEM = { "type": "object", @@ -77,35 +63,22 @@ "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"]}, "tags": {"type": "array", "items": {"type": "string"}}, - "action": {"type": "string", "enum": ["ADD", "UPDATE", "CONTRADICT"]}, - "supersedes_id": {"type": ["string", "null"]}, }, "required": [ "text", "category", - "subject", - "predicate", - "object", "confidence", "salience", "temporal_context", "tags", - "action", - "supersedes_id", ], "additionalProperties": False, } @@ -115,7 +88,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 +105,6 @@ "required": [ "scope_type", "scope_value", - "text", "situation", "action_taken", "outcome", @@ -148,33 +119,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 @@ -216,14 +173,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 @@ -258,14 +215,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", @@ -278,7 +235,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 c83a05a..21a80d9 100644 --- a/azure/cosmos/agent_memory/prompts/dedup.prompty +++ b/azure/cosmos/agent_memory/prompts/dedup.prompty @@ -1,12 +1,12 @@ --- 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: seed: 43 - maxOutputTokens: 2000 + maxOutputTokens: 16384 additionalProperties: response_format: type: json_object @@ -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 d266f1d..1ce30a9 100644 --- a/azure/cosmos/agent_memory/prompts/extract_memories.prompty +++ b/azure/cosmos/agent_memory/prompts/extract_memories.prompty @@ -1,44 +1,43 @@ --- 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: seed: 42 - maxOutputTokens: 2500 + maxOutputTokens: 16384 additionalProperties: response_format: type: json_object inputs: - existing_facts: - type: string - default: '(none)' - existing_episodics: - type: string - default: '(none)' transcript: type: string --- 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. -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 conversation — 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. -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. +## Speaker Discrimination - Where Memories May Come From -- **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. +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. ## Confidence Scoring @@ -48,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: @@ -63,135 +62,54 @@ 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 **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 -- 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. +- **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`). -- **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."` +- **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`). -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) -- **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 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. --- @@ -209,38 +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 - -### Existing facts (REFERENCE ONLY — never source ADDs from this list) - -The list below shows facts already stored for this user. Use it ONLY for these three purposes: - -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. - -{{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}} +- 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. --- @@ -260,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 @@ -282,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. Is every emitted fact grounded in a `[user]:` line in this transcript? (Existing-facts entries are reference-only; never source new ADDs from them.) +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? @@ -304,49 +191,28 @@ 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"], - "action": "ADD", - "supersedes_id": null - }, - { - "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"], - "action": "ADD", - "supersedes_id": 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", - "tags": ["topic:project", "topic:ETL"], - "action": "ADD", - "supersedes_id": null + "tags": ["topic:project", "topic:ETL"] } ], "episodic": [] } ``` +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:** @@ -358,23 +224,17 @@ 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", - "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.", @@ -402,28 +262,18 @@ 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, - "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.", "category": "requirement", - "subject": "marketing team", - "predicate": "budget", - "object": "$50,000 for the current quarter", "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": [] @@ -436,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 @@ -445,22 +295,16 @@ 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, - "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,115 +320,48 @@ 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. -``` +### Example 5: Discipline - extract only what was stated, infer nothing **Conversation:** -> User: "For this Tokyo trip, I want luxury hotels." -> User: "Normally, I prefer moderate hotels." +> User: "The dashboard took about 8 seconds to load the report. Anyway, the weather's been nice this week." -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`. +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 normally prefers moderate hotels.", - "category": "preference", - "subject": "user", - "predicate": "hotel_preference", - "object": "moderate hotels", - "confidence": 0.95, - "salience": 0.7, + "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:travel", "topic:hotels"], - "action": "ADD", - "supersedes_id": null + "tags": ["topic:performance", "topic:dashboard"] } ], "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 +### Example 6: Capturing a transition as a single memory **Conversation:** -> User: "I don't eat meat and I need wheelchair-accessible restaurants." +> User: "I moved our CI from Jenkins to GitHub Actions last quarter - the Jenkins maintenance was eating too much of my time." -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. +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 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 + "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": [] @@ -595,30 +372,24 @@ A single user turn that combines a `preference` (diet) and a `requirement` (acce ## 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. ```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", - "tags": ["topic:x"], - "action": "ADD|UPDATE|CONTRADICT", - "supersedes_id": "id 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"] } ], "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", @@ -630,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 14ddf8f..a44314a 100644 --- a/azure/cosmos/agent_memory/services/_pipeline_helpers.py +++ b/azure/cosmos/agent_memory/services/_pipeline_helpers.py @@ -16,7 +16,58 @@ from pathlib import Path from typing import Any, Iterable, Mapping, 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 @@ -267,7 +318,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. @@ -276,7 +327,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 @@ -309,47 +360,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. @@ -407,6 +417,46 @@ def format_existing_episodics(memories: list[dict[str, Any]]) -> str: "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", } ) @@ -433,7 +483,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 @@ -441,7 +491,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 @@ -453,8 +503,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: @@ -488,7 +538,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), @@ -505,7 +555,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"), @@ -531,11 +581,38 @@ 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()) + obj, end = json.JSONDecoder().raw_decode(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 + 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: + """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: @@ -604,7 +681,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 2211669..c99a2c2 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, @@ -47,17 +54,16 @@ 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 ( - format_existing_episodics as _format_existing_episodics, -) from azure.cosmos.agent_memory.services._pipeline_helpers import ( is_real_number as _is_real_number, ) @@ -206,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) @@ -248,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. """ @@ -276,6 +282,180 @@ 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). + # Flag the failure so the *destructive* in-place fold path can skip + # entirely (a defaulted cosine on a euclidean container would fold and + # overwrite unrelated memories). + self._distance_function_read_failed = True + logger.debug( + "vector dedup: could not read container vector policy; defaulting to cosine (not cached)", + exc_info=True, + ) + return "cosine" + fn = distance_function_from_container_properties(props) + self._distance_function_cache = fn + self._distance_function_read_failed = False + return fn + + def _warn_euclidean_autodrop_once(self, distance_function: str) -> None: + """One-shot WARN that the near-exact vector auto-drop is disabled. + + The ``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 _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, + *, + 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 _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) @@ -300,11 +480,11 @@ 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, "dropped_episodic_count": 0, + "inplace_updated": 0, } @staticmethod @@ -314,56 +494,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,33 +592,55 @@ 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: - existing_text = "\n".join( - f"- [ID: {mem['id']}] {mem.get('content', '')} (type=fact, salience={mem.get('salience', 'N/A')})" - for mem in existing_facts - ) - 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, - }, - ) - parsed = self._parse_llm_json(response_text) - facts = parsed.get("facts", []) - episodic = parsed.get("episodic", []) - unclassified = parsed.get("unclassified", []) + # 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: + 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 + 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]] = [] @@ -498,14 +650,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, @@ -530,10 +681,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, @@ -542,22 +690,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 +709,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 +735,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 +746,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, @@ -645,52 +764,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_facts, + processed_turns, + existing_for_hashes, user_id=user_id, thread_id=thread_id, logger=logger, @@ -700,7 +790,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", @@ -712,6 +802,234 @@ 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]]]: + """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. + + 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 + 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() + distance_function = self._vector_distance_function() + 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]]] = { + "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")] + # 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 = self._embed_batch([str(doc["content"]) for doc in missing_embeddings]) + for doc, embedding in zip(missing_embeddings, embeddings): + doc["embedding"] = embedding + + 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 = self._nearest_active_full( + user_id=user_id, + embedding=embedding, + memory_type=memory_type, + 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 + + 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 + + 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, + ) + ) + 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. + + 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. + """ + 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) + new_content = str(new_doc.get("content") or "") + 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")]) + 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 + + 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)", + neighbor.get("id"), + exc, + ) + return False + def persist_extracted_memories( self, user_id: str, @@ -752,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 @@ -764,40 +1079,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) - 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 ("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: @@ -805,7 +1092,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: @@ -839,7 +1126,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) - return self.persist_extracted_memories(user_id, extracted) + # 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) + 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, @@ -954,40 +1267,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()] @@ -1006,7 +1291,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): @@ -1049,7 +1334,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, @@ -1403,24 +1688,24 @@ def _emit_reconcile_outcome( }, ) - def reconcile_memories(self, user_id: str, n: int = 50) -> dict[str, int]: - """Reconcile a user's active facts in a single LLM pass. + 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. - 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: + 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. - * **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. + 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": int, "merged": int, "contradicted": int}`` where - ``merged`` and ``contradicted`` count the *losers* that were - soft-deleted (duplicates and contradictions respectively). + Returns ``{"kept", "merged", "contradicted"}``; ``merged`` is always 0. """ if not user_id: raise ValidationError("user_id is required") @@ -1428,54 +1713,40 @@ 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 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", user_id, n) + logger.info("reconcile_memories started user_id=%s n=%d memory_type=%s", user_id, n, memory_type) - # ---- 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 = list( - self._memories_container.query_items( - query=query, - parameters=parameters, - enable_cross_partition_query=True, - ) + facts = self._active_memories_for_reconcile(user_id, memory_type, n) + result = self._reconcile_contradictions(user_id, memory_type, facts) + self._emit_reconcile_outcome( + started_at=started_at, + user_id=user_id, + candidates=len(facts), + result=result, ) + return result + + 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 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: - logger.info( - "reconcile_memories: %d facts, nothing to reconcile", - len(facts), - ) - 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 - - # ---- 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) @@ -1487,360 +1758,80 @@ def reconcile_memories(self, user_id: str, n: int = 50) -> dict[str, int]: 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( - "dedup.prompty", - inputs={"facts_text": 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 [] - # ``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 []) 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 not in seen_tags: - seen_tags.add(t) - merged_tags.append(t) - if not merged_tags: - merged_tags = ["sys:fact"] - - # 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_id = "fact_" + hashlib.sha256(merged_id_seed.encode()).hexdigest()[:32] - - try: - merged_record = construct_internal( - FactRecord, - { - "id": merged_id, - "user_id": user_id, - "role": "system", - "type": "fact", - "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": { - "category": "preference", - "merged_via": "reconcile", - "merged_from_count": len(valid_source_ids), - }, - **self._prompt_lineage("dedup.prompty"), - "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 = 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, + # 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(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, - ) - 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, + 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, + ) + ) + def build_procedural_context(self, user_id: str) -> str: """Return the active synthesized procedural prompt for system injection.""" if not user_id: diff --git a/azure/cosmos/agent_memory/store/_search_helpers.py b/azure/cosmos/agent_memory/store/_search_helpers.py index b0f6c5e..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") @@ -94,16 +95,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..f67b532 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 ( @@ -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 @@ -840,7 +841,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 +859,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 +881,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): @@ -899,13 +904,72 @@ 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, 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 +979,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 +991,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 +1008,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..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. @@ -28,6 +28,17 @@ # 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. +# --------------------------------------------------------------------------- +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, "episodic": 7_776_000, @@ -48,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. @@ -126,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, @@ -138,6 +149,27 @@ 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_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: + """Whether Stage-3 vector deduplication is enabled (internal; on).""" + return DEDUP_VECTOR_ENABLED + + +def get_dedup_sim_high() -> float: + """Similarity at/above which a new memory is folded into its existing + canonical record in place (internal).""" + return DEDUP_SIM_HIGH + + def get_procedural_synthesis_auto() -> bool: """Whether procedural synthesis auto-fires after extract. @@ -164,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 @@ -180,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") @@ -214,6 +246,9 @@ def get_processor_owner() -> Optional[str]: "get_user_summary_every_n", "get_dedup_every_n", "get_dedup_pool_size", + "get_extraction_batch_max_tokens", + "get_dedup_vector_enabled", + "get_dedup_sim_high", "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..4952b11 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,19 +34,36 @@ def ExtractMemoriesOrchestrator(context: df.DurableOrchestrationContext): user_id = payload["user_id"] thread_id = payload["thread_id"] should_reconcile = bool(payload.get("reconcile", 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: @@ -86,11 +104,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 +123,18 @@ 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 +147,34 @@ 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"] 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") 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 f15ed35..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 @@ -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/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 1ab09cc..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 ( @@ -177,7 +177,9 @@ 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") - # Stamp the writing backend (advisory only — not enforced server-side). + 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 elif existing_doc is not None and "last_owner" in existing_doc: @@ -220,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 @@ -233,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") @@ -253,9 +255,46 @@ 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: 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/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 c7b2c8b..80d806e 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, ) @@ -97,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, @@ -109,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, @@ -206,6 +207,16 @@ 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)) + 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 +226,7 @@ async def process_changefeed_batch( "thread_id": thread_id, "count": new_count, "reconcile": should_reconcile, + "recent_k": recent_k, }, orchestration_errors, ) @@ -274,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. @@ -320,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 cf7554a..1f95742 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/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 new file mode 100644 index 0000000..2bb6347 --- /dev/null +++ b/tests/integration/test_async_full_pipeline.py @@ -0,0 +1,269 @@ +"""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_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 6e80bd9..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``). @@ -122,6 +122,67 @@ 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 # --------------------------------------------------------------------------- @@ -138,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) @@ -243,7 +304,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 +329,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 +564,60 @@ 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..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. """ @@ -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..e96b8d7 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), + ] + ) # --------------------------------------------------------------------------- @@ -158,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 2d0852b..c933e7d 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 @@ -52,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") @@ -62,8 +63,17 @@ 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", + } + assert pipeline.reconcile_memories.await_args_list[1].kwargs == { + "n": 37, + "memory_type": "episodic", + } + assert count == 10 # (merged+contradicted) x2 types @pytest.mark.asyncio @@ -75,9 +85,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..3667be0 --- /dev/null +++ b/tests/unit/aio/services/test_dedup_vector_async.py @@ -0,0 +1,354 @@ +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_dedup_extracted_folds_near_dup_in_place_and_keeps_novel(): + p = _service() + 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) + + 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_dedup_extracted_failed_inplace_update_keeps_new_doc(): + p = _service() + 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": []} + + out = await p.dedup_extracted_memories("u1", extracted) + + 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_dedup_extracted_below_threshold_is_novel(): + p = _service() + 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": []} + + out = await p.dedup_extracted_memories("u1", extracted) + + assert [doc["id"] for doc in out["facts"]] == ["f-new"] + p._apply_inplace_update.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_euclidean_disables_inplace_folding(): + p = _service() + 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 [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_apply_inplace_update_recency_wins_and_unions(): + p = _service() + p._replace_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 + + ok = await p._apply_inplace_update(neighbor, new_doc) + + assert ok is True + 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" # recency wins + 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_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 + + 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() + doc_a = _fact("a", "first") + doc_b = _fact("b", "second") + + 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_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", "content")], "episodic": [], "updates": []} + + out = await p.dedup_extracted_memories("u1", extracted) + + assert out is extracted + p._embed_batch.assert_not_awaited() + + +@pytest.mark.asyncio +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_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) + + 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": []}) + + await asyncio.gather(run("userA"), run("userB")) + assert set(seen_users) == {"userA", "userB"} + + +@pytest.mark.asyncio +async def test_reconcile_memory_type_routes_episodic_and_procedural_noop(): + p = _service() + p._run_prompty = AsyncMock() + + episodic_result = await p.reconcile_memories("u1", memory_type="episodic") + assert episodic_result == {"kept": 0, "merged": 0, "contradicted": 0} + + 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_reconcile_fact_contradiction_only(): + p = _service() + 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": "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", memory_type="fact") + + 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 f244563..668b134 100644 --- a/tests/unit/aio/test_auto_trigger.py +++ b/tests/unit/aio/test_auto_trigger.py @@ -11,11 +11,46 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from azure.cosmos.exceptions import CosmosResourceNotFoundError 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 +60,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 +98,69 @@ async def fake_upsert(body): await asyncio.gather(*list(client._background_tasks), return_exceptions=True) +class TestAsyncExtractRecentK: + @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 + 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) + 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), + ): + 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") + + @pytest.mark.asyncio + 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) + 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.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) + + stamp.assert_called_once() + + 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 2b7904a..e3dfb3e 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, @@ -216,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 == {} @@ -433,7 +435,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() @@ -455,29 +458,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 @@ -802,6 +804,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()] @@ -812,7 +845,6 @@ async def test_search_hybrid(self): results = await mem.search_cosmos( search_terms="weather Seattle", - hybrid_search=True, top_k=5, ) @@ -852,6 +884,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_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 25ae108..83d8be5 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"), + call("u", n=50, memory_type="episodic"), + ] + ) 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"} @@ -170,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 858e328..ef2fdbd 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_vector_enabled", + lambda: False, + ) + + def _make_async_pipeline() -> AsyncPipelineService: p = AsyncPipelineService.__new__(AsyncPipelineService) p._embeddings = MagicMock() 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 f5b7dd8..18b6dbb 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 @@ -409,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)) @@ -418,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) @@ -478,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. # --------------------------------------------------------------------------- @@ -546,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 @@ -565,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) @@ -609,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() @@ -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,29 @@ 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": "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_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 4f8f60f..634514b 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"} 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", "limit": 20} + 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", "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,90 @@ 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"}) + + assert pipeline.reconcile_memories.call_args_list == [ + 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_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/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 be8d841..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 @@ -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/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 d12e61b..33afa5e 100644 --- a/tests/unit/services/test_chaos_extract_persist.py +++ b/tests/unit/services/test_chaos_extract_persist.py @@ -11,6 +11,18 @@ 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.aio.services.pipeline.get_dedup_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..9cabf3f --- /dev/null +++ b/tests/unit/services/test_dedup_vector.py @@ -0,0 +1,406 @@ +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_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._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-1", "content": "same", "type": "fact"}, 0.99), + (None, 0.0), + ] + ) + p._apply_inplace_update = MagicMock(return_value=True) + extracted = { + "facts": [ + _doc("f-dup", "restatement of existing"), + _doc("f-novel", "brand new fact"), + ], + "episodic": [], + "updates": [], + } + + out = p.dedup_extracted_memories("u1", extracted) + + # 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_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() + 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": []} + + out = p.dedup_extracted_memories("u1", extracted) + + 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_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="cosine") + p._embed_batch.return_value = [[1.0, 0.0]] + 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) + + assert [doc["id"] for doc in out["facts"]] == ["f-new"] + p._apply_inplace_update.assert_not_called() + + +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() + 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": [], + } + + out = p.dedup_extracted_memories("u1", extracted) + + assert out["facts"] == [] + assert p._apply_inplace_update.call_count == 1 + assert out["updates"][-1]["inplace_updated"] == 1 + + +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() + 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": []} + + out = p.dedup_extracted_memories("u1", extracted) + + 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_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() + 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], + ) + + ok = p._apply_inplace_update(neighbor, new_doc) + + assert ok is True + # 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" # recency wins + 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_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() + 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") + 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: + 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() -> 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() + + episodic_result = p.reconcile_memories("u1", memory_type="episodic") + assert episodic_result == {"kept": 0, "merged": 0, "contradicted": 0} + + procedural_result = p.reconcile_memories("u1", memory_type="procedural") + assert procedural_result == {"kept": 0, "merged": 0, "contradicted": 0} + + p._run_prompty.assert_not_called() + + +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() + 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 == {"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" + + +def test_reconcile_skips_chained_contradiction() -> None: + # (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" diff --git a/tests/unit/services/test_extract_dry.py b/tests/unit/services/test_extract_dry.py index 32b4390..97bae44 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,6 +246,23 @@ def test_extract_memories_dry_is_byte_deterministic_for_same_llm_response() -> N ) +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([]) + 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([_turn(1)])), + ) + + service.extract_memories_dry("u1", "t1") + + assert memories_store.search_calls == [] + + @pytest.mark.asyncio async def test_async_extract_memories_dry_shape_is_small_and_has_no_embeddings() -> None: chat = _AsyncChat([_response()]) @@ -261,258 +303,63 @@ async def test_async_extract_memories_dry_is_byte_deterministic_for_same_llm_res ) -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.""" - 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), - ) - - output = 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"} - - -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()]) - 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), - ) - - service.extract_memories_dry("u1", "t1") - - for doc in turns_store.docs: - assert "extracted_at" not in doc or doc["extracted_at"] is None - - -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.""" - 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), - ) - - 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), - ) - - service.extract_memories("u1", "t1") - calls_after_first = chat.calls - - # 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 - - @pytest.mark.asyncio -async def test_async_extract_memories_marks_turns_after_successful_persist() -> None: - chat = _AsyncChat([_response()]) - memories_store = _AsyncStore([]) - turns_store = _AsyncStore([_turn(i) for i in range(3)]) +async def test_async_extract_memories_dry_does_not_call_store_search() -> None: + store = _AsyncStore([]) + + store.search = AsyncMock(return_value=[]) 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=_AsyncStore([_turn(1)])), ) - 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"], - } + await service.extract_memories_dry("u1", "t1") + store.search.assert_not_awaited() -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", - } +class _BatchChat: + """Returns a unique fact per call; optionally raises on a chosen call index.""" -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", - } - + def __init__(self, fail_on_call=None, error=None): + self.calls = 0 + self.fail_on_call = fail_on_call + self.error = error -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, - }, + 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( { - # 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), - ) + "facts": [ + { + "text": f"Fact from call {self.calls}.", + "category": "other", + "confidence": 0.9, + "salience": 0.8, + "temporal_context": None, + "tags": [], + } + ], + "episodic": [], + } + ) - with caplog.at_level("WARNING", logger="azure.cosmos.agent_memory.pipeline"): - 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 - ) +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_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": [ - { - "text": "The user normally prefers moderate hotels.", - "action": "ADD", - "category": "preference", - "confidence": 0.9, - "salience": 0.7, - } - ], - "episodic": [], - } - chat = _SyncChat([clean_response]) - memories_store = _Store(existing) - turns_store = _Store([_moderate_hotels_turn()]) +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, @@ -520,53 +367,18 @@ def test_grounding_check_silent_when_add_is_grounded_in_user_turn(caplog) -> Non 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") + out = 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]}" - ) + assert chat.calls == 3 # one LLM call per batch + assert len(out["facts"]) == 3 + assert len(out["processed_turn_docs"]) == 3 -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()]) +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, @@ -574,45 +386,21 @@ def test_grounding_check_warns_on_phantom_explicit_negation_fact(caplog) -> None 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") + out = 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() - ] - 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 - ) + # 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 -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()]) +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, @@ -620,65 +408,11 @@ def test_grounding_check_silent_on_clean_implicit_contradict(caplog) -> None: 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 == [], ( - "clean implicit-contradict must not emit grounding warnings; got: " - f"{[rec.getMessage() for rec in grounding_warnings]}" - ) - + out = service.extract_memories_dry("u1", "t1") -@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, - }, - { - "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()]) - service = AsyncPipelineService( - memories_store, - chat, - _AsyncEmbeddings(), - containers=_async_containers_for_store(memories_store, turns_store=turns_store), - ) - - with caplog.at_level("WARNING", logger="azure.cosmos.agent_memory.pipeline.aio"): - 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]}" - ) + # 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 new file mode 100644 index 0000000..2f7ddd2 --- /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) diff --git a/tests/unit/services/test_pipeline_service.py b/tests/unit/services/test_pipeline_service.py index 587f680..0800ab1 100644 --- a/tests/unit/services/test_pipeline_service.py +++ b/tests/unit/services/test_pipeline_service.py @@ -10,6 +10,14 @@ 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_vector_enabled", + lambda: False, + ) + + class FakeLLMService: """Test helper exposing chat_client + embeddings_client pair. @@ -203,7 +211,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 +233,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 +250,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 +260,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: @@ -291,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), @@ -303,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} ], @@ -314,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: @@ -339,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 45a5117..80a61c9 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() @@ -432,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 03c5420..a2891eb 100644 --- a/tests/unit/test_auto_trigger.py +++ b/tests/unit/test_auto_trigger.py @@ -10,10 +10,50 @@ from unittest.mock import MagicMock, patch +from azure.cosmos.exceptions import CosmosResourceNotFoundError + 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 +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") + pipeline.extract_memories.assert_called_once_with("u1", "t1", recent_k=None) def test_push_to_cosmos_durable_does_not_fire_trigger(monkeypatch): @@ -114,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. @@ -147,6 +187,57 @@ def test_extract_fires_independently_of_summary(self, monkeypatch): processor.process_thread_summary.assert_not_called() processor.process_user_summary.assert_not_called() + 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") + + 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=(0, 1), + ): + 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") + + 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") + + 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.stamp_failure_sync", + ) as stamp, + ): + client.add_local(user_id="u1", role="user", thread_id="t1", content="hi") + client.push_to_cosmos() + + stamp.assert_called_once() + 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") @@ -195,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 5333210..b438103 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 == {} @@ -549,31 +549,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 @@ -590,7 +590,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") @@ -636,12 +636,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) @@ -717,7 +717,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] @@ -893,7 +893,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()] @@ -902,14 +902,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_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_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" 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 fce8308..d31c3a4 100644 --- a/tests/unit/test_pipeline_confidence.py +++ b/tests/unit/test_pipeline_confidence.py @@ -13,6 +13,14 @@ 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, + ) + + def _make_pipeline(llm_response: dict): turns_container = MagicMock() memories_container = MagicMock() @@ -103,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( { @@ -139,7 +118,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 +264,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 +278,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 +326,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 +340,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 +353,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 +365,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 +383,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 +449,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 +460,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..e96ace4 100644 --- a/tests/unit/test_procedural_synthesis.py +++ b/tests/unit/test_procedural_synthesis.py @@ -15,6 +15,14 @@ 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, + ) + + def _assert_iso8601(text: str) -> None: assert text datetime.fromisoformat(text) @@ -25,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] = [] @@ -234,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) @@ -306,7 +313,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 = { @@ -322,7 +328,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 == [] @@ -450,7 +456,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 = [ @@ -635,7 +641,6 @@ def test_synthesize_procedural_retries_with_fresh_llm_call_when_winner_has_parti [prior_v1], fact_docs, [], - [], [ prior_v1, _procedural_doc( @@ -693,7 +698,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 @@ -724,7 +729,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 d81a0f7..11132bd 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"), + call("u1", n=50, memory_type="episodic"), + ] + ) 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"} @@ -103,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): @@ -130,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() @@ -151,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"} @@ -169,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() @@ -250,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 be3f63b..81c563c 100644 --- a/tests/unit/test_reconcile.py +++ b/tests/unit/test_reconcile.py @@ -29,6 +29,16 @@ from azure.cosmos.agent_memory.services.pipeline import PipelineService +@pytest.fixture(autouse=True) +def _pin_legacy_dedup_paths(monkeypatch): + """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, + ) + + def _make_pipeline() -> PipelineService: p = PipelineService.__new__(PipelineService) p._embeddings = MagicMock() @@ -140,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([]) @@ -293,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 = [ @@ -358,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() @@ -601,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", @@ -649,396 +387,88 @@ 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"). - """ +class TestExtractEarlyReturnShape: + """The no-memories early-return must include every key the success + path returns; otherwise callers using ``result["exact_dedup_skipped"]`` + KeyError on empty threads.""" - def _build(self) -> PipelineService: + def test_empty_thread_returns_full_dict_shape(self): 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", - } - ] + p._container.query_items.return_value = iter([]) # no items + out = p.extract_memories("u1", "t-empty") + for key in ( + "fact_count", + "episodic_count", + "updated_count", + "exact_dedup_skipped", + ): + assert key in out, f"missing key: {key}" + assert out[key] == 0 - 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") +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 - + otherwise contradictions get redirected to a doc that doesn't claim + the source, and ``kept`` undercounts.""" - # 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. + def test_failed_supersede_does_not_consume_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"), + _fact("f2", "alpha-restated"), + _fact("f3", "beta"), + ] + p._container.query_items.return_value = iter(facts) 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."), + "duplicate_groups": [ + { + "merged_content": "alpha (consolidated)", + "source_ids": ["f1", "f2"], + "confidence": 0.9, + "salience": 0.7, + } ], - "unclassified": [], + "contradicted_pairs": [], + "kept_ids": ["f3"], } ) ) + p._upsert_memory = MagicMock() + p._embeddings = MagicMock() + p._embeddings.generate.return_value = [0.0] + # Both supersede attempts lose the race. + p._mark_superseded = MagicMock(return_value=False) + result = p.reconcile_memories("u1") + # Sources stay active: kept counts ALL three originals. + assert result == {"kept": 3, "merged": 0, "contradicted": 0} - 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"]`` - KeyError on empty threads.""" - - def test_empty_thread_returns_full_dict_shape(self): - p = PipelineService.__new__(PipelineService) - p._container = MagicMock() - p._memories_container = p._container - p._turns_container = p._container - p._summaries_container = p._container - p._container.query_items.return_value = iter([]) # no items - out = p.extract_memories("u1", "t-empty") - for key in ( - "fact_count", - "episodic_count", - "unclassified_count", - "updated_count", - "exact_dedup_skipped", - "dropped_episodic_count", - ): - assert key in out, f"missing key: {key}" - 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 — - otherwise contradictions get redirected to a doc that doesn't claim - the source, and ``kept`` undercounts.""" - - def test_failed_supersede_does_not_consume_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"), - _fact("f2", "alpha-restated"), - _fact("f3", "beta"), - ] - 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": ["f3"], - } - ) - ) - p._upsert_memory = MagicMock() - p._embeddings = MagicMock() - p._embeddings.generate.return_value = [0.0] - # Both supersede attempts lose the race. - p._mark_superseded = MagicMock(return_value=False) - result = p.reconcile_memories("u1") - # Sources stay active: kept counts ALL three originals. - assert result == {"kept": 3, "merged": 0, "contradicted": 0} - - -class TestReconcileWinnerValidation: - """Hallucinated ``winner_id`` must be refused — never write a dangling - ``superseded_by`` that breaks the audit trail.""" - - def test_hallucinated_winner_id_skipped(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", "user is vegetarian"), - _fact("f2", "user loves ribeye"), - ] - p._container.query_items.return_value = iter(facts) +class TestReconcileWinnerValidation: + """Hallucinated ``winner_id`` must be refused - never write a dangling + ``superseded_by`` that breaks the audit trail.""" + + def test_hallucinated_winner_id_skipped(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", "user is vegetarian"), + _fact("f2", "user loves ribeye"), + ] + p._container.query_items.return_value = iter(facts) p._run_prompty = MagicMock( return_value=json.dumps( { @@ -1062,89 +492,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.""" @@ -1171,7 +518,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"] @@ -1211,183 +558,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 @@ -1427,132 +599,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 @@ -1586,81 +632,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 +651,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..d867d71 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,137 @@ 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("EXTRACTION_BATCH_MAX_TOKENS", "999") + monkeypatch.setenv("DEDUP_VECTOR_ENABLED", "false") + monkeypatch.setenv("DEDUP_SIM_HIGH", "0.50") + + assert thresholds.get_extraction_batch_max_tokens() == 7000 + assert thresholds.get_dedup_vector_enabled() is True + assert thresholds.get_dedup_sim_high() == 0.97 diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 32e6fad..6c29d95 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -5,6 +5,7 @@ from azure.cosmos.agent_memory._utils import ( COSMOS_USER_AGENT, DEFAULT_TTL_BY_TYPE, + MAX_FULLTEXT_TERMS, _build_container_kwargs, _container_policies, _make_memory, @@ -14,7 +15,11 @@ _resolve_vector_index_type, build_cosmos_user_agent, compute_content_hash, + distance_function_from_container_properties, + extract_keywords, normalize_ai_foundry_endpoint, + vector_order_direction, + vector_similarity_at_least, ) from azure.cosmos.agent_memory.exceptions import ConfigurationError, ValidationError @@ -194,65 +199,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", @@ -268,19 +340,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" # ---------------------------------------------------------------------------