diff --git a/.env.template b/.env.template index caa0c95..b8b1b82 100644 --- a/.env.template +++ b/.env.template @@ -21,9 +21,16 @@ COSMOS_DB_LEASE_CONTAINER=leases # ---- Processing thresholds (set to 0 to disable) ---- THREAD_SUMMARY_EVERY_N=10 -FACT_EXTRACTION_EVERY_N=1 +FACT_EXTRACTION_EVERY_N=2 USER_SUMMARY_EVERY_N=20 +# Episodic memory (boundary-based segmentation). EPISODE_EVAL_EVERY_N=0 disables episodic memory. +EPISODE_EVAL_EVERY_N=4 +EPISODE_IDLE_GAP_SECONDS=1800 +EPISODE_TOPIC_DRIFT=0 +EPISODE_MAX_TURNS=40 +EPISODE_MIN_TURNS=2 + # ---- Processor ownership (in-process SDK vs. Function App / Durable) ---- # Controls which side runs the auto-trigger to avoid double-firing when both # the SDK and the Function App are deployed against the same database. diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c7f2e0..4215adb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,33 @@ ## Release History +## [0.3.0b2] (Unreleased) + +#### Features Added +* Episodic memory is now a first-class memory type. Bounded experiences are segmented from the turn stream at idle-gap, topic-drift, and max-size boundaries, each captured as an `EpisodicRecord` with a summary, timeline events, an optional outcome, and first-class `lessons`. See [PR:#37](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/37) +* `search_cosmos(include_episodes=True)` blends facts and episodes into a single ranked query sharing one `top_k` budget, and `search_episodic_memories()` searches episodes directly. See [PR:#37](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/37) +* Fact extraction now defaults to the higher-recall v2 prompt (`extract_memories-v2.prompty`); set `AMT_EXTRACT_MEMORIES_PROMPT=extract_memories.prompty` to fall back to v1. See [PR:#37](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/37) +* Procedural memory is now an atomic, retrievable skill and policy library. `ProceduralRecord` stores individual procedures (behavioral policies, workflows, decision rules, tool-usage notes, recovery strategies) with scope, activation conditions, steps, status, and source provenance. See [PR:#38](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/38) +* The Durable Functions backend now runs episodic extraction, driven by the Cosmos DB change feed on a per-thread cadence set via `EPISODE_EVAL_EVERY_N` (with `EPISODE_IDLE_GAP_SECONDS`, `EPISODE_TOPIC_DRIFT`, `EPISODE_MAX_TURNS`, and `EPISODE_MIN_TURNS` mirrored on the Functions side). See [PR:#38](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/38) +* New delete helpers on both clients: `delete_turn()`, `delete_thread_summary()`, `delete_user_summary()`, and bulk `delete_thread()` (removes a thread's turns and, by default, its summary; distilled facts, episodes, and procedures are left intact). See [PR:#38](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/38) + +#### Breaking Changes +* The default fact-extraction cadence is now every 2 turns (`FACT_EXTRACTION_EVERY_N=2`) instead of every turn, across the SDK, and the Functions deploy default. See [PR:#38](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/38) +* `add_cosmos()` is renamed to `upsert_memory()` on both clients and the store; behavior is unchanged (write-or-replace by id). See [PR:#38](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/38) +* `delete_cosmos()` is renamed to `delete_memory()`. See [PR:#38](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/38) +* Procedural memory has been reshaped: `ProceduralRecord` is now an atomic procedure rather than a single compiled system-prompt document, and the compiled prompt is produced on demand by `build_procedural_context()`. Pre-existing single-prompt procedural documents from earlier betas are not migrated. See [PR:#38](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/38) +* Write-time vector deduplication (in-place fold) has been removed, along with the `DEDUP_VECTOR_ENABLED` and similarity-threshold knobs. Fact dedup is now in-batch hash plus deterministic-id create/409; contradiction reconciliation is unchanged. See [PR:#38](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/38) + +#### Bugs Fixed +* Episode extraction now isolates LLM failures: a transient error leaves the open segment un-stamped for retry, while a non-retryable error (content filter, context-length) quarantines the segment so it can neither wedge the thread nor grow it without bound. See [PR:#37](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/37) +* Out-of-range or non-numeric fact `salience` / `confidence` values are clamped instead of aborting the whole extraction batch and stalling the fact watermark. See [PR:#37](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/37) +* Segment time bounds and the open-episode-segment loader now order turns chronologically by parsed timestamp, so mixed UTC offsets and tied timestamps no longer invert episode bounds or destabilize the deterministic episode id. See [PR:#37](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/37) +* Per-turn extraction watermarks are stamped with a single-field conditional patch, so concurrent fact and episode writers no longer clobber each other's watermark field. See [PR:#38](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/38) +* `parse_llm_json` now rejects a non-object JSON root with a typed error instead of letting it surface downstream as a misclassified transient failure. See [PR:#37](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/37) +* Threshold environment values of `NaN` / `inf` are rejected instead of silently disabling the affected boundary. See [PR:#37](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/37) + +#### Other Changes +* Fact hash-dedup no longer issues a per-extraction query to preload the user's existing fact hashes; exact duplicates are caught in-batch and by the deterministic-id create (409), reducing per-turn latency. See [PR:#38](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/38) + ## [0.3.0b1] (2026-07-24) #### Features Added @@ -23,42 +51,21 @@ ## [0.2.0b3] (2026-07-08) #### Features Added -* A custom user-agent can now be supplied via the new `user_agent` constructor - argument on `CosmosMemoryClient` and `AsyncCosmosMemoryClient`. The toolkit's - own user-agent (`azsdk-python-cosmos-agent-memory/`) is always sent to - Azure Cosmos DB so usage can be tracked; when a custom value is provided it is prefixed and - the toolkit's user-agent is suffixed behind it (`" "`). See [PR:#30](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/30) -* Per-turn processing cadence can now be set in-process via the new - `cadence_thresholds` constructor argument on `CosmosMemoryClient` and - `AsyncCosmosMemoryClient`, instead of only through environment variables. Pass a - mapping keyed by the same names as the env vars (e.g. `FACT_EXTRACTION_EVERY_N`, - `DEDUP_EVERY_N`, `THREAD_SUMMARY_EVERY_N`, `USER_SUMMARY_EVERY_N`); any key not - present falls back to the environment/defaults, and `None` preserves today's - env-only behavior. See [PR:#29](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/29) +* A custom user-agent can now be supplied via the new `user_agent` constructor argument on `CosmosMemoryClient` and `AsyncCosmosMemoryClient`. The toolkit's own user-agent (`azsdk-python-cosmos-agent-memory/`) is always sent to Azure Cosmos DB so usage can be tracked; when a custom value is provided it is prefixed and the toolkit's user-agent is suffixed behind it (`" "`). See [PR:#30](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/30) +* Per-turn processing cadence can now be set in-process via the new `cadence_thresholds` constructor argument on `CosmosMemoryClient` and `AsyncCosmosMemoryClient`, instead of only through environment variables. Pass a mapping keyed by the same names as the env vars (e.g. `FACT_EXTRACTION_EVERY_N`,`DEDUP_EVERY_N`, `THREAD_SUMMARY_EVERY_N`, `USER_SUMMARY_EVERY_N`); any key not present falls back to the environment/defaults, and `None` preserves today's env-only behavior. See [PR:#29](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/29) ## [0.2.0b2] (2026-07-01) #### Features Added -* Embeddings and chat clients can now be injected via the new `embeddings_client` - and `chat_client` constructor arguments on `CosmosMemoryClient` and - `AsyncCosmosMemoryClient`. See [PR:#27](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/27) +* Embeddings and chat clients can now be injected via the new `embeddings_client` and `chat_client` constructor arguments on `CosmosMemoryClient` and `AsyncCosmosMemoryClient`. See [PR:#27](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/27) ## [0.2.0b1] (2026-06-30) #### Features Added -* Raw conversation turns can now be embedded and vector-searched. Set - `enable_turn_embeddings=True` (env `ENABLE_TURN_EMBEDDINGS`) to generate an - embedding when each turn is written, then call `search_turns()` (sync and - async, on both the client and store) to semantically search the raw turn - log. See [PR:#22](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/22/) +* Raw conversation turns can now be embedded and vector-searched. Set `enable_turn_embeddings=True` (env `ENABLE_TURN_EMBEDDINGS`) to generate an embedding when each turn is written, then call `search_turns()` (sync and async, on both the client and store) to semantically search the raw turn log. See [PR:#22](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/22/) #### Other Changes -* The memories container's vector index type is now configurable instead of being - hard-coded to `diskANN`. Set it via the `vector_index_type` argument to - `create_memory_store(...)` or the `AI_FOUNDRY_EMBEDDING_VECTOR_INDEX_TYPE` - environment variable. See [PR:#24](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/24) -* `ai_foundry_endpoint` now accepts a project-scoped Azure AI Foundry URL - (`https://.services.ai.azure.com/api/projects/`) in addition - to the account-level inference endpoint. See [PR:#23](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/23) +* The memories container's vector index type is now configurable instead of being hard-coded to `diskANN`. Set it via the `vector_index_type` argument to `create_memory_store(...)` or the `AI_FOUNDRY_EMBEDDING_VECTOR_INDEX_TYPE` environment variable. See [PR:#24](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/24) +* `ai_foundry_endpoint` now accepts a project-scoped Azure AI Foundry URL (`https://.services.ai.azure.com/api/projects/`) in addition to the account-level inference endpoint. See [PR:#23](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/23) ## [0.1.0b2] (2026-06-03) @@ -72,39 +79,21 @@ ## [0.1.0b1] - 2026-06-01 - Initial public preview release. -This is a **beta release**. The public surface may evolve in -backward-incompatible ways before the `1.0.0` general-availability cut. +This is a **beta release**. The public surface may evolve in backward-incompatible ways before the `1.0.0` general-availability cut. Pin a specific version when integrating. #### Added -- Sync (`CosmosMemoryClient`) and async (`AsyncCosmosMemoryClient`) clients - for storing, retrieving, and transforming agent memories backed by Azure - Cosmos DB. -- Typed memory record hierarchy (Pydantic): `TurnRecord`, `FactRecord`, - `EpisodicRecord`, `ProceduralRecord`, `ThreadSummaryRecord`, - `UserSummaryRecord`. -- Vector + full-text + hybrid search over memories with metadata filters, - tag filters, and per-type scoping. -- Built-in memory processing pipeline: fact extraction, thread/user - summarization, procedural-memory synthesis, contradiction handling, and - deduplication - all driven by versioned `.prompty` prompts. -- Two processor backends: `InProcessProcessor` (default, runs in your - application process) and `DurableFunctionProcessor` (offloads work to a - sibling Azure Function app via Cosmos DB change feed). -- One-command `azd up` deployment that provisions Cosmos DB (with vector + - full-text search enabled), Azure AI Foundry (chat + embedding - deployments), Azure Function app (Flex Consumption), Storage, App - Insights, and the User-Assigned Managed Identity wiring all of it - together. -- Focused exception hierarchy: `AgentMemoryError`, `ConfigurationError`, - `ValidationError`, `CosmosNotConnectedError`, `CosmosOperationError`, - `MemoryNotFoundError`, `MemoryTypeMismatchError`, `LLMError`. -- Structured JSON logging via `azure.cosmos.agent_memory.logging` - (`configure_logging`, `JsonFormatter`). +- Sync (`CosmosMemoryClient`) and async (`AsyncCosmosMemoryClient`) clients for storing, retrieving, and transforming agent memories backed by Azure Cosmos DB. +- Typed memory record hierarchy (Pydantic): `TurnRecord`, `FactRecord`, `EpisodicRecord`, `ProceduralRecord`, `ThreadSummaryRecord`, `UserSummaryRecord`. +- Vector + full-text + hybrid search over memories with metadata filters, tag filters, and per-type scoping. +- Built-in memory processing pipeline: fact extraction, thread/user summarization, procedural-memory synthesis, contradiction handling, and deduplication - all driven by versioned `.prompty` prompts. +- Two processor backends: `InProcessProcessor` (default, runs in your application process) and `DurableFunctionProcessor` (offloads work to a sibling Azure Function app via Cosmos DB change feed). +- One-command `azd up` deployment that provisions Cosmos DB (with vector + full-text search enabled), Azure AI Foundry (chat + embedding deployments), Azure Function app (Flex Consumption), Storage, App Insights, and the User-Assigned Managed Identity wiring all of it together. +- Focused exception hierarchy: `AgentMemoryError`, `ConfigurationError`, `ValidationError`, `CosmosNotConnectedError`, `CosmosOperationError`, `MemoryNotFoundError`, `MemoryTypeMismatchError`, `LLMError`. +- Structured JSON logging via `azure.cosmos.agent_memory.logging` (`configure_logging`, `JsonFormatter`). #### Package layout diff --git a/Docs/azure_testing.md b/Docs/azure_testing.md index f2ac726..c91b560 100644 --- a/Docs/azure_testing.md +++ b/Docs/azure_testing.md @@ -270,7 +270,7 @@ Bring the environment up in this order: 2. verify Cosmos DB RBAC 3. verify Azure OpenAI RBAC 4. create Cosmos resources with `create_memory_store()` -5. test `add_cosmos()` / `push_to_cosmos()` / `get_memories()` +5. test `upsert_memory()` / `push_to_cosmos()` / `get_memories()` 6. test `get_memories(user_id=..., thread_id=...)` filtering 7. test `search_cosmos()` 8. deploy the Function App (e.g., via `azd up`) so the change-feed processor is running @@ -288,7 +288,7 @@ This keeps failures isolated and easier to diagnose. ### Basic Cosmos operations ```python -memory.add_cosmos(user_id="user-1", role="user", content="Hello from Azure") +memory.upsert_memory(user_id="user-1", role="user", content="Hello from Azure") print(memory.get_memories(user_id="user-1")) ``` @@ -300,12 +300,12 @@ print(memory.search_cosmos("hello", user_id="user-1")) ### Durable processing (change-feed driven) -Processing is no longer invoked directly from the SDK - write turns with `add_cosmos()` / `push_to_cosmos()` and the deployed Function App's change-feed trigger fires the `extract_memories`, `thread_summary`, and `user_summary` orchestrators per the configured thresholds. +Processing is no longer invoked directly from the SDK - write turns with `upsert_memory()` / `push_to_cosmos()` and the deployed Function App's change-feed trigger fires the `extract_memories`, `thread_summary`, and `user_summary` orchestrators per the configured thresholds. ```python # Write enough turns to cross THREAD_SUMMARY_EVERY_N (default 10). for i in range(10): - memory.add_cosmos( + memory.upsert_memory( user_id="user-1", thread_id="thread-1", role="user", @@ -329,7 +329,7 @@ import uuid # Use a threshold of 3 (THREAD_SUMMARY_EVERY_N=3) for testing thread_id = str(uuid.uuid4()) for i in range(3): - memory.add_cosmos( + memory.upsert_memory( user_id="user-1", thread_id=thread_id, role="user", diff --git a/Docs/concepts.md b/Docs/concepts.md index 2eba614..183510f 100644 --- a/Docs/concepts.md +++ b/Docs/concepts.md @@ -28,7 +28,7 @@ Every memory uses the same base shape: **Type:** `turn` -Turn memories are raw conversation records. They are created by `add_local()`, `add_cosmos()`, and `push_to_cosmos()` (which bulk-uploads local memories to Cosmos DB). They act as the source material for summaries and facts. +Turn memories are raw conversation records. They are created by `add_local()`, `upsert_memory()`, and `push_to_cosmos()` (which bulk-uploads local memories to Cosmos DB). They act as the source material for summaries and facts. **Use for:** full conversation history and short-term context. @@ -68,7 +68,7 @@ Like thread summaries, user summaries update incrementally by merging the existi |-----------------|-----------------------------------------------------|-----------------------------------------------------------------------------| | **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()` | +| **Created by** | `add_local()` / `upsert_memory()` / `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. @@ -203,11 +203,11 @@ on_memory_change trigger | 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) | +| `THREAD_SUMMARY_EVERY_N` | Per `(user_id, thread_id)` | `10` | +| `FACT_EXTRACTION_EVERY_N` | Per `(user_id, thread_id)` | `2` | +| `USER_SUMMARY_EVERY_N` | Per `user_id` (across all threads) | `20` | -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. +These defaults are shared by both backends (`function_app/shared/config.py` imports the same constants from `azure.cosmos.agent_memory.thresholds`), so the InProcess and Durable processors fire on the same turn boundaries unless overridden. 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 diff --git a/Docs/design_patterns.md b/Docs/design_patterns.md index 0a4ed7d..527a4c5 100644 --- a/Docs/design_patterns.md +++ b/Docs/design_patterns.md @@ -27,19 +27,19 @@ await mem.connect_cosmos() THREAD_ID = "thread-abc-123" # Store user message -await mem.add_cosmos( +await mem.upsert_memory( user_id="user-1", thread_id=THREAD_ID, role="user", content="I need to migrate our PostgreSQL database to Cosmos DB.", ) # Store agent response -await mem.add_cosmos( +await mem.upsert_memory( user_id="user-1", thread_id=THREAD_ID, role="agent", content="I can help with that. What's your current schema look like?", ) # Store a tool call result with metadata -await mem.add_cosmos( +await mem.upsert_memory( user_id="user-1", thread_id=THREAD_ID, role="tool", content='{"tables": 12, "foreign_keys": 3}', @@ -63,10 +63,10 @@ await mem.push_to_cosmos() ```python # Update content of an existing memory -await mem.update_cosmos(memory_id="", content="Corrected message text") +await mem.update_cosmos(memory_id="", user_id="user-1", thread_id=THREAD_ID, memory_type="fact", content="Corrected message text") -# Delete a memory (requires all partition key values) -await mem.delete_cosmos(memory_id="", user_id="user-1", thread_id=THREAD_ID) +# Delete a memory (requires the partition keys and the memory_type) +await mem.delete_memory(memory_id="", user_id="user-1", thread_id=THREAD_ID, memory_type="fact") ``` --- @@ -206,7 +206,7 @@ New session starts ├─ Semantic search for prior facts (search_cosmos, memory_types=["fact"]) │ │ ┌── Conversation loop ──┐ - │ │ Store each turn │ (add_cosmos) + │ │ Store each turn │ (upsert_memory) │ │ Optionally extract │ (extract_facts - every N turns or on key exchanges) │ └────────────────────────┘ │ @@ -228,10 +228,10 @@ system_prompt = build_prompt(profile, relevant) # --- Conversation loop --- while not done: user_msg = get_user_input() - await mem.add_cosmos(user_id="user-1", thread_id=THREAD_ID, role="user", content=user_msg) + await mem.upsert_memory(user_id="user-1", thread_id=THREAD_ID, role="user", content=user_msg) agent_reply = call_llm(system_prompt, user_msg) - await mem.add_cosmos(user_id="user-1", thread_id=THREAD_ID, role="agent", content=agent_reply) + await mem.upsert_memory(user_id="user-1", thread_id=THREAD_ID, role="agent", content=agent_reply) # --- Session end --- await mem.generate_thread_summary(user_id="user-1", thread_id=THREAD_ID) @@ -266,7 +266,7 @@ In a multi-agent system, different agents share the same memory store but may re ```python # Research agent stores findings as turns -await mem.add_cosmos( +await mem.upsert_memory( user_id="user-1", thread_id="research-thread", role="agent", agent_id="research-agent", content="Found that the source DB has 12 tables with 3 foreign key chains.", @@ -284,7 +284,7 @@ facts = await mem.search_cosmos( ) # Planner writes its plan as a turn in its own thread -await mem.add_cosmos( +await mem.upsert_memory( user_id="user-1", thread_id="planning-thread", role="agent", agent_id="planner-agent", content=plan_text, @@ -353,10 +353,10 @@ Both approaches use the same orchestrator and activities, so the output is ident | Operation | Method | When | |-----------|--------|------| -| Store a turn | `add_cosmos` / `add_local` | Every user or agent message | +| Store a turn | `upsert_memory` / `add_local` | Every user or agent message | | Bulk upload | `push_to_cosmos` | After collecting local turns | | Update a memory | `update_cosmos` | Correct or annotate an existing record | -| Delete a memory | `delete_cosmos` | Remove incorrect or sensitive data | +| Delete a memory | `delete_memory` | Remove incorrect or sensitive data | | Get a thread | `get_thread` | Load recent conversation context | | Semantic search | `search_cosmos` | Find relevant facts or summaries for a prompt | | Summarize a thread | `generate_thread_summary` | End of conversation, periodically, or automatic via change feed | diff --git a/Docs/local_testing.md b/Docs/local_testing.md index 45307e7..c0d3ea7 100644 --- a/Docs/local_testing.md +++ b/Docs/local_testing.md @@ -182,7 +182,7 @@ memory.add_local(user_id="user-001", role="user", thread_id=thread_id, content=" memory.push_to_cosmos() # Or add directly to Cosmos -memory.add_cosmos(user_id="user-001", role="agent", thread_id=thread_id, content="Direct Cosmos write") +memory.upsert_memory(user_id="user-001", role="agent", thread_id=thread_id, content="Direct Cosmos write") # Query with filters including thread_id results = memory.get_memories(user_id="user-001", thread_id=thread_id) @@ -225,7 +225,7 @@ await memory.connect_cosmos( await memory.create_memory_store() thread_id = str(uuid.uuid4()) -await memory.add_cosmos(user_id="user-001", role="user", thread_id=thread_id, content="Async Cosmos write") +await memory.upsert_memory(user_id="user-001", role="user", thread_id=thread_id, content="Async Cosmos write") results = await memory.get_memories(user_id="user-001", thread_id=thread_id) for r in results: print(f" [{r['thread_id'][:8]}...] [{r['id'][:8]}...] role={r['role']:<6} {r['content'][:60]}") @@ -365,7 +365,7 @@ import uuid thread_id = str(uuid.uuid4()) for i in range(3): - memory.add_cosmos( + memory.upsert_memory( user_id="user-001", thread_id=thread_id, role="user", diff --git a/Docs/public_api.md b/Docs/public_api.md index 3030bba..d54362e 100644 --- a/Docs/public_api.md +++ b/Docs/public_api.md @@ -26,18 +26,18 @@ - `get_local(memory_id=None, user_id=None, role=None, memory_types=None) -> list[dict]` - filter local buffered memories. - `update_local(memory_id, content=None, role=None, memory_type=None, metadata=None) -> None` - update a local buffered memory. - `delete_local(memory_id) -> None` - remove a local buffered memory. -- `add_cosmos(user_id, role, content, memory_type='turn', metadata=None, thread_id=None, tags=None, ttl=None, salience=None, embedding=None, embed=None) -> str` - upsert one memory to Cosmos and return its id. +- `upsert_memory(user_id, role, content, memory_type='turn', metadata=None, thread_id=None, tags=None, ttl=None, salience=None, embedding=None, embed=None) -> str` - upsert one memory to Cosmos and return its id. - `push_to_cosmos(batch_size=25) -> None` - flush local buffered memories to Cosmos. - `get_memories(memory_id=None, user_id=None, thread_id=None, role=None, memory_types=None, recent_k=None, 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]` - retrieve memories from the MEMORIES container. `memory_types` defaults to `["fact", "episodic", "procedural"]` and must be a subset of those three. - `update_cosmos(memory_id, *, user_id, thread_id, memory_type, content=None, role=None, metadata=None) -> None` - point-update a memory in the container that holds `memory_type`. The `type` field itself is never mutated. -- `delete_cosmos(memory_id, *, user_id, thread_id, memory_type) -> None` - delete a memory from the container that holds `memory_type`. +- `delete_memory(memory_id, *, user_id, thread_id, memory_type) -> None` - delete a memory from the container that holds `memory_type`. - `get_thread(thread_id, user_id=None, recent_k=None, tags_all=None, tags_any=None, exclude_tags=None, include_superseded=False, created_after=None, created_before=None) -> list[dict]` - retrieve turns from the TURNS container oldest-first. - `get_thread_summary(user_id, thread_id, recent_k=None) -> list[dict]` - retrieve thread summary documents from the SUMMARIES container for a single `(user_id, thread_id)` partition. - `get_user_summary(user_id) -> Optional[dict]` - retrieve the active user-summary document. ### Retrieval -- `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_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, include_episodes=False) -> 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. @@ -46,9 +46,18 @@ - `build_procedural_context(user_id) -> str` - format procedural context for prompts. - `build_episodic_context(user_id, query, top_k=3) -> str` - format relevant episodic context. +Episodic retrieval has two scoping modes. `search_episodic(...)` on the store is +episode-only and, when called with `thread_id`, hard-filters to that thread. +`search_cosmos(..., include_episodes=True, thread_id=...)` treats episodic +memory as user-scoped and can recall episodes from other threads for the same +user. Use the episode-only surface for thread-local recall; use +`search_cosmos(include_episodes=True)` when cross-thread episodic recall should +compete with facts in one ranked result set. + ### Processing - `extract_memories(user_id, thread_id, recent_k=None) -> dict[str, int]` - extract facts/episodic memories from a thread. +- `extract_episodes(user_id, thread_id, *, flush=False) -> dict[str, int]` - segment the thread's open turn stream into episodes. Automatic processing uses `flush=False`, so the trailing open segment closes lazily on a later boundary. Call with `flush=True` on session close when the in-process backend owns processing and you need to drain the tail immediately. - `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. @@ -80,18 +89,18 @@ Local-buffer methods remain synchronous in-memory operations; Cosmos, retrieval, - `get_local(memory_id=None, user_id=None, role=None, memory_types=None) -> list[dict]` - filter local buffered memories. - `update_local(memory_id, content=None, role=None, memory_type=None, metadata=None) -> None` - update a local buffered memory. - `delete_local(memory_id) -> None` - remove a local buffered memory. -- `async add_cosmos(user_id, role, content, memory_type='turn', metadata=None, thread_id=None, tags=None, ttl=None, salience=None, embedding=None, embed=None) -> str` - upsert one memory to Cosmos and return its id. +- `async upsert_memory(user_id, role, content, memory_type='turn', metadata=None, thread_id=None, tags=None, ttl=None, salience=None, embedding=None, embed=None) -> str` - upsert one memory to Cosmos and return its id. - `async push_to_cosmos(batch_size=25) -> None` - flush local buffered memories to Cosmos. - `async get_memories(memory_id=None, user_id=None, thread_id=None, role=None, memory_types=None, recent_k=None, 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]` - retrieve memories from the MEMORIES container. `memory_types` defaults to `["fact", "episodic", "procedural"]` and must be a subset of those three. - `async update_cosmos(memory_id, *, user_id, thread_id, memory_type, content=None, role=None, metadata=None) -> None` - point-update a memory in the container that holds `memory_type`. The `type` field itself is never mutated. -- `async delete_cosmos(memory_id, *, user_id, thread_id, memory_type) -> None` - delete a memory from the container that holds `memory_type`. +- `async delete_memory(memory_id, *, user_id, thread_id, memory_type) -> None` - delete a memory from the container that holds `memory_type`. - `async get_thread(thread_id, user_id=None, recent_k=None, tags_all=None, tags_any=None, exclude_tags=None, include_superseded=False, created_after=None, created_before=None) -> list[dict]` - retrieve turns from the TURNS container oldest-first. - `async get_thread_summary(user_id, thread_id, recent_k=None) -> list[dict]` - retrieve thread summary documents from the SUMMARIES container for a single `(user_id, thread_id)` partition. - `async get_user_summary(user_id) -> Optional[dict]` - retrieve the active user-summary document. ### Retrieval -- `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_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, include_episodes=False) -> 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. @@ -100,9 +109,15 @@ Local-buffer methods remain synchronous in-memory operations; Cosmos, retrieval, - `async build_procedural_context(user_id) -> str` - format procedural context for prompts. - `async build_episodic_context(user_id, query, top_k=3) -> str` - format relevant episodic context. +Episodic retrieval has the same scoping behavior as the sync API. The store's +episode-only `search_episodic(...)` hard-filters to `thread_id` when provided, +while `search_cosmos(..., include_episodes=True, thread_id=...)` treats episodic +memory as user-scoped and may return episodes from other threads for that user. + ### Processing - `async extract_memories(user_id, thread_id, recent_k=None) -> dict[str, int]` - extract facts/episodic memories from a thread. +- `async extract_episodes(user_id, thread_id, *, flush=False) -> dict[str, int]` - segment the thread's open turn stream into episodes. Automatic processing uses `flush=False`, so call with `flush=True` on session close when the in-process backend owns processing and you need to drain the tail immediately. - `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. @@ -124,7 +139,7 @@ Use `validate_topology()` (sync) or `await validate_topology()` (async) after `c Sync extension protocols live in `azure.cosmos.agent_memory.services`; async variants live in `azure.cosmos.agent_memory.aio.services`. -- `MemoryStoreProtocol` (`azure.cosmos.agent_memory.services`): persistence primitives (`query`, `read_item`, `add_cosmos`, `mark_superseded`) consumed by the pipeline. +- `MemoryStoreProtocol` (`azure.cosmos.agent_memory.services`): persistence primitives (`query`, `read_item`, `upsert_memory`, `mark_superseded`) consumed by the pipeline. Concrete service classes are exported from their respective packages: diff --git a/README.md b/README.md index 136fec9..f37322d 100644 --- a/README.md +++ b/README.md @@ -94,8 +94,8 @@ memory.connect_cosmos() # auto-creates database + containers if missing USER, THREAD = "user-001", str(uuid.uuid4()) # Add raw turns to a conversation -memory.add_cosmos(user_id=USER, thread_id=THREAD, role="user", content="I love Cosmos DB.") -memory.add_cosmos(user_id=USER, thread_id=THREAD, role="assistant", content="It is fantastic.") +memory.upsert_memory(user_id=USER, thread_id=THREAD, role="user", content="I love Cosmos DB.") +memory.upsert_memory(user_id=USER, thread_id=THREAD, role="assistant", content="It is fantastic.") # Run the processing pipeline (thread summary + fact extraction + user summary) memory.process_now(user_id=USER, thread_id=THREAD) @@ -128,7 +128,7 @@ See [`Samples/`](Samples/) for end-to-end scenarios (chat memory, RAG, multi-age | Concept | What it is | API | |--------------------|-------------------------------------------------------------------------|-------------------------------------------------------| -| **Turn** | One message (user or assistant) - the raw conversation atom | `add_cosmos(...)`, `add_local(...)` | +| **Turn** | One message (user or assistant) - the raw conversation atom | `upsert_memory(...)`, `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(...)` | @@ -191,13 +191,13 @@ By default, the **InProcess processor** runs each pipeline step independently as | Env var | Default | Step that fires | Async behavior | |---------------------------|------------------|-------------------------------------------------------------------------------------------------------------------|-------------------------------------| -| `FACT_EXTRACTION_EVERY_N` | `1` (every turn) | `process_extract_memories` | scheduled via `asyncio.create_task` | +| `FACT_EXTRACTION_EVERY_N` | `2` (every 2 turns) | `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`. +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; bump `FACT_EXTRACTION_EVERY_N` / `DEDUP_EVERY_N` for cost-sensitive, higher-volume production traffic. 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`. The async client (`AsyncCosmosMemoryClient.push_to_cosmos`) does **not** await the auto-trigger; it schedules it as a background `asyncio.Task` so the write call returns as soon as the Cosmos upserts complete. Background failures are surfaced via `logger.warning` (search for `"Background auto-trigger task failed"`). diff --git a/Samples/Advanced/advanced_memory_lifecycle.py b/Samples/Advanced/advanced_memory_lifecycle.py index 47ec5a5..859b43a 100644 --- a/Samples/Advanced/advanced_memory_lifecycle.py +++ b/Samples/Advanced/advanced_memory_lifecycle.py @@ -75,7 +75,7 @@ def main() -> None: ("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) + mem.upsert_memory(user_id=user_id, role=role, content=content, thread_id=thread_id) _print_memories(mem, user_id, thread_id) _header(2, "Extract structured memories (facts / procedural / episodic)") @@ -93,7 +93,7 @@ def main() -> None: _header(5, "Archive: delete raw turns, keep derived memories") deleted = 0 for m in mem.get_thread(thread_id=thread_id, user_id=user_id): - mem.delete_cosmos(memory_id=m["id"], user_id=user_id, thread_id=thread_id, memory_type="turn") + mem.delete_memory(memory_id=m["id"], user_id=user_id, thread_id=thread_id, memory_type="turn") deleted += 1 print(f" deleted {deleted} raw turn(s)") diff --git a/Samples/Advanced/advanced_search_patterns.py b/Samples/Advanced/advanced_search_patterns.py index 92c0e82..0230772 100644 --- a/Samples/Advanced/advanced_search_patterns.py +++ b/Samples/Advanced/advanced_search_patterns.py @@ -74,7 +74,7 @@ def seed_memories(mem: CosmosMemoryClient, user_id: str, thread_id: str) -> None print("Seeding memories …") for entry in entries: - mem.add_cosmos( + mem.upsert_memory( user_id=user_id, role=entry["role"], content=entry["content"], diff --git a/Samples/Notebooks/Demo.ipynb b/Samples/Notebooks/Demo.ipynb index 5c7a73c..a2f587e 100644 --- a/Samples/Notebooks/Demo.ipynb +++ b/Samples/Notebooks/Demo.ipynb @@ -11,7 +11,7 @@ "\n", "1. **Setup** - Install dependencies and load environment variables\n", "2. **Local memory operations** - `add_local`, `get_local`, `update_local`, `delete_local`\n", - "3. **Cosmos DB operations** - `add_cosmos`, `get_memories`, `get_thread`\n", + "3. **Cosmos DB operations** - `upsert_memory`, `get_memories`, `get_thread`\n", "4. **Thread Summary** - `generate_thread_summary()` (in-process LLM)\n", "5. **Memory Extraction** - `extract_memories()` (facts + episodic + procedural)\n", "6. **User Summary** - `generate_user_summary()` (cross-thread profile)\n", @@ -400,7 +400,7 @@ "> - A Cosmos DB for NoSQL account with a database and container matching your `.env` values\n", "> - The container should have a [vector embedding policy](https://learn.microsoft.com/azure/cosmos-db/nosql/vector-search) configured on the `embedding` field\n", "> - Entra ID / managed identity RBAC role (e.g. *Cosmos DB Built-in Data Contributor*)\n", - "> - An Azure AI Foundry embedding model deployment for `add_cosmos` and `search_cosmos`" + "> - An Azure AI Foundry embedding model deployment for `upsert_memory` and `search_cosmos`" ] }, { @@ -430,7 +430,7 @@ "id": "4b497000", "metadata": {}, "source": [ - "### 3b. Add memories to Cosmos DB with `add_cosmos`" + "### 3b. Add memories to Cosmos DB with `upsert_memory`" ] }, { @@ -476,25 +476,25 @@ "new_thread_id = str(uuid.uuid4())\n", "print(f\"New Thread ID: {new_thread_id}\\n\")\n", "\n", - "# Add memories directly to Cosmos DB using add_cosmos\n", - "memory.add_cosmos(\n", + "# Add memories directly to Cosmos DB using upsert_memory\n", + "memory.upsert_memory(\n", " user_id=\"user-002\", role=\"user\", thread_id=new_thread_id,\n", " content=\"Can you recommend some good restaurants in New York City?\",\n", ")\n", - "memory.add_cosmos(\n", + "memory.upsert_memory(\n", " user_id=\"user-002\", role=\"tool\", thread_id=new_thread_id,\n", " content='{\"query\": \"top restaurants NYC\", \"results\": [\"Carbone\", \"Nobu\", \"Katz\\'s Deli\", \"Le Bernardin\"]}',\n", " metadata={\"tool_name\": \"restaurant_search\", \"tool_call_id\": \"call_abc123\"},\n", ")\n", - "memory.add_cosmos(\n", + "memory.upsert_memory(\n", " user_id=\"user-002\", role=\"agent\", thread_id=new_thread_id,\n", " content=\"Absolutely! NYC has incredible dining options. For Italian, try Carbone in Greenwich Village. For sushi, Nobu in Tribeca is world-class. For a classic NYC experience, Katz's Delicatessen on the Lower East Side is a must.\",\n", ")\n", - "memory.add_cosmos(\n", + "memory.upsert_memory(\n", " user_id=\"user-002\", role=\"user\", thread_id=new_thread_id,\n", " content=\"I love Italian food. Are there any options that are budget-friendly?\",\n", ")\n", - "memory.add_cosmos(\n", + "memory.upsert_memory(\n", " user_id=\"user-002\", role=\"agent\", thread_id=new_thread_id,\n", " content=\"For budget-friendly Italian in NYC, check out L'industrie Pizzeria in Williamsburg or Artichoke Basille's Pizza. Both are highly rated and won't break the bank.\",\n", ")\n", @@ -557,7 +557,7 @@ "source": [ "### 3d. Update & Delete in Cosmos DB\n", "\n", - "`update_cosmos` and `delete_cosmos` work just like their local counterparts. If the content changes, the embedding is automatically re-generated." + "`update_cosmos` and `delete_memory` work just like their local counterparts. If the content changes, the embedding is automatically re-generated." ] }, { @@ -618,7 +618,7 @@ "tool_mems = [t for t in memory.get_thread(thread_id=new_thread_id, user_id=\"user-002\") if t.get(\"role\") == \"tool\"]\n", "print(tool_mems[0])\n", "if tool_mems:\n", - " memory.delete_cosmos(\n", + " memory.delete_memory(\n", " tool_mems[0][\"id\"],\n", " user_id=tool_mems[0][\"user_id\"],\n", " thread_id=tool_mems[0][\"thread_id\"],\n", diff --git a/Samples/Notebooks/Demo_async.ipynb b/Samples/Notebooks/Demo_async.ipynb index 0658adb..4f6ac3c 100644 --- a/Samples/Notebooks/Demo_async.ipynb +++ b/Samples/Notebooks/Demo_async.ipynb @@ -397,7 +397,7 @@ "id": "5c18edf3", "metadata": {}, "source": [ - "### 3b. Add memories to Cosmos DB with `add_cosmos`" + "### 3b. Add memories to Cosmos DB with `upsert_memory`" ] }, { @@ -441,25 +441,25 @@ "new_thread_id = str(uuid.uuid4())\n", "print(f\"New Thread ID: {new_thread_id}\\n\")\n", "\n", - "# Add memories directly to Cosmos DB using add_cosmos\n", - "await memory.add_cosmos(\n", + "# Add memories directly to Cosmos DB using upsert_memory\n", + "await memory.upsert_memory(\n", " user_id=\"user-002\", role=\"user\", thread_id=new_thread_id,\n", " content=\"Can you recommend some good restaurants in New York City?\",\n", ")\n", - "await memory.add_cosmos(\n", + "await memory.upsert_memory(\n", " user_id=\"user-002\", role=\"tool\", thread_id=new_thread_id,\n", " content='{\"query\": \"top restaurants NYC\", \"results\": [\"Carbone\", \"Nobu\", \"Katz\\'s Deli\", \"Le Bernardin\"]}',\n", " metadata={\"tool_name\": \"restaurant_search\", \"tool_call_id\": \"call_abc123\"},\n", ")\n", - "await memory.add_cosmos(\n", + "await memory.upsert_memory(\n", " user_id=\"user-002\", role=\"agent\", thread_id=new_thread_id,\n", " content=\"Absolutely! NYC has incredible dining options. For Italian, try Carbone in Greenwich Village. For sushi, Nobu in Tribeca is world-class. For a classic NYC experience, Katz's Delicatessen on the Lower East Side is a must.\",\n", ")\n", - "await memory.add_cosmos(\n", + "await memory.upsert_memory(\n", " user_id=\"user-002\", role=\"user\", thread_id=new_thread_id,\n", " content=\"I love Italian food. Are there any options that are budget-friendly?\",\n", ")\n", - "await memory.add_cosmos(\n", + "await memory.upsert_memory(\n", " user_id=\"user-002\", role=\"agent\", thread_id=new_thread_id,\n", " content=\"For budget-friendly Italian in NYC, check out L'industrie Pizzeria in Williamsburg or Artichoke Basille's Pizza. Both are highly rated and won't break the bank.\",\n", ")\n", @@ -583,7 +583,7 @@ "tool_mems = [t for t in await memory.get_thread(thread_id=new_thread_id, user_id=\"user-002\") if t.get(\"role\") == \"tool\"]\n", "print(tool_mems[0])\n", "if tool_mems:\n", - " await memory.delete_cosmos(\n", + " await memory.delete_memory(\n", " tool_mems[0][\"id\"],\n", " user_id=tool_mems[0][\"user_id\"],\n", " thread_id=tool_mems[0][\"thread_id\"],\n", diff --git a/Samples/Notebooks/Demo_function_app.ipynb b/Samples/Notebooks/Demo_function_app.ipynb index 64e1481..2dcc7e9 100644 --- a/Samples/Notebooks/Demo_function_app.ipynb +++ b/Samples/Notebooks/Demo_function_app.ipynb @@ -30,7 +30,7 @@ "\n", "## What the SDK does in this mode\n", "\n", - "* `add_cosmos(..., memory_type=\"turn\")` → writes the raw turn to Cosmos.\n", + "* `upsert_memory(..., memory_type=\"turn\")` → writes the raw turn to Cosmos.\n", "* The change-feed-triggered Function App reads new turns and runs orchestrators.\n", "* `process_now()` is a **debug-logged no-op** - the Function App owns processing.\n", "* `process_now_and_wait()` polls Cosmos for the summary doc; useful for demos / tests (RU-costly)." @@ -156,7 +156,7 @@ " (\"agent\", \"Understood - no overnight bookings without your approval.\"),\n", "]\n", "for role, content in transcript:\n", - " memory.add_cosmos(\n", + " memory.upsert_memory(\n", " user_id=USER_ID,\n", " thread_id=THREAD_ID,\n", " role=role,\n", @@ -308,7 +308,7 @@ "source": [ "## 7. Going further\n", "\n", - "* **Per-turn embedding**: the SDK auto-embeds non-`turn` documents you `add_cosmos` directly. Raw `turn`\n", + "* **Per-turn embedding**: the SDK auto-embeds non-`turn` documents you `upsert_memory` directly. Raw `turn`\n", " records are intentionally not embedded - the function app does that during summary/extraction.\n", "* **Search across function-app-produced memories**: works exactly as in the in-process demo:\n", " `memory.search_cosmos(search_terms=\"…\", memory_types=[\"fact\"], user_id=USER_ID)`.\n", diff --git a/Samples/Notebooks/Demo_function_app_async.ipynb b/Samples/Notebooks/Demo_function_app_async.ipynb index 77b8167..9c135c8 100644 --- a/Samples/Notebooks/Demo_function_app_async.ipynb +++ b/Samples/Notebooks/Demo_function_app_async.ipynb @@ -125,7 +125,7 @@ " (\"agent\", \"Understood - no overnight bookings without your approval.\"),\n", "]\n", "for role, content in transcript:\n", - " await memory.add_cosmos(\n", + " await memory.upsert_memory(\n", " user_id=USER_ID,\n", " thread_id=THREAD_ID,\n", " role=role,\n", diff --git a/Samples/Processing/processing_episodic_memory.py b/Samples/Processing/processing_episodic_memory.py index 1cad3a0..8f7d7b1 100644 --- a/Samples/Processing/processing_episodic_memory.py +++ b/Samples/Processing/processing_episodic_memory.py @@ -183,7 +183,7 @@ def main() -> None: try: banner("1. Adding conversation turns") for role, content in CONVERSATION: - mem.add_cosmos(user_id=user_id, role=role, content=content, thread_id=thread_id) + mem.upsert_memory(user_id=user_id, role=role, content=content, thread_id=thread_id) print(f" [{role:>5}] {content}") banner("2. Finalizing episodes (flush the open segment)") @@ -228,7 +228,7 @@ def main() -> None: continue seen_ids.add(memory_id) try: - mem.delete_cosmos( + mem.delete_memory( memory_id=memory_id, user_id=user_id, thread_id=record.get("thread_id", thread_id), diff --git a/Samples/Processing/processing_fact_extraction.py b/Samples/Processing/processing_fact_extraction.py index 819e5ba..a85d458 100644 --- a/Samples/Processing/processing_fact_extraction.py +++ b/Samples/Processing/processing_fact_extraction.py @@ -66,7 +66,7 @@ def main() -> None: print("Adding conversation turns…") for role, content in conversations: - mem.add_cosmos(user_id=user_id, role=role, content=content, thread_id=thread_id) + mem.upsert_memory(user_id=user_id, role=role, content=content, thread_id=thread_id) print(f" [{role:>5}] {content[:80]}") print() diff --git a/Samples/Processing/processing_procedural_memory.py b/Samples/Processing/processing_procedural_memory.py new file mode 100644 index 0000000..f45608c --- /dev/null +++ b/Samples/Processing/processing_procedural_memory.py @@ -0,0 +1,334 @@ +"""Demonstrate procedural memory extraction, retrieval, and prompt projection. + +Procedural memory is a skill and policy library. Instead of storing one +mutable prompt blob, the toolkit stores many atomic ``ProceduralRecord`` items: +behavioral policies, workflows, decision rules, tool-usage notes, and recovery +strategies. Each procedure carries provenance, activation conditions, scope, +status, and executable steps where appropriate. + +The personalized system prompt is compiled on demand by +``CosmosMemoryClient.build_procedural_context(...)``. That prompt is a +deterministic projection of ACTIVE procedures only. Candidate procedures remain +in the library for inspection and later promotion, but they are not injected. + +This sample seeds both kinds of sources: + +1. An explicit user instruction: "Always ask for confirmation before deleting + cloud resources." Explicit user instructions are trusted provenance, so the + synthesized behavioral policy should become ACTIVE. +2. A realistic debugging episode about a Cosmos DB ORDER BY query. Lessons + distilled from episodes are useful skills, but episode-distilled procedures + stay CANDIDATE under the provenance gate and are excluded from prompt + injection until promoted by trusted policy or review. + +Required environment variables (.env supported via python-dotenv): + + COSMOS_DB_ENDPOINT + COSMOS_DB_DATABASE + COSMOS_DB_MEMORIES_CONTAINER + AI_FOUNDRY_ENDPOINT + AI_FOUNDRY_API_KEY + AI_FOUNDRY_CHAT_DEPLOYMENT_NAME + AI_FOUNDRY_EMBEDDING_DEPLOYMENT_NAME + AI_FOUNDRY_EMBEDDING_DIMENSIONS + +Cosmos DB authentication uses DefaultAzureCredential. Do not pass a Cosmos DB +key when local auth is disabled on the account. +""" + +from __future__ import annotations + +import json +import os +import sys +import uuid +from typing import Any + +from dotenv import load_dotenv + +from azure.cosmos.agent_memory import CosmosMemoryClient + +load_dotenv() + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +DIVIDER = "-" * 72 + + +def banner(title: str) -> None: + """Print a section banner.""" + print(f"\n{DIVIDER}") + print(f" {title}") + print(DIVIDER) + + +def _short(value: Any, width: int = 120) -> str: + """Return a single-line string trimmed for readable terminal output.""" + text = str(value or "").replace("\n", " ").strip() + return text if len(text) <= width else f"{text[: width - 3]}..." + + +def print_procedure(proc: dict[str, Any], *, index: int | None = None) -> None: + """Pretty-print the fields that make procedural memory understandable.""" + prefix = f" {index}." if index is not None else " -" + print(f"{prefix} {proc.get('name') or '(unnamed)'}") + print(f" kind: {proc.get('procedure_kind')}") + print(f" scope: {proc.get('scope_type') or '(none)'}:{proc.get('scope_value') or '*'}") + print(f" status: {proc.get('status')}") + print(f" source_kind: {proc.get('source_kind')}") + print(f" summary: {_short(proc.get('summary') or proc.get('content'))}") + conditions = proc.get("activation_conditions") or [] + if conditions: + print(" activates:") + for condition in conditions[:3]: + print(f" - {_short(condition, 96)}") + steps = proc.get("steps") or [] + if steps: + print(" steps:") + for step in sorted(steps, key=lambda item: int(item.get("sequence") or 0))[:4]: + instruction = step.get("instruction") if isinstance(step, dict) else step + print(f" {step.get('sequence', '-')}. {_short(instruction, 96)}") + + +def print_procedures(title: str, procedures: list[dict[str, Any]]) -> None: + """Print a titled procedure list.""" + print(title) + if not procedures: + print(" (none)") + return + for index, procedure in enumerate(procedures, start=1): + print_procedure(procedure, index=index) + + +# --------------------------------------------------------------------------- +# Demo data +# --------------------------------------------------------------------------- + +# Thread 1 is a direct instruction from the user. The procedural synthesizer can +# turn this into a behavioral_policy. Because the source is an explicit user +# instruction, provenance gating should mark it ACTIVE. +POLICY_TURNS = [ + ( + "user", + "Always ask for confirmation before deleting cloud resources, including " + "Cosmos DB accounts, databases, containers, or resource groups.", + ), + ( + "agent", + "Understood. I will ask for confirmation before deleting any cloud resource.", + ), +] + +# Thread 2 is an ordinary debugging conversation. Episodic extraction should +# produce an episode with a lesson, and procedural synthesis can distill that +# lesson into a reusable Cosmos DB recovery_strategy or workflow. Because this +# is episode-distilled rather than an explicit policy, it remains CANDIDATE and +# is not injected into the compiled system prompt. +DEBUGGING_TURNS = [ + ( + "user", + "My Cosmos DB query is failing: SELECT * FROM c WHERE c.customerId = " + "@customerId ORDER BY c.createdAt DESC. The portal says the ORDER BY " + "query needs a matching composite index.", + ), + ( + "agent", + "First confirm the query filters by the partition key, then inspect the " + "indexing policy for a composite index on customerId ASC and createdAt DESC.", + ), + ( + "user", + "The partition key is /customerId. The container only has the default " + "range index; there is no composite index configured.", + ), + ( + "agent", + "Add a composite index matching the equality filter and ORDER BY sort, " + "wait for indexing to finish, and retry with a small TOP literal while testing.", + ), + ( + "user", + "After adding the composite index and waiting for transformation progress, " + "the query succeeded. Lesson learned: for Cosmos DB ORDER BY failures, " + "check partition scope and composite indexes before changing application code.", + ), +] + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + required = [ + "COSMOS_DB_ENDPOINT", + "COSMOS_DB_DATABASE", + "COSMOS_DB_MEMORIES_CONTAINER", + "AI_FOUNDRY_ENDPOINT", + "AI_FOUNDRY_API_KEY", + "AI_FOUNDRY_CHAT_DEPLOYMENT_NAME", + "AI_FOUNDRY_EMBEDDING_DEPLOYMENT_NAME", + "AI_FOUNDRY_EMBEDDING_DIMENSIONS", + ] + missing = [v for v in required if not os.environ.get(v)] + if missing: + print(f"ERROR: missing env vars: {', '.join(missing)}") + print("Set the variables listed in the module docstring, then rerun this sample.") + sys.exit(1) + + mem = CosmosMemoryClient( + cosmos_endpoint=os.environ["COSMOS_DB_ENDPOINT"], + cosmos_database=os.environ["COSMOS_DB_DATABASE"], + cosmos_container=os.environ["COSMOS_DB_MEMORIES_CONTAINER"], + ai_foundry_endpoint=os.environ["AI_FOUNDRY_ENDPOINT"], + ai_foundry_api_key=os.environ["AI_FOUNDRY_API_KEY"], + embedding_deployment_name=os.environ["AI_FOUNDRY_EMBEDDING_DEPLOYMENT_NAME"], + embedding_dimensions=int(os.environ["AI_FOUNDRY_EMBEDDING_DIMENSIONS"]), + chat_deployment_name=os.environ["AI_FOUNDRY_CHAT_DEPLOYMENT_NAME"], + use_default_credential=True, + cadence_thresholds={ + "FACT_EXTRACTION_EVERY_N": 0, + "THREAD_SUMMARY_EVERY_N": 0, + "EPISODE_EVAL_EVERY_N": 0, + "USER_SUMMARY_EVERY_N": 0, + }, + ) + print("Connected to Cosmos DB with DefaultAzureCredential.") + + user_id = f"procedural-demo-{uuid.uuid4().hex[:8]}" + policy_thread_id = f"procedural-policy-{uuid.uuid4().hex[:8]}" + debug_thread_id = f"procedural-debug-{uuid.uuid4().hex[:8]}" + print(f"User ID: {user_id}") + print(f"Policy thread ID: {policy_thread_id}") + print(f"Debug thread ID: {debug_thread_id}") + + try: + banner("1. Seed explicit instruction turns") + for role, content in POLICY_TURNS: + mem.upsert_memory(user_id=user_id, role=role, content=content, thread_id=policy_thread_id) + print(f" [{role:>5}] {_short(content)}") + + banner("2. Extract facts from the explicit instruction") + # Fact extraction supplies behavioral facts to synthesize_procedural. + # The direct "Always ask..." instruction should be categorized as a + # user requirement or high-salience preference by the extraction prompt. + fact_stats = mem.extract_memories(user_id=user_id, thread_id=policy_thread_id) + print(f" stats: {json.dumps(fact_stats, indent=2)}") + + banner("3. Seed a Cosmos DB debugging episode") + for role, content in DEBUGGING_TURNS: + mem.upsert_memory(user_id=user_id, role=role, content=content, thread_id=debug_thread_id) + print(f" [{role:>5}] {_short(content)}") + + banner("4. Finalize the episode and expose its lesson") + # In production, episode boundaries are evaluated automatically by the + # processing backend. This script knows the demo conversation is over, + # so flush=True drains the trailing open segment now. + episode_stats = mem.extract_episodes(user_id, debug_thread_id, flush=True) + print(f" stats: {json.dumps(episode_stats, indent=2)}") + episodes = mem.get_episodes(user_id, thread_id=debug_thread_id) + for episode in episodes: + print(f" episode: {_short(episode.get('title'))}") + for lesson in episode.get("lessons") or []: + print(f" lesson: {_short(lesson)}") + + banner("5. Synthesize the procedural skill and policy library") + synthesis_stats = mem.synthesize_procedural(user_id) + print(f" stats: {json.dumps(synthesis_stats, indent=2)}") + print(f" procedures created: {synthesis_stats.get('procedures_created', 0)}") + + banner("6. Inspect the procedural memory library") + procedures = mem.get_procedural_memories(user_id) + active = [p for p in procedures if p.get("status") == "active"] + candidates = [p for p in procedures if p.get("status") == "candidate"] + print_procedures("ACTIVE procedures from trusted provenance:", active) + print_procedures("CANDIDATE procedures retained for review, not injection:", candidates) + + banner("7. Context-aware retrieval") + # The default retrieval API is injection-safe: status defaults to + # "active", so candidates are excluded just like they are excluded from + # prompt projection. + active_matches = mem.retrieve_procedures(user_id, "cosmos db order by query failing") + print_procedures("Default active-only retrieval:", active_matches) + + # For library inspection or review workflows, pass status=None to rank + # both active and candidate procedures. This should surface the Cosmos DB + # debugging skill while the delete-confirmation policy ranks poorly or is + # absent for this coding task. + all_matches = mem.retrieve_procedures(user_id, "cosmos db order by query failing", status=None) + print_procedures("Review retrieval including candidates:", all_matches) + + banner("8. Compile deterministic prompt projections") + # With no task, only always-on global/user behavioral policies are folded + # into the prompt. Candidate skills are deliberately excluded. + global_context = mem.build_procedural_context(user_id) + print("Global policies only:") + print(global_context or "(no active global policies)") + + # With a task, ACTIVE task procedures relevant to that task can be folded + # in after the global policies. Episode-distilled candidates still stay + # out of the prompt until promoted by trusted provenance. + task_context = mem.build_procedural_context(user_id, task="fix a failing cosmos db query") + print("\nPolicies plus relevant active task skills:") + print(task_context or "(no active policies or task skills)") + + banner("9. Closing notes") + print( + "Procedural memory keeps a two-layer model: atomic procedure " + "records form the skill/policy library, and build_procedural_context " + "compiles the deterministic prompt projection on demand. Provenance " + "gating keeps explicit user instructions, observed preferences, and " + "organization policy active, while episode-distilled, document, and " + "inferred procedures remain candidates. Episodic extraction, one " + "source of candidate procedures, also runs under the Durable Functions " + "backend in deployed processing setups." + ) + + finally: + banner("10. Cleanup") + deleted = 0 + try: + memory_records = mem.get_memories(user_id=user_id, include_superseded=True) + policy_turns = mem.get_thread(thread_id=policy_thread_id, user_id=user_id, include_superseded=True) + debug_turns = mem.get_thread(thread_id=debug_thread_id, user_id=user_id, include_superseded=True) + procedural_records = mem.get_procedural_memories(user_id=user_id, include_superseded=True) + thread_summaries = [ + *mem.get_thread_summary(user_id=user_id, thread_id=policy_thread_id), + *mem.get_thread_summary(user_id=user_id, thread_id=debug_thread_id), + ] + user_summary = mem.get_user_summary(user_id) + all_records = [*memory_records, *policy_turns, *debug_turns, *procedural_records, *thread_summaries] + if user_summary: + all_records.append(user_summary) + + seen_ids: set[str] = set() + for record in all_records: + memory_id = record.get("id") + if not memory_id or memory_id in seen_ids: + continue + seen_ids.add(memory_id) + record_type = record.get("type") + if not record_type: + continue + try: + mem.delete_memory( + memory_id=memory_id, + user_id=user_id, + thread_id=record.get("thread_id") or debug_thread_id, + memory_type=record_type, + ) + deleted += 1 + except Exception as exc: # pragma: no cover - best effort cleanup + print(f" WARN: failed to delete {memory_id}: {exc}") + print(f" Deleted {deleted} record(s) for user {user_id}") + except Exception as exc: # pragma: no cover - best effort cleanup + print(f" WARN: cleanup failed: {exc}") + finally: + mem.close() + + +if __name__ == "__main__": + main() diff --git a/Samples/Processing/processing_thread_summary.py b/Samples/Processing/processing_thread_summary.py index afd02c4..01cbb32 100644 --- a/Samples/Processing/processing_thread_summary.py +++ b/Samples/Processing/processing_thread_summary.py @@ -60,7 +60,7 @@ def main() -> None: ("agent", "10–14 days lets you spend ~5 days in each city plus some day trips."), ] for role, content in initial: - mem.add_cosmos(user_id=user_id, role=role, content=content, thread_id=thread_id) + mem.upsert_memory(user_id=user_id, role=role, content=content, thread_id=thread_id) print(f" [{role:>5}] {content[:80]}") _banner("STEP 2 – generate first summary") @@ -79,7 +79,7 @@ def main() -> None: ("agent", "Shigetsu inside Tenryu-ji temple is famous for its shojin-ryori meals."), ] for role, content in follow_up: - mem.add_cosmos(user_id=user_id, role=role, content=content, thread_id=thread_id) + mem.upsert_memory(user_id=user_id, role=role, content=content, thread_id=thread_id) print(f" [{role:>5}] {content[:80]}") doc2 = mem.generate_thread_summary(user_id=user_id, thread_id=thread_id) diff --git a/Samples/Processing/processing_user_profile.py b/Samples/Processing/processing_user_profile.py index ecb35db..d0e161f 100644 --- a/Samples/Processing/processing_user_profile.py +++ b/Samples/Processing/processing_user_profile.py @@ -59,7 +59,7 @@ def main() -> None: ("user", "Sounds great. I love simple Italian food, especially fresh ingredients."), ("agent", "Italian cuisine emphasises quality ingredients prepared simply."), ]: - mem.add_cosmos(user_id=user_id, role=role, content=content, thread_id=t1) + mem.upsert_memory(user_id=user_id, role=role, content=content, thread_id=t1) mem.extract_memories(user_id=user_id, thread_id=t1) mem.generate_thread_summary(user_id=user_id, thread_id=t1) @@ -73,7 +73,7 @@ def main() -> None: ("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) + mem.upsert_memory(user_id=user_id, role=role, content=content, thread_id=t2) mem.extract_memories(user_id=user_id, thread_id=t2) mem.generate_thread_summary(user_id=user_id, thread_id=t2) @@ -87,7 +87,7 @@ def main() -> None: ("user", "Cool. I'm a Python engineer building AI tooling."), ("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) + mem.upsert_memory(user_id=user_id, role=role, content=content, thread_id=t3) mem.extract_memories(user_id=user_id, thread_id=t3) mem.generate_thread_summary(user_id=user_id, thread_id=t3) diff --git a/Samples/Quickstarts/quickstart_cosmos.py b/Samples/Quickstarts/quickstart_cosmos.py index 4d06e96..26d8b6a 100644 --- a/Samples/Quickstarts/quickstart_cosmos.py +++ b/Samples/Quickstarts/quickstart_cosmos.py @@ -37,7 +37,7 @@ def main() -> None: # Add a memory directly to Cosmos. memory_type="turn" skips auto-embedding, # which keeps the quickstart runnable without AI Foundry credentials. - mem.add_cosmos(user_id="u1", role="user", content="Hello from quickstart!", thread_id="t1") + mem.upsert_memory(user_id="u1", role="user", content="Hello from quickstart!", thread_id="t1") print("Added memory to Cosmos") # Retrieve the thread we just wrote to @@ -61,7 +61,7 @@ def main() -> None: print("Pushed local memories to Cosmos") # Clean up – delete the memories we created - mem.delete_cosmos(memory_id=memory_id, user_id="u1", thread_id="t1", memory_type="turn") + mem.delete_memory(memory_id=memory_id, user_id="u1", thread_id="t1", memory_type="turn") print(f"Deleted memory {memory_id}") print("\nQuickstart complete!") diff --git a/Samples/Scenarios/scenario_chat_memory.py b/Samples/Scenarios/scenario_chat_memory.py index 33680e9..3dd999b 100644 --- a/Samples/Scenarios/scenario_chat_memory.py +++ b/Samples/Scenarios/scenario_chat_memory.py @@ -116,7 +116,7 @@ def run_session( print() for role, content in turns: - mem.add_cosmos( + mem.upsert_memory( user_id=user_id, role=role, content=content, @@ -223,7 +223,7 @@ def main() -> None: ), ] for role, content in new_turns: - mem.add_cosmos( + mem.upsert_memory( user_id=user_id, role=role, content=content, @@ -236,7 +236,7 @@ def main() -> None: # ── Summary ─────────────────────────────────────────────── banner("Done") print(" This sample demonstrated:") - print(" • Storing multi-turn conversations with add_cosmos") + print(" • Storing multi-turn conversations with upsert_memory") print(" • Retrieving a full thread with get_thread") print(" • Searching across sessions with search_cosmos") print(" • Using recalled context to inform new sessions") diff --git a/Samples/Scenarios/scenario_counter_tuning.py b/Samples/Scenarios/scenario_counter_tuning.py index 153d362..8c3d43a 100644 --- a/Samples/Scenarios/scenario_counter_tuning.py +++ b/Samples/Scenarios/scenario_counter_tuning.py @@ -76,7 +76,7 @@ def main() -> None: ] print(f"Writing {len(transcript)} turns to Cosmos (thread={thread_id})...") for role, content in transcript: - client.add_cosmos( + client.upsert_memory( user_id=user_id, thread_id=thread_id, role=role, diff --git a/Samples/Scenarios/scenario_customer_support.py b/Samples/Scenarios/scenario_customer_support.py index 3c7aa9c..ab0d374 100644 --- a/Samples/Scenarios/scenario_customer_support.py +++ b/Samples/Scenarios/scenario_customer_support.py @@ -35,7 +35,7 @@ def _banner(title: str) -> None: def _add_dialogue(mem: CosmosMemoryClient, user_id: str, ticket: str, dialogue: list[tuple[str, str]]) -> None: for role, content in dialogue: - mem.add_cosmos(user_id=user_id, role=role, content=content, thread_id=ticket) + mem.upsert_memory(user_id=user_id, role=role, content=content, thread_id=ticket) print(f" [{role:>5}] {content[:90]}") diff --git a/Samples/Scenarios/scenario_memory_reconciliation.py b/Samples/Scenarios/scenario_memory_reconciliation.py index b628efc..8d90bf9 100644 --- a/Samples/Scenarios/scenario_memory_reconciliation.py +++ b/Samples/Scenarios/scenario_memory_reconciliation.py @@ -100,7 +100,7 @@ def main() -> None: try: banner("1. Seeding paraphrased facts (duplicates)") for content in PARAPHRASED_FACTS: - mem.add_cosmos( + mem.upsert_memory( user_id=unique_user_id, role="user", content=content, @@ -112,7 +112,7 @@ def main() -> None: banner("2. Seeding contradicting facts") for content in CONTRADICTING_FACTS: - mem.add_cosmos( + mem.upsert_memory( user_id=unique_user_id, role="user", content=content, @@ -161,7 +161,7 @@ def main() -> None: deleted = 0 for rec in all_records: try: - mem.delete_cosmos( + mem.delete_memory( memory_id=rec["id"], user_id=unique_user_id, thread_id=rec.get("thread_id", unique_thread_id), diff --git a/Samples/Scenarios/scenario_multi_agent.py b/Samples/Scenarios/scenario_multi_agent.py index 2ea88d1..f27014e 100644 --- a/Samples/Scenarios/scenario_multi_agent.py +++ b/Samples/Scenarios/scenario_multi_agent.py @@ -5,7 +5,7 @@ thread so each can read the other's contributions via get_thread. Agent identity is tracked through metadata={"agent_id": "..."} on every -add_cosmos call. +upsert_memory call. Workflow: 1. User posts a complex question. @@ -76,7 +76,7 @@ def step1_user_question( """User asks a complex, multi-part question.""" print_header("Step 1 - User posts a complex question") - mem.add_cosmos( + mem.upsert_memory( user_id=user_id, role="user", content=( @@ -111,7 +111,7 @@ def step2_planner_creates_plan( " 3. Research end-of-life recycling challenges for both\n" " 4. Synthesise findings into a recommendation for urban fleets" ) - mem.add_cosmos( + mem.upsert_memory( user_id=user_id, role="agent", content=plan, @@ -170,7 +170,7 @@ def step3_researcher_performs_research( ] for f in findings: - mem.add_cosmos( + mem.upsert_memory( user_id=user_id, role="agent", content=f["content"], @@ -206,7 +206,7 @@ def step4_planner_synthesises_answer( "mature. Hydrogen fuel-cell vehicles may become competitive once " "green-hydrogen costs fall and recycling capacity scales." ) - mem.add_cosmos( + mem.upsert_memory( user_id=user_id, role="agent", content=recommendation, diff --git a/Samples/Scenarios/scenario_rag_with_memory.py b/Samples/Scenarios/scenario_rag_with_memory.py index 9adc8a0..2269a77 100644 --- a/Samples/Scenarios/scenario_rag_with_memory.py +++ b/Samples/Scenarios/scenario_rag_with_memory.py @@ -128,7 +128,7 @@ def run_demo() -> None: "User prefers NoSQL databases for their current project", ] for fact in facts: - mem.add_cosmos( + mem.upsert_memory( user_id=user_id, role="system", content=fact, @@ -145,7 +145,7 @@ def run_demo() -> None: ("agent", "Azure Cosmos DB with its NoSQL API would be a great fit."), ] for role, content in conversation: - mem.add_cosmos( + mem.upsert_memory( user_id=user_id, role=role, content=content, @@ -261,7 +261,7 @@ def run_demo() -> None: turns = mem.get_thread(thread_id=thread_id, user_id=user_id) for item in turns: - mem.delete_cosmos( + mem.delete_memory( memory_id=item["id"], user_id=user_id, thread_id=item.get("thread_id", thread_id), @@ -269,7 +269,7 @@ def run_demo() -> None: ) stored = mem.get_memories(user_id=user_id, thread_id=thread_id) for item in stored: - mem.delete_cosmos( + mem.delete_memory( memory_id=item["id"], user_id=user_id, thread_id=item.get("thread_id", thread_id), diff --git a/Samples/Scenarios/scenario_remote_processor.py b/Samples/Scenarios/scenario_remote_processor.py index 16f74ab..0ea0330 100644 --- a/Samples/Scenarios/scenario_remote_processor.py +++ b/Samples/Scenarios/scenario_remote_processor.py @@ -54,7 +54,7 @@ def main() -> None: ("agent", "HPK lets you co-locate related items for efficient queries."), ] for role, content in transcript: - client.add_cosmos( + client.upsert_memory( user_id=user_id, thread_id=thread_id, role=role, diff --git a/Samples/Scenarios/scenario_remote_processor_async.py b/Samples/Scenarios/scenario_remote_processor_async.py index 4ab9aca..9f624a7 100644 --- a/Samples/Scenarios/scenario_remote_processor_async.py +++ b/Samples/Scenarios/scenario_remote_processor_async.py @@ -45,7 +45,7 @@ async def main() -> None: ("agent", "HPK lets you co-locate related items for efficient queries."), ] for role, content in transcript: - await client.add_cosmos( + await client.upsert_memory( user_id=user_id, thread_id=thread_id, role=role, diff --git a/Samples/Scenarios/scenario_tagging_and_filtering.py b/Samples/Scenarios/scenario_tagging_and_filtering.py index 5bdbbb7..250baa0 100644 --- a/Samples/Scenarios/scenario_tagging_and_filtering.py +++ b/Samples/Scenarios/scenario_tagging_and_filtering.py @@ -79,7 +79,7 @@ def main() -> None: ["workflow", "ops", "important"]), ] for mem_type, content, tags in seeds: - client.add_cosmos( + client.upsert_memory( user_id=user_id, role="system", content=content, diff --git a/azure/cosmos/agent_memory/_utils.py b/azure/cosmos/agent_memory/_utils.py index 7cac875..5e89729 100644 --- a/azure/cosmos/agent_memory/_utils.py +++ b/azure/cosmos/agent_memory/_utils.py @@ -395,32 +395,6 @@ def vector_order_direction(distance_function: str) -> str: 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. @@ -585,15 +559,9 @@ def _container_policies( embedding_data_type: str, distance_function: str, full_text_language: str, - include_salience_composite: bool = True, vector_index_type: str = "quantizedFlat", ) -> tuple[dict, dict, dict]: - """Build the vector, indexing, and full-text policies for container creation. - - ``include_salience_composite`` adds the ``(salience, created_at, id)`` - composite index required by procedural synthesis on the MEMORIES container. - Turns reuse this builder with it disabled (turns are never synthesized). - """ + """Build the vector, indexing, and full-text policies for container creation.""" vector_embedding_policy = { "vectorEmbeddings": [ { @@ -616,20 +584,6 @@ def _container_policies( "fullTextIndexes": [{"path": "/content"}], } - if include_salience_composite: - # Procedural synthesis selects TOP N by (salience DESC, created_at ASC, id ASC). - # Cosmos requires a composite index for multi-property ORDER BY; without it the - # query returns a non-deterministic 50 of N when many docs share the default - # salience (0.5), which makes the source-id short-circuit in synthesize_procedural - # thrash and burn LLM calls on every reconcile. - indexing_policy["compositeIndexes"] = [ - [ - {"path": "/salience", "order": "descending"}, - {"path": "/created_at", "order": "ascending"}, - {"path": "/id", "order": "ascending"}, - ] - ] - full_text_policy = { "defaultLanguage": full_text_language, "fullTextPaths": [{"path": "/content", "language": full_text_language}], diff --git a/azure/cosmos/agent_memory/aio/cosmos_memory_client.py b/azure/cosmos/agent_memory/aio/cosmos_memory_client.py index 008e0e8..9874762 100644 --- a/azure/cosmos/agent_memory/aio/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/aio/cosmos_memory_client.py @@ -27,7 +27,12 @@ from azure.cosmos.agent_memory.aio.processors import AsyncInProcessProcessor, AsyncMemoryProcessor from azure.cosmos.agent_memory.aio.services.pipeline import AsyncPipelineService from azure.cosmos.agent_memory.aio.store import AsyncMemoryStore -from azure.cosmos.agent_memory.exceptions import CosmosNotConnectedError, CosmosOperationError, ValidationError +from azure.cosmos.agent_memory.exceptions import ( + CosmosNotConnectedError, + CosmosOperationError, + MemoryNotFoundError, + ValidationError, +) from azure.cosmos.agent_memory.logging import get_logger from azure.cosmos.agent_memory.services._pipeline_helpers import ( _normalize_cadence_thresholds, @@ -340,7 +345,6 @@ async def create_memory_store( turns_vec_policy, turns_idx_policy, turns_ft_policy = _container_policies( **{**_policy_kwargs, "vector_index_type": "quantizedFlat"}, - include_salience_composite=False, ) self._turns_container_client = await db.create_container_if_not_exists( **_build_container_kwargs( @@ -356,7 +360,6 @@ async def create_memory_store( logger.info("Created turns container: %s/%s", self._cosmos_database, self._cosmos_turns_container) summaries_vec_policy, summaries_idx_policy, summaries_ft_policy = _container_policies( **{**_policy_kwargs, "vector_index_type": "quantizedFlat"}, - include_salience_composite=False, ) summaries_idx_policy["compositeIndexes"] = [ [ @@ -562,7 +565,7 @@ def _container_for_type(self, memory_type: str) -> Any: """Return the Cosmos container client that owns ``memory_type``.""" return self._containers[container_key_for_type(memory_type)] - async def add_cosmos( + async def upsert_memory( self, user_id: str, role: str, @@ -677,7 +680,7 @@ async def update_cosmos( metadata=metadata, ) - async def delete_cosmos( + async def delete_memory( self, memory_id: str, *, @@ -692,6 +695,79 @@ async def delete_cosmos( memory_type=memory_type, ) + async def delete_turn(self, turn_id: str, *, user_id: str, thread_id: str) -> None: + """Delete a single turn document. Raises if it does not exist.""" + return await self.delete_memory( + turn_id, + user_id=user_id, + thread_id=thread_id, + memory_type="turn", + ) + + async def delete_thread_summary(self, user_id: str, thread_id: str) -> bool: + """Delete a thread's summary if present. Returns True when one was deleted. + + The summary id is deterministic, so callers need not look it up first; a + missing summary is a no-op that returns False. + """ + try: + await self.delete_memory( + f"summary_{user_id}_{thread_id}", + user_id=user_id, + thread_id=thread_id, + memory_type="thread_summary", + ) + return True + except MemoryNotFoundError: + return False + + async def delete_user_summary(self, user_id: str) -> bool: + """Delete a user's summary if present. Returns True when one was deleted. + + A missing summary is a no-op that returns False. + """ + try: + await self.delete_memory( + f"user_summary_{user_id}", + user_id=user_id, + thread_id="__user_summary__", + memory_type="user_summary", + ) + return True + except MemoryNotFoundError: + return False + + async def delete_thread(self, user_id: str, thread_id: str, *, include_summary: bool = True) -> int: + """Bulk-delete a conversation thread: all of its turns and, by default, its + thread summary. Returns the number of documents deleted. + + This targets the conversation itself (turns + summary). Durable memories + distilled from the thread - facts, episodes, procedures - are user-scoped + knowledge and are left intact. Deletion is best-effort per document: a + concurrently-removed turn is skipped rather than aborting the whole sweep. + """ + if not user_id: + raise ValidationError("user_id is required") + if not thread_id: + raise ValidationError("thread_id is required") + + deleted = 0 + turns = await self.get_thread(thread_id=thread_id, user_id=user_id, include_superseded=True) or [] + for turn in turns: + turn_id = turn.get("id") + if not turn_id: + continue + try: + await self.delete_memory(turn_id, user_id=user_id, thread_id=thread_id, memory_type="turn") + deleted += 1 + except MemoryNotFoundError: + continue + + if include_summary and await self.delete_thread_summary(user_id, thread_id): + deleted += 1 + + return deleted + async def search_cosmos( self, search_terms: str, @@ -725,6 +801,11 @@ async def search_cosmos( store = self._get_store() # Facts + episodes share one ranked query and one top_k budget: episodic # is added when include_episodes is True and stripped when it is False. + if memory_types is not None and "episodic" in memory_types and not include_episodes: + logger.warning( + "Episodic memories requested via memory_types are only returned when include_episodes=True; " + "proceeding without episodic memories and using facts or other requested memory types only." + ) if memory_types is not None: base_memory_types = [t for t in memory_types if t != "episodic"] else: @@ -967,7 +1048,8 @@ async def remove_tags( return await self._get_store().remove_tags(memory_id, user_id, thread_id, memory_type, tags) async def get_procedural_prompt(self, user_id: str) -> Optional[str]: - return await self._get_store().get_procedural_prompt(user_id=user_id) + prompt = await self._get_pipeline().build_procedural_context(user_id) + return prompt or None async def get_procedural_history(self, user_id: str, limit: int = 10) -> list[dict[str, Any]]: return await self._get_store().get_procedural_history(user_id=user_id, limit=limit) @@ -1004,8 +1086,31 @@ async def search_episodic_memories( include_superseded, ) - async def build_procedural_context(self, user_id: str) -> str: - return await self._get_pipeline().build_procedural_context(user_id) + async def retrieve_procedures( + self, + user_id: str, + search_terms: str, + top_k: int = 5, + *, + scope_type: Optional[str] = None, + scope_value: Optional[str] = None, + procedure_kind: Optional[str] = None, + status: Optional[str] = "active", + include_superseded: bool = False, + ) -> list[dict[str, Any]]: + return await self._get_store().retrieve_procedures( + user_id=user_id, + search_terms=search_terms, + top_k=top_k, + scope_type=scope_type, + scope_value=scope_value, + procedure_kind=procedure_kind, + status=status, + include_superseded=include_superseded, + ) + + async def build_procedural_context(self, user_id: str, task: Optional[str] = None) -> str: + return await self._get_pipeline().build_procedural_context(user_id, task) async def build_episodic_context( self, diff --git a/azure/cosmos/agent_memory/aio/processors/durable.py b/azure/cosmos/agent_memory/aio/processors/durable.py index 52cf32a..62c7669 100644 --- a/azure/cosmos/agent_memory/aio/processors/durable.py +++ b/azure/cosmos/agent_memory/aio/processors/durable.py @@ -12,10 +12,6 @@ logger = get_logger(__name__) -# Set once we have warned that episodic memory is inert under the durable backend, -# so the warning fires a single time per process rather than on every no-op call. -_EPISODIC_DURABLE_WARNED = False - class AsyncDurableFunctionProcessor: """Async mirror of :class:`DurableFunctionProcessor`. @@ -59,17 +55,10 @@ async def process_extract_episodes( user_id: str, thread_id: str, ) -> dict[str, int]: - # Episodic segmentation is in-process only; the Durable backend has no - # episodic path yet, so this is an explicit no-op (the auto-trigger also - # gates episode extraction to the in-process processor). Warn once so a - # durable-mode operator can see that episodic memory is not being produced. - global _EPISODIC_DURABLE_WARNED - if not _EPISODIC_DURABLE_WARNED: - _EPISODIC_DURABLE_WARNED = True - logger.warning( - "Episodic memory is not available under the Durable Functions backend " - "(no episodic write path yet); episode extraction is a no-op in durable mode." - ) + # The Durable Function app owns episodic extraction via the Cosmos DB + # Change Feed trigger: ExtractEpisodesOrchestrator -> ee_ExtractEpisodes + # -> pipeline.extract_episodes. This hook mirrors synthesize_procedural + # by leaving Durable-owned work to the orchestrator. logger.debug( "AsyncDurableFunctionProcessor.process_extract_episodes no-op user_id=%s thread_id=%s", user_id, @@ -128,11 +117,13 @@ async def synthesize_procedural( user_id: str, force: bool = False, ) -> dict[str, Any]: - raise NotImplementedError( - "Procedural synthesis runs automatically after reconcile in durable mode; " - "manual invocation via the SDK is not supported when the Durable Function " - "app is the active processor." - ) + # No-op, like the other durable hooks (mirror of the sync processor): + # procedural synthesis runs in the Durable Function app after reconcile, + # so returning instead of raising keeps the in-process auto-trigger from + # stamping a spurious failure each cadence. + del force + logger.debug("DurableFunctionProcessor.synthesize_procedural no-op user_id=%s", user_id) + return {"status": "skipped", "procedures_created": 0} async def close(self) -> None: logger.debug("AsyncDurableFunctionProcessor.close no-op") diff --git a/azure/cosmos/agent_memory/aio/services/pipeline.py b/azure/cosmos/agent_memory/aio/services/pipeline.py index 74b0520..6a3cfed 100644 --- a/azure/cosmos/agent_memory/aio/services/pipeline.py +++ b/azure/cosmos/agent_memory/aio/services/pipeline.py @@ -13,6 +13,7 @@ import hashlib import inspect import json +import re import time from collections import defaultdict from datetime import datetime, timezone @@ -22,24 +23,22 @@ CosmosResourceExistsError, CosmosResourceNotFoundError, ) +from pydantic import ValidationError as PydanticValidationError from azure.cosmos.agent_memory._container_routing import ContainerKey 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, - MemoryConflictError, ValidationError, ) from azure.cosmos.agent_memory.logging import get_logger from azure.cosmos.agent_memory.models import ( + TRUSTED_PROCEDURE_SOURCE_KINDS, EpisodicRecord, FactRecord, ProceduralRecord, @@ -73,13 +72,8 @@ from azure.cosmos.agent_memory.services._pipeline_helpers import ( is_real_number as _is_real_number, ) -from azure.cosmos.agent_memory.services._pipeline_helpers import ( - 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_episode_idle_gap_seconds, get_episode_max_turns, get_episode_min_turns, @@ -168,9 +162,44 @@ async def upsert_item(self, *, body: dict[str, Any]) -> dict[str, Any]: if inspect.isawaitable(response): response = await response return response if isinstance(response, dict) else body - response = await self._store.add_cosmos(body) + response = await self._store.upsert_memory(body) return response if isinstance(response, dict) else body + @staticmethod + def _apply_patch_operations(doc: dict[str, Any], patch_operations: list[dict[str, Any]]) -> dict[str, Any]: + patched = dict(doc) + for operation in patch_operations: + if operation.get("op") != "set": + raise ValueError(f"unsupported patch operation: {operation.get('op')!r}") + path = operation.get("path") + if not isinstance(path, str) or not path.startswith("/") or path == "/": + raise ValueError(f"unsupported patch path: {path!r}") + keys = [part.replace("~1", "/").replace("~0", "~") for part in path[1:].split("/")] + target = patched + for key in keys[:-1]: + value = target.get(key) + if not isinstance(value, dict): + value = {} + target[key] = value + target = value + target[keys[-1]] = operation.get("value") + return patched + + async def patch_item( + self, *, item: str, partition_key: Any, patch_operations: list[dict[str, Any]] + ) -> dict[str, Any]: + container = self._target_container() + patch_item = getattr(container, "patch_item", None) + if callable(patch_item): + response = patch_item(item=item, partition_key=partition_key, patch_operations=patch_operations) + if inspect.isawaitable(response): + response = await response + if isinstance(response, dict): + return response + return await self.read_item(item=item, partition_key=partition_key) + doc = await self.read_item(item=item, partition_key=partition_key) + return await self.upsert_item(body=self._apply_patch_operations(doc, patch_operations)) + async def create_item(self, *, body: dict[str, Any]) -> dict[str, Any]: container = self._target_container() if container is not None and hasattr(container, "create_item"): @@ -184,7 +213,7 @@ async def create_item(self, *, body: dict[str, Any]) -> dict[str, Any]: if inspect.isawaitable(response): response = await response return response if isinstance(response, dict) else body - response = await self._store.add_cosmos(body) + response = await self._store.upsert_memory(body) return response if isinstance(response, dict) else body async def replace_item(self, **kwargs: Any) -> Any: @@ -307,17 +336,17 @@ def _warn_distance_policy_unavailable_once(self) -> None: 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." + "vector dedup: container vector policy could not be read; skipping " + "near-exact auto-drop this run to avoid mis-calibrated drops. Memories are " + "written as-is and reconciled on a later run once the policy is readable." ) def _warn_euclidean_autodrop_once(self, distance_function: str) -> None: """One-shot WARN that the near-exact vector auto-drop is disabled. - The ``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. + The near-exact threshold is cosine-calibrated; on euclidean + the destructive auto-drop is skipped and LLM reconcile still runs. + Logged once per pipeline instance to avoid hot-path spam. """ if getattr(self, "_warned_euclidean_autodrop", False): return @@ -423,40 +452,6 @@ def _build_transcript( include_timestamp=include_timestamp, ) - async def _load_existing_memories( - self, - user_id: str, - memory_types: list[str], - limit: int = 100, - ) -> list[dict[str, Any]]: - """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 - in implementation-defined order and the dedup comparison set is - non-deterministic. - """ - type_placeholders = ", ".join(f"@mtype{i}" for i in range(len(memory_types))) - capped_limit = top_literal(limit, name="_load_existing_memories.limit") - query = ( - f"SELECT TOP {capped_limit} * FROM c " - f"WHERE c.user_id = @user_id " - f"AND c.type IN ({type_placeholders}) " - f"AND {_ACTIVE_DOC_FILTER} " - f"ORDER BY c._ts DESC" - ) - parameters: list[dict[str, Any]] = [ - {"name": "@user_id", "value": user_id}, - ] - for i, mt in enumerate(memory_types): - parameters.append({"name": f"@mtype{i}", "value": mt}) - - return await self._query_items( - self._memories_container, - query=query, - parameters=parameters, - ) - async def _upsert_memory(self, doc: dict[str, Any]) -> dict[str, Any]: """Upsert a fact, episodic, or procedural document to the memories container.""" return await self._upsert_item(self._memories_container, body=doc) @@ -547,17 +542,20 @@ async def extract_memories_durable( logger.warning("extract_memories_durable no memories found user_id=%s thread_id=%s", user_id, thread_id) return {"facts": [], "episodic": [], "updates": [], "processed_turn_docs": []} - existing_for_hash = await self._load_existing_memories(user_id, ["fact"]) - existing_fact_hashes: set[str] = { - m["content_hash"] for m in existing_for_hash if m.get("type") == "fact" and m.get("content_hash") - } + # Exact-duplicate detection is in-batch only: a content_hash seen earlier + # in THIS extraction is skipped. Cross-turn / cross-run exact duplicates + # are handled at write time by the deterministic-id create (a repeat of + # the same fact collides on id and is skipped with a 409), so there is no + # per-extract query to preload the user's existing fact hashes. + existing_fact_hashes: set[str] = set() # 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. + # go into ``processed_turns``; the in-process caller marks them + # ``extracted_at`` so they are never re-processed (the Durable backend + # instead advances a count-based watermark). Turns from batches that fail + # with a *retryable* error are left out and retried on the next run. batches = batch_turns_by_tokens(items, get_extraction_batch_max_tokens()) facts: list[dict[str, Any]] = [] processed_turns: list[dict[str, Any]] = [] @@ -681,200 +679,6 @@ async def extract_memories_durable( ) 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). - - Folds only within the same ``metadata.source`` (user vs agent); a - cross-source pair returns False so the caller keeps it as a novel ADD, - preventing tag/source desync. - """ - from azure.core import MatchConditions - from azure.cosmos.exceptions import CosmosAccessConditionFailedError - - neighbor_source = (neighbor.get("metadata") or {}).get("source") or "user" - new_source = (new_doc.get("metadata") or {}).get("source") or "user" - if neighbor_source != new_source: - logger.info( - "in-place dedup update skipped (source mismatch neighbor=%s new=%s) " - "target_id=%s; keeping new doc as novel", - neighbor_source, - new_source, - neighbor.get("id"), - ) - return False - - 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, @@ -916,7 +720,7 @@ 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) - for key in ("inplace_updated", "deferred_turn_count", "quarantined_turn_count"): + for key in ("deferred_turn_count", "quarantined_turn_count"): if key in op: result[key] = result.get(key, 0) + int(op.get(key) or 0) @@ -924,11 +728,11 @@ async def persist_extracted_memories( return result - async def _mark_turns_extracted(self, turn_docs: list[dict[str, Any]], *, field: str = "extracted_at") -> int: - """Stamp a processed-watermark field on each turn doc and upsert. Mirror of - the sync helper - ``field`` selects the independent watermark - (``extracted_at`` for facts, ``episode_extracted_at`` for episodic - segmentation). Per-turn failures are logged but never raise. + async def _mark_turns_extracted(self, turn_docs: list[dict[str, Any]]) -> int: + """Stamp the fact-extraction ``extracted_at`` watermark on each turn doc + (mirror of the sync helper). Episodic segmentation uses a separate + per-thread cursor doc and never stamps turns. Per-turn failures are + logged but never raise. """ if not turn_docs: return 0 @@ -939,14 +743,15 @@ async def _mark_turns_extracted(self, turn_docs: list[dict[str, Any]], *, field: if not turn_id: continue try: - doc_to_write = dict(turn) - doc_to_write[field] = now_iso - await self._upsert_item(self._turns_container, body=doc_to_write) + await self._turns_container.patch_item( + item=turn_id, + partition_key=[turn.get("user_id"), turn.get("thread_id")], + patch_operations=[{"op": "set", "path": "/extracted_at", "value": now_iso}], + ) marked += 1 except Exception as exc: logger.warning( - "_mark_turns_extracted(%s) failed for turn_id=%s err=%s (turn may be re-processed on next call)", - field, + "_mark_turns_extracted failed for turn_id=%s err=%s (turn may be re-processed on next call)", turn_id, exc, ) @@ -964,11 +769,8 @@ async def extract_memories( extracted = await self.extract_memories_durable(user_id, thread_id, recent_k, turns=turns) # Capture the processed turns from the compute stage 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. + # Persist uses this exact list for stamping after all creates finish. 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) @@ -1180,20 +982,101 @@ async def _build_episode_docs( def _deterministic_episode_id(segment_key: str, index: int) -> str: return deterministic_episode_id(segment_key, index) + @staticmethod + def _episode_cursor_id(user_id: str, thread_id: str) -> str: + return f"episode_cursor_{user_id}_{thread_id}" + + async def _read_episode_cursor(self, user_id: str, thread_id: str) -> tuple[str, str]: + """``(created_at, id)`` watermark of the last turn folded into an episode, + or ``("", "")`` when none. Stored as a single doc in the MEMORIES + container so advancing it never writes to the turns container and cannot + re-enter the change feed (mirror of the sync helper). ``created_at`` is + UTC-normalized, so the lexical ``>`` comparison matches chronology.""" + try: + doc = await self._read_item( + self._memories_container, + item=self._episode_cursor_id(user_id, thread_id), + partition_key=[user_id, thread_id], + ) + except CosmosResourceNotFoundError: + return "", "" + return str(doc.get("last_episode_at") or ""), str(doc.get("last_episode_id") or "") + + async def _advance_episode_cursor(self, user_id: str, thread_id: str, last_turn: dict[str, Any]) -> None: + """Advance the episodic watermark to ``last_turn``, never backwards + (mirror of the sync helper): a single-doc write that cannot partially + fail and never touches the change-feed-monitored turns container. The + advance is atomic - ETag ``IfNotModified`` with a re-read retry - so a + late, out-of-order concurrent orchestration cannot regress the cursor and + duplicate episodes under topic drift.""" + from azure.core import MatchConditions + from azure.cosmos.exceptions import CosmosAccessConditionFailedError + + last_at = str(last_turn.get("created_at") or "") + last_id = str(last_turn.get("id") or "") + if not last_at: + return + cursor_id = self._episode_cursor_id(user_id, thread_id) + partition_key = [user_id, thread_id] + for _ in range(3): + etag: Optional[str] = None + try: + existing = await self._read_item(self._memories_container, item=cursor_id, partition_key=partition_key) + except CosmosResourceNotFoundError: + existing = None + if existing is not None: + current = ( + str(existing.get("last_episode_at") or ""), + str(existing.get("last_episode_id") or ""), + ) + if (last_at, last_id) <= current: + return # monotonic: never regress + etag = existing.get("_etag") + body = { + "id": cursor_id, + "type": "episode_cursor", + "user_id": user_id, + "thread_id": thread_id, + "last_episode_at": last_at, + "last_episode_id": last_id, + "updated_at": datetime.now(tz=timezone.utc).isoformat(), + } + try: + if existing is None: + await self._create_item(self._memories_container, body=body) + else: + await self._replace_item( + self._memories_container, + item=cursor_id, + body=body, + etag=etag, + match_condition=MatchConditions.IfNotModified, + ) + return + except (CosmosResourceExistsError, CosmosAccessConditionFailedError): + continue # a concurrent run advanced first; re-read and re-check + logger.debug("episode cursor advance retries exhausted user_id=%s thread_id=%s", user_id, thread_id) + async def _load_open_episode_segment(self, user_id: str, thread_id: str) -> list[dict[str, Any]]: - """Return the open episode segment: turns not yet folded into an episode, - oldest first. Independent ``episode_extracted_at`` watermark (mirror of - the sync helper).""" + """Return the open episode segment: turns created after the episodic + watermark (not yet folded), oldest first. Uses a per-thread + ``(created_at, id)`` cursor doc instead of a per-turn stamp, so episodic + segmentation writes nothing to the turns container (mirror of the sync + helper).""" + last_at, last_id = await self._read_episode_cursor(user_id, thread_id) items = await self._query_items( self._turns_container, query=( "SELECT * FROM c WHERE c.user_id = @user_id " "AND c.thread_id = @thread_id AND c.type = 'turn' " - "AND (NOT IS_DEFINED(c.episode_extracted_at) OR IS_NULL(c.episode_extracted_at))" + "AND (c.created_at > @last_at " + "OR (c.created_at = @last_at AND c.id > @last_id))" ), parameters=[ {"name": "@user_id", "value": user_id}, {"name": "@thread_id", "value": thread_id}, + {"name": "@last_at", "value": last_at}, + {"name": "@last_id", "value": last_id}, ], partition_key=[user_id, thread_id], ) @@ -1248,17 +1131,20 @@ async def extract_episodes( ) -> dict[str, int]: """Segment the open turn stream into episodes at detected boundaries. - Mirror of the sync pipeline: the open segment is every turn without an - ``episode_extracted_at`` stamp; at each boundary (idle time-gap, topic - drift, or max-size cap) the closed segment is extracted, embedded, and - persisted, then its turns are stamped. ``flush=True`` drains the trailing - open segment. The caller never signals "session end". - - Idempotency is best-effort (not absolute): each episode's id is - deterministic in its segment key and ordinal - not the LLM summary text - - so re-running the same still-open segment skips the duplicate write (409) - while the segment's turn set is stable; a partial watermark-stamp failure - can still admit a duplicate, which episodic reconciliation does not fold. + Mirror of the sync pipeline: the open segment is every turn created after + the episodic watermark (see ``_read_episode_cursor``); at each boundary + (idle time-gap, topic drift, or max-size cap) the closed segment is + extracted, embedded, and persisted, then the watermark advances past those + turns. ``flush=True`` drains the trailing open segment. The caller never + signals "session end". + + Idempotency: each episode's id is deterministic in its segment key and + ordinal - not the LLM summary text - so re-running the same still-open + segment skips the duplicate write (409). The watermark is a single doc + advanced only after a segment's episodes are created, so a crash between + the create and the advance simply re-loads the same open segment next run + (same turn set -> same ids -> 409); there is no partial-stamp state that + could shift the boundary and admit a duplicate. """ if not user_id: raise ValidationError("user_id is required") @@ -1299,29 +1185,29 @@ async def extract_episodes( ) break # Non-retryable (e.g. content filter, context-length): quarantine the - # poison segment - stamp it so it never re-poisons future runs and the - # open segment cannot grow without bound - then advance to the next. + # poison segment - advance the watermark past it so it never + # re-poisons future runs and the open segment cannot grow without + # bound - then move to the next segment. logger.warning( "extract_episodes: quarantining %d turns after non-retryable extraction error " - "(marking episode_extracted_at so they do not re-poison future runs) " + "(advancing the episode watermark past them so they do not re-poison future runs) " "user_id=%s thread_id=%s err=%s", len(closing), user_id, thread_id, exc, ) - await self._mark_turns_extracted(closing, field="episode_extracted_at") + await self._advance_episode_cursor(user_id, thread_id, closing[-1]) segment = segment[boundary:] continue for doc, embedding in zip(docs, embeddings_for_docs): doc["embedding"] = embedding - validated = self._validate_extracted_doc(doc) try: - await self._create_memory(validated) + await self._create_memory(doc) total += 1 except CosmosResourceExistsError: - logger.info("extract_episodes idempotent skip duplicate episode id=%s", validated.get("id")) - await self._mark_turns_extracted(closing, field="episode_extracted_at") + logger.info("extract_episodes idempotent skip duplicate episode id=%s", doc.get("id")) + await self._advance_episode_cursor(user_id, thread_id, closing[-1]) segment = segment[boundary:] return {"episodes": total} @@ -1331,40 +1217,12 @@ async def synthesize_procedural( *, force: bool = False, ) -> dict[str, Any]: - """Synthesize the active procedural prompt for a user.""" + """Extract atomic procedural memories from behavioral facts and lessons.""" + del force if not user_id: raise ValidationError("user_id is required") - logger.info("synthesize_procedural started user_id=%s force=%s", user_id, force) - - async def _read_latest_procedural() -> Optional[dict[str, Any]]: - docs = await self._query_items( - self._memories_container, - query=( - "SELECT * FROM c WHERE c.user_id = @uid " - "AND c.thread_id = @thread_id " - "AND c.type = @type " - f"AND {_ACTIVE_DOC_FILTER}" - ), - parameters=[ - {"name": "@uid", "value": user_id}, - {"name": "@thread_id", "value": "__procedural__"}, - {"name": "@type", "value": "procedural"}, - ], - ) - docs.sort( - key=lambda doc: (int(doc.get("version") or 0), int(doc.get("_ts") or 0)), - reverse=True, - ) - if len(docs) > 1: - logger.warning( - "synthesize_procedural found multiple active docs user_id=%s count=%d", - user_id, - len(docs), - ) - return docs[0] if docs else None - - prior_doc = await _read_latest_procedural() + logger.info("synthesize_procedural extraction started user_id=%s", user_id) behavioral_fact_docs = await self._query_items( self._memories_container, @@ -1375,7 +1233,7 @@ async def _read_latest_procedural() -> Optional[dict[str, Any]]: "AND ((IS_DEFINED(c.metadata.category) " "AND c.metadata.category IN ('preference', 'requirement')) " "OR (IS_DEFINED(c.salience) AND c.salience >= @min_salience)) " - "ORDER BY c.salience DESC, c.created_at ASC, c.id ASC" + "ORDER BY c.created_at ASC" ), parameters=[ {"name": "@uid", "value": user_id}, @@ -1388,7 +1246,6 @@ async def _read_latest_procedural() -> Optional[dict[str, Any]]: for doc in behavioral_fact_docs if isinstance(doc.get("content"), str) and doc.get("content", "").strip() ] - behavioral_fact_ids = [doc["id"] for doc in behavioral_fact_docs] episodic_docs = await self._query_items( self._memories_container, @@ -1396,150 +1253,195 @@ async def _read_latest_procedural() -> Optional[dict[str, Any]]: "SELECT TOP 50 * FROM c WHERE c.user_id = @uid " "AND c.type = @type " f"AND {_ACTIVE_DOC_FILTER} " - "AND IS_DEFINED(c.lessons) " - "AND ARRAY_LENGTH(c.lessons) > 0 " - "ORDER BY c.salience DESC, c.created_at ASC, c.id ASC" + "AND IS_DEFINED(c.lessons) AND ARRAY_LENGTH(c.lessons) > 0 " + "ORDER BY c.created_at ASC" ), parameters=[ {"name": "@uid", "value": user_id}, {"name": "@type", "value": "episodic"}, ], ) - episodic_with_lessons = [ - doc - for doc in episodic_docs - if isinstance(doc.get("lessons"), list) - and any(isinstance(lesson, str) and lesson.strip() for lesson in doc.get("lessons", [])) - ] - source_episodic_ids = [doc["id"] for doc in episodic_with_lessons] - current_source_ids = set(behavioral_fact_ids) | set(source_episodic_ids) - - def _covered_by(prior: Optional[dict[str, Any]]) -> bool: - if prior is None: - return False - covered = set(prior.get("source_fact_ids") or []) | set(prior.get("source_episodic_ids") or []) - return current_source_ids.issubset(covered) - - if prior_doc and not force and _covered_by(prior_doc): - logger.info( - "synthesize_procedural unchanged user_id=%s fact_count=%d episodic_count=%d", - user_id, - len(behavioral_fact_ids), - len(source_episodic_ids), - ) - return {"status": "unchanged", "procedural": prior_doc} + def _episodic_lessons(doc: dict[str, Any]) -> list[str]: + lessons = doc.get("lessons") + if isinstance(lessons, list): + return [lesson.strip() for lesson in lessons if isinstance(lesson, str) and lesson.strip()] + return [] - if not current_source_ids: - logger.info( - "synthesize_procedural skipping LLM user_id=%s - no behavioral facts or episodic lessons", - user_id, - ) - return {"status": "unchanged", "procedural": prior_doc} - - user_name = "the user" - - def _render_bullets(values: list[str]) -> str: - cleaned = [value.strip() for value in values if isinstance(value, str) and value.strip()] - if not cleaned: - return "(none)" - return "\n".join(f"- {value}" for value in cleaned) - - static_prompty_inputs = { - "behavioral_facts": _render_bullets([doc.get("content", "") for doc in behavioral_fact_docs]), - "episodic_lessons": _render_bullets( - [ - lesson - for doc in episodic_with_lessons - for lesson in doc.get("lessons", []) - if isinstance(lesson, str) and lesson.strip() - ] - ), - "user_name": user_name, - } + fact_lines: list[str] = [] + fact_label_to_id: dict[str, str] = {} + for index, doc in enumerate(behavioral_fact_docs, start=1): + label = f"fact-{index}" + category = "" + metadata = doc.get("metadata") + if isinstance(metadata, dict) and isinstance(metadata.get("category"), str): + category = metadata["category"].strip() + fact_lines.append(f"{label} [{category or 'unknown'}]: {doc['content'].strip()}") + if isinstance(doc.get("id"), str): + fact_label_to_id[label] = doc["id"] + + episodic_lines: list[str] = [] + episodic_label_to_id: dict[str, str] = {} + ep_index = 0 + for doc in episodic_docs: + doc_id = doc.get("id") + for lesson in _episodic_lessons(doc): + ep_index += 1 + label = f"ep-{ep_index}" + episodic_lines.append(f"{label}: {lesson}") + if isinstance(doc_id, str): + episodic_label_to_id[label] = doc_id + + if not fact_lines and not episodic_lines: + logger.info("synthesize_procedural unchanged user_id=%s - no procedure sources", user_id) + return {"status": "unchanged", "procedures_created": 0} - # 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 - # source coverage, not just version number. - written_doc: Optional[dict[str, Any]] = None - for attempt in range(1, _PROCEDURAL_MAX_CREATE_ATTEMPTS + 1): + try: response_text = await self._run_prompty( - "synthesize_procedural.prompty", + "extract_procedure.prompty", inputs={ - "prior_prompt": (prior_doc.get("content") or "") if prior_doc else "", - **static_prompty_inputs, + "behavioral_facts": "\n".join(fact_lines), + "episodic_lessons": "\n".join(episodic_lines), }, ) - parsed = self._parse_llm_json(response_text) - system_prompt = parsed.get("system_prompt") if isinstance(parsed, dict) else None - if not isinstance(system_prompt, str) or not system_prompt.strip(): - raise LLMError("synthesize_procedural returned JSON without a non-empty 'system_prompt' string") - system_prompt = system_prompt.strip() - - new_seq = (int(prior_doc.get("version") or 0) + 1) if prior_doc else 1 - new_doc: dict[str, Any] = { - "id": f"proc_{user_id}_{new_seq}", - "user_id": user_id, - "thread_id": "__procedural__", - "type": "procedural", - "version": new_seq, - "content": system_prompt, - "source_fact_ids": behavioral_fact_ids, - "source_episodic_ids": source_episodic_ids, - "supersedes_ids": [prior_doc["id"]] if prior_doc else [], - "created_at": datetime.now(timezone.utc).isoformat(), - "updated_at": datetime.now(timezone.utc).isoformat(), - "role": "system", - "tags": ["sys:procedural", "sys:synthesized"], - **self._prompt_lineage("synthesize_procedural.prompty"), - "metadata": {}, - } - validated = construct_internal(ProceduralRecord, new_doc).to_doc() + procedures = parsed.get("procedures", []) if isinstance(parsed, dict) else [] + if not isinstance(procedures, list): + procedures = [] + except Exception as exc: # LLM/parsing quarantine: one bad call must not stop the pipeline. + if is_retryable_llm_error(exc): + logger.warning("synthesize_procedural deferred user_id=%s: %s", user_id, exc) + return {"status": "deferred", "procedures_created": 0} + logger.exception("synthesize_procedural skipped user_id=%s after non-retryable LLM error", user_id) + return {"status": "skipped", "procedures_created": 0} + + trusted_source_kinds = {kind.value for kind in TRUSTED_PROCEDURE_SOURCE_KINDS} + authority_by_source = { + "explicit_user_instruction": "high", + "organization_policy": "high", + "observed_user_preference": "medium", + "episode_distillation": "medium", + "document_content": "low", + "agent_inference": "low", + } + now = datetime.now(timezone.utc).isoformat() + created = 0 + skipped = 0 + + for proc in procedures: try: - 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", - user_id, - new_seq, - attempt, - _PROCEDURAL_MAX_CREATE_ATTEMPTS, - ) - latest = await _read_latest_procedural() - if latest is None: + if not isinstance(proc, dict): + skipped += 1 continue - prior_doc = latest - if _covered_by(prior_doc): - logger.info( - "synthesize_procedural race resolved by coverage user_id=%s winner=%s", - user_id, - prior_doc["id"], - ) - return {"status": "unchanged", "procedural": prior_doc} - if written_doc is None: - raise MemoryConflictError( - "synthesize_procedural failed after " - f"{_PROCEDURAL_MAX_CREATE_ATTEMPTS} attempts due to id collisions " - f"user_id={user_id!r}" - ) + name = proc.get("name") + if not isinstance(name, str) or not name.strip(): + skipped += 1 + continue + name = name.strip() - new_id = written_doc["id"] - if prior_doc: - await self._mark_superseded(prior_doc, new_id, reason="update") + grounded_in = proc.get("grounded_in") + if isinstance(grounded_in, str): + labels = [grounded_in] + elif isinstance(grounded_in, list): + labels = [label for label in grounded_in if isinstance(label, str)] + else: + labels = [] + source_fact_ids = sorted({fact_label_to_id[label] for label in labels if label in fact_label_to_id}) + source_episodic_ids = sorted( + {episodic_label_to_id[label] for label in labels if label in episodic_label_to_id} + ) + + source_kind = proc.get("source_kind", "agent_inference") + if not isinstance(source_kind, str): + source_kind = "agent_inference" + # Episode-only grounding cannot corroborate a user/org instruction: + # any trusted label backed solely by episodes (no behavioral fact) + # downgrades to episode_distillation - a candidate, never an + # auto-active policy. + if source_episodic_ids and not source_fact_ids and source_kind in trusted_source_kinds: + source_kind = "episode_distillation" + source_authority = authority_by_source.get(source_kind, "low") + status = "active" if source_kind in trusted_source_kinds else "candidate" + # Grounding is the trust anchor: a procedure whose ``grounded_in`` + # resolved to no persisted fact or episodic source is never + # auto-activated, regardless of the LLM's self-declared + # source_kind - this blocks an ungrounded self-labeled instruction + # from being compiled into the runtime system prompt. + if not source_fact_ids and not source_episodic_ids: + status = "candidate" + + summary = proc.get("summary") if isinstance(proc.get("summary"), str) else "" + retrieval_text = proc.get("retrieval_text") if isinstance(proc.get("retrieval_text"), str) else "" + scope_type = proc.get("scope_type") if isinstance(proc.get("scope_type"), str) else "user" + scope_value = proc.get("scope_value") if isinstance(proc.get("scope_value"), str) else None + proc_id = ( + "proc_" + + hashlib.sha256( + f"{user_id}|{scope_type}|{scope_value or ''}|{name.strip().lower()}".encode() + ).hexdigest()[:32] + ) + doc: dict[str, Any] = { + "id": proc_id, + "user_id": user_id, + "thread_id": "__procedural__", + "type": "procedural", + "role": "system", + "tags": ["sys:procedural", "sys:auto-extracted"], + "created_at": now, + "updated_at": now, + "name": name, + "summary": summary.strip() or name, + "retrieval_text": retrieval_text.strip() or summary.strip() or name, + "procedure_kind": proc.get("procedure_kind", "behavioral_policy"), + "scope_type": scope_type, + "scope_value": scope_value, + "activation_conditions": proc.get("activation_conditions", []), + "preconditions": proc.get("preconditions", []), + "steps": proc.get("steps", []), + "success_conditions": proc.get("success_conditions", []), + "failure_conditions": proc.get("failure_conditions", []), + "safety_constraints": proc.get("safety_constraints", []), + "status": status, + "priority": proc.get("priority", 0), + # Seed utility from the LLM's extraction confidence (the + # extract_procedure schema emits ``confidence``, not + # ``utility_score``). Procedures are create-only today with no + # promotion or outcome-scoring path, so this is the record's + # final utility value. + "utility_score": clamp_unit_interval(proc.get("confidence"), 0.5), + "successful_uses": proc.get("successful_uses", 0), + "failed_uses": proc.get("failed_uses", 0), + "source_kind": source_kind, + "source_authority": source_authority, + "source_fact_ids": source_fact_ids, + "source_episodic_ids": source_episodic_ids, + "source_turn_ids": proc.get("source_turn_ids", []), + "content": summary.strip() or name, + "version": proc.get("version", 1), + "metadata": {}, + **self._prompt_lineage("extract_procedure.prompty"), + } + validated = construct_internal(ProceduralRecord, doc).to_doc() + validated["embedding"] = await self._embed_one(validated["retrieval_text"]) + try: + await self._create_memory(validated) + created += 1 + except CosmosResourceExistsError: + skipped += 1 + except (ValidationError, PydanticValidationError, ValueError) as exc: + skipped += 1 + logger.warning("synthesize_procedural dropping malformed procedure user_id=%s: %s", user_id, exc) + except Exception: + skipped += 1 + logger.exception("synthesize_procedural failed to persist one procedure user_id=%s", user_id) logger.info( - "synthesize_procedural synthesized user_id=%s version=%d fact_count=%d episodic_count=%d", + "synthesize_procedural extracted user_id=%s procedures_created=%d procedures_skipped=%d", user_id, - written_doc["version"], - len(behavioral_fact_ids), - len(source_episodic_ids), + created, + skipped, ) - return {"status": "synthesized", "procedural": written_doc} + return {"status": "synthesized", "procedures_created": created, "procedures_skipped": skipped} async def generate_thread_summary_durable( self, @@ -1908,11 +1810,10 @@ async def _load_memories_by_ids( 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. + Async mirror of the sync contradiction-only reconcile. 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: @@ -2013,29 +1914,144 @@ async def _reconcile_contradictions( ) return result - async def build_procedural_context(self, user_id: str) -> str: - """Return the active synthesized procedural prompt for system injection.""" + async def build_procedural_context(self, user_id: str, task: Optional[str] = None) -> str: + """Build a deterministic system prompt projection from active procedures.""" if not user_id: raise ValidationError("user_id is required") query = ( - "SELECT TOP 1 c.content, c.version FROM c WHERE c.user_id = @user_id " - "AND c.thread_id = @thread_id AND c.type = @type " - f"AND {_ACTIVE_DOC_FILTER} " - "ORDER BY c.version DESC" + "SELECT * FROM c WHERE c.user_id=@uid AND c.type='procedural' " + "AND c.status='active' " + "AND (NOT IS_DEFINED(c.superseded_by) OR IS_NULL(c.superseded_by))" ) - items = await self._query_items( + procedures = await self._query_items( self._memories_container, query=query, - parameters=[ - {"name": "@user_id", "value": user_id}, - {"name": "@thread_id", "value": "__procedural__"}, - {"name": "@type", "value": "procedural"}, - ], + parameters=[{"name": "@uid", "value": user_id}], ) - if not items: + procedures = [ + proc + for proc in procedures + if proc.get("user_id") == user_id + and proc.get("type") == "procedural" + and proc.get("status") == "active" + and not proc.get("superseded_by") + ] + + stopwords = { + "a", + "an", + "and", + "are", + "as", + "at", + "be", + "by", + "for", + "from", + "in", + "is", + "it", + "of", + "on", + "or", + "the", + "to", + "with", + } + + def _tokens(value: str) -> set[str]: + return {token for token in re.findall(r"[a-z0-9]+", value.lower()) if token not in stopwords} + + task_tokens = _tokens(task) if isinstance(task, str) and task.strip() else set() + policies: list[dict[str, Any]] = [] + task_procedures: list[dict[str, Any]] = [] + for proc in procedures: + kind = proc.get("procedure_kind") + scope_type = proc.get("scope_type") + if kind in {"behavioral_policy", "decision_rule"} and scope_type in {"global", "user"}: + policies.append(proc) + continue + if task_tokens and kind in {"workflow", "recovery_strategy", "tool_usage"}: + searchable = " ".join( + [ + proc.get("retrieval_text") if isinstance(proc.get("retrieval_text"), str) else "", + *[ + condition + for condition in proc.get("activation_conditions", []) + if isinstance(condition, str) + ], + ] + ) + if task_tokens & _tokens(searchable): + task_procedures.append(proc) + + included = policies + task_procedures + if not included: return "" - content = items[0].get("content") - return content if isinstance(content, str) else "" + + authority_rank = {"mandatory": 3, "high": 2, "medium": 1, "low": 0} + + def _sort_key(proc: dict[str, Any]) -> tuple[int, int, str]: + try: + priority = int(proc.get("priority") or 0) + except (TypeError, ValueError): + priority = 0 + authority = proc.get("source_authority") + rank = authority_rank.get(authority if isinstance(authority, str) else "low", 0) + name = proc.get("name") if isinstance(proc.get("name"), str) else "" + return (-priority, -rank, name.lower()) + + policies.sort(key=_sort_key) + task_procedures.sort(key=_sort_key) + included = policies + task_procedures + + fingerprint_payload = sorted( + ( + str(proc.get("id", "")), + str(proc.get("version", "")), + str(proc.get("status", "")), + str(proc.get("priority", "")), + str(proc.get("scope_type", "")), + str(proc.get("scope_value", "")), + ) + for proc in included + ) + fingerprint = hashlib.sha256(json.dumps(fingerprint_payload, separators=(",", ":")).encode()).hexdigest() + logger.debug( + "build_procedural_context fingerprint=%s included_ids=%s user_id=%s", + fingerprint, + [proc.get("id") for proc in included], + user_id, + ) + + lines = ["# Learned procedures"] + if policies: + lines.append("") + lines.append("## Behavioral policies") + for proc in policies: + name = proc.get("name") if isinstance(proc.get("name"), str) else "Unnamed procedure" + summary = proc.get("summary") if isinstance(proc.get("summary"), str) else proc.get("content", "") + lines.append(f"- {name}: {summary}") + if task_procedures: + lines.append("") + lines.append("## Task procedures") + for proc in task_procedures: + name = proc.get("name") if isinstance(proc.get("name"), str) else "Unnamed procedure" + summary = proc.get("summary") if isinstance(proc.get("summary"), str) else proc.get("content", "") + lines.append(f"### {name}") + if summary: + lines.append(f"Summary: {summary}") + steps = proc.get("steps") if isinstance(proc.get("steps"), list) else [] + if steps: + sorted_steps = sorted( + [step for step in steps if isinstance(step, dict)], + key=lambda step: int(step.get("sequence") or 0), + ) + for index, step in enumerate(sorted_steps, start=1): + instruction = step.get("instruction") if isinstance(step.get("instruction"), str) else "" + if instruction: + lines.append(f"{index}. {instruction}") + return "\n".join(lines) __all__ = ["AsyncPipelineService"] diff --git a/azure/cosmos/agent_memory/aio/store/memory_store.py b/azure/cosmos/agent_memory/aio/store/memory_store.py index ceb8fc0..6c04a4d 100644 --- a/azure/cosmos/agent_memory/aio/store/memory_store.py +++ b/azure/cosmos/agent_memory/aio/store/memory_store.py @@ -137,23 +137,23 @@ async def _query_items( except Exception as exc: raise CosmosOperationError(f"{operation} failed: {exc}") from exc - async def add_cosmos(self, record: dict[str, Any]) -> dict[str, Any]: + async def upsert_memory(self, record: dict[str, Any]) -> dict[str, Any]: """Upsert a pre-built Cosmos memory document and return the stored body.""" body = self._prepare_doc(record) memory_type = body.get("type") if memory_type not in _CONTAINER_FOR_TYPE: raise ValueError( - f"add_cosmos: record id={body.get('id')!r} has invalid type={memory_type!r}. " - f"Set 'type' to one of {sorted(_CONTAINER_FOR_TYPE)} before calling add_cosmos." + f"upsert_memory: record id={body.get('id')!r} has invalid type={memory_type!r}. " + f"Set 'type' to one of {sorted(_CONTAINER_FOR_TYPE)} before calling upsert_memory." ) container = self._container_for_type(memory_type) try: response = await container.upsert_item(body=body) except Exception as exc: raise _wrap_cosmos_exception( - exc, message=f"async add_cosmos upsert failed for record {body.get('id')}: {exc}" + exc, message=f"async upsert_memory upsert failed for record {body.get('id')}: {exc}" ) from exc - logger.info("add_cosmos id=%s role=%s type=%s", body.get("id"), body.get("role"), body.get("type")) + logger.info("upsert_memory id=%s role=%s type=%s", body.get("id"), body.get("role"), body.get("type")) return response if isinstance(response, dict) else body async def add( @@ -217,7 +217,7 @@ async def add( body["embedding"] = await self._embeddings_client.generate(content) except Exception as exc: # noqa: BLE001 logger.warning( - "add_cosmos: embedding generation failed for %s (%s); proceeding without embedding", + "upsert_memory: embedding generation failed for %s (%s); proceeding without embedding", record.id, exc, ) @@ -228,7 +228,7 @@ async def add( await container.upsert_item(body=body) except Exception as exc: raise _wrap_cosmos_exception(exc, message=f"Async upsert failed for record {record.id}: {exc}") from exc - logger.info("add_cosmos id=%s role=%s type=%s", record.id, role, memory_type) + logger.info("upsert_memory id=%s role=%s type=%s", record.id, role, memory_type) return record.id async def push(self, local_memory: list[dict[str, Any]], batch_size: int = 25) -> None: @@ -1147,6 +1147,54 @@ async def build_episodic_context(self, user_id: str, query: str, top_k: int = 3) memories = await self.search_episodic(user_id, query, top_k=top_k) return format_episodic_context(memories) + async def retrieve_procedures( + self, + user_id: str, + search_terms: str, + top_k: int = 5, + *, + scope_type: Optional[str] = None, + scope_value: Optional[str] = None, + procedure_kind: Optional[str] = None, + status: Optional[str] = "active", + include_superseded: bool = False, + ) -> list[dict[str, Any]]: + """Semantic search across procedural memories for a user.""" + if not user_id: + raise ValidationError("user_id is required for retrieve_procedures") + terms = require_search_terms(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.type", "@type", "procedural") + qb.add_filter("c.user_id", "@user_id", user_id) + qb.add_filter("c.scope_type", "@scope_type", scope_type) + qb.add_filter("c.scope_value", "@scope_value", scope_value) + qb.add_filter("c.procedure_kind", "@procedure_kind", procedure_kind) + qb.add_filter("c.status", "@status", status) + + 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}) + for i, kw in enumerate(keywords): + parameters.append({"name": f"@kw{i}", "value": kw}) + + partition_key, _ = query_scope(user_id, None) + logger.debug("AsyncMemoryStore.retrieve_procedures query: %s", sql) + return await self.query( + sql, + parameters, + container_key=ContainerKey.MEMORIES, + partition_key=partition_key, + ) + async def _embed(self, text: str) -> list[float]: if self._embeddings_client is None: raise ConfigurationError( diff --git a/azure/cosmos/agent_memory/cosmos_memory_client.py b/azure/cosmos/agent_memory/cosmos_memory_client.py index 7943614..fc5390a 100644 --- a/azure/cosmos/agent_memory/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/cosmos_memory_client.py @@ -25,7 +25,7 @@ from .auto_trigger import maybe_trigger_steps from .chat import ChatClient from .embeddings import EmbeddingsClient -from .exceptions import CosmosOperationError, ValidationError +from .exceptions import CosmosOperationError, MemoryNotFoundError, ValidationError from .processors import InProcessProcessor, MemoryProcessor from .services._pipeline_helpers import _normalize_cadence_thresholds, _normalize_metadata_keys from .services.pipeline import PipelineService @@ -294,11 +294,9 @@ def create_memory_store( vector_index_type=_resolve_vector_index_type(vector_index_type), ) vec_policy, idx_policy, ft_policy = _container_policies(**_policy_kwargs) - # Turns always carry the vector index (primed for search) but skip the - # salience composite index, which only procedural synthesis needs. + # Turns always carry the vector index (primed for search). turns_vec_policy, turns_idx_policy, turns_ft_policy = _container_policies( **{**_policy_kwargs, "vector_index_type": "quantizedFlat"}, - include_salience_composite=False, ) self._memories_container_client = db.create_container_if_not_exists( **_build_container_kwargs( @@ -328,7 +326,6 @@ def create_memory_store( # composite (user_id, thread_id, version) get_*_summary relies on. summaries_vec_policy, summaries_idx_policy, summaries_ft_policy = _container_policies( **{**_policy_kwargs, "vector_index_type": "quantizedFlat"}, - include_salience_composite=False, ) summaries_idx_policy["compositeIndexes"] = [ [ @@ -521,7 +518,7 @@ def _container_for_type(self, memory_type: str) -> Any: """Return the Cosmos container client that owns ``memory_type``.""" return self._containers[container_key_for_type(memory_type)] - def add_cosmos( + def upsert_memory( self, user_id: str, role: str, @@ -569,7 +566,7 @@ def add_cosmos( try: self._maybe_auto_trigger({(user_id, thread_id): 1}) except Exception as exc: - logger.warning("Auto-trigger after add_cosmos failed: %s", exc) + logger.warning("Auto-trigger after upsert_memory failed: %s", exc) return memory_id def push_to_cosmos(self, batch_size: int = 25) -> None: @@ -638,7 +635,7 @@ def update_cosmos( metadata=metadata, ) - def delete_cosmos( + def delete_memory( self, memory_id: str, *, @@ -654,6 +651,79 @@ def delete_cosmos( memory_type=memory_type, ) + def delete_turn(self, turn_id: str, *, user_id: str, thread_id: str) -> None: + """Delete a single turn document. Raises if it does not exist.""" + return self.delete_memory( + turn_id, + user_id=user_id, + thread_id=thread_id, + memory_type="turn", + ) + + def delete_thread_summary(self, user_id: str, thread_id: str) -> bool: + """Delete a thread's summary if present. Returns True when one was deleted. + + The summary id is deterministic, so callers need not look it up first; a + missing summary is a no-op that returns False. + """ + try: + self.delete_memory( + f"summary_{user_id}_{thread_id}", + user_id=user_id, + thread_id=thread_id, + memory_type="thread_summary", + ) + return True + except MemoryNotFoundError: + return False + + def delete_user_summary(self, user_id: str) -> bool: + """Delete a user's summary if present. Returns True when one was deleted. + + A missing summary is a no-op that returns False. + """ + try: + self.delete_memory( + f"user_summary_{user_id}", + user_id=user_id, + thread_id="__user_summary__", + memory_type="user_summary", + ) + return True + except MemoryNotFoundError: + return False + + def delete_thread(self, user_id: str, thread_id: str, *, include_summary: bool = True) -> int: + """Bulk-delete a conversation thread: all of its turns and, by default, its + thread summary. Returns the number of documents deleted. + + This targets the conversation itself (turns + summary). Durable memories + distilled from the thread - facts, episodes, procedures - are user-scoped + knowledge and are left intact. Deletion is best-effort per document: a + concurrently-removed turn is skipped rather than aborting the whole sweep. + """ + if not user_id: + raise ValidationError("user_id is required") + if not thread_id: + raise ValidationError("thread_id is required") + + deleted = 0 + turns = self.get_thread(thread_id=thread_id, user_id=user_id, include_superseded=True) or [] + for turn in turns: + turn_id = turn.get("id") + if not turn_id: + continue + try: + self.delete_memory(turn_id, user_id=user_id, thread_id=thread_id, memory_type="turn") + deleted += 1 + except MemoryNotFoundError: + continue + + if include_summary and self.delete_thread_summary(user_id, thread_id): + deleted += 1 + + return deleted + def search_cosmos( self, search_terms: str, @@ -692,6 +762,11 @@ def search_cosmos( store = self._get_store() # Facts + episodes share one ranked query and one top_k budget: episodic # is added when include_episodes is True and stripped when it is False. + if memory_types is not None and "episodic" in memory_types and not include_episodes: + logger.warning( + "Episodic memories requested via memory_types are only returned when include_episodes=True; " + "proceeding without episodic memories and using facts or other requested memory types only." + ) if memory_types is not None: base_memory_types = [t for t in memory_types if t != "episodic"] else: @@ -941,8 +1016,14 @@ def remove_tags( return self._get_store().remove_tags(memory_id, user_id, thread_id, memory_type, tags) def get_procedural_prompt(self, user_id: str) -> Optional[str]: - """Return the active synthesized procedural prompt for a user.""" - return self._get_store().get_procedural_prompt(user_id=user_id) + """Return the compiled procedural system prompt for a user. + + The prompt is a deterministic projection of the user's active procedures + (compiled on demand), not a stored record. Returns None when the user has + no active procedures. + """ + prompt = self._get_pipeline().build_procedural_context(user_id) + return prompt or None def get_procedural_history(self, user_id: str, limit: int = 10) -> list[dict[str, Any]]: """Return synthesized procedural docs for a user, newest first.""" @@ -976,9 +1057,42 @@ def search_episodic_memories( 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 retrieve_procedures( + self, + user_id: str, + search_terms: str, + top_k: int = 5, + *, + scope_type: Optional[str] = None, + scope_value: Optional[str] = None, + procedure_kind: Optional[str] = None, + status: Optional[str] = "active", + include_superseded: bool = False, + ) -> list[dict[str, Any]]: + """Context-aware semantic retrieval of a user's procedures. + + Ranks active procedures by relevance to ``search_terms`` (a task or + situation), optionally narrowed by scope or kind. Pass ``status=None`` to + include non-active (e.g. candidate) procedures. + """ + return self._get_store().retrieve_procedures( + user_id=user_id, + search_terms=search_terms, + top_k=top_k, + scope_type=scope_type, + scope_value=scope_value, + procedure_kind=procedure_kind, + status=status, + include_superseded=include_superseded, + ) + + def build_procedural_context(self, user_id: str, task: Optional[str] = None) -> str: + """Compile a procedural system-prompt projection for prompt injection. + + With no ``task``, returns the user's always-on behavioral policies. With a + ``task`` string, also folds in relevant task procedures (skills). + """ + return self._get_pipeline().build_procedural_context(user_id, task) def build_episodic_context( self, diff --git a/azure/cosmos/agent_memory/models.py b/azure/cosmos/agent_memory/models.py index 7e0683c..628d858 100644 --- a/azure/cosmos/agent_memory/models.py +++ b/azure/cosmos/agent_memory/models.py @@ -12,6 +12,7 @@ from __future__ import annotations +import math import re import uuid from datetime import datetime, timezone @@ -53,6 +54,70 @@ class MemoryType(str, Enum): episodic = "episodic" +class ProcedureKind(str, Enum): + """The kind of atomic procedural knowledge a procedure captures.""" + + behavioral_policy = "behavioral_policy" + workflow = "workflow" + decision_rule = "decision_rule" + tool_usage = "tool_usage" + recovery_strategy = "recovery_strategy" + + +class ProcedureScopeType(str, Enum): + """How broadly a procedure applies. More specific scopes override broader ones.""" + + global_scope = "global" + user = "user" + agent = "agent" + domain = "domain" + project = "project" + workflow = "workflow" + tool = "tool" + + +class ProcedureStatus(str, Enum): + """Lifecycle state of a procedure. Only ``active`` procedures are compiled/injected.""" + + candidate = "candidate" + active = "active" + deprecated = "deprecated" + rejected = "rejected" + + +class ProcedureSourceKind(str, Enum): + """Where a procedure came from - gates whether it may become active policy.""" + + explicit_user_instruction = "explicit_user_instruction" + observed_user_preference = "observed_user_preference" + organization_policy = "organization_policy" + episode_distillation = "episode_distillation" + document_content = "document_content" + agent_inference = "agent_inference" + + +class ProcedureSourceAuthority(str, Enum): + """How authoritative a procedure's source is, for conflict resolution and gating.""" + + mandatory = "mandatory" + high = "high" + medium = "medium" + low = "low" + + +# Source kinds trusted enough to auto-activate a procedure as behavioral policy. +# Untrusted kinds (document_content, agent_inference) stay ``candidate`` until +# validated or explicitly approved - this stops an imperative sentence lifted +# from a document, or a one-off agent guess, from silently becoming agent policy. +TRUSTED_PROCEDURE_SOURCE_KINDS: frozenset[ProcedureSourceKind] = frozenset( + { + ProcedureSourceKind.explicit_user_instruction, + ProcedureSourceKind.observed_user_preference, + ProcedureSourceKind.organization_policy, + } +) + + def _uuid4_str() -> str: return str(uuid.uuid4()) @@ -471,18 +536,59 @@ def _validate_time_order(self) -> "EpisodicRecord": return self +class ProcedureStep(BaseModel): + """One ordered step in a procedure's execution.""" + + sequence: int + instruction: str + expected_result: Optional[str] = None + on_failure: Optional[str] = None + tool_name: Optional[str] = None + + class ProceduralRecord(MemoryRecordBase): - """Synthesized agent self-knowledge: the active personalized system prompt.""" + """An atomic, reusable procedural memory: a behavioral policy or task skill. + + Procedural memory answers "what should I do, and how should I do it" - stored + as many small, independently retrievable units (a policy/skill library), not + as one compiled system prompt. The runtime personalized system prompt is a + deterministic projection of the active, in-scope procedures (built by + ``build_procedural_context``); it is not stored as a record here. + """ memory_type: Literal[MemoryType.procedural] = Field( # type: ignore[assignment] alias="type", default=MemoryType.procedural ) + + name: str + summary: str + retrieval_text: str + + procedure_kind: ProcedureKind + scope_type: ProcedureScopeType = ProcedureScopeType.user + scope_value: Optional[str] = None + + activation_conditions: list[str] = Field(default_factory=list) + preconditions: list[str] = Field(default_factory=list) + steps: list[ProcedureStep] = Field(default_factory=list) + success_conditions: list[str] = Field(default_factory=list) + failure_conditions: list[str] = Field(default_factory=list) + safety_constraints: list[str] = Field(default_factory=list) + + status: ProcedureStatus = ProcedureStatus.candidate + priority: int = 0 + utility_score: float = 0.5 + successful_uses: int = 0 + failed_uses: int = 0 + + source_kind: ProcedureSourceKind = ProcedureSourceKind.agent_inference + source_authority: ProcedureSourceAuthority = ProcedureSourceAuthority.low + source_turn_ids: list[str] = Field(default_factory=list) + content_hash: Optional[str] = None - prompt_id: str - prompt_version: str = "v1" + prompt_id: Optional[str] = None + prompt_version: str = "v2" version: int = 1 - source_fact_ids: list[str] = Field(default_factory=list) - source_episodic_ids: list[str] = Field(default_factory=list) _ID_PREFIX: ClassVar[Optional[str]] = "proc_" @@ -495,12 +601,26 @@ def _validate_version(cls, v: Any) -> Any: raise ValueError(f"ProceduralRecord.version must be a positive integer, got {v!r}") return v + @field_validator("utility_score", mode="before") + @classmethod + def _clamp_utility(cls, v: Any) -> Any: + if v is None: + return 0.5 + try: + value = float(v) + except (TypeError, ValueError): + return 0.5 + if not math.isfinite(value): + return 0.5 + return min(1.0, max(0.0, value)) + @model_validator(mode="after") - def _require_sources(self) -> "ProceduralRecord": - if not self.source_fact_ids and not self.source_episodic_ids: - raise ValueError( - "ProceduralRecord requires at least one of source_fact_ids or source_episodic_ids to be non-empty" - ) + def _require_executable_steps(self) -> "ProceduralRecord": + # A workflow/recovery_strategy asserts a multi-step method; requiring at + # least one step keeps such procedures executable rather than empty. + # ``procedure_kind`` is a plain string here (model_config use_enum_values). + if self.procedure_kind in (ProcedureKind.workflow, ProcedureKind.recovery_strategy) and not self.steps: + raise ValueError(f"ProceduralRecord of kind {self.procedure_kind} requires at least one step") return self @@ -603,6 +723,13 @@ class OrchestrationResult(BaseModel): "EpisodeEvent", "EpisodeOutcome", "EpisodicRecord", + "ProcedureStep", + "ProcedureKind", + "ProcedureScopeType", + "ProcedureStatus", + "ProcedureSourceKind", + "ProcedureSourceAuthority", + "TRUSTED_PROCEDURE_SOURCE_KINDS", "ProceduralRecord", "TYPED_RECORD_CLASSES", "TAG_PATTERN", diff --git a/azure/cosmos/agent_memory/processors/durable.py b/azure/cosmos/agent_memory/processors/durable.py index e9aedc6..0f6d81b 100644 --- a/azure/cosmos/agent_memory/processors/durable.py +++ b/azure/cosmos/agent_memory/processors/durable.py @@ -16,10 +16,6 @@ logger = get_logger(__name__) -# Set once we have warned that episodic memory is inert under the durable backend, -# so the warning fires a single time per process rather than on every no-op call. -_EPISODIC_DURABLE_WARNED = False - class DurableFunctionProcessor: """Signals "an Azure Durable Function app is the active processor." @@ -65,17 +61,10 @@ def process_extract_episodes( user_id: str, thread_id: str, ) -> dict[str, int]: - # Episodic segmentation is in-process only; the Durable backend has no - # episodic path yet, so this is an explicit no-op (the auto-trigger also - # gates episode extraction to the in-process processor). Warn once so a - # durable-mode operator can see that episodic memory is not being produced. - global _EPISODIC_DURABLE_WARNED - if not _EPISODIC_DURABLE_WARNED: - _EPISODIC_DURABLE_WARNED = True - logger.warning( - "Episodic memory is not available under the Durable Functions backend " - "(no episodic write path yet); episode extraction is a no-op in durable mode." - ) + # The Durable Function app owns episodic extraction via the Cosmos DB + # Change Feed trigger: ExtractEpisodesOrchestrator -> ee_ExtractEpisodes + # -> pipeline.extract_episodes. This hook mirrors synthesize_procedural + # by leaving Durable-owned work to the orchestrator. logger.debug( "DurableFunctionProcessor.process_extract_episodes no-op user_id=%s thread_id=%s", user_id, @@ -134,11 +123,13 @@ def synthesize_procedural( user_id: str, force: bool = False, ) -> dict[str, Any]: - raise NotImplementedError( - "Procedural synthesis runs automatically after reconcile in durable mode; " - "manual invocation via the SDK is not supported when the Durable Function " - "app is the active processor." - ) + # No-op, like the other durable hooks: procedural synthesis runs in the + # Durable Function app after reconcile. Returning instead of raising keeps + # the in-process auto-trigger from stamping a spurious failure each cadence + # when the Durable app is the active processor. + del force + logger.debug("DurableFunctionProcessor.synthesize_procedural no-op user_id=%s", user_id) + return {"status": "skipped", "procedures_created": 0} def close(self) -> None: logger.debug("DurableFunctionProcessor.close no-op") diff --git a/azure/cosmos/agent_memory/prompts/_schemas.py b/azure/cosmos/agent_memory/prompts/_schemas.py index a0536c4..fd0a940 100644 --- a/azure/cosmos/agent_memory/prompts/_schemas.py +++ b/azure/cosmos/agent_memory/prompts/_schemas.py @@ -171,6 +171,87 @@ } +# --------------------------------------------------------------------------- +# extract_procedure.prompty - distill atomic procedural memories (skills/rules) +# --------------------------------------------------------------------------- +_PROCEDURE_STEP = { + "type": "object", + "properties": { + "sequence": {"type": "integer"}, + "instruction": {"type": "string"}, + "expected_result": {"type": ["string", "null"]}, + "on_failure": {"type": ["string", "null"]}, + "tool_name": {"type": ["string", "null"]}, + }, + "required": ["sequence", "instruction", "expected_result", "on_failure", "tool_name"], + "additionalProperties": False, +} + +_PROCEDURE_ITEM = { + "type": "object", + "properties": { + "name": {"type": "string"}, + "summary": {"type": "string"}, + "retrieval_text": {"type": "string"}, + "procedure_kind": { + "type": "string", + "enum": ["behavioral_policy", "workflow", "decision_rule", "tool_usage", "recovery_strategy"], + }, + "scope_type": { + "type": "string", + "enum": ["global", "user", "agent", "domain", "project", "workflow", "tool"], + }, + "scope_value": {"type": ["string", "null"]}, + "activation_conditions": {"type": "array", "items": {"type": "string"}}, + "preconditions": {"type": "array", "items": {"type": "string"}}, + "steps": {"type": "array", "items": _PROCEDURE_STEP}, + "success_conditions": {"type": "array", "items": {"type": "string"}}, + "failure_conditions": {"type": "array", "items": {"type": "string"}}, + "safety_constraints": {"type": "array", "items": {"type": "string"}}, + "source_kind": { + "type": "string", + "enum": [ + "explicit_user_instruction", + "observed_user_preference", + "organization_policy", + "episode_distillation", + "document_content", + "agent_inference", + ], + }, + "grounded_in": {"type": "array", "items": {"type": "string"}}, + "confidence": {"type": "number"}, + }, + "required": [ + "name", + "summary", + "retrieval_text", + "procedure_kind", + "scope_type", + "scope_value", + "activation_conditions", + "preconditions", + "steps", + "success_conditions", + "failure_conditions", + "safety_constraints", + "source_kind", + "grounded_in", + "confidence", + ], + "additionalProperties": False, +} + +EXTRACT_PROCEDURE_SCHEMA: dict[str, Any] = { + "type": "object", + "properties": { + "procedures": {"type": "array", "items": _PROCEDURE_ITEM}, + }, + "required": ["procedures"], + "additionalProperties": False, +} + + # --------------------------------------------------------------------------- # summarize.prompty - first-pass thread summary # @@ -281,6 +362,7 @@ PROMPTY_SCHEMAS: dict[str, tuple[str, dict[str, Any]]] = { "dedup.prompty": ("DedupOutput", DEDUP_SCHEMA), "extract_episode.prompty": ("ExtractEpisodesOutput", EXTRACT_EPISODE_SCHEMA), + "extract_procedure.prompty": ("ExtractProceduresOutput", EXTRACT_PROCEDURE_SCHEMA), "extract_memories.prompty": ("ExtractMemoriesOutput", EXTRACT_MEMORIES_SCHEMA), "extract_memories-v2.prompty": ("ExtractMemoriesOutput", EXTRACT_MEMORIES_SCHEMA), "summarize.prompty": ("SummarizeOutput", SUMMARIZE_SCHEMA), diff --git a/azure/cosmos/agent_memory/prompts/extract_procedure.prompty b/azure/cosmos/agent_memory/prompts/extract_procedure.prompty new file mode 100644 index 0000000..1500c30 --- /dev/null +++ b/azure/cosmos/agent_memory/prompts/extract_procedure.prompty @@ -0,0 +1,115 @@ +--- +name: extract_procedure +version: v2 +description: Distill atomic, reusable procedural memories (skills and behavioral rules) from a user's behavioral facts and episodic lessons. +model: + apiType: chat + options: + seed: 42 + maxOutputTokens: 16384 + additionalProperties: + response_format: + type: json_object +inputs: + behavioral_facts: + type: string + episodic_lessons: + type: string +--- + +system: +You are a precision procedural-memory distiller. You are given a user's durable behavioral facts and lessons distilled from their past experiences. Your job is to extract ATOMIC, REUSABLE procedures: individual behavioral rules the agent should follow, or task skills describing how to perform a recurring task. + +Procedural memory answers "what should I do, and how should I do it" - NOT "what is true" (that is a fact) and NOT "what happened once" (that is an episode). Extract a procedure only when the source states or clearly implies a reusable way to act. + +## Output +Return a JSON object `{"procedures": [ ... ]}`. Emit one procedure per distinct reusable rule or skill. If the input contains no reusable procedural knowledge, return exactly `{"procedures": []}`. Prefer fewer, higher-quality, generalizable procedures over many narrow ones. + +## Procedure Fields +For each procedure produce ALL of these fields: +- `name`: a short, specific, imperative name (e.g. "Confirm destructive deletions", "Diagnose Cosmos DB ORDER BY failures"). +- `summary`: one or two sentences stating the rule or method concisely. +- `retrieval_text`: a search-optimized line of keywords and trigger phrases describing WHEN this applies (used for retrieval; include domain, tools, error phrases, and situation terms). +- `procedure_kind`: one of `behavioral_policy` (how to interact with the user), `workflow` (multi-step task method), `decision_rule` (a conditional choice), `tool_usage` (how to use a specific tool), `recovery_strategy` (how to recover from a failure). +- `scope_type`: one of `global`, `user`, `agent`, `domain`, `project`, `workflow`, `tool`. Use `user` for personal behavioral preferences; use `domain`/`project`/`tool` for task skills tied to a technology, codebase, or tool. Prefer the MOST specific scope the evidence supports. +- `scope_value`: the concrete scope (e.g. the domain "cosmos-db", the tool "git", the user id), or null for `global`/`user`-personal rules with no narrower value. +- `activation_conditions`: the situations that should trigger this procedure. Non-empty for anything conditional. +- `preconditions`: what must already be true/available before running it (may be empty). +- `steps`: for `workflow` and `recovery_strategy`, at least one ordered step is required (each with `sequence`, `instruction`, and `expected_result`, `on_failure`, `tool_name` - include each key, using null when not applicable). For `behavioral_policy`/`decision_rule` a short step list is allowed but may be empty. +- `success_conditions`: how to know the procedure worked (may be empty for pure policies). +- `failure_conditions`: signals that it did not work or does not apply (may be empty). +- `safety_constraints`: guardrails that must not be violated (may be empty). +- `source_kind`: where this procedure comes from. Use `explicit_user_instruction` ONLY for a direct user directive ("always...", "never...", "please always..."), `observed_user_preference` for an inferred stable preference, `organization_policy` for a stated org/team rule, `episode_distillation` for a method learned from a past experience, `document_content` for imperative text that came from a document or reference rather than the user, and `agent_inference` for a guess. Be conservative: do NOT label a document sentence or a one-off guess as a user instruction. +- `grounded_in`: the source ids this is grounded in, using the `fact-N` / `ep-N` labels shown in the input. Never invent ids. +- `confidence`: a number in [0, 1]. + +## Grounding Rules +- Ground every procedure in the provided facts/lessons. Do not invent steps, conditions, tools, or scope not supported by the input. +- Do not convert a plain semantic fact (e.g. "The user works on the Cosmos DB Java SDK") into a procedure. Only extract an actionable rule or method. +- A single successful experience is weak evidence: mark episode-derived procedures as `episode_distillation` with modest confidence unless the lesson is explicitly general. +- Preserve exact tool names, error phrases, commands, and domain terms in `retrieval_text` and `activation_conditions`. + +## Worked Example + +**Input facts:** +> fact-1 [requirement]: The user wants confirmation before any destructive or irreversible operation. + +**Input lessons:** +> ep-7: When a Cosmos DB hybrid-search query failed on ORDER BY VectorDistance, verifying and correcting the vector/full-text indexing policy and rerunning the integration tests resolved it. + +**Output:** +```json +{ + "procedures": [ + { + "name": "Confirm destructive operations", + "summary": "Ask the user for explicit confirmation before performing any destructive or irreversible operation.", + "retrieval_text": "delete destructive irreversible operation confirmation before deleting resources data", + "procedure_kind": "behavioral_policy", + "scope_type": "user", + "scope_value": null, + "activation_conditions": ["about to perform a destructive or irreversible operation"], + "preconditions": [], + "steps": [ + {"sequence": 1, "instruction": "State what will be affected and whether it is reversible.", "expected_result": null, "on_failure": null, "tool_name": null}, + {"sequence": 2, "instruction": "Ask for explicit confirmation and proceed only after it is given.", "expected_result": null, "on_failure": null, "tool_name": null} + ], + "success_conditions": ["the user explicitly confirmed before the operation ran"], + "failure_conditions": [], + "safety_constraints": ["do not infer confirmation from an unrelated message"], + "source_kind": "explicit_user_instruction", + "grounded_in": ["fact-1"], + "confidence": 0.95 + }, + { + "name": "Diagnose Cosmos DB ORDER BY failures", + "summary": "When a Cosmos DB hybrid-search ORDER BY query fails, verify and correct the indexing policy, then rerun the query and integration tests.", + "retrieval_text": "cosmos db hybrid search ORDER BY VectorDistance query failure indexing policy vector full-text integration tests", + "procedure_kind": "recovery_strategy", + "scope_type": "domain", + "scope_value": "cosmos-db", + "activation_conditions": ["a Cosmos DB hybrid-search query fails while evaluating an ORDER BY expression"], + "preconditions": ["the failing query and current indexing policy are available"], + "steps": [ + {"sequence": 1, "instruction": "Identify the path and ORDER BY expression in the failure.", "expected_result": null, "on_failure": null, "tool_name": null}, + {"sequence": 2, "instruction": "Verify the required vector and full-text indexing policies are configured.", "expected_result": null, "on_failure": null, "tool_name": null}, + {"sequence": 3, "instruction": "Correct the indexing-policy mismatch and rerun the original query.", "expected_result": "the query completes successfully", "on_failure": "check SDK support and document paths", "tool_name": null}, + {"sequence": 4, "instruction": "Run the relevant integration tests.", "expected_result": "all related tests pass", "on_failure": null, "tool_name": null} + ], + "success_conditions": ["the original query completes and integration tests pass"], + "failure_conditions": ["the same error remains after the indexing policy has propagated"], + "safety_constraints": [], + "source_kind": "episode_distillation", + "grounded_in": ["ep-7"], + "confidence": 0.8 + } + ] +} +``` + +user: +## Behavioral facts +{{behavioral_facts}} + +## Episodic lessons +{{episodic_lessons}} diff --git a/azure/cosmos/agent_memory/services/__init__.py b/azure/cosmos/agent_memory/services/__init__.py index cee665a..4dbf393 100644 --- a/azure/cosmos/agent_memory/services/__init__.py +++ b/azure/cosmos/agent_memory/services/__init__.py @@ -18,7 +18,7 @@ def query( def read_item(self, item_id: str, partition_key: Any) -> dict[str, Any]: ... - def add_cosmos(self, record: dict[str, Any]) -> dict[str, Any]: ... + def upsert_memory(self, record: dict[str, Any]) -> dict[str, Any]: ... def mark_superseded( self, diff --git a/azure/cosmos/agent_memory/services/pipeline.py b/azure/cosmos/agent_memory/services/pipeline.py index 2584fbc..f308d4e 100644 --- a/azure/cosmos/agent_memory/services/pipeline.py +++ b/azure/cosmos/agent_memory/services/pipeline.py @@ -11,6 +11,7 @@ import hashlib import json +import re import time from collections import defaultdict from datetime import datetime, timezone @@ -20,6 +21,7 @@ CosmosResourceExistsError, CosmosResourceNotFoundError, ) +from pydantic import ValidationError as PydanticValidationError from azure.cosmos.agent_memory import thresholds as threshold_config from azure.cosmos.agent_memory._container_routing import ContainerKey @@ -27,17 +29,14 @@ 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, ValidationError, ) from azure.cosmos.agent_memory.logging import get_logger from azure.cosmos.agent_memory.models import ( + TRUSTED_PROCEDURE_SOURCE_KINDS, EpisodicRecord, FactRecord, ProceduralRecord, @@ -72,9 +71,6 @@ from azure.cosmos.agent_memory.services._pipeline_helpers import ( is_real_number as _is_real_number, ) -from azure.cosmos.agent_memory.services._pipeline_helpers import ( - max_or_none as _max_or_none, -) from azure.cosmos.agent_memory.store._search_helpers import top_literal logger = get_logger("azure.cosmos.agent_memory.pipeline") @@ -155,7 +151,36 @@ def upsert_item(self, *, body: dict[str, Any]) -> dict[str, Any]: if upsert is not None: response = upsert(body=body) return response if isinstance(response, dict) else body - return self._store.add_cosmos(body) + return self._store.upsert_memory(body) + + @staticmethod + def _apply_patch_operations(doc: dict[str, Any], patch_operations: list[dict[str, Any]]) -> dict[str, Any]: + patched = dict(doc) + for operation in patch_operations: + if operation.get("op") != "set": + raise ValueError(f"unsupported patch operation: {operation.get('op')!r}") + path = operation.get("path") + if not isinstance(path, str) or not path.startswith("/") or path == "/": + raise ValueError(f"unsupported patch path: {path!r}") + keys = [part.replace("~1", "/").replace("~0", "~") for part in path[1:].split("/")] + target = patched + for key in keys[:-1]: + value = target.get(key) + if not isinstance(value, dict): + value = {} + target[key] = value + target = value + target[keys[-1]] = operation.get("value") + return patched + + def patch_item(self, *, item: str, partition_key: Any, patch_operations: list[dict[str, Any]]) -> dict[str, Any]: + container = self._target_container() + patch_item = getattr(container, "patch_item", None) + if callable(patch_item): + response = patch_item(item=item, partition_key=partition_key, patch_operations=patch_operations) + return response if isinstance(response, dict) else self.read_item(item=item, partition_key=partition_key) + doc = self.read_item(item=item, partition_key=partition_key) + return self.upsert_item(body=self._apply_patch_operations(doc, patch_operations)) def create_item(self, *, body: dict[str, Any]) -> dict[str, Any]: container = self._target_container() @@ -166,7 +191,7 @@ def create_item(self, *, body: dict[str, Any]) -> dict[str, Any]: if create is not None: response = create(body=body) return response if isinstance(response, dict) else body - return self._store.add_cosmos(body) + return self._store.upsert_memory(body) def replace_item(self, **kwargs: Any) -> Any: container = self._target_container() @@ -258,43 +283,6 @@ def _build_transcript( include_timestamp=include_timestamp, ) - def _load_existing_memories( - self, - user_id: str, - memory_types: list[str], - limit: int = 100, - ) -> list[dict[str, Any]]: - """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 - in implementation-defined order and the dedup comparison set is - non-deterministic. - """ - type_placeholders = ", ".join(f"@mtype{i}" for i in range(len(memory_types))) - capped_limit = top_literal(limit, name="_load_existing_memories.limit") - query = ( - f"SELECT TOP {capped_limit} * FROM c " - f"WHERE c.user_id = @user_id " - f"AND c.type IN ({type_placeholders}) " - f"AND {_ACTIVE_DOC_FILTER} " - f"ORDER BY c._ts DESC" - ) - parameters: list[dict[str, Any]] = [ - {"name": "@user_id", "value": user_id}, - ] - for i, mt in enumerate(memory_types): - parameters.append({"name": f"@mtype{i}", "value": mt}) - - items = list( - self._memories_container.query_items( - query=query, - parameters=parameters, - enable_cross_partition_query=True, - ) - ) - return items - def _vector_distance_function(self) -> str: """Return the container's configured Cosmos ``distanceFunction`` (cached). @@ -314,7 +302,7 @@ def _vector_distance_function(self) -> str: # "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). + # 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). @@ -332,9 +320,9 @@ def _vector_distance_function(self) -> str: 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. + The near-exact threshold is cosine-calibrated; on euclidean + the destructive auto-drop is skipped and LLM reconcile still runs. + Logged once per pipeline instance to avoid hot-path spam. """ if getattr(self, "_warned_euclidean_autodrop", False): return @@ -353,9 +341,9 @@ def _warn_distance_policy_unavailable_once(self) -> None: 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." + "vector dedup: container vector policy could not be read; skipping " + "near-exact auto-drop this run to avoid mis-calibrated drops. Memories are " + "written as-is and reconciled on a later run once the policy is readable." ) def _vector_candidates( @@ -497,7 +485,6 @@ def _empty_extract_counts() -> dict[str, int]: "contradicted_count": 0, "exact_dedup_skipped": 0, "dropped_episodic_count": 0, - "inplace_updated": 0, "deferred_turn_count": 0, "quarantined_turn_count": 0, } @@ -607,18 +594,21 @@ def extract_memories_durable( logger.warning("extract_memories_durable no memories found user_id=%s thread_id=%s", user_id, thread_id) return {"facts": [], "episodic": [], "updates": [], "processed_turn_docs": []} - existing_for_hashes = self._load_existing_memories(user_id, ["fact"]) - existing_fact_hashes: set[str] = { - m["content_hash"] for m in existing_for_hashes if m.get("type") == "fact" and m.get("content_hash") - } + # Exact-duplicate detection is in-batch only: a content_hash seen earlier + # in THIS extraction is skipped. Cross-turn / cross-run exact duplicates + # are handled at write time by the deterministic-id create (a repeat of + # the same fact collides on id and is skipped with a 409), so there is no + # per-extract query to preload the user's existing fact hashes. + existing_fact_hashes: set[str] = set() # 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. + # ``processed_turns``; the in-process caller marks them ``extracted_at`` + # so they are never re-processed (the Durable backend instead advances a + # count-based watermark). Turns from batches that fail with a *retryable* + # error are left out and retried on the next run. batches = batch_turns_by_tokens(items, threshold_config.get_extraction_batch_max_tokens()) facts: list[dict[str, Any]] = [] processed_turns: list[dict[str, Any]] = [] @@ -912,20 +902,108 @@ def _build_episode_docs( def _deterministic_episode_id(segment_key: str, index: int) -> str: return deterministic_episode_id(segment_key, index) + @staticmethod + def _episode_cursor_id(user_id: str, thread_id: str) -> str: + return f"episode_cursor_{user_id}_{thread_id}" + + def _read_episode_cursor(self, user_id: str, thread_id: str) -> tuple[str, str]: + """Return the ``(created_at, id)`` watermark of the last turn folded into + an episode for this thread, or ``("", "")`` when none has been extracted. + + The cursor is a single per-thread doc in the MEMORIES container - not a + per-turn stamp - so advancing it never writes to the turns container and + therefore never re-enters the turns change feed. That is load-bearing: + the Durable cadence counter must observe each turn exactly once (at + creation); a turn mutated by episodic segmentation would otherwise be + redelivered and mis-counted. ``created_at`` is stored UTC-normalized, so + the lexical ``>`` comparison below matches chronological order. + """ + try: + doc = self._memories_container.read_item( + item=self._episode_cursor_id(user_id, thread_id), + partition_key=[user_id, thread_id], + ) + except CosmosResourceNotFoundError: + return "", "" + return str(doc.get("last_episode_at") or ""), str(doc.get("last_episode_id") or "") + + def _advance_episode_cursor(self, user_id: str, thread_id: str, last_turn: dict[str, Any]) -> None: + """Advance the episodic watermark to ``last_turn`` (newest turn just + folded), never backwards. A single-doc write: unlike the former per-turn + stamp it cannot partially fail and leave the open segment torn, and it + writes to the memories container, never the change-feed-monitored turns + container. The advance is atomic - ETag ``IfNotModified`` with a re-read + retry - so a late, out-of-order concurrent ``episode:{u}:{t}:{count}`` + orchestration cannot regress the cursor (a regressed cursor would, with + topic drift enabled, re-segment turns under a new segment key and + duplicate episodes). + """ + from azure.core import MatchConditions + from azure.cosmos.exceptions import CosmosAccessConditionFailedError + + last_at = str(last_turn.get("created_at") or "") + last_id = str(last_turn.get("id") or "") + if not last_at: + return + cursor_id = self._episode_cursor_id(user_id, thread_id) + partition_key = [user_id, thread_id] + for _ in range(3): + etag: Optional[str] = None + try: + existing = self._memories_container.read_item(item=cursor_id, partition_key=partition_key) + except CosmosResourceNotFoundError: + existing = None + if existing is not None: + current = ( + str(existing.get("last_episode_at") or ""), + str(existing.get("last_episode_id") or ""), + ) + if (last_at, last_id) <= current: + return # monotonic: never regress + etag = existing.get("_etag") + body = { + "id": cursor_id, + "type": "episode_cursor", + "user_id": user_id, + "thread_id": thread_id, + "last_episode_at": last_at, + "last_episode_id": last_id, + "updated_at": datetime.now(tz=timezone.utc).isoformat(), + } + try: + if existing is None: + self._memories_container.create_item(body=body) + else: + self._memories_container.replace_item( + item=cursor_id, + body=body, + etag=etag, + match_condition=MatchConditions.IfNotModified, + ) + return + except (CosmosResourceExistsError, CosmosAccessConditionFailedError): + continue # a concurrent run advanced first; re-read and re-check + logger.debug("episode cursor advance retries exhausted user_id=%s thread_id=%s", user_id, thread_id) + def _load_open_episode_segment(self, user_id: str, thread_id: str) -> list[dict[str, Any]]: - """Return the open episode segment: turns not yet folded into an episode, - oldest first. + """Return the open episode segment: turns created after the episodic + watermark (not yet folded into an episode), oldest first. - Mirrors the fact ``extracted_at`` watermark with an independent - ``episode_extracted_at`` cursor, so episodic segmentation neither blocks - nor is blocked by fact extraction. + Uses a per-thread ``(created_at, id)`` cursor (see + ``_read_episode_cursor``) instead of a per-turn stamp, so episodic + segmentation writes nothing to the turns container and cannot perturb + the change-feed cadence counter. """ + last_at, last_id = self._read_episode_cursor(user_id, thread_id) query = ( "SELECT * FROM c WHERE c.user_id = @user_id " "AND c.thread_id = @thread_id AND c.type = 'turn' " - "AND (NOT IS_DEFINED(c.episode_extracted_at) OR IS_NULL(c.episode_extracted_at))" + "AND (c.created_at > @last_at " + "OR (c.created_at = @last_at AND c.id > @last_id))" ) parameters: list[dict[str, Any]] = [ + {"name": "@last_at", "value": last_at}, + {"name": "@last_id", "value": last_id}, {"name": "@user_id", "value": user_id}, {"name": "@thread_id", "value": thread_id}, ] @@ -993,23 +1071,24 @@ def extract_episodes( ) -> dict[str, int]: """Segment the open turn stream into episodes at detected boundaries. - The open segment is every turn not yet folded into an episode (no - ``episode_extracted_at`` stamp). At each detected boundary - an idle - time-gap, a topic-drift shift, or the max-size cap - the closed segment is - extracted into one or more immutable episodes, embedded, and persisted, - then its turns are stamped so re-evaluation never re-extracts them. - ``flush=True`` also drains the trailing open segment even without a + The open segment is every turn created after the episodic watermark (see + ``_read_episode_cursor``). At each detected boundary - an idle time-gap, a + topic-drift shift, or the max-size cap - the closed segment is extracted + into one or more immutable episodes, embedded, and persisted, then the + watermark advances past those turns so re-evaluation never re-extracts + them. ``flush=True`` also drains the trailing open segment even without a boundary (end of conversation / explicit close); the default leaves the current, possibly-incomplete segment open. The caller never signals "session end" - boundaries are inferred from the stream itself. - Idempotency (best-effort, not absolute): each episode's id is - deterministic in its segment key and ordinal - not the LLM summary text - - so re-running the same still-open segment collides on id and the - duplicate write is skipped (409). This holds while the segment's turn set - is stable; a partial watermark-stamp failure that shifts the open - segment's boundaries can still admit a duplicate, which episodic - reconciliation does not currently fold. + Idempotency: each episode's id is deterministic in its segment key and + ordinal - not the LLM summary text - so re-running the same still-open + segment collides on id and the duplicate write is skipped (409). The + watermark is a single doc advanced only after a segment's episodes are + created, so a crash between the create and the advance simply re-loads + the same open segment next run (same turn set -> same ids -> 409); unlike + the former per-turn stamp there is no partial-stamp state that could shift + the boundary and admit a duplicate. """ if not user_id: raise ValidationError("user_id is required") @@ -1050,18 +1129,19 @@ def extract_episodes( ) break # Non-retryable (e.g. content filter, context-length): quarantine the - # poison segment - stamp it so it never re-poisons future runs and the - # open segment cannot grow without bound - then advance to the next. + # poison segment - advance the watermark past it so it never + # re-poisons future runs and the open segment cannot grow without + # bound - then move to the next segment. logger.warning( "extract_episodes: quarantining %d turns after non-retryable extraction error " - "(marking episode_extracted_at so they do not re-poison future runs) " + "(advancing the episode watermark past them so they do not re-poison future runs) " "user_id=%s thread_id=%s err=%s", len(closing), user_id, thread_id, exc, ) - self._mark_turns_extracted(closing, field="episode_extracted_at") + self._advance_episode_cursor(user_id, thread_id, closing[-1]) segment = segment[boundary:] continue for doc, embedding in zip(docs, embeddings_for_docs): @@ -1074,256 +1154,10 @@ def extract_episodes( # Advance the episode watermark so the closed turns leave the open # segment even when the segment yielded no episode (nothing episodic # to remember) - otherwise they would be re-evaluated forever. - self._mark_turns_extracted(closing, field="episode_extracted_at") + self._advance_episode_cursor(user_id, thread_id, closing[-1]) segment = segment[boundary:] return {"episodes": total} - 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. - - Folds only happen within the same ``metadata.source`` (user vs agent): an - agent action and a user statement are distinct records even when their - embeddings are near-identical, so folding across sources would corrupt - attribution (tag/source desync, content swap). Cross-source pairs return - False and the caller keeps the new doc as a novel ADD. - """ - from azure.core import MatchConditions - from azure.cosmos.exceptions import CosmosAccessConditionFailedError - - neighbor_source = (neighbor.get("metadata") or {}).get("source") or "user" - new_source = (new_doc.get("metadata") or {}).get("source") or "user" - if neighbor_source != new_source: - logger.info( - "in-place dedup update skipped (source mismatch neighbor=%s new=%s) " - "target_id=%s; keeping new doc as novel", - neighbor_source, - new_source, - neighbor.get("id"), - ) - return False - - 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, @@ -1370,7 +1204,7 @@ def persist_extracted_memories( if op.get("op") == "stats": result["exact_dedup_skipped"] += int(op.get("exact_dedup_skipped") or 0) result["dropped_episodic_count"] += int(op.get("dropped_episodic_count") or 0) - for key in ("inplace_updated", "deferred_turn_count", "quarantined_turn_count"): + for key in ("deferred_turn_count", "quarantined_turn_count"): if key in op: result[key] = result.get(key, 0) + int(op.get(key) or 0) @@ -1378,18 +1212,17 @@ def persist_extracted_memories( return result - def _mark_turns_extracted(self, turn_docs: list[dict[str, Any]], *, field: str = "extracted_at") -> int: - """Stamp a processed-watermark field on each turn doc and upsert. + def _mark_turns_extracted(self, turn_docs: list[dict[str, Any]]) -> int: + """Stamp the fact-extraction ``extracted_at`` watermark on each turn doc. - ``field`` selects the independent watermark: ``extracted_at`` for fact - extraction, ``episode_extracted_at`` for episodic segmentation. The two - cursors are independent so a turn can be fact-extracted but still belong - to the open episode segment (or vice versa). + Used by the in-process fact path so a turn is not re-extracted. (Episodic + segmentation uses a separate per-thread cursor doc - see + ``_read_episode_cursor`` - and never stamps turns, so it does not perturb + the change-feed cadence counter.) - 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-processed on the - next call, which is bounded and recoverable. + Per-turn failures are logged but do not raise - the worst case is one + turn gets re-processed on the next call, which is bounded and + recoverable. """ if not turn_docs: return 0 @@ -1400,14 +1233,15 @@ def _mark_turns_extracted(self, turn_docs: list[dict[str, Any]], *, field: str = if not turn_id: continue try: - doc_to_write = dict(turn) - doc_to_write[field] = now_iso - self._turns_container.upsert_item(body=doc_to_write) + self._turns_container.patch_item( + item=turn_id, + partition_key=[turn.get("user_id"), turn.get("thread_id")], + patch_operations=[{"op": "set", "path": "/extracted_at", "value": now_iso}], + ) marked += 1 except Exception as exc: logger.warning( - "_mark_turns_extracted(%s) failed for turn_id=%s err=%s (turn may be re-processed on next call)", - field, + "_mark_turns_extracted failed for turn_id=%s err=%s (turn may be re-processed on next call)", turn_id, exc, ) @@ -1425,11 +1259,8 @@ def extract_memories( extracted = self.extract_memories_durable(user_id, thread_id, recent_k, turns=turns) # Capture the processed turns from the compute stage 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. + # Persist uses this exact list for stamping after all creates finish. 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) @@ -1457,42 +1288,12 @@ def synthesize_procedural( *, force: bool = False, ) -> dict[str, Any]: - """Synthesize the active procedural prompt for a user.""" + """Extract atomic procedural memories from behavioral facts and lessons.""" + del force if not user_id: raise ValidationError("user_id is required") - logger.info("synthesize_procedural started user_id=%s force=%s", user_id, force) - - def _read_latest_procedural() -> Optional[dict[str, Any]]: - docs = list( - self._memories_container.query_items( - query=( - "SELECT * FROM c WHERE c.user_id = @uid " - "AND c.thread_id = @thread_id " - "AND c.type = @type " - f"AND {_ACTIVE_DOC_FILTER}" - ), - parameters=[ - {"name": "@uid", "value": user_id}, - {"name": "@thread_id", "value": "__procedural__"}, - {"name": "@type", "value": "procedural"}, - ], - enable_cross_partition_query=True, - ) - ) - docs.sort( - key=lambda doc: (int(doc.get("version") or 0), int(doc.get("_ts") or 0)), - reverse=True, - ) - if len(docs) > 1: - logger.warning( - "synthesize_procedural found multiple active docs user_id=%s count=%d", - user_id, - len(docs), - ) - return docs[0] if docs else None - - prior_doc = _read_latest_procedural() + logger.info("synthesize_procedural extraction started user_id=%s", user_id) behavioral_fact_docs = list( self._memories_container.query_items( @@ -1503,7 +1304,7 @@ def _read_latest_procedural() -> Optional[dict[str, Any]]: "AND ((IS_DEFINED(c.metadata.category) " "AND c.metadata.category IN ('preference', 'requirement')) " "OR (IS_DEFINED(c.salience) AND c.salience >= @min_salience)) " - "ORDER BY c.salience DESC, c.created_at ASC, c.id ASC" + "ORDER BY c.created_at ASC" ), parameters=[ {"name": "@uid", "value": user_id}, @@ -1518,7 +1319,6 @@ def _read_latest_procedural() -> Optional[dict[str, Any]]: for doc in behavioral_fact_docs if isinstance(doc.get("content"), str) and doc.get("content", "").strip() ] - behavioral_fact_ids = [doc["id"] for doc in behavioral_fact_docs] episodic_docs = list( self._memories_container.query_items( @@ -1526,9 +1326,8 @@ def _read_latest_procedural() -> Optional[dict[str, Any]]: "SELECT TOP 50 * FROM c WHERE c.user_id = @uid " "AND c.type = @type " f"AND {_ACTIVE_DOC_FILTER} " - "AND IS_DEFINED(c.lessons) " - "AND ARRAY_LENGTH(c.lessons) > 0 " - "ORDER BY c.salience DESC, c.created_at ASC, c.id ASC" + "AND IS_DEFINED(c.lessons) AND ARRAY_LENGTH(c.lessons) > 0 " + "ORDER BY c.created_at ASC" ), parameters=[ {"name": "@uid", "value": user_id}, @@ -1537,141 +1336,187 @@ def _read_latest_procedural() -> Optional[dict[str, Any]]: enable_cross_partition_query=True, ) ) - episodic_with_lessons = [ - doc - for doc in episodic_docs - if isinstance(doc.get("lessons"), list) - and any(isinstance(lesson, str) and lesson.strip() for lesson in doc.get("lessons", [])) - ] - source_episodic_ids = [doc["id"] for doc in episodic_with_lessons] - current_source_ids = set(behavioral_fact_ids) | set(source_episodic_ids) - - def _covered_by(prior: Optional[dict[str, Any]]) -> bool: - if prior is None: - return False - covered = set(prior.get("source_fact_ids") or []) | set(prior.get("source_episodic_ids") or []) - return current_source_ids.issubset(covered) - - if prior_doc and not force and _covered_by(prior_doc): - logger.info( - "synthesize_procedural unchanged user_id=%s fact_count=%d episodic_count=%d", - user_id, - len(behavioral_fact_ids), - len(source_episodic_ids), - ) - return {"status": "unchanged", "procedural": prior_doc} + def _episodic_lessons(doc: dict[str, Any]) -> list[str]: + lessons = doc.get("lessons") + if isinstance(lessons, list): + return [lesson.strip() for lesson in lessons if isinstance(lesson, str) and lesson.strip()] + return [] - if not current_source_ids: - logger.info( - "synthesize_procedural skipping LLM user_id=%s - no behavioral facts or episodic lessons", - user_id, - ) - return {"status": "unchanged", "procedural": prior_doc} - - user_name = "the user" - - def _render_bullets(values: list[str]) -> str: - cleaned = [value.strip() for value in values if isinstance(value, str) and value.strip()] - if not cleaned: - return "(none)" - return "\n".join(f"- {value}" for value in cleaned) - - static_prompty_inputs = { - "behavioral_facts": _render_bullets([doc.get("content", "") for doc in behavioral_fact_docs]), - "episodic_lessons": _render_bullets( - [ - lesson - for doc in episodic_with_lessons - for lesson in doc.get("lessons", []) - if isinstance(lesson, str) and lesson.strip() - ] - ), - "user_name": user_name, - } + fact_lines: list[str] = [] + fact_label_to_id: dict[str, str] = {} + for index, doc in enumerate(behavioral_fact_docs, start=1): + label = f"fact-{index}" + category = "" + metadata = doc.get("metadata") + if isinstance(metadata, dict) and isinstance(metadata.get("category"), str): + category = metadata["category"].strip() + fact_lines.append(f"{label} [{category or 'unknown'}]: {doc['content'].strip()}") + if isinstance(doc.get("id"), str): + fact_label_to_id[label] = doc["id"] + + episodic_lines: list[str] = [] + episodic_label_to_id: dict[str, str] = {} + ep_index = 0 + for doc in episodic_docs: + doc_id = doc.get("id") + for lesson in _episodic_lessons(doc): + ep_index += 1 + label = f"ep-{ep_index}" + episodic_lines.append(f"{label}: {lesson}") + if isinstance(doc_id, str): + episodic_label_to_id[label] = doc_id + + if not fact_lines and not episodic_lines: + logger.info("synthesize_procedural unchanged user_id=%s - no procedure sources", user_id) + return {"status": "unchanged", "procedures_created": 0} - # 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 - # source coverage, not just version number. - written_doc: Optional[dict[str, Any]] = None - for attempt in range(1, _PROCEDURAL_MAX_CREATE_ATTEMPTS + 1): + try: response_text = self._run_prompty( - "synthesize_procedural.prompty", + "extract_procedure.prompty", inputs={ - "prior_prompt": (prior_doc.get("content") or "") if prior_doc else "", - **static_prompty_inputs, + "behavioral_facts": "\n".join(fact_lines), + "episodic_lessons": "\n".join(episodic_lines), }, ) - parsed = self._parse_llm_json(response_text) - system_prompt = parsed.get("system_prompt") if isinstance(parsed, dict) else None - if not isinstance(system_prompt, str) or not system_prompt.strip(): - raise LLMError("synthesize_procedural returned JSON without a non-empty 'system_prompt' string") - system_prompt = system_prompt.strip() - - new_seq = (int(prior_doc.get("version") or 0) + 1) if prior_doc else 1 - new_doc: dict[str, Any] = { - "id": f"proc_{user_id}_{new_seq}", - "user_id": user_id, - "thread_id": "__procedural__", - "type": "procedural", - "version": new_seq, - "content": system_prompt, - "source_fact_ids": behavioral_fact_ids, - "source_episodic_ids": source_episodic_ids, - "supersedes_ids": [prior_doc["id"]] if prior_doc else [], - "created_at": datetime.now(timezone.utc).isoformat(), - "updated_at": datetime.now(timezone.utc).isoformat(), - "role": "system", - "tags": ["sys:procedural", "sys:synthesized"], - **self._prompt_lineage("synthesize_procedural.prompty"), - "metadata": {}, - } - validated = construct_internal(ProceduralRecord, new_doc).to_doc() + procedures = parsed.get("procedures", []) if isinstance(parsed, dict) else [] + if not isinstance(procedures, list): + procedures = [] + except Exception as exc: # LLM/parsing quarantine: one bad call must not stop the pipeline. + if is_retryable_llm_error(exc): + logger.warning("synthesize_procedural deferred user_id=%s: %s", user_id, exc) + return {"status": "deferred", "procedures_created": 0} + logger.exception("synthesize_procedural skipped user_id=%s after non-retryable LLM error", user_id) + return {"status": "skipped", "procedures_created": 0} + + trusted_source_kinds = {kind.value for kind in TRUSTED_PROCEDURE_SOURCE_KINDS} + authority_by_source = { + "explicit_user_instruction": "high", + "organization_policy": "high", + "observed_user_preference": "medium", + "episode_distillation": "medium", + "document_content": "low", + "agent_inference": "low", + } + now = datetime.now(timezone.utc).isoformat() + created = 0 + skipped = 0 + + for proc in procedures: try: - self._memories_container.create_item(body=validated) - written_doc = validated - break - except CosmosResourceExistsError: - logger.info( - "synthesize_procedural id collision user_id=%s seq=%d attempt=%d/%d - re-reading", - user_id, - new_seq, - attempt, - _PROCEDURAL_MAX_CREATE_ATTEMPTS, - ) - latest = _read_latest_procedural() - if latest is None: + if not isinstance(proc, dict): + skipped += 1 continue - prior_doc = latest - if _covered_by(prior_doc): - logger.info( - "synthesize_procedural race resolved by coverage user_id=%s winner=%s", - user_id, - prior_doc["id"], - ) - return {"status": "unchanged", "procedural": prior_doc} - if written_doc is None: - raise MemoryConflictError( - "synthesize_procedural failed after " - f"{_PROCEDURAL_MAX_CREATE_ATTEMPTS} attempts due to id collisions " - f"user_id={user_id!r}" - ) + name = proc.get("name") + if not isinstance(name, str) or not name.strip(): + skipped += 1 + continue + name = name.strip() - new_id = written_doc["id"] - if prior_doc: - self._mark_superseded(prior_doc, new_id, reason="update") + grounded_in = proc.get("grounded_in") + if isinstance(grounded_in, str): + labels = [grounded_in] + elif isinstance(grounded_in, list): + labels = [label for label in grounded_in if isinstance(label, str)] + else: + labels = [] + source_fact_ids = sorted({fact_label_to_id[label] for label in labels if label in fact_label_to_id}) + source_episodic_ids = sorted( + {episodic_label_to_id[label] for label in labels if label in episodic_label_to_id} + ) + + source_kind = proc.get("source_kind", "agent_inference") + if not isinstance(source_kind, str): + source_kind = "agent_inference" + # Episode-only grounding cannot corroborate a user/org instruction: + # any trusted label backed solely by episodes (no behavioral fact) + # downgrades to episode_distillation - a candidate, never an + # auto-active policy. + if source_episodic_ids and not source_fact_ids and source_kind in trusted_source_kinds: + source_kind = "episode_distillation" + source_authority = authority_by_source.get(source_kind, "low") + status = "active" if source_kind in trusted_source_kinds else "candidate" + # Grounding is the trust anchor: a procedure whose ``grounded_in`` + # resolved to no persisted fact or episodic source is never + # auto-activated, regardless of the LLM's self-declared + # source_kind - this blocks an ungrounded self-labeled instruction + # from being compiled into the runtime system prompt. + if not source_fact_ids and not source_episodic_ids: + status = "candidate" + + summary = proc.get("summary") if isinstance(proc.get("summary"), str) else "" + retrieval_text = proc.get("retrieval_text") if isinstance(proc.get("retrieval_text"), str) else "" + scope_type = proc.get("scope_type") if isinstance(proc.get("scope_type"), str) else "user" + scope_value = proc.get("scope_value") if isinstance(proc.get("scope_value"), str) else None + proc_id = ( + "proc_" + + hashlib.sha256( + f"{user_id}|{scope_type}|{scope_value or ''}|{name.strip().lower()}".encode() + ).hexdigest()[:32] + ) + doc: dict[str, Any] = { + "id": proc_id, + "user_id": user_id, + "thread_id": "__procedural__", + "type": "procedural", + "role": "system", + "tags": ["sys:procedural", "sys:auto-extracted"], + "created_at": now, + "updated_at": now, + "name": name, + "summary": summary.strip() or name, + "retrieval_text": retrieval_text.strip() or summary.strip() or name, + "procedure_kind": proc.get("procedure_kind", "behavioral_policy"), + "scope_type": scope_type, + "scope_value": scope_value, + "activation_conditions": proc.get("activation_conditions", []), + "preconditions": proc.get("preconditions", []), + "steps": proc.get("steps", []), + "success_conditions": proc.get("success_conditions", []), + "failure_conditions": proc.get("failure_conditions", []), + "safety_constraints": proc.get("safety_constraints", []), + "status": status, + "priority": proc.get("priority", 0), + # Seed utility from the LLM's extraction confidence (the + # extract_procedure schema emits ``confidence``, not + # ``utility_score``). Procedures are create-only today with no + # promotion or outcome-scoring path, so this is the record's + # final utility value. + "utility_score": clamp_unit_interval(proc.get("confidence"), 0.5), + "successful_uses": proc.get("successful_uses", 0), + "failed_uses": proc.get("failed_uses", 0), + "source_kind": source_kind, + "source_authority": source_authority, + "source_fact_ids": source_fact_ids, + "source_episodic_ids": source_episodic_ids, + "source_turn_ids": proc.get("source_turn_ids", []), + "content": summary.strip() or name, + "version": proc.get("version", 1), + "metadata": {}, + **self._prompt_lineage("extract_procedure.prompty"), + } + validated = construct_internal(ProceduralRecord, doc).to_doc() + validated["embedding"] = self._embed_one(validated["retrieval_text"]) + try: + self._create_memory(validated) + created += 1 + except CosmosResourceExistsError: + skipped += 1 + except (ValidationError, PydanticValidationError, ValueError) as exc: + skipped += 1 + logger.warning("synthesize_procedural dropping malformed procedure user_id=%s: %s", user_id, exc) + except Exception: + skipped += 1 + logger.exception("synthesize_procedural failed to persist one procedure user_id=%s", user_id) logger.info( - "synthesize_procedural synthesized user_id=%s version=%d fact_count=%d episodic_count=%d", + "synthesize_procedural extracted user_id=%s procedures_created=%d procedures_skipped=%d", user_id, - written_doc["version"], - len(behavioral_fact_ids), - len(source_episodic_ids), + created, + skipped, ) - return {"status": "synthesized", "procedural": written_doc} + return {"status": "synthesized", "procedures_created": created, "procedures_skipped": skipped} def generate_thread_summary_durable( self, @@ -1999,13 +1844,10 @@ def reconcile_memories(self, user_id: str, n: int = 50, *, memory_type: str = "f Each loser is soft-deleted with ``supersede_reason="contradict"`` and ``superseded_by`` set to the winner. - 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, + Near-duplicate *paraphrases* are not merged here: 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). + because they have no contradiction semantics. Returns ``{"kept", "merged", "contradicted"}``; ``merged`` is always 0. """ @@ -2139,31 +1981,143 @@ def _active_memories_for_reconcile(self, user_id: str, memory_type: str, n: int) ) ) - def build_procedural_context(self, user_id: str) -> str: - """Return the active synthesized procedural prompt for system injection.""" + def build_procedural_context(self, user_id: str, task: Optional[str] = None) -> str: + """Build a deterministic system prompt projection from active procedures.""" if not user_id: raise ValidationError("user_id is required") query = ( - "SELECT TOP 1 c.content, c.version FROM c WHERE c.user_id = @user_id " - "AND c.thread_id = @thread_id AND c.type = @type " - f"AND {_ACTIVE_DOC_FILTER} " - "ORDER BY c.version DESC" + "SELECT * FROM c WHERE c.user_id=@uid AND c.type='procedural' " + "AND c.status='active' " + "AND (NOT IS_DEFINED(c.superseded_by) OR IS_NULL(c.superseded_by))" ) - items = list( + procedures = list( self._memories_container.query_items( query=query, - parameters=[ - {"name": "@user_id", "value": user_id}, - {"name": "@thread_id", "value": "__procedural__"}, - {"name": "@type", "value": "procedural"}, - ], + parameters=[{"name": "@uid", "value": user_id}], enable_cross_partition_query=True, ) ) - if not items: + procedures = [ + proc + for proc in procedures + if proc.get("user_id") == user_id + and proc.get("type") == "procedural" + and proc.get("status") == "active" + and not proc.get("superseded_by") + ] + + stopwords = { + "a", + "an", + "and", + "are", + "as", + "at", + "be", + "by", + "for", + "from", + "in", + "is", + "it", + "of", + "on", + "or", + "the", + "to", + "with", + } + + def _tokens(value: str) -> set[str]: + return {token for token in re.findall(r"[a-z0-9]+", value.lower()) if token not in stopwords} + + task_tokens = _tokens(task) if isinstance(task, str) and task.strip() else set() + policies: list[dict[str, Any]] = [] + task_procedures: list[dict[str, Any]] = [] + for proc in procedures: + kind = proc.get("procedure_kind") + scope_type = proc.get("scope_type") + if kind in {"behavioral_policy", "decision_rule"} and scope_type in {"global", "user"}: + policies.append(proc) + continue + if task_tokens and kind in {"workflow", "recovery_strategy", "tool_usage"}: + searchable = " ".join( + [ + proc.get("retrieval_text") if isinstance(proc.get("retrieval_text"), str) else "", + *[ + condition + for condition in proc.get("activation_conditions", []) + if isinstance(condition, str) + ], + ] + ) + if task_tokens & _tokens(searchable): + task_procedures.append(proc) + + included = policies + task_procedures + if not included: return "" - content = items[0].get("content") - return content if isinstance(content, str) else "" + authority_rank = {"mandatory": 3, "high": 2, "medium": 1, "low": 0} + + def _sort_key(proc: dict[str, Any]) -> tuple[int, int, str]: + try: + priority = int(proc.get("priority") or 0) + except (TypeError, ValueError): + priority = 0 + authority = proc.get("source_authority") + rank = authority_rank.get(authority if isinstance(authority, str) else "low", 0) + name = proc.get("name") if isinstance(proc.get("name"), str) else "" + return (-priority, -rank, name.lower()) + + policies.sort(key=_sort_key) + task_procedures.sort(key=_sort_key) + included = policies + task_procedures + + fingerprint_payload = sorted( + ( + str(proc.get("id", "")), + str(proc.get("version", "")), + str(proc.get("status", "")), + str(proc.get("priority", "")), + str(proc.get("scope_type", "")), + str(proc.get("scope_value", "")), + ) + for proc in included + ) + fingerprint = hashlib.sha256(json.dumps(fingerprint_payload, separators=(",", ":")).encode()).hexdigest() + logger.debug( + "build_procedural_context fingerprint=%s included_ids=%s user_id=%s", + fingerprint, + [proc.get("id") for proc in included], + user_id, + ) -__all__ = ["PipelineService"] + lines = ["# Learned procedures"] + if policies: + lines.append("") + lines.append("## Behavioral policies") + for proc in policies: + name = proc.get("name") if isinstance(proc.get("name"), str) else "Unnamed procedure" + summary = proc.get("summary") if isinstance(proc.get("summary"), str) else proc.get("content", "") + lines.append(f"- {name}: {summary}") + if task_procedures: + lines.append("") + lines.append("## Task procedures") + for proc in task_procedures: + name = proc.get("name") if isinstance(proc.get("name"), str) else "Unnamed procedure" + summary = proc.get("summary") if isinstance(proc.get("summary"), str) else proc.get("content", "") + lines.append(f"### {name}") + if summary: + lines.append(f"Summary: {summary}") + steps = proc.get("steps") if isinstance(proc.get("steps"), list) else [] + if steps: + sorted_steps = sorted( + [step for step in steps if isinstance(step, dict)], + key=lambda step: int(step.get("sequence") or 0), + ) + for index, step in enumerate(sorted_steps, start=1): + instruction = step.get("instruction") if isinstance(step.get("instruction"), str) else "" + if instruction: + lines.append(f"{index}. {instruction}") + return "\n".join(lines) diff --git a/azure/cosmos/agent_memory/store/memory_store.py b/azure/cosmos/agent_memory/store/memory_store.py index f54833e..c4ed304 100644 --- a/azure/cosmos/agent_memory/store/memory_store.py +++ b/azure/cosmos/agent_memory/store/memory_store.py @@ -171,23 +171,23 @@ def _query_items( except Exception as exc: raise CosmosOperationError(f"{operation} failed: {exc}") from exc - def add_cosmos(self, record: dict[str, Any]) -> dict[str, Any]: + def upsert_memory(self, record: dict[str, Any]) -> dict[str, Any]: """Upsert a pre-built Cosmos memory document and return the stored body.""" body = self._prepare_doc(record) memory_type = body.get("type") if memory_type not in _CONTAINER_FOR_TYPE: raise ValueError( - f"add_cosmos: record id={body.get('id')!r} has invalid type={memory_type!r}. " - f"Set 'type' to one of {sorted(_CONTAINER_FOR_TYPE)} before calling add_cosmos." + f"upsert_memory: record id={body.get('id')!r} has invalid type={memory_type!r}. " + f"Set 'type' to one of {sorted(_CONTAINER_FOR_TYPE)} before calling upsert_memory." ) container = self._container_for_type(memory_type) try: response = container.upsert_item(body=body) except Exception as exc: raise _wrap_cosmos_exception( - exc, message=f"add_cosmos upsert failed for record {body.get('id')}: {exc}" + exc, message=f"upsert_memory upsert failed for record {body.get('id')}: {exc}" ) from exc - logger.info("add_cosmos id=%s role=%s type=%s", body.get("id"), body.get("role"), body.get("type")) + logger.info("upsert_memory id=%s role=%s type=%s", body.get("id"), body.get("role"), body.get("type")) return response if isinstance(response, dict) else body def add( @@ -251,7 +251,7 @@ def add( body["embedding"] = self._embeddings_client.generate(content) except Exception as exc: # noqa: BLE001 logger.warning( - "add_cosmos: embedding generation failed for %s (%s); proceeding without embedding", + "upsert_memory: embedding generation failed for %s (%s); proceeding without embedding", record.id, exc, ) @@ -262,7 +262,7 @@ def add( container.upsert_item(body=body) except Exception as exc: raise _wrap_cosmos_exception(exc, message=f"Upsert failed for record {record.id}: {exc}") from exc - logger.info("add_cosmos id=%s role=%s type=%s", record.id, role, memory_type) + logger.info("upsert_memory id=%s role=%s type=%s", record.id, role, memory_type) return record.id def push(self, local_memory: list[dict[str, Any]], batch_size: int = 25) -> None: @@ -1189,6 +1189,55 @@ def build_episodic_context(self, user_id: str, query: str, top_k: int = 3) -> st memories = self.search_episodic(user_id, query, top_k=top_k) return format_episodic_context(memories) + def retrieve_procedures( + self, + user_id: str, + search_terms: str, + top_k: int = 5, + *, + scope_type: Optional[str] = None, + scope_value: Optional[str] = None, + procedure_kind: Optional[str] = None, + status: Optional[str] = "active", + include_superseded: bool = False, + ) -> list[dict[str, Any]]: + """Semantic search across procedural memories for a user.""" + if not user_id: + raise ValidationError("user_id is required for retrieve_procedures") + terms = require_search_terms(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.type", "@type", "procedural") + qb.add_filter("c.user_id", "@user_id", user_id) + qb.add_filter("c.scope_type", "@scope_type", scope_type) + qb.add_filter("c.scope_value", "@scope_value", scope_value) + qb.add_filter("c.procedure_kind", "@procedure_kind", procedure_kind) + qb.add_filter("c.status", "@status", status) + + 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}) + for i, kw in enumerate(keywords): + parameters.append({"name": f"@kw{i}", "value": kw}) + + partition_key, cross_partition = query_scope(user_id, None) + logger.debug("MemoryStore.retrieve_procedures query: %s", sql) + return self.query( + sql, + parameters, + container_key=ContainerKey.MEMORIES, + partition_key=partition_key, + cross_partition=cross_partition, + ) + def _embed(self, text: str) -> list[float]: if self._embeddings_client is None: raise ConfigurationError( diff --git a/azure/cosmos/agent_memory/thresholds.py b/azure/cosmos/agent_memory/thresholds.py index 1f29056..8b58ceb 100644 --- a/azure/cosmos/agent_memory/thresholds.py +++ b/azure/cosmos/agent_memory/thresholds.py @@ -5,9 +5,6 @@ Operators override via the documented env vars; both backends read the same keys, so a single setting flips both. -Exception: the ``EPISODE_*`` knobs (boundary segmentation cadence and tuning) -are in-process only - the Durable Functions backend has no episodic path yet, -so there is no ``function_app/shared/config.py`` mirror for them. """ from __future__ import annotations @@ -20,7 +17,7 @@ logger = get_logger(__name__) -DEFAULT_FACT_EXTRACTION_EVERY_N = 1 +DEFAULT_FACT_EXTRACTION_EVERY_N = 2 DEFAULT_THREAD_SUMMARY_EVERY_N = 10 # Episodic memory is boundary-based, not turn-cadence: the turn stream is # segmented into coherent experiences and each *closed* segment becomes one @@ -63,14 +60,6 @@ # parameter of :py:meth:`ProcessingPipeline.reconcile_memories`. Hard cap # of 500 (enforced by the pipeline) bounds prompt size and LLM cost. DEFAULT_DEDUP_POOL_SIZE = 50 -# Write-time in-place near-duplicate folding. When enabled, a freshly -# extracted memory that is >= DEDUP_SIM_HIGH similar to an existing active -# memory is folded into that record in place instead of persisting as a new -# doc. Default OFF (add-only): keeping every extracted memory preserves the -# retrieval surface, which benchmarked better than folding. Operators set -# ``DEDUP_VECTOR_ENABLED=true`` to turn folding back on. -DEFAULT_DEDUP_VECTOR_ENABLED = False - # --------------------------------------------------------------------------- # INTERNAL dedup/search tuning - NOT customer-configurable. # These ship as fixed feature constants (no env vars, not in any settings @@ -78,7 +67,6 @@ # needs to become operator-facing we add the env plumbing back deliberately. # --------------------------------------------------------------------------- EXTRACTION_BATCH_MAX_TOKENS = 7000 -DEDUP_SIM_HIGH = 0.97 # >= -> fold new memory into existing canonical in place DEFAULT_TTL_BY_TYPE: dict[str, int] = { "turn": 2_592_000, @@ -254,17 +242,6 @@ def get_extraction_batch_max_tokens() -> int: return EXTRACTION_BATCH_MAX_TOKENS -def get_dedup_vector_enabled() -> bool: - """Whether write-time vector deduplication (in-place folding) is enabled.""" - return _parse_bool("DEDUP_VECTOR_ENABLED", DEFAULT_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. @@ -335,7 +312,6 @@ def get_processor_owner() -> Optional[str]: "DEFAULT_USER_SUMMARY_EVERY_N", "DEFAULT_DEDUP_EVERY_N", "DEFAULT_DEDUP_POOL_SIZE", - "DEFAULT_DEDUP_VECTOR_ENABLED", "DEFAULT_TTL_BY_TYPE", "DEFAULT_PROCEDURAL_SYNTHESIS_AUTO", "DEFAULT_ENABLE_TURN_EMBEDDINGS", @@ -353,8 +329,6 @@ def get_processor_owner() -> Optional[str]: "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/function_app.py b/function_app/function_app.py index 686cb44..8861fae 100644 --- a/function_app/function_app.py +++ b/function_app/function_app.py @@ -8,6 +8,7 @@ import azure.durable_functions as df import azure.functions as func +from orchestrators import extract_episodes as extract_episodes_bp from orchestrators import extract_memories as extract_memories_bp from orchestrators import synthesize_procedural as synthesize_procedural_bp from orchestrators import thread_summary as thread_summary_bp @@ -19,5 +20,6 @@ app.register_functions(change_feed_bp.bp) app.register_functions(thread_summary_bp.bp) app.register_functions(extract_memories_bp.bp) +app.register_functions(extract_episodes_bp.bp) app.register_functions(synthesize_procedural_bp.bp) app.register_functions(user_summary_bp.bp) diff --git a/function_app/orchestrators/extract_episodes.py b/function_app/orchestrators/extract_episodes.py new file mode 100644 index 0000000..eefbe57 --- /dev/null +++ b/function_app/orchestrators/extract_episodes.py @@ -0,0 +1,64 @@ +"""Episodic-extraction orchestrator + activities. + +Chain: ``ExtractEpisodes``. + +The pipeline writes episodes to Cosmos DB during ``ExtractEpisodes``; the +Function App returns only a slim status payload because Durable persists +activity outputs to orchestration history. +""" + +from __future__ import annotations + +import logging + +import azure.durable_functions as df +from shared.pipeline_factory import get_pipeline + +from ._retry import default_retry_options + +logger = logging.getLogger(__name__) + +bp = df.Blueprint() + + +# --------------------------------------------------------------------------- +# Orchestrator +# --------------------------------------------------------------------------- + + +@bp.orchestration_trigger(context_name="context") +def ExtractEpisodesOrchestrator(context: df.DurableOrchestrationContext): + payload = context.get_input() or {} + user_id = payload["user_id"] + thread_id = payload["thread_id"] + + retry = default_retry_options() + + result = yield context.call_activity_with_retry( + "ee_ExtractEpisodes", + retry, + {"user_id": user_id, "thread_id": thread_id}, + ) + + return result + + +# --------------------------------------------------------------------------- +# Activities +# --------------------------------------------------------------------------- + + +@bp.activity_trigger(input_name="payload") +def ee_ExtractEpisodes(payload: dict) -> dict: + user_id = payload["user_id"] + thread_id = payload["thread_id"] + pipeline = get_pipeline() + result = pipeline.extract_episodes(user_id=user_id, thread_id=thread_id, flush=False) or {} + slim = {"episodes": int(result.get("episodes", 0))} + logger.info( + "ExtractEpisodes user=%s thread=%s episodes=%s", + user_id, + thread_id, + slim["episodes"], + ) + return slim diff --git a/function_app/orchestrators/extract_memories.py b/function_app/orchestrators/extract_memories.py index dc30ef6..6fe1a66 100644 --- a/function_app/orchestrators/extract_memories.py +++ b/function_app/orchestrators/extract_memories.py @@ -1,6 +1,6 @@ """Memory-extraction orchestrator + activities. -Chain: ``Extract`` → ``Dedup`` → ``Persist`` followed by an optional +Chain: ``Extract`` -> ``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 @@ -8,7 +8,7 @@ ``reconcile`` flag on its input payload. Procedural synthesis fires only after reconcile and only when ``PROCEDURAL_SYNTHESIS_AUTO`` is enabled, so operators have a kill-switch for the extra LLM call. The prompt is always -derived from the deduped fact pool. Redundant concurrent runs across threads +derived from the extracted fact pool. Redundant concurrent runs across threads are cheap because the pipeline short-circuits with ``status="unchanged"`` when the source fact/episodic IDs have not moved. """ @@ -45,15 +45,10 @@ def ExtractMemoriesOrchestrator(context: df.DurableOrchestrationContext): retry, 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": deduped}, + {"user_id": user_id, "extracted": extracted}, ) count = payload.get("count") @@ -123,18 +118,6 @@ 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.""" @@ -149,7 +132,7 @@ def em_Persist(payload: dict) -> dict: @bp.activity_trigger(input_name="payload") async def em_AdvanceExtractWatermark(payload: dict) -> bool: - """Advance the extraction watermark after a successful extract→persist. + """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. @@ -169,7 +152,7 @@ async def em_AdvanceExtractWatermark(payload: dict) -> bool: @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→dedup→persist split handled here. + # operations are larger/more coupled than the extract/persist flow handled here. user_id = payload["user_id"] pipeline = get_pipeline() from azure.cosmos.agent_memory.thresholds import get_dedup_pool_size diff --git a/function_app/orchestrators/synthesize_procedural.py b/function_app/orchestrators/synthesize_procedural.py index 94c1c23..9713fa0 100644 --- a/function_app/orchestrators/synthesize_procedural.py +++ b/function_app/orchestrators/synthesize_procedural.py @@ -58,12 +58,12 @@ def sp_SynthesizeProcedural(payload: dict) -> dict: result = pipeline.synthesize_procedural(user_id=user_id, force=force) or {} slim = { "status": result.get("status"), - "version": (result.get("procedural") or {}).get("version"), + "procedures_created": int(result.get("procedures_created") or 0), } logger.info( - "SynthesizeProcedural user=%s status=%s version=%s", + "SynthesizeProcedural user=%s status=%s procedures_created=%s", user_id, slim["status"], - slim["version"], + slim["procedures_created"], ) return slim diff --git a/function_app/shared/config.py b/function_app/shared/config.py index 40a1b97..2ace7a6 100644 --- a/function_app/shared/config.py +++ b/function_app/shared/config.py @@ -3,16 +3,21 @@ All knobs are read from environment variables / Azure Functions app settings. Defaults: -* ``FACT_EXTRACTION_EVERY_N`` - default 1 (per-turn extraction) +* ``FACT_EXTRACTION_EVERY_N`` - default 2 (extract every 2 turns) * ``THREAD_SUMMARY_EVERY_N`` - default 10 (rolling summary cadence) +* ``EPISODE_EVAL_EVERY_N`` - default 4 (boundary-evaluation cadence) +* ``EPISODE_IDLE_GAP_SECONDS`` - default 1800 +* ``EPISODE_TOPIC_DRIFT`` - default 0.0 +* ``EPISODE_MAX_TURNS`` - default 40 +* ``EPISODE_MIN_TURNS`` - default 2 * ``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 -workloads. Summaries default to ``10`` because each summary call sees the -full recent context window and is the most expensive per-call operation. +The fact-extraction default of ``2`` extracts facts every second turn. +Operators can raise this for cost-sensitive workloads. Summaries default to +``10`` because each summary call sees the full recent context window and is +the most expensive per-call operation. Setting any ``*_EVERY_N`` env var to ``"0"`` disables that orchestrator entirely. ``PROCEDURAL_SYNTHESIS_AUTO=false`` disables the chained @@ -37,6 +42,7 @@ from __future__ import annotations import logging +import math import os logger = logging.getLogger(__name__) @@ -62,6 +68,11 @@ from azure.cosmos.agent_memory.thresholds import ( # noqa: E402 DEFAULT_DEDUP_EVERY_N, + DEFAULT_EPISODE_EVAL_EVERY_N, + DEFAULT_EPISODE_IDLE_GAP_SECONDS, + DEFAULT_EPISODE_MAX_TURNS, + DEFAULT_EPISODE_MIN_TURNS, + DEFAULT_EPISODE_TOPIC_DRIFT, DEFAULT_FACT_EXTRACTION_EVERY_N, DEFAULT_PROCEDURAL_SYNTHESIS_AUTO, DEFAULT_THREAD_SUMMARY_EVERY_N, @@ -129,6 +140,34 @@ def _parse_bool(name: str, default: bool) -> bool: return default +def _parse_threshold_float(name: str, default: float) -> float: + raw = os.environ.get(name) + if raw is None or raw == "": + return default + try: + parsed = float(raw) + except (ValueError, TypeError): + logger.warning("Invalid value for %s=%r, using default %s", name, raw, default) + return default + if not math.isfinite(parsed): + logger.warning( + "Non-finite value for %s=%r is not allowed; using default %s", + name, + raw, + default, + ) + return default + if parsed < 0: + logger.warning( + "Negative value for %s=%r is not allowed; using default %s (set to 0 to disable)", + name, + raw, + default, + ) + return default + return parsed + + def get_max_batch_size() -> int: return _parse_int("MAX_BATCH_SIZE", DEFAULT_MAX_BATCH_SIZE) @@ -149,6 +188,42 @@ def get_fact_extraction_every_n() -> int: ) +def get_episode_eval_every_n() -> int: + """Boundary-evaluation cadence in turns. ``0`` disables episodic memory.""" + return _parse_threshold( + "EPISODE_EVAL_EVERY_N", + DEFAULT_EPISODE_EVAL_EVERY_N, + ) + + +def get_episode_idle_gap_seconds() -> int: + return _parse_threshold( + "EPISODE_IDLE_GAP_SECONDS", + DEFAULT_EPISODE_IDLE_GAP_SECONDS, + ) + + +def get_episode_topic_drift() -> float: + return _parse_threshold_float( + "EPISODE_TOPIC_DRIFT", + DEFAULT_EPISODE_TOPIC_DRIFT, + ) + + +def get_episode_max_turns() -> int: + return _parse_threshold( + "EPISODE_MAX_TURNS", + DEFAULT_EPISODE_MAX_TURNS, + ) + + +def get_episode_min_turns() -> int: + return _parse_threshold( + "EPISODE_MIN_TURNS", + DEFAULT_EPISODE_MIN_TURNS, + ) + + def get_user_summary_every_n() -> int: """Threshold for triggering ``UserSummaryOrchestrator``. ``0`` disables.""" return _parse_threshold( diff --git a/function_app/shared/counters.py b/function_app/shared/counters.py index be0d14c..991787d 100644 --- a/function_app/shared/counters.py +++ b/function_app/shared/counters.py @@ -292,9 +292,27 @@ async def advance_extract_watermark( 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)}] + """Advance ``last_extract_count`` to ``count`` after a successful extract. + + Monotonic: a conditional patch (``filter_predicate``) applies the new value + only when it is strictly greater than the stored one, so an out-of-order + completion from a concurrent run can never regress the watermark below a + higher value another run already wrote (a regression would re-extract + already-covered turns - RU waste, though the deterministic-id create dedups + the facts). A 412 means a higher watermark already won; keep it. + """ + count = int(count) + patch_ops = [{"op": "add", "path": "/last_extract_count", "value": count}] try: - await container.patch_item(item=counter_id, partition_key=[user_id, thread_id], patch_operations=patch_ops) + await container.patch_item( + item=counter_id, + partition_key=[user_id, thread_id], + patch_operations=patch_ops, + filter_predicate=(f"FROM c WHERE NOT IS_DEFINED(c.last_extract_count) OR c.last_extract_count < {count}"), + ) + except CosmosHttpResponseError as exc: + if exc.status_code == 412: + return # another run already advanced past this count + logger.debug("advance_extract_watermark failed counter_id=%s: %s", counter_id, exc) except Exception as exc: # pragma: no cover - best-effort logger.debug("advance_extract_watermark failed counter_id=%s: %s", counter_id, exc) diff --git a/function_app/triggers/change_feed.py b/function_app/triggers/change_feed.py index 80d806e..8ef9c98 100644 --- a/function_app/triggers/change_feed.py +++ b/function_app/triggers/change_feed.py @@ -123,6 +123,7 @@ async def process_changefeed_batch( n_thread = config.get_thread_summary_every_n() n_facts = config.get_fact_extraction_every_n() + n_episode = config.get_episode_eval_every_n() n_user = config.get_user_summary_every_n() n_dedup = config.get_dedup_every_n() @@ -130,7 +131,7 @@ async def process_changefeed_batch( # auto-trigger contract. Disabled when either knob is 0. n_dedup_turns = n_facts * n_dedup if (n_facts > 0 and n_dedup > 0) else 0 - if n_thread == 0 and n_facts == 0 and n_user == 0: + if n_thread == 0 and n_facts == 0 and n_episode == 0 and n_user == 0: return # all orchestrators disabled # ---- Step 1: Filter to turns + group by scope ---- @@ -147,6 +148,10 @@ async def process_changefeed_batch( for doc in documents: if doc.get("type") != "turn": continue + # Turns are append-only from this trigger's perspective: the Durable + # backend never mutates a turn doc (episodic segmentation advances a + # separate cursor doc; fact extraction is count-based), so each turn is + # delivered - and counted - exactly once, at creation. user_id = doc.get("user_id") thread_id = doc.get("thread_id") if not user_id or not thread_id: @@ -167,7 +172,7 @@ async def process_changefeed_batch( thread_max_lsn[tkey] = max(thread_max_lsn.get(tkey, 0), lsn_int) user_max_lsn[user_id] = max(user_max_lsn.get(user_id, 0), lsn_int) - thread_enabled = n_thread > 0 or n_facts > 0 + thread_enabled = n_thread > 0 or n_facts > 0 or n_episode > 0 user_enabled = n_user > 0 if not thread_counts and not user_counts: @@ -231,6 +236,16 @@ async def process_changefeed_batch( orchestration_errors, ) + if n_episode > 0 and crosses_threshold(old_count, new_count, n_episode): + instance_id = f"episode:{user_id}:{thread_id}:{new_count}" + await _safe_start( + starter, + "ExtractEpisodesOrchestrator", + instance_id, + {"user_id": user_id, "thread_id": thread_id, "count": new_count}, + orchestration_errors, + ) + # ---- Step 3: User-scoped counters ---- if user_enabled: for user_id, batch_count in user_counts.items(): diff --git a/infra/README.md b/infra/README.md index 0b17a6b..2374c26 100644 --- a/infra/README.md +++ b/infra/README.md @@ -113,7 +113,7 @@ The Function app uses a counter document per `(user_id, thread_id)` to decide wh | `azd env` variable | Bicep param | Default | Effect | |---|---|---|---| | `THREAD_SUMMARY_EVERY_N` | `threadSummaryEveryN` | `10` | Run thread-summary orchestration every N turns within a `(user_id, thread_id)`. `0` disables it. | -| `FACT_EXTRACTION_EVERY_N` | `factExtractionEveryN` | `1` | Run fact / episodic / procedural extraction every N turns within a `(user_id, thread_id)`. `0` disables it. | +| `FACT_EXTRACTION_EVERY_N` | `factExtractionEveryN` | `2` | Run fact / episodic / procedural extraction every N turns within a `(user_id, thread_id)`. `0` disables it. | | `DEDUP_EVERY_N` | `dedupEveryN` | `5` | Run fact dedup every Nth fact-extraction (so dedup actually fires every `FACT_EXTRACTION_EVERY_N × DEDUP_EVERY_N` turns). | | `USER_SUMMARY_EVERY_N` | `userSummaryEveryN` | `20` | Run user-summary orchestration every N turns from a given `user_id` across all threads. `0` disables it. | | `MAX_BATCH_SIZE` | `maxBatchSize` | `20` | Maximum number of change-feed items processed per orchestration batch. | diff --git a/infra/main.bicep b/infra/main.bicep index eaf3573..986715a 100644 --- a/infra/main.bicep +++ b/infra/main.bicep @@ -69,8 +69,8 @@ param embeddingDimensions int = 1536 @description('Run thread-summary orchestration every N turns within a (user_id, thread_id). 0 = disabled.') param threadSummaryEveryN int = 10 -@description('Run extract-memories every N change-feed batches. Default 1 = run on every batch (matches SDK + local template). Bump for cost-sensitive production deployments.') -param factExtractionEveryN int = 1 +@description('Run extract-memories every N turns within a (user_id, thread_id). Default 2 = extract every 2 turns (matches SDK + local template). Bump for cost-sensitive production deployments.') +param factExtractionEveryN int = 2 @description('Run dedup once per N fact-extraction batches. Effective cadence = factExtractionEveryN * dedupEveryN turns.') param dedupEveryN int = 5 @@ -78,6 +78,21 @@ param dedupEveryN int = 5 @description('Run user-summary orchestration every N turns from a given user_id across all threads. 0 = disabled.') param userSummaryEveryN int = 20 +@description('Evaluate an episode boundary every N turns within a (user_id, thread_id). 0 = disabled (no episodic memory).') +param episodeEvalEveryN int = 4 + +@description('Idle gap (seconds) between two consecutive turns that closes the open episode.') +param episodeIdleGapSeconds int = 1800 + +@description('Cosine drift from the open segment centroid past which a new turn closes the prior episode. 0 = disabled (idle-gap + max-size only).') +param episodeTopicDrift string = '0' + +@description('Hard cap on turns in one open episode segment before a boundary is forced.') +param episodeMaxTurns int = 40 + +@description('Minimum turns before a natural (idle/drift) boundary may close an episode.') +param episodeMinTurns int = 2 + @description('Maximum number of change-feed items processed per orchestration batch.') param maxBatchSize int = 20 @@ -193,6 +208,11 @@ module functions 'modules/functions.bicep' = if (deployFunctionApp) { factExtractionEveryN: factExtractionEveryN dedupEveryN: dedupEveryN userSummaryEveryN: userSummaryEveryN + episodeEvalEveryN: episodeEvalEveryN + episodeIdleGapSeconds: episodeIdleGapSeconds + episodeTopicDrift: episodeTopicDrift + episodeMaxTurns: episodeMaxTurns + episodeMinTurns: episodeMinTurns maxBatchSize: maxBatchSize memoryProcessorOwner: memoryProcessorOwner tags: commonTags diff --git a/infra/main.parameters.json b/infra/main.parameters.json index 4207199..442d02b 100644 --- a/infra/main.parameters.json +++ b/infra/main.parameters.json @@ -51,7 +51,7 @@ "value": "${THREAD_SUMMARY_EVERY_N=10}" }, "factExtractionEveryN": { - "value": "${FACT_EXTRACTION_EVERY_N=1}" + "value": "${FACT_EXTRACTION_EVERY_N=2}" }, "dedupEveryN": { "value": "${DEDUP_EVERY_N=5}" @@ -59,6 +59,21 @@ "userSummaryEveryN": { "value": "${USER_SUMMARY_EVERY_N=20}" }, + "episodeEvalEveryN": { + "value": "${EPISODE_EVAL_EVERY_N=4}" + }, + "episodeIdleGapSeconds": { + "value": "${EPISODE_IDLE_GAP_SECONDS=1800}" + }, + "episodeTopicDrift": { + "value": "${EPISODE_TOPIC_DRIFT=0}" + }, + "episodeMaxTurns": { + "value": "${EPISODE_MAX_TURNS=40}" + }, + "episodeMinTurns": { + "value": "${EPISODE_MIN_TURNS=2}" + }, "maxBatchSize": { "value": "${MAX_BATCH_SIZE=20}" }, diff --git a/infra/modules/cosmos.bicep b/infra/modules/cosmos.bicep index 74ef0e7..82ef253 100644 --- a/infra/modules/cosmos.bicep +++ b/infra/modules/cosmos.bicep @@ -126,22 +126,6 @@ resource memoriesContainer 'Microsoft.DocumentDB/databaseAccounts/sqlDatabases/c path: '/content' } ] - compositeIndexes: [ - [ - { - path: '/salience' - order: 'descending' - } - { - path: '/created_at' - order: 'ascending' - } - { - path: '/id' - order: 'ascending' - } - ] - ] } vectorEmbeddingPolicy: { vectorEmbeddings: [ diff --git a/infra/modules/functions.bicep b/infra/modules/functions.bicep index 55e50ab..b1a77af 100644 --- a/infra/modules/functions.bicep +++ b/infra/modules/functions.bicep @@ -82,6 +82,21 @@ param dedupEveryN int @description('Run user-summary orchestration every N turns from a given user_id across all threads. 0 = disabled.') param userSummaryEveryN int +@description('Evaluate an episode boundary every N turns within a (user_id, thread_id). 0 = disabled (no episodic memory).') +param episodeEvalEveryN int + +@description('Idle gap (seconds) between two consecutive turns that closes the open episode.') +param episodeIdleGapSeconds int + +@description('Cosine drift from the open segment centroid past which a new turn closes the prior episode. 0 = disabled.') +param episodeTopicDrift string + +@description('Hard cap on turns in one open episode segment before a boundary is forced.') +param episodeMaxTurns int + +@description('Minimum turns before a natural (idle/drift) boundary may close an episode.') +param episodeMinTurns int + @description('Maximum number of change-feed items processed per orchestration batch.') param maxBatchSize int @@ -317,6 +332,26 @@ resource functionApp 'Microsoft.Web/sites@2023-12-01' = { name: 'USER_SUMMARY_EVERY_N' value: string(userSummaryEveryN) } + { + name: 'EPISODE_EVAL_EVERY_N' + value: string(episodeEvalEveryN) + } + { + name: 'EPISODE_IDLE_GAP_SECONDS' + value: string(episodeIdleGapSeconds) + } + { + name: 'EPISODE_TOPIC_DRIFT' + value: episodeTopicDrift + } + { + name: 'EPISODE_MAX_TURNS' + value: string(episodeMaxTurns) + } + { + name: 'EPISODE_MIN_TURNS' + value: string(episodeMinTurns) + } { name: 'DEDUP_EVERY_N' value: string(dedupEveryN) diff --git a/tests/integration/test_async_full_pipeline.py b/tests/integration/test_async_full_pipeline.py index 2bb6347..a1a3af4 100644 --- a/tests/integration/test_async_full_pipeline.py +++ b/tests/integration/test_async_full_pipeline.py @@ -22,7 +22,6 @@ import asyncio import time -import uuid import pytest @@ -85,7 +84,7 @@ async def _async_add_turns( turns: list[tuple[str, str]], ) -> None: for role, content in turns: - await mem.add_cosmos( + await mem.upsert_memory( user_id=user_id, role=role, content=content, @@ -134,7 +133,7 @@ async def _async_seed_fact_with_embedding( check = "SELECT c.id FROM c WHERE c.user_id = @uid AND c.content = @content AND IS_DEFINED(c.embedding)" params = [{"name": "@uid", "value": user_id}, {"name": "@content", "value": content}] for _ in range(retries): - await mem.add_cosmos( + await mem.upsert_memory( user_id=user_id, role="user", content=content, @@ -219,51 +218,3 @@ async def test_extract_reconcile_search( 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_episodic_pipeline.py b/tests/integration/test_episodic_pipeline.py index a8feebb..01c2577 100644 --- a/tests/integration/test_episodic_pipeline.py +++ b/tests/integration/test_episodic_pipeline.py @@ -140,7 +140,7 @@ def _write_hiking_thread(mem: CosmosMemoryClient, user_id: str, thread_id: str) turn_ids = set() for role, content in turns: turn_ids.add( - mem.add_cosmos( + mem.upsert_memory( user_id=user_id, role=role, content=content, diff --git a/tests/integration/test_full_pipeline.py b/tests/integration/test_full_pipeline.py index 2f9e8ac..7091c4f 100644 --- a/tests/integration/test_full_pipeline.py +++ b/tests/integration/test_full_pipeline.py @@ -78,7 +78,7 @@ def _add_turns( turns: list[tuple[str, str]], ) -> None: for role, content in turns: - mem.add_cosmos( + mem.upsert_memory( user_id=user_id, role=role, content=content, @@ -132,7 +132,7 @@ def _seed_fact_with_embedding( ) -> None: """Seed a fact and confirm it was stored *with* an embedding. - ``add_cosmos`` generates the embedding synchronously; a transient + ``upsert_memory`` 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 - @@ -141,7 +141,7 @@ def _seed_fact_with_embedding( check = "SELECT c.id FROM c WHERE c.user_id = @uid AND c.content = @content AND IS_DEFINED(c.embedding)" params = [{"name": "@uid", "value": user_id}, {"name": "@content", "value": content}] for _ in range(retries): - mem.add_cosmos( + mem.upsert_memory( user_id=user_id, role="user", content=content, @@ -169,7 +169,7 @@ def _wait_vector_searchable( ) -> None: """Poll vector search until the user's seeded fact is retrievable. - ``add_cosmos`` stores the embedding synchronously, but Cosmos's DiskANN vector + ``upsert_memory`` stores the embedding synchronously, but Cosmos's DiskANN vector index catches up asynchronously (~1-2s). Gating on a real vector search makes the subsequent ``_vector_candidates`` lookup deterministic instead of racing the index.""" @@ -342,7 +342,7 @@ def test_search_after_extraction(self, agent_memory, unique_user_id, unique_thre class TestTaggingAndSalience: def test_add_remove_tags_and_salience_filter(self, agent_memory, unique_user_id, unique_thread_id): try: - agent_memory.add_cosmos( + agent_memory.upsert_memory( user_id=unique_user_id, role="user", content="The user prefers dark mode UI and uses VS Code.", @@ -402,7 +402,7 @@ def test_dedup_near_duplicate_facts(self, agent_memory, unique_user_id, unique_t "The user resides in Seattle.", "The user works at Microsoft as an engineer.", ]: - agent_memory.add_cosmos( + agent_memory.upsert_memory( user_id=unique_user_id, role="user", content=content, @@ -443,7 +443,7 @@ def test_reconcile_resolves_contradiction(self, agent_memory, unique_user_id, un "User often orders the bone-in pork chop at steakhouses.", ] for content in contradictory_facts: - agent_memory.add_cosmos( + agent_memory.upsert_memory( user_id=unique_user_id, role="user", content=content, @@ -483,13 +483,13 @@ def test_reconcile_resolves_contradiction(self, agent_memory, unique_user_id, un def test_extract_content_hash_short_circuit(self, agent_memory, unique_user_id, unique_thread_id): try: - agent_memory.add_cosmos( + agent_memory.upsert_memory( user_id=unique_user_id, role="user", content="My favorite color is teal.", thread_id=unique_thread_id, ) - agent_memory.add_cosmos( + agent_memory.upsert_memory( user_id=unique_user_id, role="agent", content="Got it, teal is a great color.", @@ -531,7 +531,7 @@ def test_reconcile_writes_supersede_metadata(self, agent_memory, unique_user_id, "The user works as a data engineer at Microsoft in Seattle.", ] for content in paraphrases: - agent_memory.add_cosmos( + agent_memory.upsert_memory( user_id=unique_user_id, role="user", content=content, @@ -564,60 +564,3 @@ 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_procedural_pipeline.py b/tests/integration/test_procedural_pipeline.py new file mode 100644 index 0000000..c36915b --- /dev/null +++ b/tests/integration/test_procedural_pipeline.py @@ -0,0 +1,256 @@ +"""Live procedural-memory integration test. + +This module deliberately builds its Cosmos client with Microsoft Entra ID +(``DefaultAzureCredential``) rather than a key because the live integration +account disables local auth. +""" + +from __future__ import annotations + +import os +import time + +import pytest + +from azure.cosmos.agent_memory import CosmosMemoryClient + +pytestmark = [ + pytest.mark.integration, + pytest.mark.skipif( + os.getenv("AGENT_MEMORY_RUN_INTEGRATION") != "true", + reason="Set AGENT_MEMORY_RUN_INTEGRATION=true", + ), +] + + +@pytest.fixture(scope="module") +def procedural_memory( + cosmos_endpoint, + cosmos_database, + cosmos_container, + ai_foundry_endpoint, + ai_foundry_api_key, + embedding_deployment_name, + embedding_dimensions, + chat_deployment_name, +): + """Live client using AAD for Cosmos and existing containers only.""" + if not cosmos_endpoint or not ai_foundry_endpoint: + pytest.skip("COSMOS_DB_ENDPOINT / AI_FOUNDRY_ENDPOINT not set") + + mem = CosmosMemoryClient( + 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, + use_default_credential=True, + cadence_thresholds={ + "FACT_EXTRACTION_EVERY_N": 1_000_000, + "EPISODE_EVAL_EVERY_N": 1_000_000, + "THREAD_SUMMARY_EVERY_N": 1_000_000, + "USER_SUMMARY_EVERY_N": 1_000_000, + "DEDUP_EVERY_N": 1_000_000, + }, + ) + mem._maybe_auto_trigger = lambda turn_counts: None # type: ignore[method-assign] + mem.connect_cosmos( + endpoint=cosmos_endpoint, + database=cosmos_database, + container=cosmos_container, + ) + try: + yield mem + finally: + mem.close() + + +def _delete_user_records(mem: CosmosMemoryClient, user_id: str) -> None: + query = "SELECT c.id, c.thread_id FROM c WHERE c.user_id = @user_id" + params = [{"name": "@user_id", "value": user_id}] + for container in ( + mem._turns_container_client, + mem._memories_container_client, + mem._summaries_container_client, + ): + try: + docs = list( + container.query_items( + query=query, + parameters=params, + enable_cross_partition_query=True, + ) + ) + except Exception: + continue + for doc in docs: + try: + container.delete_item( + item=doc["id"], + partition_key=[user_id, doc.get("thread_id", "")], + ) + except Exception: + pass + + +def _write_procedural_seed_thread(mem: CosmosMemoryClient, user_id: str, thread_id: str) -> set[str]: + turns = [ + ( + "user", + "This is a durable requirement for future work: always confirm with me before deleting " + "files, database records, cloud resources, or any other destructive or irreversible data.", + ), + ( + "agent", + "Understood. I will ask for explicit confirmation before destructive or irreversible deletions.", + ), + ( + "user", + "Last week I debugged a Cosmos DB hybrid search failure where ORDER BY VectorDistance " + "started failing in an integration test.", + ), + ( + "agent", + "What finally resolved the Cosmos DB hybrid search failure?", + ), + ( + "user", + "The useful lesson was to inspect the vector and full-text indexing policy, fix the " + "missing vector path, wait for the policy to propagate, and rerun the focused integration test.", + ), + ( + "agent", + "That is a reusable recovery strategy for Cosmos DB hybrid search ORDER BY failures.", + ), + ] + turn_ids = set() + for role, content in turns: + turn_ids.add( + mem.upsert_memory( + user_id=user_id, + role=role, + content=content, + memory_type="turn", + thread_id=thread_id, + ) + ) + return turn_ids + + +def _wait_for_procedure_retrieval( + mem: CosmosMemoryClient, + user_id: str, + search_terms: str, + active_procedure_ids: set[str], + *, + timeout: float = 20.0, +) -> list[dict]: + deadline = time.time() + timeout + last_results: list[dict] = [] + while time.time() < deadline: + try: + last_results = mem.retrieve_procedures(user_id, search_terms, top_k=5) + if any(result.get("id") in active_procedure_ids for result in last_results): + return last_results + except Exception: + pass + time.sleep(1) + return last_results + + +def _assert_context_excludes_candidates(context: str, candidates: list[dict]) -> None: + for candidate in candidates: + for field in ("name", "summary", "retrieval_text", "content"): + value = candidate.get(field) + if isinstance(value, str) and value.strip(): + assert value not in context, f"Candidate {field} leaked into procedural context: {value!r}" + + +def test_live_procedural_synthesis_retrieval_and_context( + procedural_memory, + unique_user_id, + unique_thread_id, +): + try: + _write_procedural_seed_thread(procedural_memory, unique_user_id, unique_thread_id) + time.sleep(1) + + memory_stats = procedural_memory.extract_memories(unique_user_id, unique_thread_id) + facts = procedural_memory.get_memories(user_id=unique_user_id, memory_types=["fact"]) + behavioral_facts = [ + fact + for fact in facts + if fact.get("metadata", {}).get("category") in {"preference", "requirement"} + or float(fact.get("salience") or 0.0) >= 0.8 + ] + assert memory_stats.get("fact_count", 0) >= 1 or behavioral_facts, ( + f"Expected at least one behavioral fact extracted from the deletion requirement, got {memory_stats}" + ) + + episode_stats = procedural_memory.extract_episodes(unique_user_id, unique_thread_id, flush=True) + assert episode_stats.get("episodes", 0) >= 1, ( + f"Expected at least one extracted episode with a debugging lesson, got {episode_stats}" + ) + episodes = procedural_memory.get_episodes(unique_user_id) + lesson_bearing = [ + episode + for episode in episodes + if any(isinstance(lesson, str) and lesson.strip() for lesson in (episode.get("lessons") or [])) + ] + assert lesson_bearing, f"Expected at least one episode with first-class lessons, got {episodes}" + + result = procedural_memory.synthesize_procedural(unique_user_id) + assert result.get("status") == "synthesized", result + assert result.get("procedures_created", 0) >= 1, result + + procedures = procedural_memory.get_procedural_memories(unique_user_id) + assert len(procedures) >= result.get("procedures_created", 0) + for procedure in procedures: + assert procedure.get("type") == "procedural" + assert procedure.get("name") + assert procedure.get("procedure_kind") + assert procedure.get("status") in {"active", "candidate"} + assert procedure.get("retrieval_text") + + active = [procedure for procedure in procedures if procedure.get("status") == "active"] + candidates = [procedure for procedure in procedures if procedure.get("status") == "candidate"] + assert active, f"Expected at least one active procedure from explicit requirements, got {procedures}" + + fact_only = [ + procedure + for procedure in procedures + if procedure.get("source_fact_ids") and not procedure.get("source_episodic_ids") + ] + episode_only = [ + procedure + for procedure in procedures + if procedure.get("source_episodic_ids") and not procedure.get("source_fact_ids") + ] + if fact_only: + assert any(procedure.get("status") == "active" for procedure in fact_only), fact_only + if episode_only: + assert all(procedure.get("status") == "candidate" for procedure in episode_only), episode_only + + task = "delete destructive irreversible confirmation before deleting data" + active_ids = {procedure["id"] for procedure in active} + retrieved = _wait_for_procedure_retrieval(procedural_memory, unique_user_id, task, active_ids) + assert retrieved, "Expected retrieve_procedures to return an active procedure" + assert retrieved[0].get("id") in active_ids, retrieved + + context = procedural_memory.build_procedural_context(unique_user_id, task=task) + assert context.strip() + assert any( + (procedure.get("name") and procedure["name"] in context) + or (procedure.get("summary") and procedure["summary"] in context) + for procedure in active + ), context + _assert_context_excludes_candidates(context, candidates) + + prompt = procedural_memory.get_procedural_prompt(unique_user_id) + assert prompt is None or prompt.strip() + if prompt: + _assert_context_excludes_candidates(prompt, candidates) + finally: + _delete_user_records(procedural_memory, unique_user_id) diff --git a/tests/integration/test_ttl_lifecycle.py b/tests/integration/test_ttl_lifecycle.py index 1413d9b..e7d4a23 100644 --- a/tests/integration/test_ttl_lifecycle.py +++ b/tests/integration/test_ttl_lifecycle.py @@ -80,7 +80,7 @@ def _delete_if_present( client: CosmosMemoryClient, memory_id: str, user_id: str, thread_id: str, memory_type: str ) -> None: try: - client.delete_cosmos(memory_id=memory_id, user_id=user_id, thread_id=thread_id, memory_type=memory_type) + client.delete_memory(memory_id=memory_id, user_id=user_id, thread_id=thread_id, memory_type=memory_type) except Exception: pass @@ -91,7 +91,7 @@ def test_turn_ttl_expires_while_episodic_persists(ttl_client: CosmosMemoryClient turn_id = "" episodic_id = "" try: - turn_id = ttl_client.add_cosmos( + turn_id = ttl_client.upsert_memory( user_id=user_id, role="user", content="temporary turn", @@ -99,7 +99,7 @@ def test_turn_ttl_expires_while_episodic_persists(ttl_client: CosmosMemoryClient thread_id=thread_id, ttl=60, ) - episodic_id = ttl_client.add_cosmos( + episodic_id = ttl_client.upsert_memory( user_id=user_id, role="system", content="durable episodic memory", diff --git a/tests/unit/aio/processors/test_durable.py b/tests/unit/aio/processors/test_durable.py index 9a61b97..d4df5ad 100644 --- a/tests/unit/aio/processors/test_durable.py +++ b/tests/unit/aio/processors/test_durable.py @@ -2,14 +2,20 @@ from __future__ import annotations +import logging +from unittest.mock import MagicMock + import pytest +from azure.cosmos.agent_memory.aio.cosmos_memory_client import AsyncCosmosMemoryClient from azure.cosmos.agent_memory.aio.processors import ( AsyncDurableFunctionProcessor, ProcessThreadResult, UserSummaryResult, ) +OLD_EPISODIC_DURABLE_WARNING = "Episodic memory is not available under the Durable Functions backend" + @pytest.mark.asyncio async def test_process_thread_returns_empty_result(): @@ -30,6 +36,37 @@ async def test_generate_user_summary_returns_empty_result(): assert result.summary is None +@pytest.mark.asyncio +async def test_process_extract_episodes_returns_empty_result_without_old_warning(caplog): + proc = AsyncDurableFunctionProcessor() + caplog.set_level(logging.WARNING) + + result = await proc.process_extract_episodes(user_id="u1", thread_id="t1") + + assert result == {} + assert OLD_EPISODIC_DURABLE_WARNING not in caplog.text + + +@pytest.mark.asyncio +async def test_synthesize_procedural_is_noop(): + # Mirror of the sync processor: no-op (not raise) so the auto-trigger does + # not stamp a spurious failure each cadence. + proc = AsyncDurableFunctionProcessor() + result = await proc.synthesize_procedural(user_id="u1") + assert result == {"status": "skipped", "procedures_created": 0} + + +@pytest.mark.asyncio +async def test_client_extract_episodes_raises_for_durable_processor(): + client = AsyncCosmosMemoryClient(use_default_credential=False, processor=AsyncDurableFunctionProcessor()) + client._pipeline = MagicMock() + + with pytest.raises(NotImplementedError, match="Durable Function app"): + await client.extract_episodes("u1", "t1", flush=True) + + client._pipeline.extract_episodes.assert_not_called() + + @pytest.mark.asyncio async def test_close_is_noop(): assert await AsyncDurableFunctionProcessor().close() is None diff --git a/tests/unit/aio/services/test_dedup_vector_async.py b/tests/unit/aio/services/test_dedup_vector_async.py deleted file mode 100644 index ed679ed..0000000 --- a/tests/unit/aio/services/test_dedup_vector_async.py +++ /dev/null @@ -1,385 +0,0 @@ -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 - - -@pytest.fixture(autouse=True) -def _enable_vector_folding(monkeypatch: pytest.MonkeyPatch) -> None: - # DEDUP_VECTOR_ENABLED now defaults to False (add-only); this suite exercises - # the in-place folding path, so enable it. Tests that assert the flag-off - # behavior patch the getter directly and override this. - monkeypatch.setenv("DEDUP_VECTOR_ENABLED", "true") - - -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": {}, - "title": "CI episode", - "started_at": None, - "ended_at": None, - "participants": [], - "events": [{"sequence": 1, "description": content, "occurred_at": None, "source_turn_ids": []}], - "outcome": {"status": "successful", "description": content}, - "lessons": [content], - "source_turn_ids": [], - "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_apply_inplace_update_skips_cross_source_fold(): - p = _service() - p._replace_item = AsyncMock() - neighbor = _fact( - "existing-1", "same content", tags=["sys:fact"], metadata={"category": "preference", "source": "user"} - ) - neighbor["_etag"] = "etag-xyz" - new_doc = _fact( - "f-new", - "same content", - embedding=[0.5, 0.5], - tags=["sys:fact", "sys:agent-fact"], - metadata={"category": "other", "source": "agent"}, - ) - - assert await p._apply_inplace_update(neighbor, new_doc) is False - p._replace_item.assert_not_awaited() - - -@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/services/test_episode_boundary_async.py b/tests/unit/aio/services/test_episode_boundary_async.py index 5961637..a08a55c 100644 --- a/tests/unit/aio/services/test_episode_boundary_async.py +++ b/tests/unit/aio/services/test_episode_boundary_async.py @@ -1,4 +1,8 @@ -"""Boundary-based episodic segmentation (async mirror of test_episode_boundary).""" +"""Boundary-based episodic segmentation (async mirror of test_episode_boundary). + +A per-thread ``(created_at, id)`` cursor doc (not a per-turn stamp) advances past +folded turns, so episodic segmentation writes nothing to the turns container and +cannot perturb the change-feed cadence counter.""" from __future__ import annotations @@ -64,8 +68,24 @@ def _episodes(store: _AsyncTrackingStore) -> list[dict[str, Any]]: return [doc for doc in store.docs if doc.get("type") == "episodic"] -def _stamped(turns_store: _AsyncStore) -> list[str]: - return sorted(t["id"] for t in turns_store.docs if t.get("episode_extracted_at")) +def _folded(memories: _AsyncTrackingStore, turns_store: _AsyncStore) -> list[str]: + """Turn ids folded into an episode, derived from the per-thread episodic + cursor doc (the watermark model advances a single cursor instead of stamping + each turn, so nothing is written to the turns container).""" + cursor = next((d for d in memories.docs if d.get("type") == "episode_cursor"), None) + if cursor is None: + return [] + last_at = str(cursor.get("last_episode_at") or "") + last_id = str(cursor.get("last_episode_id") or "") + return sorted( + str(t.get("id")) + for t in turns_store.docs + if t.get("type") == "turn" + and ( + (str(t.get("created_at") or "") < last_at) + or (str(t.get("created_at") or "") == last_at and str(t.get("id") or "") <= last_id) + ) + ) @pytest.mark.asyncio @@ -79,7 +99,7 @@ async def test_time_gap_closes_prior_episode_and_leaves_tail_open(monkeypatch) - assert result == {"episodes": 1} assert len(_episodes(memories)) == 1 - assert _stamped(turns_store) == ["turn-1", "turn-2"] + assert _folded(memories, turns_store) == ["turn-1", "turn-2"] @pytest.mark.asyncio @@ -109,7 +129,7 @@ async def test_flush_drains_open_tail(monkeypatch) -> None: assert flushed == {"episodes": 1} assert len(_episodes(memories)) == 2 - assert _stamped(turns_store) == ["turn-1", "turn-2", "turn-3", "turn-4"] + assert _folded(memories, turns_store) == ["turn-1", "turn-2", "turn-3", "turn-4"] assert await service.extract_episodes("u1", "t1", flush=True) == {"episodes": 0} @@ -124,7 +144,7 @@ async def test_max_turns_forces_a_boundary(monkeypatch) -> None: result = await service.extract_episodes("u1", "t1") assert result == {"episodes": 2} - assert _stamped(turns_store) == ["turn-1", "turn-2", "turn-3", "turn-4"] + assert _folded(memories, turns_store) == ["turn-1", "turn-2", "turn-3", "turn-4"] @pytest.mark.asyncio @@ -143,7 +163,7 @@ async def test_topic_drift_closes_episode(monkeypatch) -> None: result = await service.extract_episodes("u1", "t1") assert result == {"episodes": 1} - assert _stamped(turns_store) == ["turn-1", "turn-2"] + assert _folded(memories, turns_store) == ["turn-1", "turn-2"] @pytest.mark.asyncio @@ -158,7 +178,7 @@ async def test_no_boundary_keeps_segment_open_without_calling_the_llm(monkeypatc assert result == {"episodes": 0} assert _episodes(memories) == [] - assert _stamped(turns_store) == [] + assert _folded(memories, turns_store) == [] assert chat.calls == 0 @@ -175,7 +195,7 @@ async def test_idle_gap_below_min_turns_does_not_close_episode(monkeypatch) -> N assert result == {"episodes": 0} assert _episodes(memories) == [] - assert _stamped(turns_store) == [] + assert _folded(memories, turns_store) == [] @pytest.mark.asyncio @@ -190,7 +210,7 @@ async def test_idle_gap_below_min_turns_still_flushes_as_one_episode(monkeypatch result = await service.extract_episodes("u1", "t1", flush=True) assert result == {"episodes": 1} - assert _stamped(turns_store) == ["turn-1", "turn-2", "turn-3"] + assert _folded(memories, turns_store) == ["turn-1", "turn-2", "turn-3"] @pytest.mark.asyncio @@ -210,7 +230,7 @@ async def _boom(*a: Any, **k: Any) -> str: assert result == {"episodes": 0} assert _episodes(memories) == [] - assert _stamped(turns_store) == [] # un-stamped -> retried next run + assert _folded(memories, turns_store) == [] # un-stamped -> retried next run @pytest.mark.asyncio @@ -230,4 +250,49 @@ async def _boom(*a: Any, **k: Any) -> str: assert result == {"episodes": 0} assert _episodes(memories) == [] - assert _stamped(turns_store) == ["turn-1", "turn-2"] # quarantined + advanced + assert _folded(memories, turns_store) == ["turn-1", "turn-2"] # quarantined + advanced + + +@pytest.mark.asyncio +async def test_advance_episode_cursor_is_monotonic(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_IDLE_GAP_SECONDS", "120") + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + turns = [_turn_at(1, 1), _turn_at(2, 2), _turn_at(3, 30), _turn_at(4, 31)] + service, memories, turns_store, _ = _service(turns) + + await service.extract_episodes("u1", "t1") # gap closes [turn-1, turn-2] -> cursor at turn-2 + assert _folded(memories, turns_store) == ["turn-1", "turn-2"] + + # A late, out-of-order concurrent run advancing to an OLDER turn is a no-op. + await service._advance_episode_cursor("u1", "t1", _turn_at(1, 1)) + assert _folded(memories, turns_store) == ["turn-1", "turn-2"] + + # A genuinely newer turn still advances the cursor. + await service._advance_episode_cursor("u1", "t1", _turn_at(4, 31)) + assert _folded(memories, turns_store) == ["turn-1", "turn-2", "turn-3", "turn-4"] + + +@pytest.mark.asyncio +async def test_episode_cursor_is_per_thread(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_IDLE_GAP_SECONDS", "120") + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + + def _t(i: int, minute: int, thread: str) -> dict[str, Any]: + return {**_turn_at(i, minute), "thread_id": thread} + + turns = [ + _t(1, 1, "t1"), + _t(2, 2, "t1"), + _t(3, 30, "t1"), + _t(4, 1, "t2"), + _t(5, 2, "t2"), + _t(6, 30, "t2"), + ] + service, memories, turns_store, _ = _service(turns) + + await service.extract_episodes("u1", "t1") # only thread t1 + + cursors = sorted(d["thread_id"] for d in memories.docs if d.get("type") == "episode_cursor") + assert cursors == ["t1"] + t2_segment = await service._load_open_episode_segment("u1", "t2") + assert [t["id"] for t in t2_segment] == ["turn-4", "turn-5", "turn-6"] diff --git a/tests/unit/aio/services/test_episodic_retrieval_async.py b/tests/unit/aio/services/test_episodic_retrieval_async.py index 4b0973b..beabc80 100644 --- a/tests/unit/aio/services/test_episodic_retrieval_async.py +++ b/tests/unit/aio/services/test_episodic_retrieval_async.py @@ -1,5 +1,6 @@ from __future__ import annotations +import logging from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock @@ -7,6 +8,8 @@ from azure.cosmos.agent_memory.aio.cosmos_memory_client import AsyncCosmosMemoryClient from azure.cosmos.agent_memory.aio.store import AsyncMemoryStore +EPISODIC_OPT_IN_WARNING = "Episodic memories requested via memory_types are only returned when include_episodes=True" + class AsyncIterator: def __init__(self, items): @@ -95,6 +98,37 @@ async def test_async_search_cosmos_base_is_facts_only_no_episodes_without_optin( store.search_episodic.assert_not_awaited() +async def test_async_search_cosmos_warns_when_episodic_requested_without_optin(caplog): + mem, _ = _connected_client() + store = MagicMock() + store.search = AsyncMock(return_value=[]) + mem._get_store = MagicMock(return_value=store) + caplog.set_level(logging.WARNING) + + results = await mem.search_cosmos("weather", user_id="u1", memory_types=["episodic"]) + + assert results == [] + assert EPISODIC_OPT_IN_WARNING in caplog.text + assert store.search.call_args.kwargs["memory_types"] == ["fact"] + + +async def test_async_search_cosmos_does_not_warn_for_episodic_optin_or_other_types(caplog): + mem, _ = _connected_client() + store = MagicMock() + store.search = AsyncMock(return_value=[]) + mem._get_store = MagicMock(return_value=store) + caplog.set_level(logging.WARNING) + + await mem.search_cosmos("weather", user_id="u1", memory_types=["episodic"], include_episodes=True) + assert EPISODIC_OPT_IN_WARNING not in caplog.text + assert store.search.call_args.kwargs["memory_types"] == ["episodic"] + + caplog.clear() + await mem.search_cosmos("weather", user_id="u1", memory_types=["fact"]) + assert EPISODIC_OPT_IN_WARNING not in caplog.text + assert store.search.call_args.kwargs["memory_types"] == ["fact"] + + async def test_async_search_cosmos_include_episodes_combines_facts_and_episodes_in_base_query(): mem, _ = _connected_client() store = MagicMock() diff --git a/tests/unit/aio/services/test_extract_episodes_async.py b/tests/unit/aio/services/test_extract_episodes_async.py index 2815a83..a312343 100644 --- a/tests/unit/aio/services/test_extract_episodes_async.py +++ b/tests/unit/aio/services/test_extract_episodes_async.py @@ -98,12 +98,13 @@ async def test_extract_episodes_embeds_content_persists_append_only(monkeypatch) "The user planned a vacation.", ] ] - assert [doc["content"] for doc in store.docs] == [ + episodes = [doc for doc in store.docs if doc.get("type") == "episodic"] + assert [doc["content"] for doc in episodes] == [ "The user fixed flaky CI retries and the tests passed.", "The user planned a vacation.", ] - assert all(doc["id"].startswith("ep_") for doc in store.docs) - assert all(doc["embedding"] == [1.0] for doc in store.docs) + assert all(doc["id"].startswith("ep_") for doc in episodes) + assert all(doc["embedding"] == [1.0] for doc in episodes) assert store.supersede_calls == [] assert store.search_calls == [] @@ -116,7 +117,7 @@ async def test_extract_episodes_empty_window_persists_nothing(monkeypatch) -> No result = await service.extract_episodes("u1", "t1", flush=True) assert result == {"episodes": 0} - assert store.docs == [] + assert [doc for doc in store.docs if doc.get("type") == "episodic"] == [] assert embeddings.calls == [] @@ -138,7 +139,7 @@ async def test_extract_episodes_skips_malformed_episode_with_warning(caplog, mon result = await service.extract_episodes("u1", "t1", flush=True) assert result == {"episodes": 1} - assert [doc["title"] for doc in store.docs] == ["Valid episode"] + assert [doc["title"] for doc in store.docs if doc.get("type") == "episodic"] == ["Valid episode"] assert "dropping malformed episode" in caplog.text @@ -202,8 +203,8 @@ async def test_extract_episodes_skips_duplicate_when_segment_reprocessed(monkeyp ) assert await service.extract_episodes("u1", "t1", flush=True) == {"episodes": 1} - for turn in turns.docs: - turn.pop("episode_extracted_at", None) + # Simulate a crash before the cursor advanced: the watermark never moved. + store.docs = [doc for doc in store.docs if doc.get("type") != "episode_cursor"] assert await service.extract_episodes("u1", "t1", flush=True) == {"episodes": 0} episodic = [doc for doc in store.docs if doc.get("type") == "episodic"] assert len(episodic) == 1 diff --git a/tests/unit/aio/services/test_mark_turns_extracted_async.py b/tests/unit/aio/services/test_mark_turns_extracted_async.py new file mode 100644 index 0000000..4369dfc --- /dev/null +++ b/tests/unit/aio/services/test_mark_turns_extracted_async.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +from typing import Any + +import pytest + +from azure.cosmos.agent_memory._container_routing import ContainerKey +from azure.cosmos.agent_memory.aio.services.pipeline import AsyncPipelineService, _AsyncStoreContainerAdapter +from tests.unit.services.test_extract_dry import _AsyncChat, _AsyncEmbeddings, _AsyncStore + + +class _AsyncPatchTurnsContainer: + def __init__(self, docs: list[dict[str, Any]]): + self.docs = [dict(doc) for doc in docs] + self.patch_calls: list[dict[str, Any]] = [] + self.upsert_calls: list[dict[str, Any]] = [] + + async def query_items(self, **kwargs: Any) -> list[dict[str, Any]]: + del kwargs + return [dict(doc) for doc in self.docs] + + async def read_item(self, *, item: str, partition_key: Any) -> dict[str, Any]: + del partition_key + for doc in self.docs: + if doc.get("id") == item: + return dict(doc) + raise KeyError(item) + + async def upsert_item(self, *, body: dict[str, Any]) -> dict[str, Any]: + body = dict(body) + self.upsert_calls.append(body) + for index, doc in enumerate(self.docs): + if doc.get("id") == body.get("id"): + self.docs[index] = body + return body + self.docs.append(body) + return body + + async def patch_item( + self, *, item: str, partition_key: Any, patch_operations: list[dict[str, Any]] + ) -> dict[str, Any]: + self.patch_calls.append( + {"item": item, "partition_key": partition_key, "patch_operations": [dict(op) for op in patch_operations]} + ) + for doc in self.docs: + if doc.get("id") == item: + for operation in patch_operations: + assert operation["op"] == "set" + assert operation["path"].startswith("/") + doc[operation["path"][1:]] = operation["value"] + return dict(doc) + raise KeyError(item) + + +class _AsyncNoPatchTurnsContainer(_AsyncPatchTurnsContainer): + patch_item = None + + +class _AsyncContainerBackedStore(_AsyncStore): + def __init__(self, container: Any): + super().__init__([]) + self._containers = {ContainerKey.TURNS: container} + + +def _turn_doc(**overrides: Any) -> dict[str, Any]: + doc = { + "id": "turn-1", + "user_id": "u1", + "thread_id": "t1", + "type": "turn", + "content": "hello", + "created_at": "2025-01-01T00:00:00+00:00", + } + doc.update(overrides) + return doc + + +def _service(turns_container: Any) -> AsyncPipelineService: + memories_store = _AsyncStore([]) + summaries_store = _AsyncStore([]) + turns_adapter = _AsyncStoreContainerAdapter(_AsyncContainerBackedStore(turns_container), ContainerKey.TURNS) + return AsyncPipelineService( + memories_store, + _AsyncChat([]), + _AsyncEmbeddings(), + containers={ + ContainerKey.TURNS: turns_adapter, + ContainerKey.MEMORIES: _AsyncStoreContainerAdapter(memories_store, ContainerKey.MEMORIES), + ContainerKey.SUMMARIES: _AsyncStoreContainerAdapter(summaries_store, ContainerKey.SUMMARIES), + }, + ) + + +@pytest.mark.asyncio +async def test_mark_turns_extracted_patches_extracted_at() -> None: + turns_container = _AsyncPatchTurnsContainer([_turn_doc()]) + service = _service(turns_container) + + marked = await service._mark_turns_extracted([_turn_doc()]) + + assert marked == 1 + assert turns_container.patch_calls == [ + { + "item": "turn-1", + "partition_key": ["u1", "t1"], + "patch_operations": [ + {"op": "set", "path": "/extracted_at", "value": turns_container.docs[0]["extracted_at"]} + ], + } + ] + assert turns_container.docs[0]["extracted_at"] + assert turns_container.upsert_calls == [] + + +@pytest.mark.asyncio +async def test_mark_turns_extracted_falls_back_to_read_modify_upsert() -> None: + turns_container = _AsyncNoPatchTurnsContainer([_turn_doc()]) + service = _service(turns_container) + + marked = await service._mark_turns_extracted([_turn_doc()]) + + assert marked == 1 + assert turns_container.docs[0]["extracted_at"] + assert turns_container.upsert_calls == [turns_container.docs[0]] diff --git a/tests/unit/aio/services/test_procedural_retrieval_async.py b/tests/unit/aio/services/test_procedural_retrieval_async.py new file mode 100644 index 0000000..00b4b64 --- /dev/null +++ b/tests/unit/aio/services/test_procedural_retrieval_async.py @@ -0,0 +1,115 @@ +from __future__ import annotations + +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +from azure.cosmos.agent_memory._container_routing import ContainerKey +from azure.cosmos.agent_memory.aio.store import AsyncMemoryStore + + +class AsyncIterator: + def __init__(self, items: list[dict[str, Any]]) -> None: + self._items = iter(items) + + def __aiter__(self) -> AsyncIterator: + return self + + async def __anext__(self) -> dict[str, Any]: + try: + return next(self._items) + except StopIteration: + raise StopAsyncIteration + + +def _containers(*, memories: Any = None, turns: Any = None, summaries: Any = None) -> dict[ContainerKey, Any]: + return { + ContainerKey.MEMORIES: memories if memories is not None else MagicMock(), + ContainerKey.TURNS: turns if turns is not None else MagicMock(), + ContainerKey.SUMMARIES: summaries if summaries is not None else MagicMock(), + } + + +def _params_by_name(call_kwargs: dict[str, Any]) -> dict[str, Any]: + return {param["name"]: param["value"] for param in call_kwargs["parameters"]} + + +async def test_async_retrieve_procedures_filters_active_procedures_by_default() -> None: + ranked_docs = [ + {"id": "proc-1", "type": "procedural", "status": "active", "similarity_score": 0.1}, + {"id": "proc-2", "type": "procedural", "status": "active", "similarity_score": 0.2}, + ] + memories = MagicMock() + memories.query_items.return_value = AsyncIterator(ranked_docs) + embeddings = MagicMock() + embeddings.generate = AsyncMock(return_value=[0.1, 0.2]) + store = AsyncMemoryStore(containers=_containers(memories=memories), embeddings_client=embeddings) + + result = await store.retrieve_procedures("u1", "cosmos db retry", top_k=2) + + assert result == ranked_docs + call_kwargs = memories.query_items.call_args.kwargs + assert "TOP 2" in call_kwargs["query"] + assert "c.type = @type" in call_kwargs["query"] + assert "c.user_id = @user_id" in call_kwargs["query"] + assert "c.status = @status" in call_kwargs["query"] + assert "VectorDistance(c.embedding, @embedding)" in call_kwargs["query"] + assert "(NOT IS_DEFINED(c.superseded_by) OR IS_NULL(c.superseded_by))" in call_kwargs["query"] + params = _params_by_name(call_kwargs) + assert params["@type"] == "procedural" + assert params["@user_id"] == "u1" + assert params["@status"] == "active" + assert params["@embedding"] == [0.1, 0.2] + assert params["@kw0"] == "cosmos" + + +async def test_async_retrieve_procedures_status_none_drops_status_and_adds_scope_filters() -> None: + ranked_docs = [{"id": "proc-domain", "type": "procedural", "similarity_score": 0.1}] + memories = MagicMock() + memories.query_items.return_value = AsyncIterator(ranked_docs) + embeddings = MagicMock() + embeddings.generate = AsyncMock(return_value=[0.3, 0.4]) + store = AsyncMemoryStore(containers=_containers(memories=memories), embeddings_client=embeddings) + + result = await store.retrieve_procedures( + "u1", + "partition key", + scope_type="domain", + scope_value="cosmos-db", + status=None, + ) + + assert result == ranked_docs + call_kwargs = memories.query_items.call_args.kwargs + assert "c.scope_type = @scope_type" in call_kwargs["query"] + assert "c.scope_value = @scope_value" in call_kwargs["query"] + assert "c.status = @status" not in call_kwargs["query"] + params = _params_by_name(call_kwargs) + assert params["@type"] == "procedural" + assert params["@user_id"] == "u1" + assert params["@scope_type"] == "domain" + assert params["@scope_value"] == "cosmos-db" + assert "@status" not in params + + +async def test_async_retrieve_procedures_filters_by_kind_and_can_include_superseded() -> None: + ranked_docs = [{"id": "proc-workflow", "type": "procedural", "similarity_score": 0.1}] + memories = MagicMock() + memories.query_items.return_value = AsyncIterator(ranked_docs) + embeddings = MagicMock() + embeddings.generate = AsyncMock(return_value=[0.5, 0.6]) + store = AsyncMemoryStore(containers=_containers(memories=memories), embeddings_client=embeddings) + + result = await store.retrieve_procedures( + "u1", + "workflow retry", + procedure_kind="workflow", + include_superseded=True, + ) + + assert result == ranked_docs + call_kwargs = memories.query_items.call_args.kwargs + assert "c.procedure_kind = @procedure_kind" in call_kwargs["query"] + assert "(NOT IS_DEFINED(c.superseded_by) OR IS_NULL(c.superseded_by))" not in call_kwargs["query"] + params = _params_by_name(call_kwargs) + assert params["@procedure_kind"] == "workflow" + assert params["@status"] == "active" diff --git a/tests/unit/aio/services/test_procedural_synthesis_and_context_async.py b/tests/unit/aio/services/test_procedural_synthesis_and_context_async.py new file mode 100644 index 0000000..4add436 --- /dev/null +++ b/tests/unit/aio/services/test_procedural_synthesis_and_context_async.py @@ -0,0 +1,441 @@ +from __future__ import annotations + +from typing import Any + +import pytest +from azure.cosmos.exceptions import CosmosResourceExistsError + +from azure.cosmos.agent_memory.aio.services.pipeline import AsyncPipelineService +from tests.unit.services.test_extract_dry import ( + _async_containers_for_store, + _AsyncChat, + _AsyncEmbeddings, + _AsyncStore, +) + + +class _AsyncProceduralStore(_AsyncStore): + async def query(self, sql: str, parameters=None, partition_key=None, cross_partition: bool = False): + del partition_key, cross_partition + params = {p["name"]: p["value"] for p in (parameters or [])} + user_id = params.get("@uid", params.get("@user_id")) + memory_type = params.get("@type", params.get("@memory_type")) + docs = [dict(doc) for doc in self.docs] + if user_id is not None: + docs = [doc for doc in docs if doc.get("user_id") == user_id] + if memory_type is not None: + docs = [doc for doc in docs if doc.get("type") == memory_type] + if "c.status='active'" in sql: + docs = [doc for doc in docs if doc.get("status") == "active"] + if "superseded_by" in sql: + docs = [doc for doc in docs if not doc.get("superseded_by")] + return docs + + async def create_item(self, *, body: dict[str, Any]) -> dict[str, Any]: + if any(doc.get("id") == body.get("id") for doc in self.docs): + raise CosmosResourceExistsError(message="conflict") + self.docs.append(dict(body)) + return dict(body) + + +def _service( + store: _AsyncProceduralStore, + responses: list[dict[str, Any]] | None = None, +) -> AsyncPipelineService: + return AsyncPipelineService( + store, + _AsyncChat(responses or []), + _AsyncEmbeddings(), + containers=_async_containers_for_store(store), + ) + + +def _fact() -> dict[str, Any]: + return { + "id": "fact-raw-1", + "user_id": "u1", + "type": "fact", + "content": "The user explicitly said to run targeted tests before reporting success.", + "metadata": {"category": "preference"}, + "salience": 0.9, + "created_at": "2025-01-01T00:00:00+00:00", + } + + +def _episode() -> dict[str, Any]: + return { + "id": "episode-raw-1", + "user_id": "u1", + "type": "episodic", + "content": "A retry investigation succeeded.", + "lessons": ["Retry transient CI failures once before escalating."], + "salience": 0.8, + "created_at": "2025-01-01T00:01:00+00:00", + } + + +def _procedure( + name: str, + *, + grounded_in: list[str], + source_kind: str, + summary: str = "Run targeted tests before reporting success.", +) -> dict[str, Any]: + return { + "name": name, + "summary": summary, + "retrieval_text": summary, + "procedure_kind": "behavioral_policy", + "scope_type": "user", + "scope_value": None, + "activation_conditions": [], + "preconditions": [], + "steps": [], + "success_conditions": [], + "failure_conditions": [], + "safety_constraints": [], + "source_kind": source_kind, + "grounded_in": grounded_in, + "confidence": 0.8, + } + + +@pytest.mark.asyncio +async def test_synthesize_procedural_applies_provenance_gate() -> None: + store = _AsyncProceduralStore([_fact(), _episode()]) + service = _service( + store, + [ + { + "procedures": [ + _procedure( + "Targeted testing", + grounded_in=["fact-1"], + source_kind="explicit_user_instruction", + ), + _procedure( + "Retry CI failures", + grounded_in=["ep-1"], + source_kind="explicit_user_instruction", + summary="Retry transient CI failures once before escalating.", + ), + ] + } + ], + ) + + result = await service.synthesize_procedural("u1") + + assert result == {"status": "synthesized", "procedures_created": 2, "procedures_skipped": 0} + procedures = [doc for doc in store.docs if doc.get("type") == "procedural"] + assert {doc["name"]: doc["status"] for doc in procedures} == { + "Targeted testing": "active", + "Retry CI failures": "candidate", + } + assert {doc["name"]: doc["source_kind"] for doc in procedures}["Retry CI failures"] == "episode_distillation" + assert procedures[0]["thread_id"] == "__procedural__" + assert procedures[0]["embedding"] == [1.0] + + +@pytest.mark.asyncio +async def test_build_procedural_context_uses_active_procedures_only() -> None: + active = { + "id": "proc-active", + "user_id": "u1", + "type": "procedural", + "status": "active", + "name": "Targeted testing", + "summary": "Run targeted tests before reporting success.", + "retrieval_text": "tests success", + "procedure_kind": "behavioral_policy", + "scope_type": "user", + "scope_value": None, + "priority": 10, + "source_authority": "high", + "version": 1, + } + candidate = { + **active, + "id": "proc-candidate", + "status": "candidate", + "name": "Candidate policy", + "summary": "Do not include this candidate procedure.", + } + service = _service(_AsyncProceduralStore([active, candidate])) + + context = await service.build_procedural_context("u1") + + assert "Run targeted tests before reporting success." in context + assert "Do not include this candidate procedure." not in context + assert await service.build_procedural_context("missing-user") == "" + + +def _active_procedure( + name: str, + *, + summary: str, + procedure_kind: str = "behavioral_policy", + scope_type: str = "user", + retrieval_text: str | None = None, + activation_conditions: list[str] | None = None, + priority: int = 0, + source_authority: str = "low", + steps: list[dict[str, Any]] | None = None, +) -> dict[str, Any]: + return { + "id": f"proc-{name.lower().replace(' ', '-')}", + "user_id": "u1", + "type": "procedural", + "status": "active", + "name": name, + "summary": summary, + "retrieval_text": retrieval_text if retrieval_text is not None else summary, + "procedure_kind": procedure_kind, + "scope_type": scope_type, + "scope_value": None, + "activation_conditions": activation_conditions or [], + "steps": steps or [], + "priority": priority, + "source_authority": source_authority, + "version": 1, + } + + +@pytest.mark.asyncio +async def test_build_procedural_context_includes_task_matching_workflow_only() -> None: + workflow = _active_procedure( + "Partition workflow", + summary="Use partition key diagnostics.", + procedure_kind="workflow", + scope_type="domain", + retrieval_text="cosmos partition routing diagnostics", + activation_conditions=["when debugging partition fanout"], + steps=[{"sequence": 1, "instruction": "Inspect partition routing."}], + ) + service = _service(_AsyncProceduralStore([workflow])) + + matching_context = await service.build_procedural_context("u1", task="debug partition latency") + + assert "Partition workflow" in matching_context + assert "Inspect partition routing." in matching_context + assert await service.build_procedural_context("u1") == "" + assert await service.build_procedural_context("u1", task="summarize billing invoices") == "" + + +@pytest.mark.asyncio +async def test_build_procedural_context_always_includes_global_and_user_policies() -> None: + global_policy = _active_procedure( + "Global reporting", + summary="Always report validation status.", + scope_type="global", + ) + user_rule = _active_procedure( + "User test rule", + summary="Run targeted tests before reporting success.", + procedure_kind="decision_rule", + scope_type="user", + ) + domain_policy = _active_procedure( + "Domain policy", + summary="Do not include domain policies without task matching.", + scope_type="domain", + ) + service = _service(_AsyncProceduralStore([global_policy, user_rule, domain_policy])) + + context = await service.build_procedural_context("u1") + + assert "Global reporting" in context + assert "User test rule" in context + assert "Domain policy" not in context + + +@pytest.mark.asyncio +async def test_build_procedural_context_orders_policies_by_priority_then_authority() -> None: + lower_priority = _active_procedure( + "Lower priority", + summary="Lower priority policy.", + priority=1, + source_authority="mandatory", + ) + higher_priority = _active_procedure( + "Higher priority", + summary="Higher priority policy.", + priority=10, + source_authority="low", + ) + mandatory_authority = _active_procedure( + "Mandatory authority", + summary="Mandatory authority policy.", + priority=10, + source_authority="mandatory", + ) + service = _service(_AsyncProceduralStore([lower_priority, higher_priority, mandatory_authority])) + + context = await service.build_procedural_context("u1") + + assert context.index("Mandatory authority") < context.index("Higher priority") + assert context.index("Higher priority") < context.index("Lower priority") + + +@pytest.mark.asyncio +async def test_synthesize_procedural_downgrades_episode_only_explicit_claim() -> None: + store = _AsyncProceduralStore([_episode()]) + service = _service( + store, + [ + { + "procedures": [ + _procedure( + "Episode claim", + grounded_in=["ep-1"], + source_kind="explicit_user_instruction", + summary="Retry transient CI failures once before escalating.", + ) + ] + } + ], + ) + + result = await service.synthesize_procedural("u1") + + assert result == {"status": "synthesized", "procedures_created": 1, "procedures_skipped": 0} + procedure = next(doc for doc in store.docs if doc.get("type") == "procedural") + assert procedure["status"] == "candidate" + assert procedure["source_kind"] == "episode_distillation" + + +@pytest.mark.asyncio +async def test_synthesize_procedural_downgrades_episode_only_organization_policy() -> None: + store = _AsyncProceduralStore([_episode()]) + service = _service( + store, + [ + { + "procedures": [ + _procedure( + "Org policy from episode", + grounded_in=["ep-1"], + source_kind="organization_policy", + summary="Retry transient CI failures once before escalating.", + ) + ] + } + ], + ) + + result = await service.synthesize_procedural("u1") + + assert result == {"status": "synthesized", "procedures_created": 1, "procedures_skipped": 0} + procedure = next(doc for doc in store.docs if doc.get("type") == "procedural") + assert procedure["status"] == "candidate" + assert procedure["source_kind"] == "episode_distillation" + + +@pytest.mark.asyncio +async def test_synthesize_procedural_ungrounded_claim_is_candidate() -> None: + store = _AsyncProceduralStore([_fact()]) + service = _service( + store, + [ + { + "procedures": [ + _procedure( + "Ungrounded org policy", + grounded_in=["nonexistent"], + source_kind="organization_policy", + ) + ] + } + ], + ) + + result = await service.synthesize_procedural("u1") + + assert result == {"status": "synthesized", "procedures_created": 1, "procedures_skipped": 0} + procedure = next(doc for doc in store.docs if doc.get("type") == "procedural") + assert procedure["status"] == "candidate" + + +@pytest.mark.asyncio +async def test_synthesize_procedural_skips_malformed_workflow_and_creates_valid_sibling() -> None: + malformed = _procedure( + "Empty workflow", + grounded_in=["fact-1"], + source_kind="explicit_user_instruction", + summary="Malformed workflow with no steps.", + ) + malformed["procedure_kind"] = "workflow" + valid = _procedure( + "Valid policy", + grounded_in=["fact-1"], + source_kind="explicit_user_instruction", + ) + store = _AsyncProceduralStore([_fact()]) + service = _service(store, [{"procedures": [malformed, valid]}]) + + result = await service.synthesize_procedural("u1") + + assert result == {"status": "synthesized", "procedures_created": 1, "procedures_skipped": 1} + procedures = [doc for doc in store.docs if doc.get("type") == "procedural"] + assert [doc["name"] for doc in procedures] == ["Valid policy"] + + +@pytest.mark.asyncio +async def test_synthesize_procedural_quarantines_retryable_and_non_retryable_llm_errors() -> None: + retryable_service = _service(_AsyncProceduralStore([_fact()])) + + async def raise_retryable(filename: str, inputs: dict[str, Any]) -> str: + del filename, inputs + raise RuntimeError("rate limit") + + retryable_service._run_prompty = raise_retryable # type: ignore[method-assign] + + retryable_result = await retryable_service.synthesize_procedural("u1") + + assert retryable_result == {"status": "deferred", "procedures_created": 0} + + non_retryable_service = _service(_AsyncProceduralStore([_fact()])) + + async def raise_non_retryable(filename: str, inputs: dict[str, Any]) -> str: + del filename, inputs + raise RuntimeError("content_filter") + + non_retryable_service._run_prompty = raise_non_retryable # type: ignore[method-assign] + + non_retryable_result = await non_retryable_service.synthesize_procedural("u1") + + assert non_retryable_result == {"status": "skipped", "procedures_created": 0} + + +@pytest.mark.asyncio +async def test_synthesize_procedural_maps_multi_procedure_lineage() -> None: + store = _AsyncProceduralStore([_fact(), _episode()]) + service = _service( + store, + [ + { + "procedures": [ + _procedure( + "Fact lineage", + grounded_in=["fact-1"], + source_kind="explicit_user_instruction", + ), + _procedure( + "Episode lineage", + grounded_in=["ep-1"], + source_kind="episode_distillation", + summary="Retry transient CI failures once before escalating.", + ), + ] + } + ], + ) + + result = await service.synthesize_procedural("u1") + + assert result == {"status": "synthesized", "procedures_created": 2, "procedures_skipped": 0} + procedures = {doc["name"]: doc for doc in store.docs if doc.get("type") == "procedural"} + assert procedures["Fact lineage"]["source_fact_ids"] == ["fact-raw-1"] + assert procedures["Fact lineage"]["source_episodic_ids"] == [] + assert procedures["Episode lineage"]["source_fact_ids"] == [] + assert procedures["Episode lineage"]["source_episodic_ids"] == ["episode-raw-1"] diff --git a/tests/unit/aio/store/test_memory_store.py b/tests/unit/aio/store/test_memory_store.py index e0626de..4f8fc27 100644 --- a/tests/unit/aio/store/test_memory_store.py +++ b/tests/unit/aio/store/test_memory_store.py @@ -324,7 +324,7 @@ async def test_search_adds_created_time_range_filters(): assert params["@created_before"] == "2026-03-01T00:00:00+00:00" -async def test_add_cosmos_routes_by_type(): +async def test_upsert_memory_routes_by_type(): turns = MagicMock() memories = MagicMock() summaries = MagicMock() @@ -333,7 +333,7 @@ async def test_add_cosmos_routes_by_type(): store = AsyncMemoryStore(containers=_containers(turns=turns, memories=memories, summaries=summaries)) for memory_type in ("turn", "fact", "episodic", "procedural", "thread_summary", "user_summary"): - await store.add_cosmos(_doc(id=f"{memory_type}_id", type=memory_type)) + await store.upsert_memory(_doc(id=f"{memory_type}_id", type=memory_type)) assert turns.upsert_item.await_count == 1 assert memories.upsert_item.await_count == 3 diff --git a/tests/unit/aio/test_cosmos_memory_client.py b/tests/unit/aio/test_cosmos_memory_client.py index ae58496..ff223dc 100644 --- a/tests/unit/aio/test_cosmos_memory_client.py +++ b/tests/unit/aio/test_cosmos_memory_client.py @@ -546,11 +546,11 @@ async def test_validate_topology_raises_when_not_connected(self): class TestAddCosmos: - async def test_add_cosmos(self): + async def test_upsert_memory(self): mem, container = _connected_client() # Suppress the background cadence task to keep the test focused on the CRUD write. mem._maybe_auto_trigger = AsyncMock() - await mem.add_cosmos(user_id="u1", role="user", content="hello", thread_id="t1") + await mem.upsert_memory(user_id="u1", role="user", content="hello", thread_id="t1") # Drain any pending background tasks (none expected since we stubbed the trigger). await asyncio.gather(*list(mem._background_tasks), return_exceptions=True) @@ -560,37 +560,37 @@ async def test_add_cosmos(self): assert body["content"] == "hello" assert body["user_id"] == "u1" - async def test_add_cosmos_not_connected(self): + async def test_upsert_memory_not_connected(self): mem = _make_client() with pytest.raises(CosmosNotConnectedError): - await mem.add_cosmos(user_id="u1", role="user", content="hi", thread_id="t1") + await mem.upsert_memory(user_id="u1", role="user", content="hi", thread_id="t1") - async def test_add_cosmos_turn_requires_thread_id(self): + async def test_upsert_memory_turn_requires_thread_id(self): """Turn writes must declare a thread_id so the auto-trigger counter can group them.""" mem, _ = _connected_client() with pytest.raises(ValidationError, match="thread_id is required"): - await mem.add_cosmos(user_id="u1", role="user", content="hi") # memory_type='turn' default + await mem.upsert_memory(user_id="u1", role="user", content="hi") # memory_type='turn' default - async def test_add_cosmos_non_turn_does_not_require_thread_id(self): + async def test_upsert_memory_non_turn_does_not_require_thread_id(self): """Non-turn writes (facts, episodics, etc.) work without thread_id and skip cadence.""" mem, container = _connected_client() trigger = AsyncMock() mem._maybe_auto_trigger = trigger - await mem.add_cosmos(user_id="u1", role="user", content="prefers dark mode", memory_type="fact") + await mem.upsert_memory(user_id="u1", role="user", content="prefers dark mode", memory_type="fact") await asyncio.gather(*list(mem._background_tasks), return_exceptions=True) container.upsert_item.assert_awaited_once() trigger.assert_not_awaited() - async def test_add_cosmos_turn_schedules_cadence(self): + async def test_upsert_memory_turn_schedules_cadence(self): """A turn write must schedule the auto-trigger as a background task so cadence env vars apply whether the caller uses the local buffer or writes through directly.""" mem, _ = _connected_client() trigger = AsyncMock() mem._maybe_auto_trigger = trigger - await mem.add_cosmos(user_id="u1", role="user", content="hello", thread_id="t1") + await mem.upsert_memory(user_id="u1", role="user", content="hello", thread_id="t1") # Drain the background task so the AsyncMock records the call. await asyncio.gather(*list(mem._background_tasks), return_exceptions=True) @@ -723,7 +723,7 @@ async def test_success(self): container.read_item = AsyncMock(return_value=_make_doc(id="m1", type="fact")) container.delete_item = AsyncMock() - await mem.delete_cosmos(memory_id="m1", user_id="u1", thread_id="t1", memory_type="fact") + await mem.delete_memory(memory_id="m1", user_id="u1", thread_id="t1", memory_type="fact") container.delete_item.assert_awaited_once_with(item="m1", partition_key=["u1", "t1"]) @@ -735,7 +735,7 @@ async def test_not_found(self): container.delete_item = AsyncMock() with pytest.raises(MemoryNotFoundError): - await mem.delete_cosmos(memory_id="x", user_id="u1", thread_id="t1", memory_type="fact") + await mem.delete_memory(memory_id="x", user_id="u1", thread_id="t1", memory_type="fact") container.delete_item.assert_not_awaited() @@ -781,7 +781,7 @@ async def test_cosmos_ops_without_connect(self): with pytest.raises(CosmosNotConnectedError): await mem.update_cosmos(memory_id="m1", user_id="u1", thread_id="t1", memory_type="fact") with pytest.raises(CosmosNotConnectedError): - await mem.delete_cosmos(memory_id="m1", user_id="u1", thread_id="t1", memory_type="fact") + await mem.delete_memory(memory_id="m1", user_id="u1", thread_id="t1", memory_type="fact") # =================================================================== @@ -1079,3 +1079,60 @@ async def test_search_summaries_delegates_to_store(self): assert out == [{"content": "s", "type": "user_summary"}] store.search_summaries.assert_awaited_once() + + +class TestAsyncDeleteHelpers: + @pytest.mark.asyncio + async def test_delete_turn_delegates_with_turn_type(self): + mem = _make_client() + mem.delete_memory = AsyncMock() + await mem.delete_turn("turn-1", user_id="u1", thread_id="t1") + mem.delete_memory.assert_awaited_once_with("turn-1", user_id="u1", thread_id="t1", memory_type="turn") + + @pytest.mark.asyncio + async def test_delete_thread_summary_returns_true_then_false(self): + mem = _make_client() + mem.delete_memory = AsyncMock() + assert await mem.delete_thread_summary("u1", "t1") is True + mem.delete_memory.assert_awaited_once_with( + "summary_u1_t1", user_id="u1", thread_id="t1", memory_type="thread_summary" + ) + mem.delete_memory = AsyncMock(side_effect=MemoryNotFoundError(memory_id="x", user_id="u1", thread_id="t1")) + assert await mem.delete_thread_summary("u1", "t1") is False + + @pytest.mark.asyncio + async def test_delete_user_summary_uses_deterministic_id_and_scope(self): + mem = _make_client() + mem.delete_memory = AsyncMock() + assert await mem.delete_user_summary("u1") is True + mem.delete_memory.assert_awaited_once_with( + "user_summary_u1", user_id="u1", thread_id="__user_summary__", memory_type="user_summary" + ) + + @pytest.mark.asyncio + async def test_delete_thread_deletes_turns_and_summary_and_counts(self): + mem = _make_client() + mem.get_thread = AsyncMock(return_value=[{"id": "turn-1"}, {"id": "turn-2"}]) + mem.delete_memory = AsyncMock() + deleted = await mem.delete_thread("u1", "t1") + assert deleted == 3 # 2 turns + 1 summary + assert mem.get_thread.call_args.kwargs["include_superseded"] is True + + @pytest.mark.asyncio + async def test_delete_thread_skips_missing_turn(self): + mem = _make_client() + mem.get_thread = AsyncMock(return_value=[{"id": "turn-1"}, {"id": "turn-2"}]) + + async def _delete(memory_id, **kwargs): + if memory_id == "turn-1": + raise MemoryNotFoundError(memory_id="turn-1", user_id="u1", thread_id="t1") + + mem.delete_memory = AsyncMock(side_effect=_delete) + deleted = await mem.delete_thread("u1", "t1", include_summary=False) + assert deleted == 1 + + @pytest.mark.asyncio + async def test_delete_thread_requires_ids(self): + mem = _make_client() + with pytest.raises(ValidationError): + await mem.delete_thread("", "t1") diff --git a/tests/unit/aio/test_procedural_synthesis.py b/tests/unit/aio/test_procedural_synthesis.py index 513208f..21fa892 100644 --- a/tests/unit/aio/test_procedural_synthesis.py +++ b/tests/unit/aio/test_procedural_synthesis.py @@ -1,178 +1,206 @@ -"""Async tests for procedural synthesis and procedural prompt retrieval. - -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 -correctly, that the durable-processor branch short-circuits, and that -the store-backed procedural reads work over async iterators. -""" +"""Async tests for atomic procedural synthesis and compiled procedural context.""" from __future__ import annotations -from unittest.mock import AsyncMock, MagicMock +from typing import Any import pytest - -from azure.cosmos.agent_memory.aio.cosmos_memory_client import AsyncCosmosMemoryClient -from azure.cosmos.agent_memory.aio.processors import AsyncDurableFunctionProcessor - - -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 +from azure.cosmos.exceptions import CosmosResourceExistsError + +from azure.cosmos.agent_memory.aio.services.pipeline import AsyncPipelineService +from tests.unit.services.test_extract_dry import ( + _async_containers_for_store, + _AsyncChat, + _AsyncEmbeddings, + _AsyncStore, +) + + +class _AsyncProceduralStore(_AsyncStore): + async def query(self, sql: str, parameters=None, partition_key=None, cross_partition: bool = False): + del partition_key, cross_partition + params = {p["name"]: p["value"] for p in (parameters or [])} + user_id = params.get("@uid", params.get("@user_id")) + memory_type = params.get("@type", params.get("@memory_type")) + docs = [dict(doc) for doc in self.docs] + if user_id is not None: + docs = [doc for doc in docs if doc.get("user_id") == user_id] + if memory_type is not None: + docs = [doc for doc in docs if doc.get("type") == memory_type] + if "c.status='active'" in sql: + docs = [doc for doc in docs if doc.get("status") == "active"] + if "superseded_by" in sql: + docs = [doc for doc in docs if not doc.get("superseded_by")] + return docs + + async def create_item(self, *, body: dict[str, Any]) -> dict[str, Any]: + if any(doc.get("id") == body.get("id") for doc in self.docs): + raise CosmosResourceExistsError(message="conflict") + self.docs.append(dict(body)) + return dict(body) + + +def _service( + store: _AsyncProceduralStore, + responses: list[dict[str, Any]] | None = None, +) -> AsyncPipelineService: + return AsyncPipelineService( + store, + _AsyncChat(responses or []), + _AsyncEmbeddings(), + containers=_async_containers_for_store(store), + ) -def _procedural_doc( - doc_id: str, - *, - version: int, - content: str, - source_fact_ids: list[str], - source_episodic_ids: list[str], - superseded_by: str | None = None, - ts: int = 0, - etag: str = "etag-1", -) -> dict: - doc = { - "id": doc_id, +def _fact() -> dict[str, Any]: + return { + "id": "fact-raw-1", "user_id": "u1", - "thread_id": "__procedural__", - "type": "procedural", - "version": version, - "content": content, - "source_fact_ids": list(source_fact_ids), - "source_episodic_ids": list(source_episodic_ids), - "supersedes_ids": [], - "created_at": f"2025-01-0{version}T00:00:00+00:00", - "role": "system", - "tags": ["sys:procedural", "sys:synthesized"], - "_etag": etag, - "_ts": ts, + "type": "fact", + "content": "The user explicitly said to run targeted tests before reporting success.", + "metadata": {"category": "preference"}, + "salience": 0.9, + "created_at": "2025-01-01T00:00:00+00:00", } - if superseded_by is not None: - doc["superseded_by"] = superseded_by - return doc - - -def _make_client(*, processor=None) -> AsyncCosmosMemoryClient: - client = AsyncCosmosMemoryClient(use_default_credential=False, processor=processor) - client._memories_container_client = MagicMock() - client._turns_container_client = client._memories_container_client - client._summaries_container_client = client._memories_container_client - return client - -@pytest.mark.asyncio -async def test_async_synthesize_procedural_awaits_async_pipeline(): - """The client must ``await`` ``AsyncPipelineService.synthesize_procedural`` - directly (no ``asyncio.to_thread`` indirection) and forward force=True.""" - client = _make_client() - pipeline = AsyncMock() - expected = {"status": "synthesized", "procedural": {"id": "proc_u1_1", "version": 1}} - pipeline.synthesize_procedural.return_value = expected - pipeline._store = client._get_store() - pipeline._containers = dict(client._containers) - client._pipeline = pipeline - - result = await client.synthesize_procedural("u1", force=True) - - assert result == expected - pipeline.synthesize_procedural.assert_awaited_once_with("u1", force=True) +def _episode() -> dict[str, Any]: + return { + "id": "episode-raw-1", + "user_id": "u1", + "type": "episodic", + "content": "A retry investigation succeeded.", + "lessons": ["Retry transient CI failures once before escalating."], + "salience": 0.8, + "created_at": "2025-01-01T00:01:00+00:00", + } -@pytest.mark.asyncio -async def test_async_get_procedural_prompt_returns_none_when_missing(): - client = _make_client() - client._memories_container_client.query_items = MagicMock(return_value=AsyncIterator([])) - assert await client.get_procedural_prompt("u1") is None +def _procedure( + name: str, + *, + grounded_in: list[str], + source_kind: str, + summary: str = "Run targeted tests before reporting success.", +) -> dict[str, Any]: + return { + "name": name, + "summary": summary, + "retrieval_text": summary, + "procedure_kind": "behavioral_policy", + "scope_type": "user", + "scope_value": None, + "activation_conditions": [], + "preconditions": [], + "steps": [], + "success_conditions": [], + "failure_conditions": [], + "safety_constraints": [], + "source_kind": source_kind, + "grounded_in": grounded_in, + "confidence": 0.8, + } @pytest.mark.asyncio -async def test_async_get_procedural_prompt_returns_active_content(): - active_doc = _procedural_doc( - "proc_u1_2", - version=2, - content="Active prompt", - source_fact_ids=["f1"], - source_episodic_ids=["e1"], - ts=2, +async def test_synthesize_procedural_extracts_atomic_procedures_and_gates_provenance() -> None: + store = _AsyncProceduralStore([_fact(), _episode()]) + service = _service( + store, + [ + { + "procedures": [ + _procedure( + "Targeted testing", + grounded_in=["fact-1"], + source_kind="explicit_user_instruction", + ), + _procedure( + "Retry CI failures", + grounded_in=["ep-1"], + source_kind="episode_distillation", + summary="Retry transient CI failures once before escalating.", + ), + ] + } + ], ) - superseded_doc = _procedural_doc( - "proc_u1_1", - version=1, - content="Old prompt", - source_fact_ids=["f1"], - source_episodic_ids=["e1"], - superseded_by="proc_u1_2", - ts=1, - ) - docs = [superseded_doc, active_doc] - client = _make_client() - - def _query_items(**kwargs): - query = kwargs["query"] - if "superseded_by" in query: - return AsyncIterator([doc for doc in docs if not doc.get("superseded_by")]) - return AsyncIterator(docs) - client._memories_container_client.query_items = MagicMock(side_effect=_query_items) + result = await service.synthesize_procedural("u1") - assert await client.get_procedural_prompt("u1") == "Active prompt" + assert result == {"status": "synthesized", "procedures_created": 2, "procedures_skipped": 0} + procedures = [doc for doc in store.docs if doc.get("type") == "procedural"] + assert len(procedures) == 2 + assert {doc["name"]: doc["status"] for doc in procedures} == { + "Targeted testing": "active", + "Retry CI failures": "candidate", + } + assert {doc["name"]: doc["source_kind"] for doc in procedures} == { + "Targeted testing": "explicit_user_instruction", + "Retry CI failures": "episode_distillation", + } + for doc in procedures: + assert doc["id"].startswith("proc_") + assert doc["type"] == "procedural" + assert doc["name"] + assert doc["retrieval_text"] + assert doc["thread_id"] == "__procedural__" + assert doc["embedding"] == [1.0] + # utility_score is seeded from the LLM confidence (0.8), not hard-wired to 0.5. + assert doc["utility_score"] == 0.8 @pytest.mark.asyncio -async def test_async_get_procedural_history_orders_active_first_then_newest_versions(): - v1 = _procedural_doc( - "proc_u1_1", - version=1, - content="v1", - source_fact_ids=["f1"], - source_episodic_ids=["e1"], - superseded_by="proc_u1_2", - ts=1, - ) - v2 = _procedural_doc( - "proc_u1_2", - version=2, - content="v2", - source_fact_ids=["f1", "f2"], - source_episodic_ids=["e1"], - superseded_by="proc_u1_3", - ts=2, - ) - v3 = _procedural_doc( - "proc_u1_3", - version=3, - content="v3", - source_fact_ids=["f1", "f2", "f3"], - source_episodic_ids=["e1"], - ts=3, - ) - client = _make_client() - client._memories_container_client.query_items = MagicMock(return_value=AsyncIterator([v1, v3, v2])) +async def test_synthesize_procedural_is_idempotent_by_scope_and_name() -> None: + store = _AsyncProceduralStore([_fact()]) + response = { + "procedures": [ + _procedure( + "Targeted testing", + grounded_in=["fact-1"], + source_kind="explicit_user_instruction", + ) + ] + } + service = _service(store, [response, response]) - history = await client.get_procedural_history("u1", limit=10) + first = await service.synthesize_procedural("u1") + second = await service.synthesize_procedural("u1") - assert [doc["id"] for doc in history] == ["proc_u1_3", "proc_u1_2", "proc_u1_1"] + assert first == {"status": "synthesized", "procedures_created": 1, "procedures_skipped": 0} + assert second == {"status": "synthesized", "procedures_created": 0, "procedures_skipped": 1} + procedures = [doc for doc in store.docs if doc.get("type") == "procedural"] + assert len(procedures) == 1 @pytest.mark.asyncio -async def test_async_client_synthesize_procedural_raises_for_remote_processors(): - client = _make_client(processor=AsyncDurableFunctionProcessor()) - client._pipeline = AsyncMock() +async def test_build_procedural_context_uses_active_procedures_only() -> None: + active = { + "id": "proc-active", + "user_id": "u1", + "type": "procedural", + "status": "active", + "name": "Targeted testing", + "summary": "Run targeted tests before reporting success.", + "retrieval_text": "tests success", + "procedure_kind": "behavioral_policy", + "scope_type": "user", + "scope_value": None, + "priority": 10, + "source_authority": "high", + "version": 1, + } + candidate = { + **active, + "id": "proc-candidate", + "status": "candidate", + "name": "Candidate policy", + "summary": "Do not include this candidate procedure.", + } + service = _service(_AsyncProceduralStore([active, candidate])) - with pytest.raises(NotImplementedError, match="durable mode"): - await client.synthesize_procedural("u1") + context = await service.build_procedural_context("u1") - client._pipeline.synthesize_procedural.assert_not_called() + assert "Run targeted tests before reporting success." in context + assert "Do not include this candidate procedure." not in context + assert await service.build_procedural_context("missing-user") == "" diff --git a/tests/unit/aio/test_process_now.py b/tests/unit/aio/test_process_now.py index 83d8be5..7f4a536 100644 --- a/tests/unit/aio/test_process_now.py +++ b/tests/unit/aio/test_process_now.py @@ -31,7 +31,7 @@ def _patch_get_thread(client, turns): async def test_process_now_with_inprocess_invokes_full_pipeline(): """process_now must fire ALL FIVE steps for AsyncInProcess: thread_summary, extract, reconcile, procedural, user_summary. Pre-fix this was only the first 3, so - procedural + user_summary never ran when callers used add_cosmos + process_now.""" + procedural + user_summary never ran when callers used upsert_memory + process_now.""" client = _connected() pipeline = AsyncMock() pipeline.generate_thread_summary.return_value = {"id": "s"} diff --git a/tests/unit/aio/test_reconcile_telemetry.py b/tests/unit/aio/test_reconcile_telemetry.py index ef2fdbd..858e328 100644 --- a/tests/unit/aio/test_reconcile_telemetry.py +++ b/tests/unit/aio/test_reconcile_telemetry.py @@ -17,14 +17,6 @@ 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/function_app/test_change_feed.py b/tests/unit/function_app/test_change_feed.py index 18b6dbb..899babb 100644 --- a/tests/unit/function_app/test_change_feed.py +++ b/tests/unit/function_app/test_change_feed.py @@ -13,7 +13,7 @@ import pytest import triggers.change_feed as change_feed_module -from azure.cosmos.exceptions import CosmosResourceNotFoundError +from azure.cosmos.exceptions import CosmosHttpResponseError, CosmosResourceNotFoundError from triggers.change_feed import process_changefeed_batch @@ -73,8 +73,18 @@ async def upsert_item(*, body, **_kwargs): state[body["id"]] = dict(body) return body - async def patch_item(*, item, partition_key, patch_operations): + async def patch_item(*, item, partition_key, patch_operations, filter_predicate=None): doc = state.setdefault(item, {"id": item}) + if filter_predicate is not None: + # Honor the monotonic watermark guard: apply only when the new + # last_extract_count exceeds the stored one (else Cosmos returns 412). + new_val = next( + (op["value"] for op in patch_operations if op.get("path") == "/last_extract_count"), + None, + ) + cur = doc.get("last_extract_count") + if new_val is not None and isinstance(cur, int) and cur >= new_val: + raise CosmosHttpResponseError(status_code=412, message="precondition failed") for op in patch_operations: doc[op["path"].lstrip("/")] = op["value"] return dict(doc) @@ -164,6 +174,7 @@ async def run_twice(): { "THREAD_SUMMARY_EVERY_N": "0", "FACT_EXTRACTION_EVERY_N": "0", + "EPISODE_EVAL_EVERY_N": "0", "USER_SUMMARY_EVERY_N": "0", }, clear=False, @@ -186,7 +197,7 @@ def test_all_disabled_skips_everything(): def test_unset_thresholds_apply_documented_defaults(monkeypatch): """When env vars are unset, thresholds fall back to documented defaults - (fact=1, thread=10, user=20), not 0. + (fact=2, thread=10, user=20), not 0. Regression for the silent-no-op out-of-the-box deploy bug: a missing setting should NOT disable the orchestrator (only an explicit "0" does). @@ -197,10 +208,10 @@ def test_unset_thresholds_apply_documented_defaults(monkeypatch): starter = _make_starter() container = _make_counter_container_starting_at() - # 1 turn => crosses fact extraction (n=1) only. + # 2 turns => crosses fact extraction (n=2) only. asyncio.run( process_changefeed_batch( - [_turn()], + [_turn(), _turn()], starter, counter_container=container, ) @@ -308,6 +319,7 @@ def test_turn_doc_missing_ids_is_skipped(): { "THREAD_SUMMARY_EVERY_N": "4", "FACT_EXTRACTION_EVERY_N": "4", + "EPISODE_EVAL_EVERY_N": "4", "USER_SUMMARY_EVERY_N": "20", }, clear=False, @@ -329,10 +341,40 @@ def test_thread_threshold_crossing_starts_summary_and_extract(): started = {(call.args[0], call.kwargs["instance_id"]) for call in starter.start_new.await_args_list} assert ("ThreadSummaryOrchestrator", "thread_summary:u1:t1:4") in started assert ("ExtractMemoriesOrchestrator", "extract:u1:t1:4") in started + assert ("ExtractEpisodesOrchestrator", "episode:u1:t1:4") in started # User threshold 20 not crossed by 4 turns. assert not any(name == "UserSummaryOrchestrator" for name, _ in started) +@patch.dict( + os.environ, + { + "THREAD_SUMMARY_EVERY_N": "0", + "FACT_EXTRACTION_EVERY_N": "0", + "EPISODE_EVAL_EVERY_N": "4", + "USER_SUMMARY_EVERY_N": "0", + }, + clear=False, +) +def test_episode_threshold_crossing_starts_extract_episodes(): + starter = _make_starter() + container = _make_counter_container_starting_at() + + asyncio.run( + process_changefeed_batch( + [_turn() for _ in range(4)], + starter, + counter_container=container, + ) + ) + + episode_calls = [c for c in starter.start_new.await_args_list if c.args[0] == "ExtractEpisodesOrchestrator"] + assert len(episode_calls) == 1 + call = episode_calls[0] + assert call.kwargs["instance_id"] == "episode:u1:t1:4" + assert call.kwargs["client_input"] == {"user_id": "u1", "thread_id": "t1", "count": 4} + + @patch.dict( os.environ, { @@ -374,6 +416,7 @@ def test_user_threshold_crossing_starts_user_summary(): { "THREAD_SUMMARY_EVERY_N": "0", # disabled "FACT_EXTRACTION_EVERY_N": "4", # enabled + "EPISODE_EVAL_EVERY_N": "0", # disabled "USER_SUMMARY_EVERY_N": "0", # disabled }, clear=False, @@ -395,6 +438,33 @@ def test_disabled_thread_summary_does_not_start_summary_orchestrator(): assert started == ["ExtractMemoriesOrchestrator"] +@patch.dict( + os.environ, + { + "THREAD_SUMMARY_EVERY_N": "0", + "FACT_EXTRACTION_EVERY_N": "4", + "EPISODE_EVAL_EVERY_N": "0", + "USER_SUMMARY_EVERY_N": "0", + }, + clear=False, +) +def test_disabled_episode_does_not_start_extract_episodes(): + starter = _make_starter() + container = _make_counter_container_starting_at() + + asyncio.run( + process_changefeed_batch( + [_turn() for _ in range(4)], + starter, + counter_container=container, + ) + ) + + started = [call.args[0] for call in starter.start_new.await_args_list] + assert "ExtractMemoriesOrchestrator" in started + assert "ExtractEpisodesOrchestrator" not in started + + @patch.dict( os.environ, { @@ -648,6 +718,26 @@ def test_reconcile_flag_set_only_when_n_facts_times_n_dedup_threshold_crosses(): assert payload.get("recent_k") == 1 +def test_advance_extract_watermark_is_monotonic(): + """A later, out-of-order run must not regress the extract watermark below a + higher value a concurrent run already wrote (that would re-extract turns).""" + from shared.counters import advance_extract_watermark, thread_counter_id + + container = _make_counter_container_starting_at() + cid = thread_counter_id("u1", "t1") + + asyncio.run(advance_extract_watermark(container, cid, "u1", "t1", 8)) + assert container._state[cid]["last_extract_count"] == 8 + + # An out-of-order run with a lower count is a no-op (Cosmos 412, swallowed). + asyncio.run(advance_extract_watermark(container, cid, "u1", "t1", 4)) + assert container._state[cid]["last_extract_count"] == 8 + + # A higher count still advances. + asyncio.run(advance_extract_watermark(container, cid, "u1", "t1", 12)) + assert container._state[cid]["last_extract_count"] == 12 + + @patch.dict( os.environ, { diff --git a/tests/unit/function_app/test_config_episode.py b/tests/unit/function_app/test_config_episode.py new file mode 100644 index 0000000..572ed1b --- /dev/null +++ b/tests/unit/function_app/test_config_episode.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +import pytest +from shared import config + + +@pytest.mark.parametrize( + ("env_name", "getter_name", "expected"), + [ + ("EPISODE_EVAL_EVERY_N", "get_episode_eval_every_n", 4), + ("EPISODE_IDLE_GAP_SECONDS", "get_episode_idle_gap_seconds", 1800), + ("EPISODE_TOPIC_DRIFT", "get_episode_topic_drift", 0.0), + ("EPISODE_MAX_TURNS", "get_episode_max_turns", 40), + ("EPISODE_MIN_TURNS", "get_episode_min_turns", 2), + ], +) +def test_episode_config_getters_defaults( + monkeypatch: pytest.MonkeyPatch, + env_name: str, + getter_name: str, + expected: int | float, +) -> None: + monkeypatch.delenv(env_name, raising=False) + + assert getattr(config, getter_name)() == expected + + +@pytest.mark.parametrize( + ("env_name", "getter_name", "raw", "expected"), + [ + ("EPISODE_EVAL_EVERY_N", "get_episode_eval_every_n", "6", 6), + ("EPISODE_IDLE_GAP_SECONDS", "get_episode_idle_gap_seconds", "2400", 2400), + ("EPISODE_TOPIC_DRIFT", "get_episode_topic_drift", "0.35", 0.35), + ("EPISODE_MAX_TURNS", "get_episode_max_turns", "50", 50), + ("EPISODE_MIN_TURNS", "get_episode_min_turns", "3", 3), + ], +) +def test_episode_config_getters_parse_env( + monkeypatch: pytest.MonkeyPatch, + env_name: str, + getter_name: str, + raw: str, + expected: int | float, +) -> None: + monkeypatch.setenv(env_name, raw) + + assert getattr(config, getter_name)() == expected diff --git a/tests/unit/function_app/test_extract_episodes_orchestrator.py b/tests/unit/function_app/test_extract_episodes_orchestrator.py new file mode 100644 index 0000000..abcbc7e --- /dev/null +++ b/tests/unit/function_app/test_extract_episodes_orchestrator.py @@ -0,0 +1,73 @@ +"""Unit tests for the episodic-extraction Durable orchestrator.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +from orchestrators import extract_episodes as ee_mod + + +def _user_function(builder): + if hasattr(builder, "_function"): + return builder._function.get_user_function().orchestrator_function + return builder + + +def _make_context(payload): + ctx = MagicMock() + ctx.get_input.return_value = payload + + yielded_calls: list[tuple] = [] + + def call_activity_with_retry(name, retry, activity_payload): + yielded_calls.append((name, retry, activity_payload)) + return ("__call__", name, activity_payload) + + ctx.call_activity_with_retry.side_effect = call_activity_with_retry + ctx._yielded_calls = yielded_calls + return ctx + + +def _drive(gen, activity_results): + yields = [] + iterator = iter(activity_results) + try: + sent = None + while True: + value = gen.send(sent) + yields.append(value) + sent = next(iterator) + except StopIteration as stop: + return stop.value, yields + + +class TestExtractEpisodesOrchestrator: + def _orchestrator(self): + return _user_function(ee_mod.ExtractEpisodesOrchestrator) + + @patch.object(ee_mod, "default_retry_options", return_value=MagicMock(name="retry")) + def test_calls_activity_once_with_user_and_thread_and_returns_result(self, _retry): + ctx = _make_context({"user_id": "u1", "thread_id": "t1"}) + gen = self._orchestrator()(ctx) + + result, _ = _drive(gen, [{"episodes": 2}]) + + assert [call[0] for call in ctx._yielded_calls] == ["ee_ExtractEpisodes"] + assert ctx._yielded_calls[0][2] == {"user_id": "u1", "thread_id": "t1"} + assert result == {"episodes": 2} + + +@patch.object(ee_mod, "get_pipeline") +def test_activity_calls_pipeline_and_returns_slim_payload(mock_get_pipeline): + pipeline = MagicMock() + pipeline.extract_episodes.return_value = {"episodes": 3} + mock_get_pipeline.return_value = pipeline + + result = ee_mod.ee_ExtractEpisodes({"user_id": "u1", "thread_id": "t1"}) + + pipeline.extract_episodes.assert_called_once_with( + user_id="u1", + thread_id="t1", + flush=False, + ) + assert result == {"episodes": 3} diff --git a/tests/unit/function_app/test_orchestrators.py b/tests/unit/function_app/test_orchestrators.py index 14b149e..582ba18 100644 --- a/tests/unit/function_app/test_orchestrators.py +++ b/tests/unit/function_app/test_orchestrators.py @@ -193,7 +193,6 @@ 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, @@ -202,15 +201,11 @@ def test_extract_only_when_reconcile_flag_absent(self, _retry): ], ) - assert [c[0] for c in ctx._yielded_calls] == ["em_Extract", "em_Dedup", "em_Persist"] + assert [c[0] for c in ctx._yielded_calls] == ["em_Extract", "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 @@ -222,7 +217,6 @@ 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}, { @@ -234,8 +228,8 @@ def test_chains_reconcile_when_flag_true(self, _retry): ) names = [c[0] for c in ctx._yielded_calls] - assert names == ["em_Extract", "em_Dedup", "em_Persist", "em_ReconcileMemories"] - assert ctx._yielded_calls[3][2] == {"user_id": "u1"} + assert names == ["em_Extract", "em_Persist", "em_ReconcileMemories"] + assert ctx._yielded_calls[2][2] == {"user_id": "u1"} assert [s[0] for s in ctx._yielded_sub_orchestrators] == [ "SynthesizeProceduralOrchestrator", ] @@ -261,13 +255,11 @@ def boom_after_sub(name, retry, sub_payload, *args, **kwargs): gen = self._orchestrator()(ctx) # Yield 1: em_Extract gen.send(None) - # Yield 2: em_Dedup - gen.send({"facts": [{"id": "f1"}], "episodic": [], "updates": []}) - # Yield 3: em_Persist + # Yield 2: em_Persist gen.send({"facts": [{"id": "f1"}], "episodic": [], "updates": []}) - # Yield 4: em_ReconcileMemories + # Yield 3: em_ReconcileMemories gen.send({"fact_count": 2, "episodic_count": 0, "updated_count": 0}) - # Yield 5: SynthesizeProceduralOrchestrator - throw an exception + # Yield 4: SynthesizeProceduralOrchestrator - throw an exception gen.send( { "fact": {"kept": 0, "merged": 1, "contradicted": 0}, @@ -295,13 +287,12 @@ 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_Dedup", "em_Persist"] + assert [c[0] for c in ctx._yielded_calls] == ["em_Extract", "em_Persist"] assert ctx._yielded_sub_orchestrators == [] assert result["procedural"] is None @@ -309,7 +300,7 @@ def test_procedural_not_called_when_reconcile_skipped(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": []}, {"facts": []}, {"fact_count": 0}]) + _drive(gen, [{"facts": []}, {"fact_count": 0}]) extract_payload = ctx._yielded_calls[0][2] assert extract_payload == {"user_id": "u", "thread_id": "t"} @@ -318,21 +309,19 @@ def test_extract_payload_carries_user_thread_without_recent_k_when_absent(self, 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}]) + _drive(gen, [{"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): + def test_extract_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}]) + _drive(gen, [extracted, {"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): @@ -346,17 +335,17 @@ def test_activity_failure_propagates(self, _retry): 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]) + _drive(gen, [{"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} + assert names == ["em_Extract", "em_Persist", "em_AdvanceExtractWatermark"] + assert ctx._yielded_calls[2][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}]) + _drive(gen, [{"facts": []}, {"fact_count": 0}]) names = [c[0] for c in ctx._yielded_calls] assert "em_AdvanceExtractWatermark" not in names @@ -393,28 +382,6 @@ def test_em_extract_falls_back_to_max_batch_size_when_recent_k_absent(self): pipeline.extract_memories_durable.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 = [ diff --git a/tests/unit/function_app/test_synthesize_procedural_orchestrator.py b/tests/unit/function_app/test_synthesize_procedural_orchestrator.py index 29c49ea..6cb9ebe 100644 --- a/tests/unit/function_app/test_synthesize_procedural_orchestrator.py +++ b/tests/unit/function_app/test_synthesize_procedural_orchestrator.py @@ -51,11 +51,11 @@ def test_calls_activity_once_with_user_and_force_and_returns_result(self, _retry ctx = _make_context({"user_id": "u1", "force": True}) gen = self._orchestrator()(ctx) - result, _ = _drive(gen, [{"status": "synthesized", "version": 3}]) + result, _ = _drive(gen, [{"status": "synthesized", "procedures_created": 2}]) assert [call[0] for call in ctx._yielded_calls] == ["sp_SynthesizeProcedural"] assert ctx._yielded_calls[0][2] == {"user_id": "u1", "force": True} - assert result == {"status": "synthesized", "version": 3} + assert result == {"status": "synthesized", "procedures_created": 2} @pytest.mark.parametrize( @@ -63,13 +63,13 @@ def test_calls_activity_once_with_user_and_force_and_returns_result(self, _retry [ ( {"user_id": "u1", "force": True}, - {"status": "synthesized", "procedural": {"id": "proc_u1_3", "version": 3, "content": "Prompt"}}, - {"status": "synthesized", "version": 3}, + {"status": "synthesized", "procedures_created": 2}, + {"status": "synthesized", "procedures_created": 2}, ), ( {"user_id": "u2", "force": False}, - {"status": "unchanged", "procedural": None}, - {"status": "unchanged", "version": None}, + {"status": "unchanged", "procedures_created": 0}, + {"status": "unchanged", "procedures_created": 0}, ), ], ) @@ -86,4 +86,4 @@ def test_activity_calls_pipeline_and_returns_slim_payload(mock_get_pipeline, pay force=payload.get("force", False), ) assert result == expected - assert "procedural" not in result + assert "version" not in result diff --git a/tests/unit/processors/test_durable.py b/tests/unit/processors/test_durable.py index cf65591..1cf6426 100644 --- a/tests/unit/processors/test_durable.py +++ b/tests/unit/processors/test_durable.py @@ -2,12 +2,20 @@ from __future__ import annotations +import logging +from unittest.mock import MagicMock + +import pytest + +from azure.cosmos.agent_memory.cosmos_memory_client import CosmosMemoryClient from azure.cosmos.agent_memory.processors import ( DurableFunctionProcessor, ProcessThreadResult, UserSummaryResult, ) +OLD_EPISODIC_DURABLE_WARNING = "Episodic memory is not available under the Durable Functions backend" + def test_process_thread_returns_empty_result(): proc = DurableFunctionProcessor() @@ -30,6 +38,35 @@ def test_generate_user_summary_returns_empty_result(): assert result.summary is None +def test_process_extract_episodes_returns_empty_result_without_old_warning(caplog): + proc = DurableFunctionProcessor() + caplog.set_level(logging.WARNING) + + result = proc.process_extract_episodes(user_id="u1", thread_id="t1") + + assert result == {} + assert OLD_EPISODIC_DURABLE_WARNING not in caplog.text + + +def test_synthesize_procedural_is_noop(): + # Durable procedural synthesis runs in the Function app after reconcile; the + # processor no-ops (rather than raising) so the in-process auto-trigger does + # not stamp a spurious failure each cadence. + proc = DurableFunctionProcessor() + result = proc.synthesize_procedural(user_id="u1") + assert result == {"status": "skipped", "procedures_created": 0} + + +def test_client_extract_episodes_raises_for_durable_processor(): + client = CosmosMemoryClient(use_default_credential=False, processor=DurableFunctionProcessor()) + client._pipeline = MagicMock() + + with pytest.raises(NotImplementedError, match="Durable Function app"): + client.extract_episodes("u1", "t1", flush=True) + + client._pipeline.extract_episodes.assert_not_called() + + def test_close_is_noop(): assert DurableFunctionProcessor().close() is None diff --git a/tests/unit/services/test_chaos_extract_persist.py b/tests/unit/services/test_chaos_extract_persist.py index 9a40976..a0df1b7 100644 --- a/tests/unit/services/test_chaos_extract_persist.py +++ b/tests/unit/services/test_chaos_extract_persist.py @@ -11,18 +11,6 @@ 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]] = {} @@ -65,7 +53,7 @@ def read_item(self, item_id: str, partition_key: Any): del partition_key return self.container.docs[item_id] - def add_cosmos(self, record: dict[str, Any]) -> dict[str, Any]: + def upsert_memory(self, record: dict[str, Any]) -> dict[str, Any]: self.container.docs[record["id"]] = dict(record) return record @@ -81,8 +69,8 @@ async def query(self, sql: str, parameters=None, partition_key=None, cross_parti async def read_item(self, item_id: str, partition_key: Any): return super().read_item(item_id, partition_key) - async def add_cosmos(self, record: dict[str, Any]) -> dict[str, Any]: - return super().add_cosmos(record) + async def upsert_memory(self, record: dict[str, Any]) -> dict[str, Any]: + return super().upsert_memory(record) async def mark_superseded(self, old_doc: dict[str, Any], superseder_id: str, *, reason: str) -> bool: return super().mark_superseded(old_doc, superseder_id, reason=reason) diff --git a/tests/unit/services/test_dedup_vector.py b/tests/unit/services/test_dedup_vector.py deleted file mode 100644 index e2d58f6..0000000 --- a/tests/unit/services/test_dedup_vector.py +++ /dev/null @@ -1,438 +0,0 @@ -from __future__ import annotations - -import json -from typing import Any -from unittest.mock import MagicMock - -import pytest - -from azure.cosmos.agent_memory.services.pipeline import PipelineService - - -@pytest.fixture(autouse=True) -def _enable_vector_folding(monkeypatch: pytest.MonkeyPatch) -> None: - # DEDUP_VECTOR_ENABLED now defaults to False (add-only); this suite exercises - # the in-place folding path, so enable it. Tests that assert the flag-off - # behavior patch the getter directly and override this. - monkeypatch.setenv("DEDUP_VECTOR_ENABLED", "true") - - -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 {}, - ) - doc = { - "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, - } - if memory_type == "episodic": - doc.setdefault("title", content) - doc.setdefault("events", []) - doc.setdefault("participants", []) - doc.setdefault("lessons", []) - doc.setdefault("source_turn_ids", []) - return doc - - -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_apply_inplace_update_skips_cross_source_fold() -> None: - p = _make_pipeline() - neighbor = _doc( - "existing-1", - "same content", - tags=["sys:fact"], - metadata={"category": "preference", "source": "user"}, - ) - neighbor["_etag"] = "etag-xyz" - new_doc = _doc( - "f-new", - "same content", - embedding=[0.5, 0.5], - tags=["sys:fact", "sys:agent-fact"], - metadata={"category": "other", "source": "agent"}, - ) - - assert p._apply_inplace_update(neighbor, new_doc) is False - p._memories_container.replace_item.assert_not_called() - p._memories_container.upsert_item.assert_not_called() - - -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_episode_boundary.py b/tests/unit/services/test_episode_boundary.py index 97b9a4e..5140e5c 100644 --- a/tests/unit/services/test_episode_boundary.py +++ b/tests/unit/services/test_episode_boundary.py @@ -2,9 +2,10 @@ Episodes are finalized at detected boundaries in the open turn stream - an idle time-gap, a topic-drift shift, or a max-size cap - never on a fixed turn cadence -and never by the caller signaling "session end". Turns folded into an episode are -stamped ``episode_extracted_at`` (an independent watermark) so re-evaluation is -idempotent and episodes never duplicate. +and never by the caller signaling "session end". A per-thread ``(created_at, id)`` +cursor (stored as a single doc, not a per-turn stamp) advances past folded turns +so re-evaluation is idempotent and episodes never duplicate; because it never +writes to the turns container it cannot perturb the change-feed cadence counter. """ from __future__ import annotations @@ -67,8 +68,24 @@ def _episodes(store: _TrackingStore) -> list[dict[str, Any]]: return [doc for doc in store.docs if doc.get("type") == "episodic"] -def _stamped(turns_store: _Store) -> list[str]: - return sorted(t["id"] for t in turns_store.docs if t.get("episode_extracted_at")) +def _folded(memories: _TrackingStore, turns_store: _Store) -> list[str]: + """Turn ids folded into an episode, derived from the per-thread episodic + cursor doc (the watermark model advances a single cursor instead of stamping + each turn, so nothing is written to the turns container).""" + cursor = next((d for d in memories.docs if d.get("type") == "episode_cursor"), None) + if cursor is None: + return [] + last_at = str(cursor.get("last_episode_at") or "") + last_id = str(cursor.get("last_episode_id") or "") + return sorted( + str(t.get("id")) + for t in turns_store.docs + if t.get("type") == "turn" + and ( + (str(t.get("created_at") or "") < last_at) + or (str(t.get("created_at") or "") == last_at and str(t.get("id") or "") <= last_id) + ) + ) def test_time_gap_closes_prior_episode_and_leaves_tail_open(monkeypatch) -> None: @@ -84,7 +101,7 @@ def test_time_gap_closes_prior_episode_and_leaves_tail_open(monkeypatch) -> None assert result == {"episodes": 1} assert len(_episodes(memories)) == 1 # Only the pre-gap segment is closed; the tail stays open. - assert _stamped(turns_store) == ["turn-1", "turn-2"] + assert _folded(memories, turns_store) == ["turn-1", "turn-2"] def test_reevaluation_is_idempotent_via_watermark(monkeypatch) -> None: @@ -113,7 +130,7 @@ def test_flush_drains_open_tail(monkeypatch) -> None: assert flushed == {"episodes": 1} assert len(_episodes(memories)) == 2 - assert _stamped(turns_store) == ["turn-1", "turn-2", "turn-3", "turn-4"] + assert _folded(memories, turns_store) == ["turn-1", "turn-2", "turn-3", "turn-4"] # Everything stamped: a further flush is a no-op. assert service.extract_episodes("u1", "t1", flush=True) == {"episodes": 0} @@ -129,7 +146,7 @@ def test_max_turns_forces_a_boundary(monkeypatch) -> None: # Two forced segments: [turn-1, turn-2] and [turn-3, turn-4]. assert result == {"episodes": 2} - assert _stamped(turns_store) == ["turn-1", "turn-2", "turn-3", "turn-4"] + assert _folded(memories, turns_store) == ["turn-1", "turn-2", "turn-3", "turn-4"] def test_topic_drift_closes_episode(monkeypatch) -> None: @@ -149,7 +166,7 @@ def test_topic_drift_closes_episode(monkeypatch) -> None: # Topic A closes when topic B arrives at turn-3; the B tail stays open. assert result == {"episodes": 1} - assert _stamped(turns_store) == ["turn-1", "turn-2"] + assert _folded(memories, turns_store) == ["turn-1", "turn-2"] def test_no_boundary_keeps_segment_open_without_calling_the_llm(monkeypatch) -> None: @@ -163,7 +180,7 @@ def test_no_boundary_keeps_segment_open_without_calling_the_llm(monkeypatch) -> assert result == {"episodes": 0} assert _episodes(memories) == [] - assert _stamped(turns_store) == [] # nothing closed + assert _folded(memories, turns_store) == [] # nothing closed assert chat.calls == 0 # extraction LLM only runs at a boundary @@ -194,7 +211,7 @@ def test_idle_gap_below_min_turns_does_not_close_episode(monkeypatch) -> None: assert result == {"episodes": 0} assert _episodes(memories) == [] - assert _stamped(turns_store) == [] + assert _folded(memories, turns_store) == [] def test_idle_gap_below_min_turns_still_flushes_as_one_episode(monkeypatch) -> None: @@ -208,7 +225,7 @@ def test_idle_gap_below_min_turns_still_flushes_as_one_episode(monkeypatch) -> N result = service.extract_episodes("u1", "t1", flush=True) assert result == {"episodes": 1} - assert _stamped(turns_store) == ["turn-1", "turn-2", "turn-3"] + assert _folded(memories, turns_store) == ["turn-1", "turn-2", "turn-3"] def test_closed_segment_with_no_episode_still_stamps_turns(monkeypatch) -> None: @@ -224,7 +241,7 @@ def test_closed_segment_with_no_episode_still_stamps_turns(monkeypatch) -> None: assert result == {"episodes": 0} assert _episodes(memories) == [] # The closed pre-gap turns are still watermarked despite yielding no episode. - assert _stamped(turns_store) == ["turn-1", "turn-2"] + assert _folded(memories, turns_store) == ["turn-1", "turn-2"] def test_extract_episodes_defers_segment_on_retryable_error(monkeypatch) -> None: @@ -243,7 +260,7 @@ def _boom(*a, **k): assert result == {"episodes": 0} assert _episodes(memories) == [] - assert _stamped(turns_store) == [] # un-stamped -> retried next run + assert _folded(memories, turns_store) == [] # un-stamped -> retried next run def test_extract_episodes_quarantines_segment_on_non_retryable_error(monkeypatch) -> None: @@ -262,4 +279,50 @@ def _boom(*a, **k): assert result == {"episodes": 0} assert _episodes(memories) == [] - assert _stamped(turns_store) == ["turn-1", "turn-2"] # quarantined + advanced + assert _folded(memories, turns_store) == ["turn-1", "turn-2"] # quarantined + advanced + + +def test_advance_episode_cursor_is_monotonic(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_IDLE_GAP_SECONDS", "120") + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + turns = [_turn_at(1, 1), _turn_at(2, 2), _turn_at(3, 30), _turn_at(4, 31)] + service, memories, turns_store, _ = _service(turns) + + service.extract_episodes("u1", "t1") # gap closes [turn-1, turn-2] -> cursor at turn-2 + assert _folded(memories, turns_store) == ["turn-1", "turn-2"] + + # A late, out-of-order concurrent run advancing to an OLDER turn is a no-op: + # the watermark must never regress (a regression would re-open a wider + # segment and, with drift on, duplicate episodes). + service._advance_episode_cursor("u1", "t1", _turn_at(1, 1)) + assert _folded(memories, turns_store) == ["turn-1", "turn-2"] + + # A genuinely newer turn still advances the cursor. + service._advance_episode_cursor("u1", "t1", _turn_at(4, 31)) + assert _folded(memories, turns_store) == ["turn-1", "turn-2", "turn-3", "turn-4"] + + +def test_episode_cursor_is_per_thread(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_IDLE_GAP_SECONDS", "120") + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + + def _t(i: int, minute: int, thread: str) -> dict[str, Any]: + return {**_turn_at(i, minute), "thread_id": thread} + + turns = [ + _t(1, 1, "t1"), + _t(2, 2, "t1"), + _t(3, 30, "t1"), # t1: gap closes [turn-1, turn-2] + _t(4, 1, "t2"), + _t(5, 2, "t2"), + _t(6, 30, "t2"), # t2: never processed + ] + service, memories, turns_store, _ = _service(turns) + + service.extract_episodes("u1", "t1") # only thread t1 + + # Only t1 has a cursor; t2's stream is untouched. + cursors = sorted(d["thread_id"] for d in memories.docs if d.get("type") == "episode_cursor") + assert cursors == ["t1"] + t2_segment = service._load_open_episode_segment("u1", "t2") + assert [t["id"] for t in t2_segment] == ["turn-4", "turn-5", "turn-6"] diff --git a/tests/unit/services/test_episodic_retrieval.py b/tests/unit/services/test_episodic_retrieval.py index d72eaf1..fca2130 100644 --- a/tests/unit/services/test_episodic_retrieval.py +++ b/tests/unit/services/test_episodic_retrieval.py @@ -1,5 +1,6 @@ from __future__ import annotations +import logging from typing import Any from unittest.mock import MagicMock @@ -7,6 +8,8 @@ from azure.cosmos.agent_memory.cosmos_memory_client import CosmosMemoryClient from azure.cosmos.agent_memory.store import MemoryStore +EPISODIC_OPT_IN_WARNING = "Episodic memories requested via memory_types are only returned when include_episodes=True" + def _containers(*, memories: Any = None, turns: Any = None, summaries: Any = None) -> dict[ContainerKey, Any]: return { @@ -83,6 +86,35 @@ def test_search_cosmos_base_is_facts_only_no_episodes_without_optin() -> None: store.search_episodic.assert_not_called() +def test_search_cosmos_warns_when_episodic_requested_without_optin(caplog) -> None: + store = MagicMock() + store.search.return_value = [] + client = _client_with_store(store) + caplog.set_level(logging.WARNING) + + result = client.search_cosmos("ci retries", user_id="u1", memory_types=["episodic"]) + + assert result == [] + assert EPISODIC_OPT_IN_WARNING in caplog.text + assert store.search.call_args.kwargs["memory_types"] == ["fact"] + + +def test_search_cosmos_does_not_warn_for_episodic_optin_or_other_types(caplog) -> None: + store = MagicMock() + store.search.return_value = [] + client = _client_with_store(store) + caplog.set_level(logging.WARNING) + + client.search_cosmos("ci retries", user_id="u1", memory_types=["episodic"], include_episodes=True) + assert EPISODIC_OPT_IN_WARNING not in caplog.text + assert store.search.call_args.kwargs["memory_types"] == ["episodic"] + + caplog.clear() + client.search_cosmos("ci retries", user_id="u1", memory_types=["fact"]) + assert EPISODIC_OPT_IN_WARNING not in caplog.text + assert store.search.call_args.kwargs["memory_types"] == ["fact"] + + def test_search_cosmos_include_episodes_combines_facts_and_episodes_in_base_query() -> None: store = MagicMock() # A single combined query returns facts + episodes ranked together. diff --git a/tests/unit/services/test_extract_dry.py b/tests/unit/services/test_extract_dry.py index 6e86486..3104f9d 100644 --- a/tests/unit/services/test_extract_dry.py +++ b/tests/unit/services/test_extract_dry.py @@ -5,6 +5,7 @@ from unittest.mock import AsyncMock import pytest +from azure.cosmos.exceptions import CosmosResourceNotFoundError from azure.cosmos.agent_memory._container_routing import ContainerKey from azure.cosmos.agent_memory.aio.services.pipeline import AsyncPipelineService, _AsyncStoreContainerAdapter @@ -77,8 +78,16 @@ def query(self, sql: str, parameters=None, partition_key=None, cross_partition: docs = [doc for doc in docs if doc.get("type") in types] if "superseded_by" in sql: docs = [doc for doc in docs if not doc.get("superseded_by")] - if "episode_extracted_at" in sql: - docs = [doc for doc in docs if not doc.get("episode_extracted_at")] + if "@last_at" in params: + # Episodic open-segment cursor: turns after the (created_at, id) watermark. + last_at = params["@last_at"] + last_id = params.get("@last_id", "") + docs = [ + doc + for doc in docs + if (str(doc.get("created_at") or "") > last_at) + or (str(doc.get("created_at") or "") == last_at and str(doc.get("id") or "") > last_id) + ] elif "extracted_at" in sql: docs = [doc for doc in docs if not doc.get("extracted_at")] return docs @@ -98,9 +107,9 @@ def read_item(self, item_id: str, partition_key: Any): for doc in self.docs: if doc.get("id") == item_id: return dict(doc) - raise KeyError(item_id) + raise CosmosResourceNotFoundError(message=item_id) - def add_cosmos(self, record: dict[str, Any]) -> dict[str, Any]: + def upsert_memory(self, record: dict[str, Any]) -> dict[str, Any]: self.docs.append(dict(record)) return record @@ -136,8 +145,8 @@ async def query(self, sql: str, parameters=None, partition_key=None, cross_parti async def read_item(self, item_id: str, partition_key: Any): return super().read_item(item_id, partition_key) - async def add_cosmos(self, record: dict[str, Any]) -> dict[str, Any]: - return super().add_cosmos(record) + async def upsert_memory(self, record: dict[str, Any]) -> dict[str, Any]: + return super().upsert_memory(record) async def mark_superseded(self, old_doc: dict[str, Any], superseder_id: str, *, reason: str) -> bool: return super().mark_superseded(old_doc, superseder_id, reason=reason) diff --git a/tests/unit/services/test_extract_episodes.py b/tests/unit/services/test_extract_episodes.py index c4b9abd..bef72d0 100644 --- a/tests/unit/services/test_extract_episodes.py +++ b/tests/unit/services/test_extract_episodes.py @@ -127,12 +127,13 @@ def test_extract_episodes_embeds_content_persists_append_only(monkeypatch) -> No "The user planned a vacation.", ] ] - assert [doc["content"] for doc in store.docs] == [ + episodes = [doc for doc in store.docs if doc.get("type") == "episodic"] + assert [doc["content"] for doc in episodes] == [ "The user fixed flaky CI retries and the tests passed.", "The user planned a vacation.", ] - assert all(doc["id"].startswith("ep_") for doc in store.docs) - assert all(doc["embedding"] == [1.0] for doc in store.docs) + assert all(doc["id"].startswith("ep_") for doc in episodes) + assert all(doc["embedding"] == [1.0] for doc in episodes) assert store.supersede_calls == [] assert store.search_calls == [] @@ -144,7 +145,7 @@ def test_extract_episodes_empty_window_persists_nothing(monkeypatch) -> None: result = service.extract_episodes("u1", "t1", flush=True) assert result == {"episodes": 0} - assert store.docs == [] + assert [doc for doc in store.docs if doc.get("type") == "episodic"] == [] assert embeddings.calls == [] @@ -165,7 +166,7 @@ def test_extract_episodes_skips_malformed_episode_with_warning(caplog, monkeypat result = service.extract_episodes("u1", "t1", flush=True) assert result == {"episodes": 1} - assert [doc["title"] for doc in store.docs] == ["Valid episode"] + assert [doc["title"] for doc in store.docs if doc.get("type") == "episodic"] == ["Valid episode"] assert "dropping malformed episode" in caplog.text @@ -235,9 +236,9 @@ def test_extract_episodes_skips_duplicate_when_segment_reprocessed(monkeypatch) assert service.extract_episodes("u1", "t1", flush=True) == {"episodes": 1} - # Simulate crash-before-stamp: the closing turns never got watermarked. - for turn in turns.docs: - turn.pop("episode_extracted_at", None) + # Simulate a crash before the cursor advanced: the episodic watermark never + # moved, so the same open segment is re-loaded and re-segmented next run. + store.docs = [doc for doc in store.docs if doc.get("type") != "episode_cursor"] assert service.extract_episodes("u1", "t1", flush=True) == {"episodes": 0} episodic = [doc for doc in store.docs if doc.get("type") == "episodic"] diff --git a/tests/unit/services/test_mark_turns_extracted.py b/tests/unit/services/test_mark_turns_extracted.py new file mode 100644 index 0000000..ca07c65 --- /dev/null +++ b/tests/unit/services/test_mark_turns_extracted.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +from typing import Any + +from azure.cosmos.agent_memory._container_routing import ContainerKey +from azure.cosmos.agent_memory.services.pipeline import PipelineService, _StoreContainerAdapter +from tests.unit.services.test_extract_dry import _Store, _SyncChat, _SyncEmbeddings + + +class _PatchTurnsContainer: + def __init__(self, docs: list[dict[str, Any]]): + self.docs = [dict(doc) for doc in docs] + self.patch_calls: list[dict[str, Any]] = [] + self.upsert_calls: list[dict[str, Any]] = [] + + def query_items(self, **kwargs: Any) -> list[dict[str, Any]]: + del kwargs + return [dict(doc) for doc in self.docs] + + def read_item(self, *, item: str, partition_key: Any) -> dict[str, Any]: + del partition_key + for doc in self.docs: + if doc.get("id") == item: + return dict(doc) + raise KeyError(item) + + def upsert_item(self, *, body: dict[str, Any]) -> dict[str, Any]: + body = dict(body) + self.upsert_calls.append(body) + for index, doc in enumerate(self.docs): + if doc.get("id") == body.get("id"): + self.docs[index] = body + return body + self.docs.append(body) + return body + + def patch_item(self, *, item: str, partition_key: Any, patch_operations: list[dict[str, Any]]) -> dict[str, Any]: + self.patch_calls.append( + {"item": item, "partition_key": partition_key, "patch_operations": [dict(op) for op in patch_operations]} + ) + for doc in self.docs: + if doc.get("id") == item: + for operation in patch_operations: + assert operation["op"] == "set" + assert operation["path"].startswith("/") + doc[operation["path"][1:]] = operation["value"] + return dict(doc) + raise KeyError(item) + + +class _NoPatchTurnsContainer(_PatchTurnsContainer): + patch_item = None + + +class _ContainerBackedStore(_Store): + def __init__(self, container: Any): + super().__init__([]) + self._containers = {ContainerKey.TURNS: container} + + +def _turn_doc(**overrides: Any) -> dict[str, Any]: + doc = { + "id": "turn-1", + "user_id": "u1", + "thread_id": "t1", + "type": "turn", + "content": "hello", + "created_at": "2025-01-01T00:00:00+00:00", + } + doc.update(overrides) + return doc + + +def _service(turns_container: Any) -> PipelineService: + memories_store = _Store([]) + summaries_store = _Store([]) + turns_adapter = _StoreContainerAdapter(_ContainerBackedStore(turns_container), ContainerKey.TURNS) + return PipelineService( + memories_store, + _SyncChat([]), + _SyncEmbeddings(), + containers={ + ContainerKey.TURNS: turns_adapter, + ContainerKey.MEMORIES: _StoreContainerAdapter(memories_store, ContainerKey.MEMORIES), + ContainerKey.SUMMARIES: _StoreContainerAdapter(summaries_store, ContainerKey.SUMMARIES), + }, + ) + + +def test_mark_turns_extracted_patches_extracted_at() -> None: + turns_container = _PatchTurnsContainer([_turn_doc()]) + service = _service(turns_container) + + marked = service._mark_turns_extracted([_turn_doc()]) + + assert marked == 1 + assert turns_container.patch_calls == [ + { + "item": "turn-1", + "partition_key": ["u1", "t1"], + "patch_operations": [ + {"op": "set", "path": "/extracted_at", "value": turns_container.docs[0]["extracted_at"]} + ], + } + ] + assert turns_container.docs[0]["extracted_at"] + assert turns_container.upsert_calls == [] + + +def test_mark_turns_extracted_falls_back_to_read_modify_upsert() -> None: + turns_container = _NoPatchTurnsContainer([_turn_doc()]) + service = _service(turns_container) + + marked = service._mark_turns_extracted([_turn_doc()]) + + assert marked == 1 + assert turns_container.docs[0]["extracted_at"] + assert turns_container.upsert_calls == [turns_container.docs[0]] diff --git a/tests/unit/services/test_persist_extracted.py b/tests/unit/services/test_persist_extracted.py index 366537b..928f2f4 100644 --- a/tests/unit/services/test_persist_extracted.py +++ b/tests/unit/services/test_persist_extracted.py @@ -48,7 +48,7 @@ def read_item(self, item_id: str, partition_key: Any): del partition_key return self.container.docs[item_id] - def add_cosmos(self, record: dict[str, Any]) -> dict[str, Any]: + def upsert_memory(self, record: dict[str, Any]) -> dict[str, Any]: self.upserts.append(dict(record)) self.container.docs[record["id"]] = dict(record) return record @@ -66,8 +66,8 @@ async def query(self, *args, **kwargs): async def read_item(self, item_id: str, partition_key: Any): return super().read_item(item_id, partition_key) - async def add_cosmos(self, record: dict[str, Any]) -> dict[str, Any]: - return super().add_cosmos(record) + async def upsert_memory(self, record: dict[str, Any]) -> dict[str, Any]: + return super().upsert_memory(record) async def mark_superseded(self, old_doc: dict[str, Any], superseder_id: str, *, reason: str) -> bool: return super().mark_superseded(old_doc, superseder_id, reason=reason) diff --git a/tests/unit/services/test_pipeline_service.py b/tests/unit/services/test_pipeline_service.py index bbec77e..0904acc 100644 --- a/tests/unit/services/test_pipeline_service.py +++ b/tests/unit/services/test_pipeline_service.py @@ -5,19 +5,12 @@ from unittest.mock import MagicMock import pytest +from azure.cosmos.exceptions import CosmosResourceExistsError from azure.cosmos.agent_memory._container_routing import ContainerKey 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. @@ -92,6 +85,8 @@ def query( docs = [doc for doc in docs if doc.get("metadata", {}).get("category") == params["@category"]] if "@predicate" in params: docs = [doc for doc in docs if doc.get("metadata", {}).get("predicate") == params["@predicate"]] + if "c.status='active'" in sql: + docs = [doc for doc in docs if doc.get("status") == "active"] if "superseded_by" in sql: docs = [doc for doc in docs if not doc.get("superseded_by")] if "IS_DEFINED(c.lessons)" in sql: @@ -118,13 +113,21 @@ def read_item(self, item_id: str, partition_key: Any) -> dict[str, Any]: raise CosmosResourceNotFoundError(message=f"not found: {item_id}") - def add_cosmos(self, record: dict[str, Any]) -> dict[str, Any]: + def upsert_memory(self, record: dict[str, Any]) -> dict[str, Any]: body = dict(record) self.upserts.append(body) self.docs = [doc for doc in self.docs if doc.get("id") != body.get("id")] self.docs.append(body) return body + def create_item(self, *, body: dict[str, Any]) -> dict[str, Any]: + if any(doc.get("id") == body.get("id") for doc in self.docs): + raise CosmosResourceExistsError(message="conflict") + created = dict(body) + self.upserts.append(created) + self.docs.append(created) + return created + def mark_superseded(self, old_doc: dict[str, Any], superseder_id: str, *, reason: str) -> bool: self.supersede_calls.append((old_doc["id"], superseder_id, reason)) for doc in self.docs: @@ -273,17 +276,43 @@ def test_synthesize_procedural_produces_procedural_memory() -> None: }, ] ) - llm = FakeLLMService([{"system_prompt": "Use concise bullet points."}]) + llm = FakeLLMService( + [ + { + "procedures": [ + { + "name": "Concise bullet responses", + "summary": "Use concise bullet points.", + "retrieval_text": "Use concise bullet points.", + "procedure_kind": "behavioral_policy", + "scope_type": "user", + "scope_value": None, + "activation_conditions": [], + "preconditions": [], + "steps": [], + "success_conditions": [], + "failure_conditions": [], + "safety_constraints": [], + "source_kind": "explicit_user_instruction", + "grounded_in": ["fact-1"], + "confidence": 0.8, + } + ] + } + ] + ) result = _pipeline(store, llm).synthesize_procedural("u1") assert result["status"] == "synthesized" - proc = result["procedural"] + assert result["procedures_created"] >= 1 + proc = next(doc for doc in store.upserts if doc["type"] == "procedural") assert proc["type"] == "procedural" - assert proc["content"] == "Use concise bullet points." + assert proc["id"].startswith("proc_") + assert proc["name"] == "Concise bullet responses" + assert proc["status"] == "active" + assert proc["retrieval_text"] == "Use concise bullet points." assert proc["source_fact_ids"] == ["f1"] - assert proc["source_episodic_ids"] == ["e1"] - assert store.upserts == [proc] def test_reconcile_memories_returns_contradiction_counts() -> None: @@ -413,28 +442,46 @@ def test_build_procedural_context_returns_active_procedural() -> None: store = FakeStore( [ { - "id": "proc_u1_1", + "id": "proc_active", "user_id": "u1", "thread_id": "__procedural__", "type": "procedural", + "status": "active", + "name": "Targeted testing", + "summary": "Run targeted tests before reporting success.", + "retrieval_text": "tests success", + "procedure_kind": "behavioral_policy", + "scope_type": "user", + "scope_value": None, + "priority": 10, + "source_authority": "high", "version": 1, - "content": "Old prompt", - "superseded_by": "proc_u1_2", }, { - "id": "proc_u1_2", + "id": "proc_candidate", "user_id": "u1", "thread_id": "__procedural__", "type": "procedural", + "status": "candidate", + "name": "Candidate policy", + "summary": "Do not include this candidate procedure.", + "retrieval_text": "candidate policy", + "procedure_kind": "behavioral_policy", + "scope_type": "user", + "scope_value": None, + "priority": 10, + "source_authority": "low", "version": 2, - "content": "Active prompt", }, ] ) fake = FakeLLMService([]) result = _pipeline(store, fake).build_procedural_context("u1") - assert result == "Active prompt" + assert "Targeted testing" in result + assert "Run targeted tests before reporting success." in result + assert "Candidate policy" not in result + assert "Do not include this candidate procedure." not in result def test_build_procedural_context_requires_user_id() -> None: diff --git a/tests/unit/services/test_procedural_retrieval.py b/tests/unit/services/test_procedural_retrieval.py new file mode 100644 index 0000000..a439758 --- /dev/null +++ b/tests/unit/services/test_procedural_retrieval.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock + +from azure.cosmos.agent_memory._container_routing import ContainerKey +from azure.cosmos.agent_memory.store import MemoryStore + + +def _containers(*, memories: Any = None, turns: Any = None, summaries: Any = None) -> dict[ContainerKey, Any]: + return { + ContainerKey.MEMORIES: memories if memories is not None else MagicMock(), + ContainerKey.TURNS: turns if turns is not None else MagicMock(), + ContainerKey.SUMMARIES: summaries if summaries is not None else MagicMock(), + } + + +def _params_by_name(call_kwargs: dict[str, Any]) -> dict[str, Any]: + return {param["name"]: param["value"] for param in call_kwargs["parameters"]} + + +def test_retrieve_procedures_filters_active_procedures_by_default() -> None: + ranked_docs = [ + {"id": "proc-1", "type": "procedural", "status": "active", "similarity_score": 0.1}, + {"id": "proc-2", "type": "procedural", "status": "active", "similarity_score": 0.2}, + ] + memories = MagicMock() + memories.query_items.return_value = ranked_docs + embeddings = MagicMock() + embeddings.generate.return_value = [0.1, 0.2] + store = MemoryStore(containers=_containers(memories=memories), embeddings_client=embeddings) + + result = store.retrieve_procedures("u1", "cosmos db retry", top_k=2) + + assert result == ranked_docs + call_kwargs = memories.query_items.call_args.kwargs + assert "TOP 2" in call_kwargs["query"] + assert "c.type = @type" in call_kwargs["query"] + assert "c.user_id = @user_id" in call_kwargs["query"] + assert "c.status = @status" in call_kwargs["query"] + assert "VectorDistance(c.embedding, @embedding)" in call_kwargs["query"] + assert "(NOT IS_DEFINED(c.superseded_by) OR IS_NULL(c.superseded_by))" in call_kwargs["query"] + params = _params_by_name(call_kwargs) + assert params["@type"] == "procedural" + assert params["@user_id"] == "u1" + assert params["@status"] == "active" + assert params["@embedding"] == [0.1, 0.2] + assert params["@kw0"] == "cosmos" + + +def test_retrieve_procedures_status_none_drops_status_and_adds_scope_filters() -> None: + ranked_docs = [{"id": "proc-domain", "type": "procedural", "similarity_score": 0.1}] + memories = MagicMock() + memories.query_items.return_value = ranked_docs + embeddings = MagicMock() + embeddings.generate.return_value = [0.3, 0.4] + store = MemoryStore(containers=_containers(memories=memories), embeddings_client=embeddings) + + result = store.retrieve_procedures( + "u1", + "partition key", + scope_type="domain", + scope_value="cosmos-db", + status=None, + ) + + assert result == ranked_docs + call_kwargs = memories.query_items.call_args.kwargs + assert "c.scope_type = @scope_type" in call_kwargs["query"] + assert "c.scope_value = @scope_value" in call_kwargs["query"] + assert "c.status = @status" not in call_kwargs["query"] + params = _params_by_name(call_kwargs) + assert params["@type"] == "procedural" + assert params["@user_id"] == "u1" + assert params["@scope_type"] == "domain" + assert params["@scope_value"] == "cosmos-db" + assert "@status" not in params + + +def test_retrieve_procedures_filters_by_kind_and_can_include_superseded() -> None: + ranked_docs = [{"id": "proc-workflow", "type": "procedural", "similarity_score": 0.1}] + memories = MagicMock() + memories.query_items.return_value = ranked_docs + embeddings = MagicMock() + embeddings.generate.return_value = [0.5, 0.6] + store = MemoryStore(containers=_containers(memories=memories), embeddings_client=embeddings) + + result = store.retrieve_procedures( + "u1", + "workflow retry", + procedure_kind="workflow", + include_superseded=True, + ) + + assert result == ranked_docs + call_kwargs = memories.query_items.call_args.kwargs + assert "c.procedure_kind = @procedure_kind" in call_kwargs["query"] + assert "(NOT IS_DEFINED(c.superseded_by) OR IS_NULL(c.superseded_by))" not in call_kwargs["query"] + params = _params_by_name(call_kwargs) + assert params["@procedure_kind"] == "workflow" + assert params["@status"] == "active" diff --git a/tests/unit/services/test_procedural_synthesis_and_context.py b/tests/unit/services/test_procedural_synthesis_and_context.py new file mode 100644 index 0000000..79584f8 --- /dev/null +++ b/tests/unit/services/test_procedural_synthesis_and_context.py @@ -0,0 +1,450 @@ +from __future__ import annotations + +from typing import Any + +from azure.cosmos.exceptions import CosmosResourceExistsError + +from azure.cosmos.agent_memory.services.pipeline import PipelineService +from tests.unit.services.test_extract_dry import ( + _containers_for_store, + _Store, + _SyncChat, + _SyncEmbeddings, +) + + +class _ProceduralStore(_Store): + def query(self, sql: str, parameters=None, partition_key=None, cross_partition: bool = False): + del partition_key, cross_partition + params = {p["name"]: p["value"] for p in (parameters or [])} + user_id = params.get("@uid", params.get("@user_id")) + memory_type = params.get("@type", params.get("@memory_type")) + docs = [dict(doc) for doc in self.docs] + if user_id is not None: + docs = [doc for doc in docs if doc.get("user_id") == user_id] + if memory_type is not None: + docs = [doc for doc in docs if doc.get("type") == memory_type] + if "c.status='active'" in sql: + docs = [doc for doc in docs if doc.get("status") == "active"] + if "superseded_by" in sql: + docs = [doc for doc in docs if not doc.get("superseded_by")] + return docs + + def create_item(self, *, body: dict[str, Any]) -> dict[str, Any]: + if any(doc.get("id") == body.get("id") for doc in self.docs): + raise CosmosResourceExistsError(message="conflict") + self.docs.append(dict(body)) + return dict(body) + + +def _service(store: _ProceduralStore, responses: list[dict[str, Any]] | None = None) -> PipelineService: + return PipelineService( + store, + _SyncChat(responses or []), + _SyncEmbeddings(), + containers=_containers_for_store(store), + ) + + +def _fact() -> dict[str, Any]: + return { + "id": "fact-raw-1", + "user_id": "u1", + "type": "fact", + "content": "The user explicitly said to run targeted tests before reporting success.", + "metadata": {"category": "preference"}, + "salience": 0.9, + "created_at": "2025-01-01T00:00:00+00:00", + } + + +def _episode() -> dict[str, Any]: + return { + "id": "episode-raw-1", + "user_id": "u1", + "type": "episodic", + "content": "A retry investigation succeeded.", + "lessons": ["Retry transient CI failures once before escalating."], + "salience": 0.8, + "created_at": "2025-01-01T00:01:00+00:00", + } + + +def _procedure( + name: str, + *, + grounded_in: list[str], + source_kind: str, + summary: str = "Run targeted tests before reporting success.", +) -> dict[str, Any]: + return { + "name": name, + "summary": summary, + "retrieval_text": summary, + "procedure_kind": "behavioral_policy", + "scope_type": "user", + "scope_value": None, + "activation_conditions": [], + "preconditions": [], + "steps": [], + "success_conditions": [], + "failure_conditions": [], + "safety_constraints": [], + "source_kind": source_kind, + "grounded_in": grounded_in, + "confidence": 0.8, + } + + +def test_synthesize_procedural_applies_provenance_gate() -> None: + store = _ProceduralStore([_fact(), _episode()]) + service = _service( + store, + [ + { + "procedures": [ + _procedure( + "Targeted testing", + grounded_in=["fact-1"], + source_kind="explicit_user_instruction", + ), + _procedure( + "Retry CI failures", + grounded_in=["ep-1"], + source_kind="episode_distillation", + summary="Retry transient CI failures once before escalating.", + ), + ] + } + ], + ) + + result = service.synthesize_procedural("u1") + + assert result == {"status": "synthesized", "procedures_created": 2, "procedures_skipped": 0} + procedures = [doc for doc in store.docs if doc.get("type") == "procedural"] + assert {doc["name"]: doc["status"] for doc in procedures} == { + "Targeted testing": "active", + "Retry CI failures": "candidate", + } + assert procedures[0]["thread_id"] == "__procedural__" + assert procedures[0]["embedding"] == [1.0] + + +def test_synthesize_procedural_is_idempotent_by_scope_and_name() -> None: + store = _ProceduralStore([_fact()]) + response = { + "procedures": [ + _procedure( + "Targeted testing", + grounded_in=["fact-1"], + source_kind="explicit_user_instruction", + ) + ] + } + service = _service(store, [response, response]) + + first = service.synthesize_procedural("u1") + second = service.synthesize_procedural("u1") + + assert first["procedures_created"] == 1 + assert second == {"status": "synthesized", "procedures_created": 0, "procedures_skipped": 1} + procedures = [doc for doc in store.docs if doc.get("type") == "procedural"] + assert len(procedures) == 1 + + +def test_build_procedural_context_uses_active_procedures_only() -> None: + active = { + "id": "proc-active", + "user_id": "u1", + "type": "procedural", + "status": "active", + "name": "Targeted testing", + "summary": "Run targeted tests before reporting success.", + "retrieval_text": "tests success", + "procedure_kind": "behavioral_policy", + "scope_type": "user", + "scope_value": None, + "priority": 10, + "source_authority": "high", + "version": 1, + } + candidate = { + **active, + "id": "proc-candidate", + "status": "candidate", + "name": "Candidate policy", + "summary": "Do not include this candidate procedure.", + } + service = _service(_ProceduralStore([active, candidate])) + + context = service.build_procedural_context("u1") + + assert "Run targeted tests before reporting success." in context + assert "Do not include this candidate procedure." not in context + assert service.build_procedural_context("missing-user") == "" + + +def _active_procedure( + name: str, + *, + summary: str, + procedure_kind: str = "behavioral_policy", + scope_type: str = "user", + retrieval_text: str | None = None, + activation_conditions: list[str] | None = None, + priority: int = 0, + source_authority: str = "low", + steps: list[dict[str, Any]] | None = None, +) -> dict[str, Any]: + return { + "id": f"proc-{name.lower().replace(' ', '-')}", + "user_id": "u1", + "type": "procedural", + "status": "active", + "name": name, + "summary": summary, + "retrieval_text": retrieval_text if retrieval_text is not None else summary, + "procedure_kind": procedure_kind, + "scope_type": scope_type, + "scope_value": None, + "activation_conditions": activation_conditions or [], + "steps": steps or [], + "priority": priority, + "source_authority": source_authority, + "version": 1, + } + + +def test_build_procedural_context_includes_task_matching_workflow_only() -> None: + workflow = _active_procedure( + "Partition workflow", + summary="Use partition key diagnostics.", + procedure_kind="workflow", + scope_type="domain", + retrieval_text="cosmos partition routing diagnostics", + activation_conditions=["when debugging partition fanout"], + steps=[{"sequence": 1, "instruction": "Inspect partition routing."}], + ) + service = _service(_ProceduralStore([workflow])) + + matching_context = service.build_procedural_context("u1", task="debug partition latency") + + assert "Partition workflow" in matching_context + assert "Inspect partition routing." in matching_context + assert service.build_procedural_context("u1") == "" + assert service.build_procedural_context("u1", task="summarize billing invoices") == "" + + +def test_build_procedural_context_always_includes_global_and_user_policies() -> None: + global_policy = _active_procedure( + "Global reporting", + summary="Always report validation status.", + scope_type="global", + ) + user_rule = _active_procedure( + "User test rule", + summary="Run targeted tests before reporting success.", + procedure_kind="decision_rule", + scope_type="user", + ) + domain_policy = _active_procedure( + "Domain policy", + summary="Do not include domain policies without task matching.", + scope_type="domain", + ) + service = _service(_ProceduralStore([global_policy, user_rule, domain_policy])) + + context = service.build_procedural_context("u1") + + assert "Global reporting" in context + assert "User test rule" in context + assert "Domain policy" not in context + + +def test_build_procedural_context_orders_policies_by_priority_then_authority() -> None: + lower_priority = _active_procedure( + "Lower priority", + summary="Lower priority policy.", + priority=1, + source_authority="mandatory", + ) + higher_priority = _active_procedure( + "Higher priority", + summary="Higher priority policy.", + priority=10, + source_authority="low", + ) + mandatory_authority = _active_procedure( + "Mandatory authority", + summary="Mandatory authority policy.", + priority=10, + source_authority="mandatory", + ) + service = _service(_ProceduralStore([lower_priority, higher_priority, mandatory_authority])) + + context = service.build_procedural_context("u1") + + assert context.index("Mandatory authority") < context.index("Higher priority") + assert context.index("Higher priority") < context.index("Lower priority") + + +def test_synthesize_procedural_downgrades_episode_only_explicit_claim() -> None: + store = _ProceduralStore([_episode()]) + service = _service( + store, + [ + { + "procedures": [ + _procedure( + "Episode claim", + grounded_in=["ep-1"], + source_kind="explicit_user_instruction", + summary="Retry transient CI failures once before escalating.", + ) + ] + } + ], + ) + + result = service.synthesize_procedural("u1") + + assert result == {"status": "synthesized", "procedures_created": 1, "procedures_skipped": 0} + procedure = next(doc for doc in store.docs if doc.get("type") == "procedural") + assert procedure["status"] == "candidate" + assert procedure["source_kind"] == "episode_distillation" + + +def test_synthesize_procedural_downgrades_episode_only_organization_policy() -> None: + store = _ProceduralStore([_episode()]) + service = _service( + store, + [ + { + "procedures": [ + _procedure( + "Org policy from episode", + grounded_in=["ep-1"], + source_kind="organization_policy", + summary="Retry transient CI failures once before escalating.", + ) + ] + } + ], + ) + + result = service.synthesize_procedural("u1") + + assert result == {"status": "synthesized", "procedures_created": 1, "procedures_skipped": 0} + procedure = next(doc for doc in store.docs if doc.get("type") == "procedural") + # An organization_policy backed only by an episode cannot auto-activate. + assert procedure["status"] == "candidate" + assert procedure["source_kind"] == "episode_distillation" + + +def test_synthesize_procedural_ungrounded_claim_is_candidate() -> None: + store = _ProceduralStore([_fact()]) + service = _service( + store, + [ + { + "procedures": [ + _procedure( + "Ungrounded org policy", + grounded_in=["nonexistent"], + source_kind="organization_policy", + ) + ] + } + ], + ) + + result = service.synthesize_procedural("u1") + + assert result == {"status": "synthesized", "procedures_created": 1, "procedures_skipped": 0} + procedure = next(doc for doc in store.docs if doc.get("type") == "procedural") + # No resolved fact/episodic grounding -> never auto-activated, even though the + # LLM self-labeled it a trusted organization_policy. + assert procedure["status"] == "candidate" + + +def test_synthesize_procedural_skips_malformed_workflow_and_creates_valid_sibling() -> None: + malformed = _procedure( + "Empty workflow", + grounded_in=["fact-1"], + source_kind="explicit_user_instruction", + summary="Malformed workflow with no steps.", + ) + malformed["procedure_kind"] = "workflow" + valid = _procedure( + "Valid policy", + grounded_in=["fact-1"], + source_kind="explicit_user_instruction", + ) + store = _ProceduralStore([_fact()]) + service = _service(store, [{"procedures": [malformed, valid]}]) + + result = service.synthesize_procedural("u1") + + assert result == {"status": "synthesized", "procedures_created": 1, "procedures_skipped": 1} + procedures = [doc for doc in store.docs if doc.get("type") == "procedural"] + assert [doc["name"] for doc in procedures] == ["Valid policy"] + + +def test_synthesize_procedural_quarantines_retryable_and_non_retryable_llm_errors() -> None: + retryable_service = _service(_ProceduralStore([_fact()])) + + def raise_retryable(filename: str, inputs: dict[str, Any]) -> str: + del filename, inputs + raise RuntimeError("rate limit") + + retryable_service._run_prompty = raise_retryable # type: ignore[method-assign] + + retryable_result = retryable_service.synthesize_procedural("u1") + + assert retryable_result == {"status": "deferred", "procedures_created": 0} + + non_retryable_service = _service(_ProceduralStore([_fact()])) + + def raise_non_retryable(filename: str, inputs: dict[str, Any]) -> str: + del filename, inputs + raise RuntimeError("content_filter") + + non_retryable_service._run_prompty = raise_non_retryable # type: ignore[method-assign] + + non_retryable_result = non_retryable_service.synthesize_procedural("u1") + + assert non_retryable_result == {"status": "skipped", "procedures_created": 0} + + +def test_synthesize_procedural_maps_multi_procedure_lineage() -> None: + store = _ProceduralStore([_fact(), _episode()]) + service = _service( + store, + [ + { + "procedures": [ + _procedure( + "Fact lineage", + grounded_in=["fact-1"], + source_kind="explicit_user_instruction", + ), + _procedure( + "Episode lineage", + grounded_in=["ep-1"], + source_kind="episode_distillation", + summary="Retry transient CI failures once before escalating.", + ), + ] + } + ], + ) + + result = service.synthesize_procedural("u1") + + assert result == {"status": "synthesized", "procedures_created": 2, "procedures_skipped": 0} + procedures = {doc["name"]: doc for doc in store.docs if doc.get("type") == "procedural"} + assert procedures["Fact lineage"]["source_fact_ids"] == ["fact-raw-1"] + assert procedures["Fact lineage"]["source_episodic_ids"] == [] + assert procedures["Episode lineage"]["source_fact_ids"] == [] + assert procedures["Episode lineage"]["source_episodic_ids"] == ["episode-raw-1"] diff --git a/tests/unit/store/test_memory_store.py b/tests/unit/store/test_memory_store.py index a54bf0e..d6093aa 100644 --- a/tests/unit/store/test_memory_store.py +++ b/tests/unit/store/test_memory_store.py @@ -370,14 +370,14 @@ def test_build_episodic_context_forwards_search_options(): store.search_episodic.assert_called_once_with("u1", "weather", top_k=3) -def test_add_cosmos_routes_by_type(): +def test_upsert_memory_routes_by_type(): turns = MagicMock() memories = MagicMock() summaries = MagicMock() store = MemoryStore(containers=_containers(turns=turns, memories=memories, summaries=summaries)) for memory_type in ("turn", "fact", "episodic", "procedural", "thread_summary", "user_summary"): - store.add_cosmos(_doc(id=f"{memory_type}_id", type=memory_type)) + store.upsert_memory(_doc(id=f"{memory_type}_id", type=memory_type)) assert turns.upsert_item.call_count == 1 assert memories.upsert_item.call_count == 3 diff --git a/tests/unit/test_auto_trigger.py b/tests/unit/test_auto_trigger.py index 5ed82aa..ddab91a 100644 --- a/tests/unit/test_auto_trigger.py +++ b/tests/unit/test_auto_trigger.py @@ -1,8 +1,8 @@ """Tests for the InProcess push_to_cosmos auto-trigger. -Per-turn fact extraction is the new default (FACT_EXTRACTION_EVERY_N=1): -each turn flushed to Cosmos should immediately fire `process_thread` for -the in-process backend. The durable backend must remain a no-op (the +These tests set `FACT_EXTRACTION_EVERY_N=1` for per-turn extraction: each +turn flushed to Cosmos should immediately fire `process_thread` for the +in-process backend. The durable backend must remain a no-op (the change-feed function app handles it). """ diff --git a/tests/unit/test_cosmos_memory_client.py b/tests/unit/test_cosmos_memory_client.py index 61b0142..7a2e458 100644 --- a/tests/unit/test_cosmos_memory_client.py +++ b/tests/unit/test_cosmos_memory_client.py @@ -591,12 +591,12 @@ def test_constructor_rejects_invalid_throughput_mode(self): class TestAddCosmos: - def test_add_cosmos(self): + def test_upsert_memory(self): mem, container = _connected_client() # 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") + mem.upsert_memory(user_id="u1", role="user", content="hello", thread_id="t1") turns = mem._turns_container_client turns.upsert_item.assert_called_once() @@ -605,10 +605,10 @@ def test_add_cosmos(self): assert body["user_id"] == "u1" assert body["role"] == "user" - def test_add_cosmos_threads_explicit_created_at(self): + def test_upsert_memory_threads_explicit_created_at(self): mem, container = _connected_client() mem._maybe_auto_trigger = MagicMock() - mem.add_cosmos( + mem.upsert_memory( user_id="u1", role="user", content="hello", @@ -619,47 +619,47 @@ def test_add_cosmos_threads_explicit_created_at(self): body = mem._turns_container_client.upsert_item.call_args.kwargs["body"] assert body["created_at"] == "2024-03-01T12:00:00+00:00" - def test_add_cosmos_not_connected(self): + def test_upsert_memory_not_connected(self): mem = _make_client() with pytest.raises(CosmosNotConnectedError): - mem.add_cosmos(user_id="u1", role="user", content="hi", thread_id="t1") + mem.upsert_memory(user_id="u1", role="user", content="hi", thread_id="t1") - def test_add_cosmos_turn_requires_thread_id(self): + def test_upsert_memory_turn_requires_thread_id(self): """Turn writes must declare a thread_id so the auto-trigger counter can group them.""" mem, _ = _connected_client() with pytest.raises(ValidationError, match="thread_id is required"): - mem.add_cosmos(user_id="u1", role="user", content="hi") # memory_type='turn' default + mem.upsert_memory(user_id="u1", role="user", content="hi") # memory_type='turn' default - def test_add_cosmos_non_turn_does_not_require_thread_id(self): + def test_upsert_memory_non_turn_does_not_require_thread_id(self): """Non-turn writes (facts, episodics, etc.) work without thread_id and skip cadence.""" mem, container = _connected_client() trigger = MagicMock() mem._maybe_auto_trigger = trigger - mem.add_cosmos(user_id="u1", role="user", content="prefers dark mode", memory_type="fact") + mem.upsert_memory(user_id="u1", role="user", content="prefers dark mode", memory_type="fact") container.upsert_item.assert_called_once() trigger.assert_not_called() - def test_add_cosmos_turn_triggers_cadence(self): + def test_upsert_memory_turn_triggers_cadence(self): """A turn write must bump the auto-trigger counter so cadence env vars apply whether the caller uses the local buffer or writes through directly.""" mem, _ = _connected_client() trigger = MagicMock() mem._maybe_auto_trigger = trigger - mem.add_cosmos(user_id="u1", role="user", content="hello", thread_id="t1") + mem.upsert_memory(user_id="u1", role="user", content="hello", thread_id="t1") 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 - + def test_upsert_memory_swallows_cadence_failure(self): + """If the cadence trigger raises, the upsert_memory 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. - result_id = mem.add_cosmos(user_id="u1", role="user", content="hi", thread_id="t1") + result_id = mem.upsert_memory(user_id="u1", role="user", content="hi", thread_id="t1") assert isinstance(result_id, str) mem._turns_container_client.upsert_item.assert_called_once() @@ -951,7 +951,7 @@ def test_success(self): container.read_item = MagicMock(return_value=_make_doc(id="m1", type="fact")) container.delete_item = MagicMock() - mem.delete_cosmos(memory_id="m1", user_id="u1", thread_id="t1", memory_type="fact") + mem.delete_memory(memory_id="m1", user_id="u1", thread_id="t1", memory_type="fact") container.delete_item.assert_called_once_with(item="m1", partition_key=["u1", "t1"]) @@ -963,7 +963,7 @@ def test_not_found(self): container.delete_item = MagicMock() with pytest.raises(MemoryNotFoundError): - mem.delete_cosmos(memory_id="nope", user_id="u1", thread_id="t1", memory_type="fact") + mem.delete_memory(memory_id="nope", user_id="u1", thread_id="t1", memory_type="fact") container.delete_item.assert_not_called() @@ -1171,7 +1171,7 @@ def test_cosmos_ops_without_connect(self): with pytest.raises(CosmosNotConnectedError): mem.update_cosmos(memory_id="m1", user_id="u1", thread_id="t1", memory_type="fact") with pytest.raises(CosmosNotConnectedError): - mem.delete_cosmos(memory_id="m1", user_id="u1", thread_id="t1", memory_type="fact") + mem.delete_memory(memory_id="m1", user_id="u1", thread_id="t1", memory_type="fact") # =================================================================== @@ -1271,3 +1271,68 @@ def test_non_int_value_rejected(self): def test_non_mapping_rejected(self): with pytest.raises(TypeError): CosmosMemoryClient(use_default_credential=False, cadence_thresholds=[("DEDUP_EVERY_N", 5)]) + + +class TestDeleteHelpers: + def test_delete_turn_delegates_with_turn_type(self): + mem = _make_client() + mem.delete_memory = MagicMock() + mem.delete_turn("turn-1", user_id="u1", thread_id="t1") + mem.delete_memory.assert_called_once_with("turn-1", user_id="u1", thread_id="t1", memory_type="turn") + + def test_delete_thread_summary_uses_deterministic_id_and_returns_true(self): + mem = _make_client() + mem.delete_memory = MagicMock() + assert mem.delete_thread_summary("u1", "t1") is True + mem.delete_memory.assert_called_once_with( + "summary_u1_t1", user_id="u1", thread_id="t1", memory_type="thread_summary" + ) + + def test_delete_thread_summary_missing_returns_false(self): + mem = _make_client() + mem.delete_memory = MagicMock(side_effect=MemoryNotFoundError(memory_id="x", user_id="u1", thread_id="t1")) + assert mem.delete_thread_summary("u1", "t1") is False + + def test_delete_user_summary_uses_deterministic_id_and_scope(self): + mem = _make_client() + mem.delete_memory = MagicMock() + assert mem.delete_user_summary("u1") is True + mem.delete_memory.assert_called_once_with( + "user_summary_u1", user_id="u1", thread_id="__user_summary__", memory_type="user_summary" + ) + + def test_delete_thread_deletes_turns_and_summary_and_counts(self): + mem = _make_client() + mem.get_thread = MagicMock(return_value=[{"id": "turn-1"}, {"id": "turn-2"}]) + mem.delete_memory = MagicMock() + deleted = mem.delete_thread("u1", "t1") + assert deleted == 3 # 2 turns + 1 summary + assert mem.get_thread.call_args.kwargs["include_superseded"] is True + turn_deletes = [c for c in mem.delete_memory.call_args_list if c.kwargs.get("memory_type") == "turn"] + assert len(turn_deletes) == 2 + + def test_delete_thread_without_summary(self): + mem = _make_client() + mem.get_thread = MagicMock(return_value=[{"id": "turn-1"}]) + mem.delete_memory = MagicMock() + deleted = mem.delete_thread("u1", "t1", include_summary=False) + assert deleted == 1 + + def test_delete_thread_skips_missing_turn(self): + mem = _make_client() + mem.get_thread = MagicMock(return_value=[{"id": "turn-1"}, {"id": "turn-2"}]) + + def _delete(memory_id, **kwargs): + if memory_id == "turn-1": + raise MemoryNotFoundError(memory_id="turn-1", user_id="u1", thread_id="t1") + + mem.delete_memory = MagicMock(side_effect=_delete) + deleted = mem.delete_thread("u1", "t1", include_summary=False) + assert deleted == 1 # turn-2 only + + def test_delete_thread_requires_ids(self): + mem = _make_client() + with pytest.raises(ValidationError): + mem.delete_thread("", "t1") + with pytest.raises(ValidationError): + mem.delete_thread("u1", "") diff --git a/tests/unit/test_models.py b/tests/unit/test_models.py index 550c1d2..178d70b 100644 --- a/tests/unit/test_models.py +++ b/tests/unit/test_models.py @@ -26,6 +26,11 @@ MemoryType, OrchestrationResult, ProceduralRecord, + ProcedureKind, + ProcedureScopeType, + ProcedureSourceKind, + ProcedureStatus, + ProcedureStep, SearchResult, ThreadSummaryRecord, TurnRecord, @@ -112,8 +117,11 @@ def _procedural_kwargs(**overrides: Any) -> dict[str, Any]: "id": "proc_u1_1", "user_id": "u1", "content": "Be concise.", + "name": "Be concise", + "summary": "Prefer concise responses.", + "retrieval_text": "Respond concisely when answering the user.", + "procedure_kind": "behavioral_policy", "version": 1, - "source_fact_ids": ["fact_" + _HEX32], "prompt_id": "synthesize_procedural.prompty", } base.update(overrides) @@ -405,22 +413,43 @@ class TestProceduralRecord: def test_minimal_valid(self): rec = ProceduralRecord(**_procedural_kwargs()) assert rec.memory_type == "procedural" + assert rec.procedure_kind == "behavioral_policy" + assert rec.status == "candidate" + assert rec.utility_score == 0.5 + assert rec.name == "Be concise" + assert rec.summary == "Prefer concise responses." + assert rec.retrieval_text == "Respond concisely when answering the user." assert rec.version == 1 - def test_requires_non_empty_source_fact_ids(self): - with pytest.raises(pydantic.ValidationError, match="source"): - ProceduralRecord(**_procedural_kwargs(source_fact_ids=[])) - - def test_accepts_episodic_only_sources(self): - """Procedural records driven purely off episodic lessons must be valid; - the validator should accept either source set being non-empty.""" - rec = ProceduralRecord(**_procedural_kwargs(source_fact_ids=[], source_episodic_ids=["ep_abc"])) + def test_accepts_no_source_ids(self): + rec = ProceduralRecord(**_procedural_kwargs(source_fact_ids=[], source_episodic_ids=[])) assert rec.source_fact_ids == [] - assert rec.source_episodic_ids == ["ep_abc"] + assert rec.source_episodic_ids == [] + + def test_workflow_without_steps_raises(self): + with pytest.raises(pydantic.ValidationError, match="requires at least one step"): + ProceduralRecord(**_procedural_kwargs(procedure_kind=ProcedureKind.workflow)) + + def test_workflow_with_step_is_valid(self): + rec = ProceduralRecord( + **_procedural_kwargs( + procedure_kind=ProcedureKind.workflow, + steps=[ + ProcedureStep( + sequence=1, + instruction="Check current context.", + expected_result="Context is understood.", + ) + ], + ) + ) + assert rec.procedure_kind == "workflow" + assert rec.steps[0].instruction == "Check current context." - def test_rejects_when_both_source_sets_empty(self): - with pytest.raises(pydantic.ValidationError, match="source"): - ProceduralRecord(**_procedural_kwargs(source_fact_ids=[], source_episodic_ids=[])) + @pytest.mark.parametrize(("input_score", "expected"), [(5.0, 1.0), (-1, 0.0), ("not a number", 0.5)]) + def test_utility_score_is_clamped(self, input_score, expected): + rec = ProceduralRecord(**_procedural_kwargs(utility_score=input_score)) + assert rec.utility_score == expected def test_id_must_start_with_proc_prefix(self): with pytest.raises(pydantic.ValidationError, match="id must start with 'proc_'"): @@ -430,6 +459,28 @@ def test_version_must_be_positive(self): with pytest.raises(pydantic.ValidationError): ProceduralRecord(**_procedural_kwargs(version=0)) + def test_round_trip_preserves_new_fields(self): + rec = ProceduralRecord( + **_procedural_kwargs( + procedure_kind=ProcedureKind.tool_usage, + scope_type=ProcedureScopeType.project, + scope_value="agent-memory-toolkit", + status=ProcedureStatus.active, + source_kind=ProcedureSourceKind.explicit_user_instruction, + source_turn_ids=["turn-1"], + ) + ) + + restored = MemoryRecordBase.from_doc(rec.to_doc()) + + assert isinstance(restored, ProceduralRecord) + assert restored.name == "Be concise" + assert restored.procedure_kind == "tool_usage" + assert restored.scope_type == "project" + assert restored.status == "active" + assert restored.source_kind == "explicit_user_instruction" + assert restored.source_turn_ids == ["turn-1"] + # --------------------------------------------------------------------------- # Shared field validators on the base diff --git a/tests/unit/test_pipeline_confidence.py b/tests/unit/test_pipeline_confidence.py index 61dbe21..e21ef48 100644 --- a/tests/unit/test_pipeline_confidence.py +++ b/tests/unit/test_pipeline_confidence.py @@ -13,14 +13,6 @@ 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() @@ -56,7 +48,6 @@ def _make_pipeline(llm_response: dict): pipeline = PipelineService(store, chat, embeddings, containers=containers) # Avoid real LLM/prompty calls. pipeline._run_prompty = MagicMock(return_value=json.dumps(llm_response)) - pipeline._load_existing_memories = MagicMock(return_value=[]) return pipeline, upserted diff --git a/tests/unit/test_procedural_synthesis.py b/tests/unit/test_procedural_synthesis.py index 2a47dcb..215efe8 100644 --- a/tests/unit/test_procedural_synthesis.py +++ b/tests/unit/test_procedural_synthesis.py @@ -1,801 +1,199 @@ -"""Tests for procedural synthesis and procedural prompt retrieval.""" +"""Tests for atomic procedural synthesis and compiled procedural context.""" from __future__ import annotations -import json -from datetime import datetime -from unittest.mock import MagicMock +from typing import Any -import pytest +from azure.cosmos.exceptions import CosmosResourceExistsError -from azure.cosmos.agent_memory._container_routing import ContainerKey -from azure.cosmos.agent_memory.cosmos_memory_client import CosmosMemoryClient -from azure.cosmos.agent_memory.processors import DurableFunctionProcessor from azure.cosmos.agent_memory.services.pipeline import PipelineService -from azure.cosmos.agent_memory.store import MemoryStore +from tests.unit.services.test_extract_dry import ( + _containers_for_store, + _Store, + _SyncChat, + _SyncEmbeddings, +) + + +class _ProceduralStore(_Store): + def query(self, sql: str, parameters=None, partition_key=None, cross_partition: bool = False): + del partition_key, cross_partition + params = {p["name"]: p["value"] for p in (parameters or [])} + user_id = params.get("@uid", params.get("@user_id")) + memory_type = params.get("@type", params.get("@memory_type")) + docs = [dict(doc) for doc in self.docs] + if user_id is not None: + docs = [doc for doc in docs if doc.get("user_id") == user_id] + if memory_type is not None: + docs = [doc for doc in docs if doc.get("type") == memory_type] + if "c.status='active'" in sql: + docs = [doc for doc in docs if doc.get("status") == "active"] + if "superseded_by" in sql: + docs = [doc for doc in docs if not doc.get("superseded_by")] + return docs + + def create_item(self, *, body: dict[str, Any]) -> dict[str, Any]: + if any(doc.get("id") == body.get("id") for doc in self.docs): + raise CosmosResourceExistsError(message="conflict") + self.docs.append(dict(body)) + return dict(body) -@pytest.fixture(autouse=True) -def _pin_legacy_extract_dedup(monkeypatch): - monkeypatch.setattr( - "azure.cosmos.agent_memory.thresholds.get_dedup_vector_enabled", - lambda: False, +def _service(store: _ProceduralStore, responses: list[dict[str, Any]] | None = None) -> PipelineService: + return PipelineService( + store, + _SyncChat(responses or []), + _SyncEmbeddings(), + containers=_containers_for_store(store), ) -def _assert_iso8601(text: str) -> None: - assert text - datetime.fromisoformat(text) - - -def _capture_upserts(): - """Capture documents written via upsert_item OR create_item. - - 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 - helper now wires the same capture to both side_effects. - """ - upserted: list[dict] = [] - - def _capture(*, body): - upserted.append(body) - return body - - return upserted, _capture - - -def _make_extract_pipeline(llm_response: dict): - turns_container = MagicMock() - memories_container = MagicMock() - summaries_container = MagicMock() - turns_container.query_items.return_value = [ - { - "id": "turn1", - "user_id": "u1", - "thread_id": "t1", - "role": "user", - "type": "turn", - "content": "Always use bullet points.", - "created_at": "2025-01-01T00:00:00+00:00", - } - ] - upserted, capture = _capture_upserts() - memories_container.upsert_item.side_effect = capture - memories_container.create_item.side_effect = capture - - embeddings = MagicMock() - embeddings.generate_batch.side_effect = lambda texts: [[0.0] * 4 for _ in texts] - - containers = { - ContainerKey.TURNS: turns_container, - ContainerKey.MEMORIES: memories_container, - ContainerKey.SUMMARIES: summaries_container, - } - store = MemoryStore(containers=containers, embeddings_client=embeddings) - pipeline = PipelineService(store, MagicMock(), embeddings, containers=containers) - pipeline._run_prompty = MagicMock(return_value=json.dumps(llm_response)) - pipeline._load_existing_memories = MagicMock(return_value=[]) - return pipeline, memories_container, upserted - - -def _fact_doc( - doc_id: str, - content: str, - *, - category: str = "preference", - salience: float = 0.9, - created_at: str = "2025-01-01T00:00:00+00:00", - predicate: str | None = None, - obj: str | None = None, -) -> dict: - metadata = {"category": category} - if predicate is not None: - metadata["predicate"] = predicate - if obj is not None: - metadata["object"] = obj +def _fact() -> dict[str, Any]: return { - "id": doc_id, + "id": "fact-raw-1", "user_id": "u1", - "thread_id": "t-source", - "role": "system", "type": "fact", - "content": content, - "metadata": metadata, - "salience": salience, - "created_at": created_at, + "content": "The user explicitly said to run targeted tests before reporting success.", + "metadata": {"category": "preference"}, + "salience": 0.9, + "created_at": "2025-01-01T00:00:00+00:00", } -def _episodic_doc( - doc_id: str, - *, - lesson: str | None = None, - lessons: list[str] | None = None, - salience: float = 0.7, - created_at: str = "2025-01-02T00:00:00+00:00", -) -> dict: - if lessons is None: - lessons = [lesson] if lesson else [] +def _episode() -> dict[str, Any]: return { - "id": doc_id, + "id": "episode-raw-1", "user_id": "u1", - "thread_id": "t-source", - "role": "system", "type": "episodic", - "content": f"Episode {doc_id}", - "lessons": lessons, - "salience": salience, - "created_at": created_at, + "content": "A retry investigation succeeded.", + "lessons": ["Retry transient CI failures once before escalating."], + "salience": 0.8, + "created_at": "2025-01-01T00:01:00+00:00", } -def _procedural_doc( - doc_id: str, +def _procedure( + name: str, *, - version: int, - content: str, - source_fact_ids: list[str], - source_episodic_ids: list[str], - superseded_by: str | None = None, - ts: int = 0, - etag: str = "etag-1", -) -> dict: - doc = { - "id": doc_id, - "user_id": "u1", - "thread_id": "__procedural__", - "type": "procedural", - "version": version, - "content": content, - "source_fact_ids": list(source_fact_ids), - "source_episodic_ids": list(source_episodic_ids), - "supersedes_ids": [], - "created_at": f"2025-01-0{version}T00:00:00+00:00", - "role": "system", - "tags": ["sys:procedural", "sys:synthesized"], - "_etag": etag, - "_ts": ts, - } - if superseded_by is not None: - doc["superseded_by"] = superseded_by - return doc - - -def _make_synthesis_pipeline( - *, - prior_docs: list[dict] | None = None, - fact_docs: list[dict] | None = None, - episodic_docs: list[dict] | None = None, - name_docs: list[dict] | None = None, - llm_output: str = "Follow the user's preferences.", -): - turns_container = MagicMock() - memories_container = MagicMock() - summaries_container = MagicMock() - memories_container.query_items.side_effect = [ - list(prior_docs or []), - list(fact_docs or []), - list(episodic_docs or []), - list(name_docs or []), - ] - upserted, capture = _capture_upserts() - memories_container.upsert_item.side_effect = capture - memories_container.create_item.side_effect = capture - - mock_embeddings = MagicMock() - containers = { - ContainerKey.TURNS: turns_container, - ContainerKey.MEMORIES: memories_container, - ContainerKey.SUMMARIES: summaries_container, - } - store = MemoryStore(containers=containers, embeddings_client=mock_embeddings) - pipeline = PipelineService(store, MagicMock(), mock_embeddings, containers=containers) - pipeline._run_prompty = MagicMock(return_value=json.dumps({"system_prompt": llm_output})) - return pipeline, memories_container, upserted - - -def _make_client(*, processor=None) -> CosmosMemoryClient: - client = CosmosMemoryClient.__new__(CosmosMemoryClient) - memories_container = MagicMock() - turns_container = MagicMock() - summaries_container = MagicMock() - containers = { - ContainerKey.TURNS: turns_container, - ContainerKey.MEMORIES: memories_container, - ContainerKey.SUMMARIES: summaries_container, - } - client._memories_container_client = memories_container - client._turns_container_client = turns_container - client._summaries_container_client = summaries_container - client._embeddings_client = MagicMock() - client._store = MemoryStore(containers=containers, embeddings_client=client._embeddings_client) - client._pipeline = None - client._processor = processor - client._processor_explicit = processor is not None - return client - - -def test_extract_memories_returns_count_shape_and_ignores_legacy_episodic_payload(): - pipeline, _, upserted = _make_extract_pipeline( - { - "facts": [ - { - "text": "Always use bullet points.", - "category": "preference", - "action": "ADD", - } - ], - "episodic": [ - { - "scope_type": "task", - "scope_value": "refactoring tests", - "text": "Refactored test suite using focused helpers to keep it readable.", - "situation": "Refactoring tests", - "action_taken": "Used focused helpers", - "outcome": "The suite stayed readable", - } - ], - } - ) - - result = pipeline.extract_memories("u1", "t1") - legacy_fact_count_key = "_".join(("facts", "count")) - legacy_proc_key = "_".join(("procedural", "count")) - - assert result["fact_count"] == 1 - assert result["episodic_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) - - -def test_extract_memories_ignores_legacy_procedural_bucket_in_llm_payload(): - pipeline, _, upserted = _make_extract_pipeline( - { - "facts": [ - { - "text": "Never use var in TypeScript.", - "category": "requirement", - "action": "ADD", - } - ], - "procedural": [ - { - "instruction": "Use bullet points", - "action": "ADD", - } - ], - } - ) - - result = pipeline.extract_memories("u1", "t1") - legacy_fact_count_key = "_".join(("facts", "count")) - legacy_proc_key = "_".join(("procedural", "count")) - - assert result["fact_count"] == 1 - assert legacy_fact_count_key not in result - assert legacy_proc_key not in result - assert [doc["type"] for doc in upserted] == ["fact"] - - -def test_synthesize_procedural_first_synthesis_from_empty_prior(): - fact_docs = [ - _fact_doc("f1", "Always use bullet points.", category="preference", salience=0.95), - _fact_doc("f2", "Never use var in TypeScript.", category="requirement", salience=0.9), - ] - episodic_docs = [ - _episodic_doc("e1", lesson="When the user asks for brevity, keep the answer terse.", salience=0.8), - _episodic_doc("e2", lesson="", salience=0.2), - ] - pipeline, container, upserted = _make_synthesis_pipeline( - fact_docs=fact_docs, - episodic_docs=episodic_docs, - llm_output="Be concise and prefer bullet points.", - ) - - result = pipeline.synthesize_procedural("u1", force=False) - - assert pipeline._run_prompty.call_count == 1 - assert result["status"] == "synthesized" - doc = result["procedural"] - assert doc["version"] == 1 - assert doc["content"] == "Be concise and prefer bullet points." - assert set(doc["source_fact_ids"]) == {"f1", "f2"} - assert set(doc["source_episodic_ids"]) == {"e1"} - assert doc["supersedes_ids"] == [] - assert upserted == [doc] - container.replace_item.assert_not_called() - - -def test_synthesize_procedural_flattens_multiple_lessons_per_episode(): - fact_docs = [ - _fact_doc("f1", "Always use bullet points.", category="preference", salience=0.95), - ] - episodic_docs = [ - _episodic_doc( - "e1", - lessons=[ - "When the user asks for brevity, keep the answer terse.", - "Confirm cancellations explicitly before acting.", - ], - salience=0.8, - ), - _episodic_doc("e2", lessons=[], salience=0.2), - ] - pipeline, _, _ = _make_synthesis_pipeline( - fact_docs=fact_docs, - episodic_docs=episodic_docs, - llm_output="Be concise and confirm before acting.", - ) - - result = pipeline.synthesize_procedural("u1", force=False) - - assert result["status"] == "synthesized" - doc = result["procedural"] - assert set(doc["source_episodic_ids"]) == {"e1"} - - rendered = pipeline._run_prompty.call_args.kwargs["inputs"]["episodic_lessons"] - assert "When the user asks for brevity, keep the answer terse." in rendered - assert "Confirm cancellations explicitly before acting." in rendered - - -def test_synthesize_procedural_only_touches_memories_container(): - turns_container = MagicMock() - memories_container = MagicMock() - summaries_container = MagicMock() - memories_container.query_items.side_effect = [ - [], - [_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 = { - ContainerKey.TURNS: turns_container, - ContainerKey.MEMORIES: memories_container, - ContainerKey.SUMMARIES: summaries_container, + grounded_in: list[str], + source_kind: str, + summary: str = "Run targeted tests before reporting success.", +) -> dict[str, Any]: + return { + "name": name, + "summary": summary, + "retrieval_text": summary, + "procedure_kind": "behavioral_policy", + "scope_type": "user", + "scope_value": None, + "activation_conditions": [], + "preconditions": [], + "steps": [], + "success_conditions": [], + "failure_conditions": [], + "safety_constraints": [], + "source_kind": source_kind, + "grounded_in": grounded_in, + "confidence": 0.8, } - embeddings = MagicMock() - store = MemoryStore(containers=containers, embeddings_client=embeddings) - pipeline = PipelineService(store, MagicMock(), embeddings, containers=containers) - pipeline._run_prompty = MagicMock(return_value=json.dumps({"system_prompt": "Use concise bullets."})) - - result = pipeline.synthesize_procedural("u1") - - assert result["status"] == "synthesized" - assert memories_container.query_items.call_count == 3 - memories_container.create_item.assert_called_once() - turns_container.method_calls == [] - summaries_container.method_calls == [] - - -def test_synthesize_procedural_resynthesis_supersedes_prior_with_update_reason(): - prior_doc = _procedural_doc( - "proc_u1_1", - version=1, - content="Old prompt", - source_fact_ids=["f1", "f2"], - source_episodic_ids=["e1"], - ts=1, - ) - fact_docs = [ - _fact_doc("f1", "Always use bullet points.", category="preference"), - _fact_doc("f2", "Never use var in TypeScript.", category="requirement"), - _fact_doc("f3", "Lead with the final answer.", category="preference", salience=0.85), - ] - episodic_docs = [_episodic_doc("e1", lesson="Keep examples small.")] - pipeline, container, upserted = _make_synthesis_pipeline( - prior_docs=[prior_doc], - fact_docs=fact_docs, - episodic_docs=episodic_docs, - llm_output="New prompt", - ) - - result = pipeline.synthesize_procedural("u1") - - assert result["status"] == "synthesized" - new_doc = result["procedural"] - assert new_doc["id"] == "proc_u1_2" - assert new_doc["version"] == 2 - assert new_doc["supersedes_ids"] == [prior_doc["id"]] - assert upserted == [new_doc] - body = container.replace_item.call_args.kwargs["body"] - assert body["id"] == prior_doc["id"] - assert body["superseded_by"] == new_doc["id"] - _assert_iso8601(body["superseded_at"]) - assert body["supersede_reason"] == "update" - - -def test_synthesize_procedural_noop_when_source_ids_are_unchanged(): - prior_doc = _procedural_doc( - "proc_u1_1", - version=1, - content="Existing prompt", - source_fact_ids=["f1", "f2"], - source_episodic_ids=["e1"], - ts=1, - ) - fact_docs = [ - _fact_doc("f2", "Never use var in TypeScript.", category="requirement"), - _fact_doc("f1", "Always use bullet points.", category="preference"), - ] - episodic_docs = [_episodic_doc("e1", lesson="Keep examples small.")] - pipeline, container, _ = _make_synthesis_pipeline( - prior_docs=[prior_doc], - fact_docs=fact_docs, - episodic_docs=episodic_docs, - ) - - result = pipeline.synthesize_procedural("u1", force=False) - - assert result == {"status": "unchanged", "procedural": prior_doc} - pipeline._run_prompty.assert_not_called() - container.upsert_item.assert_not_called() - container.replace_item.assert_not_called() - - -def test_synthesize_procedural_force_true_reruns_when_source_ids_are_unchanged(): - prior_doc = _procedural_doc( - "proc_u1_1", - version=1, - content="Existing prompt", - source_fact_ids=["f1", "f2"], - source_episodic_ids=["e1"], - ts=1, - ) - fact_docs = [ - _fact_doc("f1", "Always use bullet points.", category="preference"), - _fact_doc("f2", "Never use var in TypeScript.", category="requirement"), - ] - episodic_docs = [_episodic_doc("e1", lesson="Keep examples small.")] - pipeline, container, upserted = _make_synthesis_pipeline( - prior_docs=[prior_doc], - fact_docs=fact_docs, - episodic_docs=episodic_docs, - llm_output="Refreshed prompt", - ) - - result = pipeline.synthesize_procedural("u1", force=True) - - assert pipeline._run_prompty.call_count == 1 - assert result["status"] == "synthesized" - new_doc = result["procedural"] - assert new_doc["version"] == 2 - assert upserted == [new_doc] - body = container.replace_item.call_args.kwargs["body"] - assert body["superseded_by"] == new_doc["id"] - assert body["supersede_reason"] == "update" - - -def test_synthesize_procedural_short_circuits_for_cold_user_with_no_sources(): - """B2 regression: a user with no facts and no episodics must not consume an - LLM call. Without this guard, ``synthesize_procedural`` for cold users - would invoke the chat client and then discard the result, wasting tokens - on every auto-trigger. - """ - pipeline, container, upserted = _make_synthesis_pipeline() - - result = pipeline.synthesize_procedural("u1") - - pipeline._run_prompty.assert_not_called() - container.upsert_item.assert_not_called() - container.replace_item.assert_not_called() - assert result == {"status": "unchanged", "procedural": None} - assert upserted == [] - - -def test_synthesize_procedural_short_circuits_on_second_call_with_tied_salience(): - """B1 regression: 60 facts at the default salience must yield a deterministic - selection so the source-id short-circuit fires on the second call. - - 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 - must be deterministic. - """ - fact_docs = [ - _fact_doc( - f"f{i:03}", - f"Fact {i}", - category="preference", - salience=0.5, - created_at=f"2025-01-{(i % 28) + 1:02}T00:00:00+00:00", - ) - for i in range(60) - ] - pipeline, container, upserted = _make_synthesis_pipeline( - fact_docs=fact_docs, - llm_output="Generated prompt", - ) - - first = pipeline.synthesize_procedural("u1") - assert first["status"] == "synthesized" - assert pipeline._run_prompty.call_count == 1 - - synthesized = first["procedural"] - container.query_items.side_effect = [ - [synthesized], - list(fact_docs), - [], - [], - ] - - second = pipeline.synthesize_procedural("u1") - - assert pipeline._run_prompty.call_count == 1 - assert second["status"] == "unchanged" - assert second["procedural"]["id"] == synthesized["id"] - - -def test_get_procedural_prompt_returns_none_when_missing(): - client = _make_client() - client._memories_container_client.query_items.return_value = [] - - assert client.get_procedural_prompt("u1") is None - - -def test_get_procedural_prompt_returns_active_content(): - active_doc = _procedural_doc( - "proc_u1_2", - version=2, - content="Active prompt", - source_fact_ids=["f1"], - source_episodic_ids=["e1"], - ts=2, - ) - superseded_doc = _procedural_doc( - "proc_u1_1", - version=1, - content="Old prompt", - source_fact_ids=["f1"], - source_episodic_ids=["e1"], - superseded_by="proc_u1_2", - ts=1, - ) - docs = [superseded_doc, active_doc] - client = _make_client() - - def _query_items(**kwargs): - query = kwargs["query"] - if "superseded_by" in query: - return [doc for doc in docs if not doc.get("superseded_by")] - return list(docs) - - client._memories_container_client.query_items.side_effect = _query_items - - assert client.get_procedural_prompt("u1") == "Active prompt" - - -def test_get_procedural_history_returns_active_first_then_newest_versions(): - v1 = _procedural_doc( - "proc_u1_1", - version=1, - content="v1", - source_fact_ids=["f1"], - source_episodic_ids=["e1"], - superseded_by="proc_u1_2", - ts=1, - ) - v2 = _procedural_doc( - "proc_u1_2", - version=2, - content="v2", - source_fact_ids=["f1", "f2"], - source_episodic_ids=["e1"], - superseded_by="proc_u1_3", - ts=2, - ) - v3 = _procedural_doc( - "proc_u1_3", - version=3, - content="v3", - source_fact_ids=["f1", "f2", "f3"], - source_episodic_ids=["e1"], - ts=3, - ) - client = _make_client() - client._memories_container_client.query_items.return_value = [v1, v3, v2] - - history = client.get_procedural_history("u1", limit=10) - - assert [doc["id"] for doc in history] == ["proc_u1_3", "proc_u1_2", "proc_u1_1"] -def test_get_procedural_history_respects_limit(): - v1 = _procedural_doc( - "proc_u1_1", - version=1, - content="v1", - source_fact_ids=["f1"], - source_episodic_ids=["e1"], - superseded_by="proc_u1_2", - ts=1, - ) - v2 = _procedural_doc( - "proc_u1_2", - version=2, - content="v2", - source_fact_ids=["f1", "f2"], - source_episodic_ids=["e1"], - superseded_by="proc_u1_3", - ts=2, - ) - v3 = _procedural_doc( - "proc_u1_3", - version=3, - content="v3", - source_fact_ids=["f1", "f2", "f3"], - source_episodic_ids=["e1"], - ts=3, - ) - client = _make_client() - client._memories_container_client.query_items.return_value = [v1, v2, v3] - - history = client.get_procedural_history("u1", limit=2) - - assert [doc["id"] for doc in history] == ["proc_u1_3", "proc_u1_2"] - assert len(history) == 2 - - -def test_client_synthesize_procedural_raises_for_remote_processors(): - client = _make_client(processor=DurableFunctionProcessor()) - client._pipeline = MagicMock() - - with pytest.raises(NotImplementedError, match="durable mode"): - client.synthesize_procedural("u1") - - client._pipeline.synthesize_procedural.assert_not_called() - - -def test_synthesize_procedural_retries_with_fresh_llm_call_when_winner_has_partial_coverage(): - """Two concurrent writers race on ``proc_u1_2``; loser must re-read the - winner, see that its source set does NOT cover the loser's current set, - re-call the LLM with the winner as the new prior, and write at the - bumped version. Verifies content monotonicity in source coverage, not - just version number. - """ - from azure.cosmos.exceptions import CosmosResourceExistsError - - prior_v1 = _procedural_doc( - "proc_u1_1", - version=1, - content="v1 prompt", - source_fact_ids=["f1"], - source_episodic_ids=[], - ts=1, - ) - # Loser sees f1 and f2; winner only saw f1 → not covered → must re-call LLM. - fact_docs = [ - _fact_doc("f1", "Always use bullet points.", category="preference"), - _fact_doc("f2", "Never use var in TypeScript.", category="requirement"), - ] - - container = MagicMock() - container.query_items.side_effect = [ - [prior_v1], - fact_docs, - [], +def test_synthesize_procedural_extracts_atomic_procedures_and_gates_provenance() -> None: + store = _ProceduralStore([_fact(), _episode()]) + service = _service( + store, [ - prior_v1, - _procedural_doc( - "proc_u1_2", - version=2, - content="v2 by winner", - source_fact_ids=["f1"], - source_episodic_ids=[], - ts=2, - ), + { + "procedures": [ + _procedure( + "Targeted testing", + grounded_in=["fact-1"], + source_kind="explicit_user_instruction", + ), + _procedure( + "Retry CI failures", + grounded_in=["ep-1"], + source_kind="episode_distillation", + summary="Retry transient CI failures once before escalating.", + ), + ] + } ], - ] - upserted, capture = _capture_upserts() - container.upsert_item.side_effect = capture - - write_log: list[str] = [] - - def _create(*, body): - write_log.append(body["id"]) - if body["id"] == "proc_u1_2": - raise CosmosResourceExistsError(message="conflict") - capture(body=body) - return body + ) - container.create_item.side_effect = _create + result = service.synthesize_procedural("u1") - mock_embeddings = MagicMock() - containers = { - ContainerKey.TURNS: MagicMock(), - ContainerKey.MEMORIES: container, - ContainerKey.SUMMARIES: MagicMock(), + assert result == {"status": "synthesized", "procedures_created": 2, "procedures_skipped": 0} + procedures = [doc for doc in store.docs if doc.get("type") == "procedural"] + assert len(procedures) == 2 + assert {doc["name"]: doc["status"] for doc in procedures} == { + "Targeted testing": "active", + "Retry CI failures": "candidate", + } + assert {doc["name"]: doc["source_kind"] for doc in procedures} == { + "Targeted testing": "explicit_user_instruction", + "Retry CI failures": "episode_distillation", } - store = MemoryStore(containers=containers, embeddings_client=mock_embeddings) - pipeline = PipelineService(store, MagicMock(), mock_embeddings, containers=containers) - # Two LLM calls expected: once against v1, once against v2-winner after 409. - pipeline._run_prompty = MagicMock( - side_effect=[ - json.dumps({"system_prompt": "v2 by us (stale, will lose race)"}), - json.dumps({"system_prompt": "v3 by us (fresh, prior=winner)"}), + for doc in procedures: + assert doc["id"].startswith("proc_") + assert doc["type"] == "procedural" + assert doc["name"] + assert doc["retrieval_text"] + assert doc["thread_id"] == "__procedural__" + assert doc["embedding"] == [1.0] + # utility_score is seeded from the LLM confidence (0.8), not hard-wired to 0.5. + assert doc["utility_score"] == 0.8 + + +def test_synthesize_procedural_is_idempotent_by_scope_and_name() -> None: + store = _ProceduralStore([_fact()]) + response = { + "procedures": [ + _procedure( + "Targeted testing", + grounded_in=["fact-1"], + source_kind="explicit_user_instruction", + ) ] - ) - - result = pipeline.synthesize_procedural("u1", force=False) - - assert result["status"] == "synthesized" - assert result["procedural"]["id"] == "proc_u1_3" - assert result["procedural"]["version"] == 3 - assert write_log == ["proc_u1_2", "proc_u1_3"] - # LLM called twice: stale content discarded, fresh content reflects winner as prior. - assert pipeline._run_prompty.call_count == 2 - assert "v3 by us" in result["procedural"]["content"] - assert result["procedural"]["supersedes_ids"] == ["proc_u1_2"] - - -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. - """ - from azure.cosmos.exceptions import CosmosResourceExistsError - - prior_v1 = _procedural_doc( - "proc_u1_1", - version=1, - content="v1 prompt", - source_fact_ids=["f1"], - source_episodic_ids=[], - ts=1, - ) - fact_docs = [ - _fact_doc("f1", "Always use bullet points.", category="preference"), - _fact_doc("f2", "Never use var in TypeScript.", category="requirement"), - ] - # Winner already wrote v2 with the same source set we see now. - winner_v2 = _procedural_doc( - "proc_u1_2", - version=2, - content="v2 by winner", - source_fact_ids=["f1", "f2"], - source_episodic_ids=[], - ts=2, - ) - - container = MagicMock() - container.query_items.side_effect = [ - [prior_v1], - fact_docs, - [], - [prior_v1, winner_v2], - ] - upserted, capture = _capture_upserts() - container.upsert_item.side_effect = capture + } + service = _service(store, [response, response]) - write_log: list[str] = [] + first = service.synthesize_procedural("u1") + second = service.synthesize_procedural("u1") - def _create(*, body): - write_log.append(body["id"]) - if body["id"] == "proc_u1_2": - raise CosmosResourceExistsError(message="conflict") - capture(body=body) - return body + assert first == {"status": "synthesized", "procedures_created": 1, "procedures_skipped": 0} + assert second == {"status": "synthesized", "procedures_created": 0, "procedures_skipped": 1} + procedures = [doc for doc in store.docs if doc.get("type") == "procedural"] + assert len(procedures) == 1 - container.create_item.side_effect = _create - mock_embeddings = MagicMock() - containers = { - ContainerKey.TURNS: MagicMock(), - ContainerKey.MEMORIES: container, - ContainerKey.SUMMARIES: MagicMock(), +def test_build_procedural_context_uses_active_procedures_only() -> None: + active = { + "id": "proc-active", + "user_id": "u1", + "type": "procedural", + "status": "active", + "name": "Targeted testing", + "summary": "Run targeted tests before reporting success.", + "retrieval_text": "tests success", + "procedure_kind": "behavioral_policy", + "scope_type": "user", + "scope_value": None, + "priority": 10, + "source_authority": "high", + "version": 1, + } + candidate = { + **active, + "id": "proc-candidate", + "status": "candidate", + "name": "Candidate policy", + "summary": "Do not include this candidate procedure.", } - store = MemoryStore(containers=containers, embeddings_client=mock_embeddings) - pipeline = PipelineService(store, MagicMock(), mock_embeddings, containers=containers) - pipeline._run_prompty = MagicMock(return_value=json.dumps({"system_prompt": "stale v2 by us"})) + service = _service(_ProceduralStore([active, candidate])) - result = pipeline.synthesize_procedural("u1", force=False) + context = service.build_procedural_context("u1") - # Loser short-circuits: returns winner's doc, never retries the LLM. - assert result["status"] == "unchanged" - assert result["procedural"]["id"] == "proc_u1_2" - assert result["procedural"]["content"] == "v2 by winner" - # Only the losing write attempt; no retry write at proc_u1_3. - assert write_log == ["proc_u1_2"] - assert pipeline._run_prompty.call_count == 1 + assert "Run targeted tests before reporting success." in context + assert "Do not include this candidate procedure." not in context + assert service.build_procedural_context("missing-user") == "" diff --git a/tests/unit/test_process_now.py b/tests/unit/test_process_now.py index 11132bd..56d1231 100644 --- a/tests/unit/test_process_now.py +++ b/tests/unit/test_process_now.py @@ -31,7 +31,7 @@ def _patch_get_thread(client, turns): def test_process_now_with_inprocess_invokes_full_pipeline(): """process_now must fire ALL FIVE steps for InProcess: thread_summary, extract, reconcile, procedural, user_summary. Pre-fix this was only the first 3, so - procedural + user_summary never ran when callers used add_cosmos + process_now.""" + procedural + user_summary never ran when callers used upsert_memory + process_now.""" client = _connected() # default → InProcessProcessor lazily built pipeline = MagicMock() pipeline.generate_thread_summary.return_value = {"id": "s", "type": "thread_summary"} diff --git a/tests/unit/test_reconcile.py b/tests/unit/test_reconcile.py index cd94c3a..74279d5 100644 --- a/tests/unit/test_reconcile.py +++ b/tests/unit/test_reconcile.py @@ -29,16 +29,6 @@ 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() @@ -226,21 +216,9 @@ def _build(self) -> PipelineService: p._mark_superseded = MagicMock(return_value=True) return p - def test_extract_skips_when_content_hash_matches_existing(self): - from azure.cosmos.agent_memory._utils import compute_content_hash - + def test_extract_skips_in_batch_duplicate_facts(self): p = self._build() - existing_text = "User likes coffee" - existing = [ - { - "id": "fact_existing", - "type": "fact", - "content": existing_text, - "content_hash": compute_content_hash(existing_text), - "thread_id": "t1", - "tags": ["sys:fact"], - } - ] + dup_text = "User likes coffee" # extract_memories pulls turns directly from the container. turns = [ { @@ -252,19 +230,28 @@ def test_extract_skips_when_content_hash_matches_existing(self): } ] p._container.query_items.return_value = iter(turns) - p._load_existing_memories = MagicMock(return_value=existing) - # Stub the LLM extraction to emit a duplicate fact (same text). + # Stub the LLM extraction to emit the SAME fact text twice in one batch; + # the second is an in-batch exact duplicate and must be skipped. There is + # no store-side preload query - cross-run duplicates are handled by the + # deterministic-id create (409) at write time instead. p._run_prompty = MagicMock( return_value=json.dumps( { "facts": [ { - "text": existing_text, + "text": dup_text, "confidence": 0.9, "salience": 0.6, "action": "ADD", "tags": ["sys:fact"], - } + }, + { + "text": dup_text, + "confidence": 0.9, + "salience": 0.6, + "action": "ADD", + "tags": ["sys:fact"], + }, ], "procedural": [], "episodic": [], @@ -275,15 +262,13 @@ def test_extract_skips_when_content_hash_matches_existing(self): out = p.extract_memories("u1", "t1") assert out["exact_dedup_skipped"] >= 1 - assert out["fact_count"] == 0 - # No new fact upserted (the only ADD got short-circuited). - assert all(call.args[0].get("type") != "fact" for call in p._upsert_memory.call_args_list) + # Exactly one fact survives the in-batch dedup. + assert out["fact_count"] == 1 def test_extract_writes_content_hash_on_new_facts(self): from azure.cosmos.agent_memory._utils import compute_content_hash p = self._build() - p._load_existing_memories = MagicMock(return_value=[]) turns = [ { "id": "turn-1", @@ -337,19 +322,11 @@ def _build(self) -> PipelineService: return p def test_fact_not_dropped_when_only_procedural_has_same_hash(self): + # With no store-side preload, extraction never cross-checks other types, + # so a procedural doc sharing a fact's content_hash cannot affect the fact + # bucket. A normal fact ADD is created and nothing is exact-dedup skipped. p = self._build() text = "Always reply in Spanish" - # Existing PROCEDURAL with that text - must NOT poison the FACT bucket. - existing = [ - { - "id": "proc_existing", - "type": "procedural", - "content": text, - "content_hash": compute_content_hash(text), - "thread_id": "__procedural__", - "tags": ["sys:procedural"], - } - ] p._container.query_items.return_value = iter( [ { @@ -361,7 +338,6 @@ def test_fact_not_dropped_when_only_procedural_has_same_hash(self): } ] ) - p._load_existing_memories = MagicMock(return_value=existing) p._run_prompty = MagicMock( return_value=json.dumps( { @@ -657,7 +633,6 @@ def _build(self) -> PipelineService: p._turns_container = p._container p._summaries_container = p._container p._chat = MagicMock() - p._load_existing_memories = MagicMock(return_value=[]) return p def test_procedural_update_with_self_referential_id_is_skipped(self): diff --git a/tests/unit/test_thresholds.py b/tests/unit/test_thresholds.py index 43a5872..0c48e4b 100644 --- a/tests/unit/test_thresholds.py +++ b/tests/unit/test_thresholds.py @@ -56,7 +56,7 @@ def test_enable_turn_embeddings_falsy_values(monkeypatch, raw) -> None: @pytest.mark.parametrize( ("env_name", "getter_name", "expected"), [ - ("FACT_EXTRACTION_EVERY_N", "get_fact_extraction_every_n", 1), + ("FACT_EXTRACTION_EVERY_N", "get_fact_extraction_every_n", 2), ("THREAD_SUMMARY_EVERY_N", "get_thread_summary_every_n", 10), ("EPISODE_EVAL_EVERY_N", "get_episode_eval_every_n", 4), ("USER_SUMMARY_EVERY_N", "get_user_summary_every_n", 20), @@ -103,7 +103,7 @@ def test_env_config_getters_parse_env( @pytest.mark.parametrize( ("env_name", "getter_name", "expected"), [ - ("FACT_EXTRACTION_EVERY_N", "get_fact_extraction_every_n", 1), + ("FACT_EXTRACTION_EVERY_N", "get_fact_extraction_every_n", 2), ("THREAD_SUMMARY_EVERY_N", "get_thread_summary_every_n", 10), ("EPISODE_EVAL_EVERY_N", "get_episode_eval_every_n", 4), ("USER_SUMMARY_EVERY_N", "get_user_summary_every_n", 20), @@ -125,7 +125,7 @@ def test_int_getters_reject_negative( @pytest.mark.parametrize( ("env_name", "getter_name", "expected"), [ - ("FACT_EXTRACTION_EVERY_N", "get_fact_extraction_every_n", 1), + ("FACT_EXTRACTION_EVERY_N", "get_fact_extraction_every_n", 2), ("THREAD_SUMMARY_EVERY_N", "get_thread_summary_every_n", 10), ("EPISODE_EVAL_EVERY_N", "get_episode_eval_every_n", 4), ("USER_SUMMARY_EVERY_N", "get_user_summary_every_n", 20), @@ -196,18 +196,4 @@ def test_processor_owner_invalid_uses_default(monkeypatch: pytest.MonkeyPatch) - def test_internalized_getters_return_fixed_constants_and_ignore_env(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("EXTRACTION_BATCH_MAX_TOKENS", "999") - monkeypatch.setenv("DEDUP_SIM_HIGH", "0.50") - assert thresholds.get_extraction_batch_max_tokens() == 7000 - assert thresholds.get_dedup_sim_high() == 0.97 - - -def test_dedup_vector_enabled_defaults_false_and_reads_env(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.delenv("DEDUP_VECTOR_ENABLED", raising=False) - assert thresholds.get_dedup_vector_enabled() is False - - monkeypatch.setenv("DEDUP_VECTOR_ENABLED", "true") - assert thresholds.get_dedup_vector_enabled() is True - - monkeypatch.setenv("DEDUP_VECTOR_ENABLED", "false") - assert thresholds.get_dedup_vector_enabled() is False diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 8b462af..7d3af33 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -19,7 +19,6 @@ extract_keywords, normalize_ai_foundry_endpoint, vector_order_direction, - vector_similarity_at_least, ) from azure.cosmos.agent_memory.exceptions import ConfigurationError, ValidationError @@ -235,21 +234,6 @@ def test_vector_order_direction_per_function(): 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_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",