diff --git a/.env.template b/.env.template index a264e6e..caa0c95 100644 --- a/.env.template +++ b/.env.template @@ -29,7 +29,7 @@ USER_SUMMARY_EVERY_N=20 # the SDK and the Function App are deployed against the same database. # # * Unset / blank -> SDK auto-trigger fires; FA change-feed SKIPS. -# (Pure SDK deployments — no env config needed.) +# (Pure SDK deployments - no env config needed.) # * "inprocess" -> SDK auto-trigger fires; FA change-feed SKIPS. # * "durable" -> SDK auto-trigger SKIPS; FA change-feed fires. # diff --git a/.github/workflows/_test_release.yml b/.github/workflows/_test_release.yml index e712c7d..a821c89 100644 --- a/.github/workflows/_test_release.yml +++ b/.github/workflows/_test_release.yml @@ -85,7 +85,7 @@ jobs: verbose: true print-hash: true repository-url: https://test.pypi.org/legacy/ - # CI-only — overwrites a same-version file if a re-run is needed. + # CI-only - overwrites a same-version file if a re-run is needed. # https://github.com/pypa/gh-action-pypi-publish#tolerating-release-package-file-duplicates skip-existing: true # Attestations default-on in v1.11.0+ and require additional diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 2f773de..0756b4e 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -105,7 +105,7 @@ jobs: env: PKG_NAME: ${{ needs.build.outputs.pkg-name }} VERSION: ${{ needs.build.outputs.version }} - # Two-pass retry — TestPyPI mirror propagation can lag by a few seconds. + # Two-pass retry - TestPyPI mirror propagation can lag by a few seconds. # Primary index is real PyPI so all the dependencies resolve normally; # the just-published preview only lives on TestPyPI. run: | @@ -164,7 +164,7 @@ jobs: packages-dir: dist/ verbose: true print-hash: true - # Attestations default-on in v1.11.0+ — opt out until configured. + # Attestations default-on in v1.11.0+ - opt out until configured. attestations: false mark-release: diff --git a/CHANGELOG.md b/CHANGELOG.md index 1e37c75..5c7f2e0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -70,7 +70,7 @@ * `ProcessThreadResult` gains `procedural` and `user_summary` fields. `extract_memories` returns a `dropped_episodic_count` for monitoring LLM-extraction quality.See [PR:#20](https://github.com/aayush3011/AgentMemoryToolkit/pull/20) -## [0.1.0b1] — 2026-06-01 +## [0.1.0b1] - 2026-06-01 Initial public preview release. @@ -91,7 +91,7 @@ Pin a specific version when integrating. 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. + 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). diff --git a/Docs/RELEASING.md b/Docs/RELEASING.md index 23b7876..0e69bb8 100644 --- a/Docs/RELEASING.md +++ b/Docs/RELEASING.md @@ -19,7 +19,7 @@ This project uses [PEP 440](https://peps.python.org/pep-0440/) versioning: Before cutting a release: -1. **CI green on `main`** — every workflow in `.github/workflows/` passes. +1. **CI green on `main`** - every workflow in `.github/workflows/` passes. 2. **Unit tests pass locally** in a fresh venv: ```bash python -m venv /tmp/release-venv @@ -32,21 +32,21 @@ Before cutting a release: ```bash AGENT_MEMORY_RUN_INTEGRATION=true pytest tests/integration -q ``` -4. **Samples and notebooks work** — every script under `Samples/` runs to +4. **Samples and notebooks work** - every script under `Samples/` runs to completion against the live environment. -5. **No uncommitted local changes** — `git status` clean. +5. **No uncommitted local changes** - `git status` clean. ## Cutting a release 1. **Bump the version** in `pyproject.toml`: -2. **Update `CHANGELOG.md`** — add a new section with the version, the +2. **Update `CHANGELOG.md`** - add a new section with the version, the date, and a summary of changes. Move entries from the unreleased section if you keep one. 3. **Bump the Function app's SDK pin** in `function_app/requirements.txt` - to match — `azure-cosmos-agent-memory==`. The FA installs + to match - `azure-cosmos-agent-memory==`. The FA installs the SDK from PyPI, so the pin must move in lockstep with the SDK release. (If the release workflow fails after merge, the FA will be - pinned to a non-existent version until you cut a follow-up patch — + pinned to a non-existent version until you cut a follow-up patch - coordinate the merge + release-workflow run together.) 4. **Open a PR** with the version bump + CHANGELOG + updated FA pin. Suggested title: `Release v`. Get it reviewed and merged @@ -55,7 +55,7 @@ Before cutting a release: - Navigate to **Actions → release → Run workflow** - Pick the `main` branch - Click **Run workflow** -6. **The workflow does the rest** — see `.github/workflows/release.yml`: +6. **The workflow does the rest** - see `.github/workflows/release.yml`: - Builds `dist/*.whl` + `dist/*.tar.gz` - Asserts no namespace `__init__.py` shadows are in the wheel - Publishes to **TestPyPI** via trusted publishing @@ -69,7 +69,7 @@ Before cutting a release: ## Namespace package note `azure-cosmos-agent-memory` installs files under `azure/cosmos/agent_memory/`. -It MUST NOT ship `azure/__init__.py` or `azure/cosmos/__init__.py` — those +It MUST NOT ship `azure/__init__.py` or `azure/cosmos/__init__.py` - those are owned by the `azure-cosmos` package. The wheel build is configured (`[tool.setuptools.packages.find]` with `include = ["azure.cosmos.agent_memory*"]` and `namespaces = true`) so that only the `agent_memory` subtree is packaged. @@ -94,4 +94,4 @@ If a release ships a regression: 3. Add a deprecation note to the GitHub Release pointing at the replacement version. -Never delete a PyPI release — yank it and ship a fixed version instead. +Never delete a PyPI release - yank it and ship a fixed version instead. diff --git a/Docs/azure_testing.md b/Docs/azure_testing.md index 6b1523b..f2ac726 100644 --- a/Docs/azure_testing.md +++ b/Docs/azure_testing.md @@ -147,7 +147,7 @@ Set any threshold to `"0"` to disable that processing type. The `leases` container is provisioned by `create_memory_store()` alongside the `memories` and `counter` containers, so the Function App should be configured to use that existing lease container. -The Function App authenticates to Cosmos DB and Azure OpenAI via its managed identity — there's no shared key or function-key handoff between the SDK and the Function App. +The Function App authenticates to Cosmos DB and Azure OpenAI via its managed identity - there's no shared key or function-key handoff between the SDK and the Function App. --- @@ -300,7 +300,7 @@ 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 `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. ```python # Write enough turns to cross THREAD_SUMMARY_EVERY_N (default 10). diff --git a/Docs/concepts.md b/Docs/concepts.md index 7484ec0..2eba614 100644 --- a/Docs/concepts.md +++ b/Docs/concepts.md @@ -94,11 +94,11 @@ Memories stored in Cosmos DB include embeddings generated by Microsoft AI Foundr Facts work especially well for vector search because each fact is stored as a small, self-contained document. -By default raw conversation turns are *not* embedded — only derived memories (facts, episodic, procedural, summaries) carry vectors. Set `enable_turn_embeddings=True` (env `ENABLE_TURN_EMBEDDINGS`) to also embed turns on write, then call `search_turns()` to vector-search the raw conversation log. The turns container is always provisioned with a `quantizedFlat` vector index, so this flag only toggles embedding generation and can be turned on or off at any time without recreating the container. +By default raw conversation turns are *not* embedded - only derived memories (facts, episodic, procedural, summaries) carry vectors. Set `enable_turn_embeddings=True` (env `ENABLE_TURN_EMBEDDINGS`) to also embed turns on write, then call `search_turns()` to vector-search the raw conversation log. The turns container is always provisioned with a `quantizedFlat` vector index, so this flag only toggles embedding generation and can be turned on or off at any time without recreating the container. ### Unified retrieval (`include_turns`) -`search_cosmos(..., include_turns=True)` returns extracted memories plus raw conversation turns in one call — useful for recovering detail that extraction dropped. Up to `turn_top_k` turns (default `top_k`) are **appended after** the memory hits, so memory results keep priority; the two sets are not score-fused. A turn is skipped only when its content is an **exact** string match of a returned memory — paraphrased overlaps are not de-duplicated. The turn search requires `enable_turn_embeddings` and is best-effort: if it fails, the memory results are returned unchanged. +`search_cosmos(..., include_turns=True)` returns extracted memories plus raw conversation turns in one call - useful for recovering detail that extraction dropped. Up to `turn_top_k` turns (default `top_k`) are **appended after** the memory hits, so memory results keep priority; the two sets are not score-fused. A turn is skipped only when its content is an **exact** string match of a returned memory - paraphrased overlaps are not de-duplicated. The turn search requires `enable_turn_embeddings` and is best-effort: if it fails, the memory results are returned unchanged. --- @@ -122,7 +122,7 @@ Prompts for summarization and fact extraction live in `azure_functions/prompts/` ## Memory Reconciliation -Two independent mechanisms keep the fact pool clean and convergent: a cheap, LLM-free **write-time in-place dedup** that folds near-duplicate restatements into their existing record before they persist, and a periodic **LLM contradiction pass** that resolves opposing claims. Paraphrases are handled entirely at write time, so the LLM pass never merges duplicates — it only adjudicates contradictions. +Two independent mechanisms keep the fact pool clean and convergent: a cheap, LLM-free **write-time in-place dedup** that folds near-duplicate restatements into their existing record before they persist, and a periodic **LLM contradiction pass** that resolves opposing claims. Paraphrases are handled entirely at write time, so the LLM pass never merges duplicates - it only adjudicates contradictions. ### Write-time in-place dedup (LLM-free) @@ -130,19 +130,19 @@ Between extraction and persist, `dedup_extracted_memories` compares each newly e | condition (cosine) | action | |-----------------------------|-----------------------------------------------------------------------------------| -| `content_hash` hit | skip — identical re-extraction, no vector query, no write (`exact_dedup_skipped`) | +| `content_hash` hit | skip - identical re-extraction, no vector query, no write (`exact_dedup_skipped`) | | `s ≥ DEDUP_SIM_HIGH` (0.97) | **fold in place** into the existing neighbor (no LLM) | | `s < DEDUP_SIM_HIGH` | persist as a novel record | -A **fold** refreshes the existing neighbor rather than minting a new doc: it keeps the neighbor's `id` / `created_at` / partition, unions tags, takes the max salience/confidence, and bumps `updated_at`. Content and embedding are recency-wins **except** that a shorter restatement never overwrites longer content, so specifics captured by the richer record aren't dropped. The write is applied with ETag optimistic concurrency (`IfNotModified`); on an ETag conflict the fold is abandoned and the new doc is added as novel, so a concurrent supersede/refresh is never clobbered. This makes the write path convergent — a restatement updates one document instead of creating a duplicate that a later sweep must merge and supersede. +A **fold** refreshes the existing neighbor rather than minting a new doc: it keeps the neighbor's `id` / `created_at` / partition, unions tags, takes the max salience/confidence, and bumps `updated_at`. Content and embedding are recency-wins **except** that a shorter restatement never overwrites longer content, so specifics captured by the richer record aren't dropped. The write is applied with ETag optimistic concurrency (`IfNotModified`); on an ETag conflict the fold is abandoned and the new doc is added as novel, so a concurrent supersede/refresh is never clobbered. This makes the write path convergent - a restatement updates one document instead of creating a duplicate that a later sweep must merge and supersede. -The threshold is calibrated for **cosine/dotproduct** on normalized embeddings. On a container whose `distanceFunction` is **euclidean** — or when the distance policy can't be read — the destructive in-place fold is **disabled** (one-shot warning) and every extracted doc persists as novel, because cosine thresholds don't translate to unbounded euclidean distances. +The threshold is calibrated for **cosine/dotproduct** on normalized embeddings. On a container whose `distanceFunction` is **euclidean** - or when the distance policy can't be read - the destructive in-place fold is **disabled** (one-shot warning) and every extracted doc persists as novel, because cosine thresholds don't translate to unbounded euclidean distances. ### LLM contradiction reconcile -`reconcile_memories(user_id, n=50, *, memory_type="fact")` loads up to `n` active (non-superseded) facts, most recent first, and asks the dedup prompt to identify **contradicted pairs** — opposing claims about the same subject (e.g. "deadline March 1" vs "March 15"). Each loser is soft-deleted with `supersede_reason="contradict"` and `superseded_by` set to the winner (more recent wins, higher confidence as tiebreaker). Chained contradictions are guarded, so a fact already superseded in this pass can't be used to tombstone a third. +`reconcile_memories(user_id, n=50, *, memory_type="fact")` loads up to `n` active (non-superseded) facts, most recent first, and asks the dedup prompt to identify **contradicted pairs** - opposing claims about the same subject (e.g. "deadline March 1" vs "March 15"). Each loser is soft-deleted with `supersede_reason="contradict"` and `superseded_by` set to the winner (more recent wins, higher confidence as tiebreaker). Chained contradictions are guarded, so a fact already superseded in this pass can't be used to tombstone a third. -Near-duplicate **paraphrases are not merged here** — write-time in-place dedup already folds them before they land — so reconcile is a single, bounded, convergent pass: no clustering, no candidate/full-pool modes, no synthesized merged documents, and no re-merge churn. `episodic` and `procedural` types are no-ops (episodic has no contradiction semantics; its near-dups fold at write time). The pass runs over one flat pool of the most-recent active facts and returns `{"kept": int, "merged": int, "contradicted": int}` — `merged` is always `0`, retained for backward-compatible callers. +Near-duplicate **paraphrases are not merged here** - write-time in-place dedup already folds them before they land - so reconcile is a single, bounded, convergent pass: no clustering, no candidate/full-pool modes, no synthesized merged documents, and no re-merge churn. `episodic` and `procedural` types are no-ops (episodic has no contradiction semantics; its near-dups fold at write time). The pass runs over one flat pool of the most-recent active facts and returns `{"kept": int, "merged": int, "contradicted": int}` - `merged` is always `0`, retained for backward-compatible callers. ### Loser preservation @@ -154,20 +154,20 @@ Each fact written by `extract_memories` carries a `content_hash` (SHA-256 of nor ### Extraction gating (`extracted_at` + `recent_k`) -Extraction only ever reads turns not yet stamped `extracted_at`; `persist` stamps every turn it consumed once the extracted memories are durably written. This `extracted_at` gate is the authoritative "what's left to extract" signal and is shared by both backends — a failed or lagging extract leaves its turns unstamped, so the full backlog is retried on the next run and no turns are skipped. +Extraction only ever reads turns not yet stamped `extracted_at`; `persist` stamps every turn it consumed once the extracted memories are durably written. This `extracted_at` gate is the authoritative "what's left to extract" signal and is shared by both backends - a failed or lagging extract leaves its turns unstamped, so the full backlog is retried on the next run and no turns are skipped. -- **In-process SDK auto-trigger** — calls extraction with `recent_k=None`, so it drains the entire unextracted backlog each run and relies purely on the `extracted_at` gate. -- **Change-feed / Durable Function App** — additionally **sizes** `recent_k` from a per-thread **watermark** (`last_extract_count` on the counter doc): `recent_k = current_count − last_extract_count` (with `last_extract_count` treated as `0` before the first successful extract), then still applies the `extracted_at` filter. The watermark advances **only after a successful extract**. The window is deliberately **not** capped by `DEDUP_POOL_SIZE` (that knob governs the reconcile pool, not the extraction window) — capping would extract only the newest N and silently strand the oldest backlog turns. +- **In-process SDK auto-trigger** - calls extraction with `recent_k=None`, so it drains the entire unextracted backlog each run and relies purely on the `extracted_at` gate. +- **Change-feed / Durable Function App** - additionally **sizes** `recent_k` from a per-thread **watermark** (`last_extract_count` on the counter doc): `recent_k = current_count − last_extract_count` (with `last_extract_count` treated as `0` before the first successful extract), then still applies the `extracted_at` filter. The watermark advances **only after a successful extract**. The window is deliberately **not** capped by `DEDUP_POOL_SIZE` (that knob governs the reconcile pool, not the extraction window) - capping would extract only the newest N and silently strand the oldest backlog turns. -> **Caveat (rare, change-feed path):** the counter increment is best-effort — under sustained optimistic-concurrency contention it can drop an increment rather than block the user's write path (see `increment_counter_sync`). A dropped increment leaves `current_count` lagging the true turn count, which can under-size the watermark-derived `recent_k`. The Function App backend avoids stranding turns by raising to force change-feed redelivery; the in-process path is unaffected because it uses `recent_k=None` and the `extracted_at` gate, not the counter, to decide what to extract. +> **Caveat (rare, change-feed path):** the counter increment is best-effort - under sustained optimistic-concurrency contention it can drop an increment rather than block the user's write path (see `increment_counter_sync`). A dropped increment leaves `current_count` lagging the true turn count, which can under-size the watermark-derived `recent_k`. The Function App backend avoids stranding turns by raising to force change-feed redelivery; the in-process path is unaffected because it uses `recent_k=None` and the `extracted_at` gate, not the counter, to decide what to extract. ### Tunable Only three reconcile knobs are operator-configurable: -- `DEDUP_EVERY_N` (default `5`) — how often reconcile runs in the auto-trigger path (every Nth **extract**, not every Nth turn). Set to `0` to disable. -- `DEDUP_POOL_SIZE` (default `50`, hard cap `500`) — the pool size `n` passed to `reconcile_memories`; also overridable per call. Larger values give the LLM a wider view at higher token cost. -- `DEDUP_VECTOR_ENABLED` (default `false`) — write-time in-place near-duplicate folding. Default off = **add-only**. Set to `true` to fold near-duplicate restatements into their canonical record at write time. +- `DEDUP_EVERY_N` (default `5`) - how often reconcile runs in the auto-trigger path (every Nth **extract**, not every Nth turn). Set to `0` to disable. +- `DEDUP_POOL_SIZE` (default `50`, hard cap `500`) - the pool size `n` passed to `reconcile_memories`; also overridable per call. Larger values give the LLM a wider view at higher token cost. +- `DEDUP_VECTOR_ENABLED` (default `false`) - write-time in-place near-duplicate folding. Default off = **add-only**. Set to `true` to fold near-duplicate restatements into their canonical record at write time. @@ -214,7 +214,7 @@ Set any value to `0` to disable that processing type. For example, setting `THRE | Container | Partition Key | Purpose | |----------------------|-----------------------------------------|------------------------------------------------------------------------| | `memories` | `/user_id`, `/thread_id` (hierarchical) | Durable derived memories (`fact`, `episodic`, `procedural`) | -| `memories_turns` | `/user_id`, `/thread_id` (hierarchical) | Raw conversation turns (`turn`) — append-only, TTL-pruned | +| `memories_turns` | `/user_id`, `/thread_id` (hierarchical) | Raw conversation turns (`turn`) - append-only, TTL-pruned | | `memories_summaries` | `/user_id`, `/thread_id` (hierarchical) | Thread + user summaries (`thread_summary`, `user_summary`) | | `counter` | `/user_id`, `/thread_id` (hierarchical) | Message count tracking for automatic processing | | `leases` | `/id` | Change feed checkpointing container created by `create_memory_store()` | @@ -233,7 +233,7 @@ This keeps the change feed dependencies aligned with the main memory store inste | Mode | Trigger | Use case | |----------------------|----------------------------------------------|-----------------------------------------------------------------------------| | **On-demand (pull)** | SDK call (`generate_thread_summary()`, etc.) | Explicit control over when processing happens | -| **Automatic (push)** | Change feed trigger | Fire-and-forget — processing happens in the background as turns are written | +| **Automatic (push)** | Change feed trigger | Fire-and-forget - processing happens in the background as turns are written | Both modes use the same Durable Functions orchestrator and activities, so prompts, incremental update logic, and stored outputs are identical. diff --git a/Docs/design_patterns.md b/Docs/design_patterns.md index 0e520c3..0a4ed7d 100644 --- a/Docs/design_patterns.md +++ b/Docs/design_patterns.md @@ -75,10 +75,10 @@ await mem.delete_cosmos(memory_id="", user_id="user-1", thread_id=THREAD_ID) ### When to call -- **End of conversation** — after the user closes a session or a support ticket is resolved. -- **Long-running thread** — when a thread exceeds a token budget (e.g. > 50 turns) and you need a compact representation for context. -- **Periodic background job** — on a schedule to keep summaries up to date for active threads. -- **Automatic (change feed)** — set `THREAD_SUMMARY_EVERY_N` and the change feed trigger handles it. See [Section 8](#8-automatic-processing-with-change-feed). +- **End of conversation** - after the user closes a session or a support ticket is resolved. +- **Long-running thread** - when a thread exceeds a token budget (e.g. > 50 turns) and you need a compact representation for context. +- **Periodic background job** - on a schedule to keep summaries up to date for active threads. +- **Automatic (change feed)** - set `THREAD_SUMMARY_EVERY_N` and the change feed trigger handles it. See [Section 8](#8-automatic-processing-with-change-feed). Summaries are incremental: if one already exists for the thread, only newer turns are merged in. @@ -101,10 +101,10 @@ The summary is stored automatically in Cosmos with id `summary_user-1_thread-abc ### When to call -- **After each meaningful exchange** — extract facts from the latest turns so they are available for retrieval immediately. -- **End of conversation** — capture all discrete preferences, decisions, and requirements from the thread. -- **Before a planning step** — in multi-agent workflows, extract facts before handing context to a planner agent. -- **Automatic (change feed)** — set `FACT_EXTRACTION_EVERY_N` and the change feed trigger handles it. See [Section 8](#8-automatic-processing-with-change-feed). +- **After each meaningful exchange** - extract facts from the latest turns so they are available for retrieval immediately. +- **End of conversation** - capture all discrete preferences, decisions, and requirements from the thread. +- **Before a planning step** - in multi-agent workflows, extract facts before handing context to a planner agent. +- **Automatic (change feed)** - set `FACT_EXTRACTION_EVERY_N` and the change feed trigger handles it. See [Section 8](#8-automatic-processing-with-change-feed). Each fact is stored as its own document with its own embedding, making it ideal for fine-grained semantic search. @@ -124,10 +124,10 @@ result = await mem.extract_facts( ### When to call -- **Cross-session onboarding** — at the start of a new thread, generate (or update) the user summary so the agent has context from all prior conversations. -- **After a thread summary is created** — chain it: summarize the thread, then update the user summary. -- **On a schedule** — for users with many threads, run periodically to keep the profile current. -- **Automatic (change feed)** — set `USER_SUMMARY_EVERY_N` and the change feed trigger handles it. See [Section 8](#8-automatic-processing-with-change-feed). +- **Cross-session onboarding** - at the start of a new thread, generate (or update) the user summary so the agent has context from all prior conversations. +- **After a thread summary is created** - chain it: summarize the thread, then update the user summary. +- **On a schedule** - for users with many threads, run periodically to keep the profile current. +- **Automatic (change feed)** - set `USER_SUMMARY_EVERY_N` and the change feed trigger handles it. See [Section 8](#8-automatic-processing-with-change-feed). User summaries are also incremental. The pipeline merges only new thread data into the existing profile. @@ -207,7 +207,7 @@ New session starts │ │ ┌── Conversation loop ──┐ │ │ Store each turn │ (add_cosmos) - │ │ Optionally extract │ (extract_facts — every N turns or on key exchanges) + │ │ Optionally extract │ (extract_facts - every N turns or on key exchanges) │ └────────────────────────┘ │ ├─ Summarize the thread (generate_thread_summary) @@ -333,8 +333,8 @@ Set any value to `0` to disable that processing type. All three default to `0` ( The change feed trigger needs two additional Cosmos DB containers beyond the existing `memories` container: -- **`counter`** — stores lightweight per-thread and per-user message counters used for threshold checks -- **`leases`** — auto-created by the Azure Functions runtime for change feed checkpointing +- **`counter`** - stores lightweight per-thread and per-user message counters used for threshold checks +- **`leases`** - auto-created by the Azure Functions runtime for change feed checkpointing The `COSMOS_DB__accountEndpoint` setting must also be configured for the identity-based change feed binding. diff --git a/Docs/local_testing.md b/Docs/local_testing.md index 703e50f..45307e7 100644 --- a/Docs/local_testing.md +++ b/Docs/local_testing.md @@ -263,7 +263,7 @@ Expected functions include: - `extract_facts` - `generate_user_summary` - `http_start` -- `on_memory_change` (change feed trigger — only active when `COSMOS_DB__accountEndpoint` is set) +- `on_memory_change` (change feed trigger - only active when `COSMOS_DB__accountEndpoint` is set) ### Function keys @@ -358,7 +358,7 @@ If you have configured the change feed settings above, you can test automatic pr 1. Set a low threshold for testing, e.g. `THREAD_SUMMARY_EVERY_N=3`. 2. Write turns to Cosmos (via the SDK or `curl`) until the threshold is crossed. -3. Watch the Functions host logs — you should see the orchestrator being started automatically. +3. Watch the Functions host logs - you should see the orchestrator being started automatically. ```python import uuid diff --git a/Docs/operations.md b/Docs/operations.md index e6d9d26..91ba69b 100644 --- a/Docs/operations.md +++ b/Docs/operations.md @@ -1,6 +1,6 @@ # Operations -Runtime knobs for an Agent-Memory-Toolkit deployment. Most ops levers live in `.env` / Function-app App Settings — change them, restart the consumer, and you're done. Deployment-time knobs (Bicep params bound to `azd env set ...`) live in [`infra/README.md`](../infra/README.md). +Runtime knobs for an Agent-Memory-Toolkit deployment. Most ops levers live in `.env` / Function-app App Settings - change them, restart the consumer, and you're done. Deployment-time knobs (Bicep params bound to `azd env set ...`) live in [`infra/README.md`](../infra/README.md). ## Memory lifecycle (TTL) diff --git a/Docs/public_api.md b/Docs/public_api.md index 96337d3..3030bba 100644 --- a/Docs/public_api.md +++ b/Docs/public_api.md @@ -4,7 +4,7 @@ `CosmosMemoryClient` and `AsyncCosmosMemoryClient` are thin orchestrators. They keep local-buffer state and Cosmos connection lifecycle, then delegate persistence to `MemoryStore` / `AsyncMemoryStore` and higher-level behavior to: -- `ChatClient` / `EmbeddingsClient` (sync) and `AsyncEmbeddingsClient` (async) — Azure OpenAI wrappers. +- `ChatClient` / `EmbeddingsClient` (sync) and `AsyncEmbeddingsClient` (async) - Azure OpenAI wrappers. - `RetrievalService` / `AsyncRetrievalService` for filtering, vector search, and episodic context. - `PipelineService` for extraction, summaries, procedural synthesis, and reconciliation. - `InProcessProcessor` / `AsyncInProcessProcessor` / `DurableFunctionProcessor` for immediate or change-feed-driven processing. @@ -14,53 +14,53 @@ ### Connection -- `__init__(cosmos_endpoint=None, cosmos_credential=None, cosmos_key=None, cosmos_database=None, cosmos_container=None, cosmos_turns_container='memories_turns', cosmos_summaries_container='memories_summaries', cosmos_counter_container=None, cosmos_lease_container=None, cosmos_throughput_mode=None, cosmos_autoscale_max_ru=None, ai_foundry_endpoint=None, ai_foundry_credential=None, ai_foundry_api_key=None, embedding_deployment_name='text-embedding-3-large', embedding_dimensions=None, chat_deployment_name='gpt-4o-mini', use_default_credential=True, enable_turn_embeddings=None, processor=None) -> None` — configure local state, model clients, optional Cosmos auto-connect, and optional processing backend. The SDK uses a hard 3-container topology: turns in `memories_turns`, facts/episodic/procedural in `memories`, and summaries in `memories_summaries` (or the names you pass). `enable_turn_embeddings` (default `False`, env `ENABLE_TURN_EMBEDDINGS`) embeds raw turns on write so they can be vector-searched via `search_turns()`; the turns container is always provisioned with a vector index, so toggling this never requires recreating it. -- `close() -> None` — close Cosmos/model clients and owned credentials. -- `connect_cosmos(endpoint=None, credential=None, key=None, database=None, container=None, turns_container=None, summaries_container=None) -> None` — connect to existing memory, turns, and summaries containers. -- `create_memory_store(database=None, container=None, turns_container=None, summaries_container=None, counter_container=None, lease_container=None, endpoint=None, credential=None, key=None, embedding_dimensions=None, embedding_data_type=None, distance_function=None, full_text_language=None, throughput_mode=None, autoscale_max_ru=None) -> None` — create/connect the memory, turns, summaries, counter, and lease containers. -- `validate_topology() -> None` — read metadata for all three memory containers and raise `RuntimeError` if any is missing or unreachable; call after connecting to catch infrastructure/config drift before writes. +- `__init__(cosmos_endpoint=None, cosmos_credential=None, cosmos_key=None, cosmos_database=None, cosmos_container=None, cosmos_turns_container='memories_turns', cosmos_summaries_container='memories_summaries', cosmos_counter_container=None, cosmos_lease_container=None, cosmos_throughput_mode=None, cosmos_autoscale_max_ru=None, ai_foundry_endpoint=None, ai_foundry_credential=None, ai_foundry_api_key=None, embedding_deployment_name='text-embedding-3-large', embedding_dimensions=None, chat_deployment_name='gpt-4o-mini', use_default_credential=True, enable_turn_embeddings=None, processor=None) -> None` - configure local state, model clients, optional Cosmos auto-connect, and optional processing backend. The SDK uses a hard 3-container topology: turns in `memories_turns`, facts/episodic/procedural in `memories`, and summaries in `memories_summaries` (or the names you pass). `enable_turn_embeddings` (default `False`, env `ENABLE_TURN_EMBEDDINGS`) embeds raw turns on write so they can be vector-searched via `search_turns()`; the turns container is always provisioned with a vector index, so toggling this never requires recreating it. +- `close() -> None` - close Cosmos/model clients and owned credentials. +- `connect_cosmos(endpoint=None, credential=None, key=None, database=None, container=None, turns_container=None, summaries_container=None) -> None` - connect to existing memory, turns, and summaries containers. +- `create_memory_store(database=None, container=None, turns_container=None, summaries_container=None, counter_container=None, lease_container=None, endpoint=None, credential=None, key=None, embedding_dimensions=None, embedding_data_type=None, distance_function=None, full_text_language=None, throughput_mode=None, autoscale_max_ru=None) -> None` - create/connect the memory, turns, summaries, counter, and lease containers. +- `validate_topology() -> None` - read metadata for all three memory containers and raise `RuntimeError` if any is missing or unreachable; call after connecting to catch infrastructure/config drift before writes. ### Memory CRUD -- `add_local(user_id, role, content, memory_type='turn', agent_id=None, metadata=None, thread_id=None, tags=None, ttl=None, salience=None) -> None` — append a memory to the local buffer. -- `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. -- `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`. -- `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. +- `add_local(user_id, role, content, memory_type='turn', agent_id=None, metadata=None, thread_id=None, tags=None, ttl=None, salience=None) -> None` - append a memory to the local buffer. +- `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. +- `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`. +- `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_turns(search_terms, user_id, thread_id=None, role=None, top_k=5, tags_all=None, tags_any=None, exclude_tags=None, created_after=None, created_before=None) -> list[dict]` — hybrid vector/full-text search the raw conversation log instead of facts/episodic/procedural (requires turn embeddings; see `enable_turn_embeddings`). `user_id` is required so the search is scoped to one partition instead of scanning every user's turns. -- `get_procedural_prompt(user_id) -> Optional[str]` — read the active procedural prompt. -- `get_procedural_history(user_id, limit=10) -> list[dict]` — read procedural prompt history. -- `get_procedural_memories(user_id, priority=None, category=None, min_salience=None, include_superseded=False) -> list[dict]` — retrieve procedural memory documents. -- `search_episodic_memories(user_id, search_terms, top_k=5, min_salience=None, include_superseded=False) -> list[dict]` — search episodic memories. -- `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. +- `search_cosmos(search_terms, memory_id=None, user_id=None, role=None, memory_types=None, thread_id=None, top_k=5, tags_all=None, tags_any=None, exclude_tags=None, include_superseded=False, min_salience=None, min_confidence=None, created_after=None, created_before=None) -> list[dict]` - hybrid vector/full-text search memories, falling back to vector-only for all-stopword queries. +- `search_turns(search_terms, user_id, thread_id=None, role=None, top_k=5, tags_all=None, tags_any=None, exclude_tags=None, created_after=None, created_before=None) -> list[dict]` - hybrid vector/full-text search the raw conversation log instead of facts/episodic/procedural (requires turn embeddings; see `enable_turn_embeddings`). `user_id` is required so the search is scoped to one partition instead of scanning every user's turns. +- `get_procedural_prompt(user_id) -> Optional[str]` - read the active procedural prompt. +- `get_procedural_history(user_id, limit=10) -> list[dict]` - read procedural prompt history. +- `get_procedural_memories(user_id, priority=None, category=None, min_salience=None, include_superseded=False) -> list[dict]` - retrieve procedural memory documents. +- `search_episodic_memories(user_id, search_terms, top_k=5, min_salience=None, include_superseded=False) -> list[dict]` - search episodic memories. +- `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. ### Processing -- `extract_memories(user_id, thread_id, recent_k=None) -> dict[str, int]` — extract facts/episodic memories from a thread. -- `synthesize_procedural(user_id, *, force=False) -> dict` — synthesize the procedural prompt. -- `generate_thread_summary(user_id, thread_id, recent_k=None, **kwargs) -> dict` — generate and persist a thread summary. -- `generate_user_summary(user_id, thread_ids=None, recent_k=None, **kwargs) -> dict` — generate and persist a user summary. -- `reconcile(user_id, n=None) -> dict[str, int]` — resolve contradictory facts (paraphrases fold at write time). -- `process_now(*, user_id, thread_id) -> ProcessThreadResult` — run the configured processor immediately. -- `process_now_and_wait(*, user_id, thread_id, timeout=30.0) -> bool` — process and wait for a summary. +- `extract_memories(user_id, thread_id, recent_k=None) -> dict[str, int]` - extract facts/episodic memories from a thread. +- `synthesize_procedural(user_id, *, force=False) -> dict` - synthesize the procedural prompt. +- `generate_thread_summary(user_id, thread_id, recent_k=None, **kwargs) -> dict` - generate and persist a thread summary. +- `generate_user_summary(user_id, thread_ids=None, recent_k=None, **kwargs) -> dict` - generate and persist a user summary. +- `reconcile(user_id, n=None) -> dict[str, int]` - resolve contradictory facts (paraphrases fold at write time). +- `process_now(*, user_id, thread_id) -> ProcessThreadResult` - run the configured processor immediately. +- `process_now_and_wait(*, user_id, thread_id, timeout=30.0) -> bool` - process and wait for a summary. ### Tagging -- `add_tags(memory_id, user_id, thread_id, memory_type, tags) -> None` — add tags to a memory. `memory_type` must be one of `fact`, `episodic`, `procedural`. -- `remove_tags(memory_id, user_id, thread_id, memory_type, tags) -> None` — remove tags from a memory. `memory_type` must be one of `fact`, `episodic`, `procedural`. -- `list_tags(user_id, *, thread_id=None, prefix=None, include_sys=False) -> list[str]` — list sorted, deduped tags for a user; omits `sys:*` by default. +- `add_tags(memory_id, user_id, thread_id, memory_type, tags) -> None` - add tags to a memory. `memory_type` must be one of `fact`, `episodic`, `procedural`. +- `remove_tags(memory_id, user_id, thread_id, memory_type, tags) -> None` - remove tags from a memory. `memory_type` must be one of `fact`, `episodic`, `procedural`. +- `list_tags(user_id, *, thread_id=None, prefix=None, include_sys=False) -> list[str]` - list sorted, deduped tags for a user; omits `sys:*` by default. ## AsyncCosmosMemoryClient @@ -68,53 +68,53 @@ Local-buffer methods remain synchronous in-memory operations; Cosmos, retrieval, ### Connection -- `__init__(cosmos_endpoint=None, cosmos_credential=None, cosmos_key=None, cosmos_database=None, cosmos_container=None, cosmos_turns_container='memories_turns', cosmos_summaries_container='memories_summaries', cosmos_counter_container=None, cosmos_lease_container=None, cosmos_throughput_mode=None, cosmos_autoscale_max_ru=None, ai_foundry_endpoint=None, ai_foundry_credential=None, ai_foundry_api_key=None, embedding_deployment_name='text-embedding-3-large', embedding_dimensions=None, chat_deployment_name='gpt-4o-mini', use_default_credential=True, enable_turn_embeddings=None, processor=None) -> None` — configure async local state, model clients, and optional processing backend. The async SDK uses the same hard 3-container topology as the sync client. `enable_turn_embeddings` (default `False`, env `ENABLE_TURN_EMBEDDINGS`) embeds raw turns on write so they can be vector-searched via `search_turns()`. -- `async close() -> None` — close async/sync resources and owned credentials. -- `async connect_cosmos(endpoint=None, credential=None, key=None, database=None, container=None, turns_container=None, summaries_container=None) -> None` — connect to existing memory, turns, and summaries containers. -- `async create_memory_store(database=None, container=None, turns_container=None, summaries_container=None, counter_container=None, lease_container=None, endpoint=None, credential=None, key=None, embedding_dimensions=None, embedding_data_type=None, distance_function=None, full_text_language=None, throughput_mode=None, autoscale_max_ru=None) -> None` — create/connect memory, turns, summaries, counter, and lease containers. -- `async validate_topology() -> None` — read metadata for all three memory containers and raise `RuntimeError` if any is missing or unreachable; call after connecting to catch infrastructure/config drift before writes. +- `__init__(cosmos_endpoint=None, cosmos_credential=None, cosmos_key=None, cosmos_database=None, cosmos_container=None, cosmos_turns_container='memories_turns', cosmos_summaries_container='memories_summaries', cosmos_counter_container=None, cosmos_lease_container=None, cosmos_throughput_mode=None, cosmos_autoscale_max_ru=None, ai_foundry_endpoint=None, ai_foundry_credential=None, ai_foundry_api_key=None, embedding_deployment_name='text-embedding-3-large', embedding_dimensions=None, chat_deployment_name='gpt-4o-mini', use_default_credential=True, enable_turn_embeddings=None, processor=None) -> None` - configure async local state, model clients, and optional processing backend. The async SDK uses the same hard 3-container topology as the sync client. `enable_turn_embeddings` (default `False`, env `ENABLE_TURN_EMBEDDINGS`) embeds raw turns on write so they can be vector-searched via `search_turns()`. +- `async close() -> None` - close async/sync resources and owned credentials. +- `async connect_cosmos(endpoint=None, credential=None, key=None, database=None, container=None, turns_container=None, summaries_container=None) -> None` - connect to existing memory, turns, and summaries containers. +- `async create_memory_store(database=None, container=None, turns_container=None, summaries_container=None, counter_container=None, lease_container=None, endpoint=None, credential=None, key=None, embedding_dimensions=None, embedding_data_type=None, distance_function=None, full_text_language=None, throughput_mode=None, autoscale_max_ru=None) -> None` - create/connect memory, turns, summaries, counter, and lease containers. +- `async validate_topology() -> None` - read metadata for all three memory containers and raise `RuntimeError` if any is missing or unreachable; call after connecting to catch infrastructure/config drift before writes. ### Memory CRUD -- `add_local(user_id, role, content, memory_type='turn', agent_id=None, metadata=None, thread_id=None, tags=None, ttl=None, salience=None) -> None` — append a memory to the local buffer. -- `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 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 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. +- `add_local(user_id, role, content, memory_type='turn', agent_id=None, metadata=None, thread_id=None, tags=None, ttl=None, salience=None) -> None` - append a memory to the local buffer. +- `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 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 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_turns(search_terms, user_id, thread_id=None, role=None, top_k=5, tags_all=None, tags_any=None, exclude_tags=None, created_after=None, created_before=None) -> list[dict]` — hybrid vector/full-text search the raw conversation log instead of facts/episodic/procedural (requires turn embeddings; see `enable_turn_embeddings`). `user_id` is required so the search is scoped to one partition instead of scanning every user's turns. -- `async get_procedural_prompt(user_id) -> Optional[str]` — read the active procedural prompt. -- `async get_procedural_history(user_id, limit=10) -> list[dict]` — read procedural prompt history. -- `async get_procedural_memories(user_id, priority=None, category=None, min_salience=None, include_superseded=False) -> list[dict]` — retrieve procedural memory documents. -- `async search_episodic_memories(user_id, search_terms, top_k=5, min_salience=None, include_superseded=False) -> list[dict]` — search episodic memories. -- `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. +- `async search_cosmos(search_terms, memory_id=None, user_id=None, role=None, memory_types=None, thread_id=None, top_k=5, tags_all=None, tags_any=None, exclude_tags=None, include_superseded=False, min_salience=None, min_confidence=None, created_after=None, created_before=None) -> list[dict]` - hybrid vector/full-text search memories, falling back to vector-only for all-stopword queries. +- `async search_turns(search_terms, user_id, thread_id=None, role=None, top_k=5, tags_all=None, tags_any=None, exclude_tags=None, created_after=None, created_before=None) -> list[dict]` - hybrid vector/full-text search the raw conversation log instead of facts/episodic/procedural (requires turn embeddings; see `enable_turn_embeddings`). `user_id` is required so the search is scoped to one partition instead of scanning every user's turns. +- `async get_procedural_prompt(user_id) -> Optional[str]` - read the active procedural prompt. +- `async get_procedural_history(user_id, limit=10) -> list[dict]` - read procedural prompt history. +- `async get_procedural_memories(user_id, priority=None, category=None, min_salience=None, include_superseded=False) -> list[dict]` - retrieve procedural memory documents. +- `async search_episodic_memories(user_id, search_terms, top_k=5, min_salience=None, include_superseded=False) -> list[dict]` - search episodic memories. +- `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. ### Processing -- `async extract_memories(user_id, thread_id, recent_k=None) -> dict[str, int]` — extract facts/episodic memories from a thread. -- `async synthesize_procedural(user_id, *, force=False) -> dict` — synthesize the procedural prompt. -- `async generate_thread_summary(user_id, thread_id, recent_k=None, **kwargs) -> dict` — generate and persist a thread summary. -- `async generate_user_summary(user_id, thread_ids=None, recent_k=None, **kwargs) -> dict` — generate and persist a user summary. -- `async reconcile(user_id, n=None) -> dict[str, int]` — resolve contradictory facts (paraphrases fold at write time). -- `async process_now(*, user_id, thread_id) -> ProcessThreadResult` — run the configured processor immediately. -- `async process_now_and_wait(*, user_id, thread_id, timeout=30.0) -> bool` — process and wait for a summary. +- `async extract_memories(user_id, thread_id, recent_k=None) -> dict[str, int]` - extract facts/episodic memories from a thread. +- `async synthesize_procedural(user_id, *, force=False) -> dict` - synthesize the procedural prompt. +- `async generate_thread_summary(user_id, thread_id, recent_k=None, **kwargs) -> dict` - generate and persist a thread summary. +- `async generate_user_summary(user_id, thread_ids=None, recent_k=None, **kwargs) -> dict` - generate and persist a user summary. +- `async reconcile(user_id, n=None) -> dict[str, int]` - resolve contradictory facts (paraphrases fold at write time). +- `async process_now(*, user_id, thread_id) -> ProcessThreadResult` - run the configured processor immediately. +- `async process_now_and_wait(*, user_id, thread_id, timeout=30.0) -> bool` - process and wait for a summary. ### Tagging -- `async add_tags(memory_id, user_id, thread_id, memory_type, tags) -> None` — add tags to a memory. `memory_type` must be one of `fact`, `episodic`, `procedural`. -- `async remove_tags(memory_id, user_id, thread_id, memory_type, tags) -> None` — remove tags from a memory. `memory_type` must be one of `fact`, `episodic`, `procedural`. -- `async list_tags(user_id, *, thread_id=None, prefix=None, include_sys=False) -> list[str]` — list sorted, deduped tags for a user; omits `sys:*` by default. +- `async add_tags(memory_id, user_id, thread_id, memory_type, tags) -> None` - add tags to a memory. `memory_type` must be one of `fact`, `episodic`, `procedural`. +- `async remove_tags(memory_id, user_id, thread_id, memory_type, tags) -> None` - remove tags from a memory. `memory_type` must be one of `fact`, `episodic`, `procedural`. +- `async list_tags(user_id, *, thread_id=None, prefix=None, include_sys=False) -> list[str]` - list sorted, deduped tags for a user; omits `sys:*` by default. ## Topology validation @@ -129,5 +129,5 @@ Sync extension protocols live in `azure.cosmos.agent_memory.services`; async var Concrete service classes are exported from their respective packages: - Sync: `RetrievalService`, `PipelineService` from `azure.cosmos.agent_memory.services` (sub-modules `retrieval`, `pipeline`). -- Async: `AsyncRetrievalService` and `AsyncPipelineService` from `azure.cosmos.agent_memory.aio.services` (sub-modules `retrieval`, `pipeline`). The async pipeline is a fully-native asyncio implementation — not an `asyncio.to_thread` shim over the sync pipeline. +- Async: `AsyncRetrievalService` and `AsyncPipelineService` from `azure.cosmos.agent_memory.aio.services` (sub-modules `retrieval`, `pipeline`). The async pipeline is a fully-native asyncio implementation - not an `asyncio.to_thread` shim over the sync pipeline. - Threshold-driven auto-trigger: `maybe_trigger_steps` from `azure.cosmos.agent_memory.auto_trigger` (sync) and `azure.cosmos.agent_memory.aio.auto_trigger` (async). diff --git a/Docs/troubleshooting.md b/Docs/troubleshooting.md index a64d1cd..5459938 100644 --- a/Docs/troubleshooting.md +++ b/Docs/troubleshooting.md @@ -157,7 +157,7 @@ Automatic processing requires these settings in the Functions app or `local.sett "USER_SUMMARY_EVERY_N": "10" ``` -Set a threshold to `"0"` to disable that processing type. `MEMORY_PROCESSOR_OWNER` must be `"durable"` for the Function App's change-feed trigger to actually fire — leave it unset (or set to `"inprocess"`) for SDK-only deployments. +Set a threshold to `"0"` to disable that processing type. `MEMORY_PROCESSOR_OWNER` must be `"durable"` for the Function App's change-feed trigger to actually fire - leave it unset (or set to `"inprocess"`) for SDK-only deployments. Cosmos DB memory documents store their category in the JSON `type` field. Only documents with `type: "turn"` increment counters. Derived memories with `type: "thread_summary"`, `type: "fact"`, `type: "episodic"`, `type: "procedural"`, or `type: "user_summary"` do not trigger threshold counts. diff --git a/README.md b/README.md index 1927225..136fec9 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ [![YouTube](https://img.shields.io/badge/YouTube-Azure%20Cosmos%20DB-FF0000?logo=youtube&logoColor=white)](https://www.youtube.com/@AzureCosmosDB) -Agent Memory Toolkit is a Python SDK for storing, retrieving, and transforming agent memories on Azure Cosmos DB. It gives your agent both raw conversation history and higher-value derived memory — thread summaries, extracted facts, and cross-thread user profiles — all searchable semantically. The processing pipeline can run **in-process** (zero infra) or in a sibling **Azure Durable Function app** that watches the Cosmos DB change feed. Sync (`CosmosMemoryClient`) and async (`AsyncCosmosMemoryClient`) APIs are mirror-images of each other. +Agent Memory Toolkit is a Python SDK for storing, retrieving, and transforming agent memories on Azure Cosmos DB. It gives your agent both raw conversation history and higher-value derived memory - thread summaries, extracted facts, and cross-thread user profiles - all searchable semantically. The processing pipeline can run **in-process** (zero infra) or in a sibling **Azure Durable Function app** that watches the Cosmos DB change feed. Sync (`CosmosMemoryClient`) and async (`AsyncCosmosMemoryClient`) APIs are mirror-images of each other. --- @@ -31,7 +31,7 @@ pip install ".[dev]" The toolkit needs a Cosmos DB account, an Azure OpenAI / AI Foundry deployment, and (optionally for the remote processor) an Azure Function app. Pick whichever path matches your situation: -**Option A — One-command provision (`azd up`).** Creates everything from scratch — Cosmos + AI Foundry + Function app (Flex Consumption, idle cost ≈ $0) + UAMI + RBAC — and writes a working `.env` to `.azure//.env`: +**Option A - One-command provision (`azd up`).** Creates everything from scratch - Cosmos + AI Foundry + Function app (Flex Consumption, idle cost ≈ $0) + UAMI + RBAC - and writes a working `.env` to `.azure//.env`: ```bash # Prereqs: az + azd installed; subscription with quota for gpt-4o-mini @@ -51,7 +51,7 @@ azd up # are provisioned. Outputs are written to .azure/memorytoolkit-dev/.env ``` -The Function app is always provisioned but only used when you opt into `DurableFunctionProcessor` — it sits idle (and bills nothing) for in-process workloads. +The Function app is always provisioned but only used when you opt into `DurableFunctionProcessor` - it sits idle (and bills nothing) for in-process workloads. Load the generated env vars and you're ready to use the SDK: @@ -61,14 +61,14 @@ set -a && . ./.azure/memorytoolkit-dev/.env && set +a To tear everything down later: `azd down --purge` (the `--purge` flag skips Cosmos / AI Foundry soft-delete so names are immediately reusable). -**Option B — Bring your own resources.** If you already have a Cosmos DB account and an AI Foundry / Azure OpenAI deployment, copy the env template and fill in the endpoints: +**Option B - Bring your own resources.** If you already have a Cosmos DB account and an AI Foundry / Azure OpenAI deployment, copy the env template and fill in the endpoints: ```bash cp .env.template .env # edit COSMOS_DB_ENDPOINT, AI_FOUNDRY_ENDPOINT, AI_FOUNDRY_EMBEDDING_DEPLOYMENT_NAME, AI_FOUNDRY_CHAT_DEPLOYMENT_NAME ``` -> For the Durable Function app counter-trigger settings, Bicep module reference, RBAC scopes, and the SDK-only escape hatch (`DEPLOY_FUNCTION_APP=false`) — see **[`infra/README.md`](infra/README.md)**. +> For the Durable Function app counter-trigger settings, Bicep module reference, RBAC scopes, and the SDK-only escape hatch (`DEPLOY_FUNCTION_APP=false`) - see **[`infra/README.md`](infra/README.md)**. ### 3. Use the SDK @@ -109,7 +109,7 @@ for h in hits: print(memory.get_user_summary(user_id=USER)) ``` -> Async API is identical — just `await` each call: +> Async API is identical - just `await` each call: > ```python > from azure.cosmos.agent_memory.aio import AsyncCosmosMemoryClient > ``` @@ -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 | `add_cosmos(...)`, `add_local(...)` | | **Thread summary** | LLM-generated, incrementally updated rollup of a single thread | `generate_thread_summary(...)` | | **Fact** | Discrete, independently searchable assertion extracted from turns | `extract_memories(...)` | | **Procedural** | Behavioral rule / instruction the user wants followed | `extract_memories(...)` | @@ -162,7 +162,7 @@ The `extract_memories` pipeline classifies each item it pulls from the conversat |---------|----------------------------------------------------| | 0.9–1.0 | Directly stated and unambiguous | | 0.7–0.9 | Clearly implied, no contradicting evidence | -| 0.5–0.7 | Inferred from context — plausible but not explicit | +| 0.5–0.7 | Inferred from context - plausible but not explicit | | < 0.5 | Should be in `unclassified` instead | Filter at retrieval time: @@ -174,9 +174,9 @@ high_conf_facts = memory.get_memories(user_id="u1", memory_types=["fact"], min_c ### Memory Reconciliation -`reconcile(user_id, n=50)` (on the public client; underlying pipeline method is `ProcessingPipeline.reconcile_memories`) resolves **semantic contradictions** in a single LLM pass over the N most-recent active facts, soft-deleting each loser with `supersede_reason="contradict"`. Paraphrased duplicates are *not* handled here — they are folded in place at write time by the LLM-free vector dedup (see below), so reconcile stays a bounded, convergent contradiction pass. See [Docs/concepts.md](Docs/concepts.md#memory-reconciliation) for details. +`reconcile(user_id, n=50)` (on the public client; underlying pipeline method is `ProcessingPipeline.reconcile_memories`) resolves **semantic contradictions** in a single LLM pass over the N most-recent active facts, soft-deleting each loser with `supersede_reason="contradict"`. Paraphrased duplicates are *not* handled here - they are folded in place at write time by the LLM-free vector dedup (see below), so reconcile stays a bounded, convergent contradiction pass. See [Docs/concepts.md](Docs/concepts.md#memory-reconciliation) for details. -> **Cost note.** Each reconciliation makes one LLM call covering up to `n` facts (default 50, hard cap 500). With auto-trigger, this fires every `FACT_EXTRACTION_EVERY_N × DEDUP_EVERY_N` turns per user, with `n` taken from `DEDUP_POOL_SIZE`. The previous cosine-cluster pre-filter was removed deliberately — it could not catch semantic contradictions like "vegetarian" vs "ribeye steak" — so the LLM is now invoked whenever there are ≥ 2 active facts. To bound LLM cost more tightly: raise `DEDUP_EVERY_N` (lower frequency — reconcile fires every Nth extraction, so a *higher* N means *less often*), lower `DEDUP_POOL_SIZE` (smaller per-call pool), or override `n` per call when invoking `reconcile()` directly. +> **Cost note.** Each reconciliation makes one LLM call covering up to `n` facts (default 50, hard cap 500). With auto-trigger, this fires every `FACT_EXTRACTION_EVERY_N × DEDUP_EVERY_N` turns per user, with `n` taken from `DEDUP_POOL_SIZE`. The previous cosine-cluster pre-filter was removed deliberately - it could not catch semantic contradictions like "vegetarian" vs "ribeye steak" - so the LLM is now invoked whenever there are ≥ 2 active facts. To bound LLM cost more tightly: raise `DEDUP_EVERY_N` (lower frequency - reconcile fires every Nth extraction, so a *higher* N means *less often*), lower `DEDUP_POOL_SIZE` (smaller per-call pool), or override `n` per call when invoking `reconcile()` directly. | New `MemoryRecord` field | Meaning | |--------------------------|-----------------------------------------------------------------------------| @@ -197,13 +197,13 @@ By default, the **InProcess processor** runs each pipeline step independently as | `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 (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`. 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"`). #### Backend exclusivity (`MEMORY_PROCESSOR_OWNER`) -Both the SDK auto-trigger and the function-app change-feed processor write into the same `counter` container. If you accidentally point an `InProcessProcessor` at a Cosmos container that already has a function app attached, both backends will run the pipeline on the same writes — double extraction, double dedup, double counters. +Both the SDK auto-trigger and the function-app change-feed processor write into the same `counter` container. If you accidentally point an `InProcessProcessor` at a Cosmos container that already has a function app attached, both backends will run the pipeline on the same writes - double extraction, double dedup, double counters. Set the env var on **both sides** to make ownership explicit: @@ -215,7 +215,7 @@ Set the env var on **both sides** to make ownership explicit: The default (unset) preserves backward compatibility. For any production deployment we recommend setting it on both sides so a misconfiguration produces a loud log line instead of silent double-work. -> **Advisory, not enforced.** `MEMORY_PROCESSOR_OWNER` is operator-configured exclusivity, not a server-side lock. Each backend reads its own env var; if the SDK is set to `inprocess` but the FA forgets to set `durable` (or vice versa), both still run. As a backstop, every counter write stamps `last_owner=` on the doc — when the SDK observes a counter previously written by `durable` (or vice versa), it logs a one-shot `WARN` so misconfiguration surfaces in logs without spamming. Treat this as a configuration audit signal, not a hard guarantee. +> **Advisory, not enforced.** `MEMORY_PROCESSOR_OWNER` is operator-configured exclusivity, not a server-side lock. Each backend reads its own env var; if the SDK is set to `inprocess` but the FA forgets to set `durable` (or vice versa), both still run. As a backstop, every counter write stamps `last_owner=` on the doc - when the SDK observes a counter previously written by `durable` (or vice versa), it logs a one-shot `WARN` so misconfiguration surfaces in logs without spamming. Treat this as a configuration audit signal, not a hard guarantee. --- @@ -225,7 +225,7 @@ Pick at construction time via the `processor=` kwarg. | | `InProcessProcessor` (default) | `DurableFunctionProcessor` | |--------------------------|-----------------------------------|------------------------------------------------------| -| Infra | None — just `pip install` | Sibling Azure Function app | +| Infra | None - just `pip install` | Sibling Azure Function app | | Best for | Prototypes, low TPS, single-agent | Fleet / multi-agent / high TPS | | `process_now()` | Synchronous, returns when done | No-op (work runs async on change feed) | | `process_now_and_wait()` | Returns immediately after flush | Polls until summary visible (RU-costly; tests/demos) | @@ -236,7 +236,7 @@ from azure.cosmos.agent_memory import CosmosMemoryClient, DurableFunctionProcess memory = CosmosMemoryClient(..., processor=DurableFunctionProcessor()) ``` -`DurableFunctionProcessor` is a thin marker — there is no SDK→Function HTTP call. The SDK just writes turns; the deployed Function app picks them up via the Cosmos change feed. Counter-based trigger configuration and Bicep module reference live in [`infra/README.md`](infra/README.md). +`DurableFunctionProcessor` is a thin marker - there is no SDK→Function HTTP call. The SDK just writes turns; the deployed Function app picks them up via the Cosmos change feed. Counter-based trigger configuration and Bicep module reference live in [`infra/README.md`](infra/README.md). --- @@ -267,13 +267,13 @@ memory = CosmosMemoryClient(..., processor=DurableFunctionProcessor()) | Symbol | Module | Purpose | |--------------------------------------|---------------------------------|----------------------------------------------------------------------------------------| -| `CosmosMemoryClient` | `azure.cosmos.agent_memory` | Sync client — local CRUD, Cosmos DB I/O, processing | +| `CosmosMemoryClient` | `azure.cosmos.agent_memory` | Sync client - local CRUD, Cosmos DB I/O, processing | | `AsyncCosmosMemoryClient` | `azure.cosmos.agent_memory.aio` | Async mirror | | `MemoryProcessor` | `azure.cosmos.agent_memory` | Protocol that any processor backend implements | -| `InProcessProcessor` | `azure.cosmos.agent_memory` | Default backend — runs the pipeline in-process | -| `DurableFunctionProcessor` | `azure.cosmos.agent_memory` | Marker backend — work runs in sibling Function app via change feed | -| `client.process_now()` | — | Run the pipeline for recent turns (in-process) or no-op (remote) | -| `client.process_now_and_wait()` | — | Opt-in poll until processing completes; useful for tests/demos with the remote backend | +| `InProcessProcessor` | `azure.cosmos.agent_memory` | Default backend - runs the pipeline in-process | +| `DurableFunctionProcessor` | `azure.cosmos.agent_memory` | Marker backend - work runs in sibling Function app via change feed | +| `client.process_now()` | - | Run the pipeline for recent turns (in-process) or no-op (remote) | +| `client.process_now_and_wait()` | - | Opt-in poll until processing completes; useful for tests/demos with the remote backend | | `MemoryRecord`, `MemoryType`, `Role` | `azure.cosmos.agent_memory` | Pydantic models / enums | Async equivalents (`AsyncInProcessProcessor`, `AsyncDurableFunctionProcessor`) live in `azure.cosmos.agent_memory.aio`. @@ -282,12 +282,12 @@ Async equivalents (`AsyncInProcessProcessor`, `AsyncDurableFunctionProcessor`) l ## Documentation -- **[Docs/concepts.md](Docs/concepts.md)** — Memory types, threads, roles, embeddings, processing pipeline -- **[Docs/design_patterns.md](Docs/design_patterns.md)** — Integration patterns for chat apps and multi-agent systems -- **[Docs/local_testing.md](Docs/local_testing.md)** — Prerequisites, environment setup, running locally, debugging -- **[Docs/azure_testing.md](Docs/azure_testing.md)** — Azure deployment, RBAC, cloud validation -- **[infra/README.md](infra/README.md)** — `azd` deployment, Bicep modules, RBAC, counter-trigger tuning, SDK-only mode -- **[Docs/troubleshooting.md](Docs/troubleshooting.md)** — Common issues and resolutions for setup, auth, Cosmos DB, embeddings, Durable Functions, vector search, change feed, etc. +- **[Docs/concepts.md](Docs/concepts.md)** - Memory types, threads, roles, embeddings, processing pipeline +- **[Docs/design_patterns.md](Docs/design_patterns.md)** - Integration patterns for chat apps and multi-agent systems +- **[Docs/local_testing.md](Docs/local_testing.md)** - Prerequisites, environment setup, running locally, debugging +- **[Docs/azure_testing.md](Docs/azure_testing.md)** - Azure deployment, RBAC, cloud validation +- **[infra/README.md](infra/README.md)** - `azd` deployment, Bicep modules, RBAC, counter-trigger tuning, SDK-only mode +- **[Docs/troubleshooting.md](Docs/troubleshooting.md)** - Common issues and resolutions for setup, auth, Cosmos DB, embeddings, Durable Functions, vector search, change feed, etc. --- @@ -298,7 +298,7 @@ azure/cosmos/agent_memory/ Python SDK (sync + aio mirror) processors/ MemoryProcessor Protocol + InProcess/Durable backends function_app/ Sibling Azure Durable Function app infra/ Bicep modules + main.bicep for `azd up` -azure.yaml `azd` config — provisions Cosmos + AI Foundry + Function app +azure.yaml `azd` config - provisions Cosmos + AI Foundry + Function app Samples/ Categorized demo notebooks + sample scripts Docs/ Conceptual + operational docs tests/ Unit + integration tests (pytest) @@ -309,7 +309,7 @@ tests/ Unit + integration tests (pytest) ## Migration notes - **`azure.cosmos.agent_memory.processing.ProcessingClient` is removed.** Drop the import and call `client.process_now()` (or `client.process_now_and_wait()`) instead. Same for the async `AsyncProcessingClient`. -- **New `processor=` kwarg.** Defaults to `InProcessProcessor()` — existing code keeps its current behavior with no edits. +- **New `processor=` kwarg.** Defaults to `InProcessProcessor()` - existing code keeps its current behavior with no edits. - **`adf_endpoint` / `adf_key` constructor kwargs are gone.** The SDK no longer makes HTTP calls to the Function app at runtime; the Function app reads from the Cosmos change feed. ## Trademark notice diff --git a/Samples/Notebooks/Demo.ipynb b/Samples/Notebooks/Demo.ipynb index d3178b9..5c7a73c 100644 --- a/Samples/Notebooks/Demo.ipynb +++ b/Samples/Notebooks/Demo.ipynb @@ -200,7 +200,7 @@ "\n", "memory.add_local(\n", " user_id=USER_ID, role=\"user\", thread_id=THREAD_ID,\n", - " content=\"Whenever you book a flight for me, always book an aisle seat — never a window or middle.\",\n", + " content=\"Whenever you book a flight for me, always book an aisle seat - never a window or middle.\",\n", ")\n", "memory.add_local(\n", " user_id=USER_ID, role=\"agent\", thread_id=THREAD_ID,\n", @@ -212,7 +212,7 @@ ")\n", "memory.add_local(\n", " user_id=USER_ID, role=\"agent\", thread_id=THREAD_ID,\n", - " content=\"Noted — I'll follow that order: weather, then flights, then hotel.\",\n", + " content=\"Noted - I'll follow that order: weather, then flights, then hotel.\",\n", ")\n", "memory.add_local(\n", " user_id=USER_ID, role=\"user\", thread_id=THREAD_ID,\n", @@ -220,7 +220,7 @@ ")\n", "memory.add_local(\n", " user_id=USER_ID, role=\"agent\", thread_id=THREAD_ID,\n", - " content=\"Will do — no overnight bookings without your explicit approval.\",\n", + " content=\"Will do - no overnight bookings without your explicit approval.\",\n", ")\n", "\n", "print(f\"Added {len(memory.local_memory)} memories\")\n", @@ -231,14 +231,14 @@ "# focused on rules/workflows rather than mixed with factual booking specifics.\n", "RULES_THREAD_ID = str(uuid.uuid4())\n", "for role, content in [\n", - " (\"user\", \"Whenever you book a flight for me, always book an aisle seat — never a window or middle.\"),\n", + " (\"user\", \"Whenever you book a flight for me, always book an aisle seat - never a window or middle.\"),\n", " (\"agent\", \"Got it. I'll always select an aisle seat for your bookings.\"),\n", " (\"user\", \"For trip planning, my workflow is: first check the weather, then check flights, and book the hotel last after everything else is confirmed.\"),\n", - " (\"agent\", \"Noted — I'll follow that order: weather, then flights, then hotel.\"),\n", + " (\"agent\", \"Noted - I'll follow that order: weather, then flights, then hotel.\"),\n", " (\"user\", \"Never book me into anything that departs or arrives between midnight and 6am unless I explicitly approve it.\"),\n", - " (\"agent\", \"Will do — no overnight bookings without your explicit approval.\"),\n", + " (\"agent\", \"Will do - no overnight bookings without your explicit approval.\"),\n", " (\"user\", \"When picking a hotel, only recommend ones that include complimentary breakfast.\"),\n", - " (\"agent\", \"Understood — only hotels with complimentary breakfast.\"),\n", + " (\"agent\", \"Understood - only hotels with complimentary breakfast.\"),\n", "]:\n", " memory.add_local(user_id=USER_ID, role=role, thread_id=RULES_THREAD_ID, content=content)\n", "\n", @@ -421,7 +421,7 @@ }, "outputs": [], "source": [ - "# Already connected via constructor — call connect_cosmos() only if you need to reconnect\n", + "# Already connected via constructor - call connect_cosmos() only if you need to reconnect\n", "print(f\"Connected: {memory._memories_container_client is not None}\")" ] }, @@ -542,7 +542,7 @@ "for r in results:\n", " print(f\" [{r['thread_id'][:8]}...] [{r['id'][:8]}...] type={r['type']:<10} {r['content'][:50]}\")\n", "\n", - "# Turns live in a separate container — fetch via get_thread and filter in-process.\n", + "# Turns live in a separate container - fetch via get_thread and filter in-process.\n", "turns = memory.get_thread(thread_id=THREAD_ID, user_id=USER_ID)\n", "agent_turns = [t for t in turns if t.get(\"role\") == \"agent\"]\n", "print(f\"\\nAgent turns in seed thread: {len(agent_turns)}\")\n", @@ -693,7 +693,7 @@ "source": [ "## 4. Thread Summary (in-process)\n", "\n", - "`generate_thread_summary()` runs the summarisation pipeline **in-process** — no Azure Functions\n", + "`generate_thread_summary()` runs the summarisation pipeline **in-process** - no Azure Functions\n", "required. It will:\n", "\n", "1. Query Cosmos DB for memories matching the given `user_id` + `thread_id`.\n", @@ -845,7 +845,7 @@ "- **`thread_ids`** (optional) – limit to specific threads; omit for all threads\n", "- **`recent_k`** (optional) – per-thread recency limit\n", "\n", - "Retrieve the latest stored profile at any time with `get_user_summary(user_id)` — useful\n", + "Retrieve the latest stored profile at any time with `get_user_summary(user_id)` - useful\n", "for priming new conversations.\n" ] }, @@ -898,7 +898,7 @@ " print(\"User Summary for\", user_id)\n", " print(stored[\"content\"])\n", "else:\n", - " print(\"No user summary found — run the generate_user_summary cell first.\")\n" + " print(\"No user summary found - run the generate_user_summary cell first.\")\n" ] }, { diff --git a/Samples/Notebooks/Demo_async.ipynb b/Samples/Notebooks/Demo_async.ipynb index 3aca2d1..0658adb 100644 --- a/Samples/Notebooks/Demo_async.ipynb +++ b/Samples/Notebooks/Demo_async.ipynb @@ -185,7 +185,7 @@ "\n", "memory.add_local(\n", " user_id=USER_ID, role=\"user\", thread_id=THREAD_ID,\n", - " content=\"Whenever you book a flight for me, always book an aisle seat — never a window or middle.\",\n", + " content=\"Whenever you book a flight for me, always book an aisle seat - never a window or middle.\",\n", ")\n", "memory.add_local(\n", " user_id=USER_ID, role=\"agent\", thread_id=THREAD_ID,\n", @@ -197,7 +197,7 @@ ")\n", "memory.add_local(\n", " user_id=USER_ID, role=\"agent\", thread_id=THREAD_ID,\n", - " content=\"Noted — I'll follow that order: weather, then flights, then hotel.\",\n", + " content=\"Noted - I'll follow that order: weather, then flights, then hotel.\",\n", ")\n", "memory.add_local(\n", " user_id=USER_ID, role=\"user\", thread_id=THREAD_ID,\n", @@ -205,7 +205,7 @@ ")\n", "memory.add_local(\n", " user_id=USER_ID, role=\"agent\", thread_id=THREAD_ID,\n", - " content=\"Will do — no overnight bookings without your explicit approval.\",\n", + " content=\"Will do - no overnight bookings without your explicit approval.\",\n", ")\n", "\n", "print(f\"Added {len(memory.local_memory)} memories\")\n", @@ -216,14 +216,14 @@ "# focused on rules/workflows rather than mixed with factual booking specifics.\n", "RULES_THREAD_ID = str(uuid.uuid4())\n", "for role, content in [\n", - " (\"user\", \"Whenever you book a flight for me, always book an aisle seat — never a window or middle.\"),\n", + " (\"user\", \"Whenever you book a flight for me, always book an aisle seat - never a window or middle.\"),\n", " (\"agent\", \"Got it. I'll always select an aisle seat for your bookings.\"),\n", " (\"user\", \"For trip planning, my workflow is: first check the weather, then check flights, and book the hotel last after everything else is confirmed.\"),\n", - " (\"agent\", \"Noted — I'll follow that order: weather, then flights, then hotel.\"),\n", + " (\"agent\", \"Noted - I'll follow that order: weather, then flights, then hotel.\"),\n", " (\"user\", \"Never book me into anything that departs or arrives between midnight and 6am unless I explicitly approve it.\"),\n", - " (\"agent\", \"Will do — no overnight bookings without your explicit approval.\"),\n", + " (\"agent\", \"Will do - no overnight bookings without your explicit approval.\"),\n", " (\"user\", \"When picking a hotel, only recommend ones that include complimentary breakfast.\"),\n", - " (\"agent\", \"Understood — only hotels with complimentary breakfast.\"),\n", + " (\"agent\", \"Understood - only hotels with complimentary breakfast.\"),\n", "]:\n", " memory.add_local(user_id=USER_ID, role=role, thread_id=RULES_THREAD_ID, content=content)\n", "\n", @@ -343,7 +343,7 @@ ")\n", "print(f\"After update:\\n{json.dumps(memory.get_local(memory_id=target_id)[0], indent=2)}\")\n", "\n", - "# Delete the third memory (index 2 — the user's booking request)\n", + "# Delete the third memory (index 2 - the user's booking request)\n", "del_target_id = memory.local_memory[2][\"id\"]\n", "print(f\"\\nDeleting memory {del_target_id[:8]}...\")\n", "memory.delete_local(del_target_id)\n", @@ -507,7 +507,7 @@ "for r in results:\n", " print(f\" [{r['thread_id'][:8]}...] [{r['id'][:8]}...] type={r['type']:<10} {r['content'][:50]}\")\n", "\n", - "# Turns live in a separate container — fetch via get_thread and filter in-process.\n", + "# Turns live in a separate container - fetch via get_thread and filter in-process.\n", "turns = await memory.get_thread(thread_id=THREAD_ID, user_id=USER_ID)\n", "agent_turns = [t for t in turns if t.get(\"role\") == \"agent\"]\n", "print(f\"\\nAgent turns in seed thread: {len(agent_turns)}\")\n", @@ -652,7 +652,7 @@ "source": [ "## 4. Thread Summary (in-process)\n", "\n", - "`AsyncCosmosMemoryClient.generate_thread_summary()` runs the summarisation pipeline **in-process** — no Azure Functions required. The async client uses the AI Foundry async openai client and the async Cosmos SDK, so calls don't block the event loop.\n", + "`AsyncCosmosMemoryClient.generate_thread_summary()` runs the summarisation pipeline **in-process** - no Azure Functions required. The async client uses the AI Foundry async openai client and the async Cosmos SDK, so calls don't block the event loop.\n", "\n", "If a prior summary exists, the call performs an **incremental update** that preserves the metadata-tracked `source_count`.\n" ] @@ -838,7 +838,7 @@ " print(\"User Summary for\", user_id)\n", " print(stored[\"content\"])\n", "else:\n", - " print(\"No user summary found — run the generate_user_summary cell first.\")\n" + " print(\"No user summary found - run the generate_user_summary cell first.\")\n" ] }, { @@ -848,7 +848,7 @@ "source": [ "### 7. Vector search with `search_cosmos`\n", "\n", - "Same as the sync version — embeds the query and runs a `VectorDistance` similarity search." + "Same as the sync version - embeds the query and runs a `VectorDistance` similarity search." ] }, { diff --git a/Samples/Notebooks/Demo_function_app.ipynb b/Samples/Notebooks/Demo_function_app.ipynb index a9f031b..64e1481 100644 --- a/Samples/Notebooks/Demo_function_app.ipynb +++ b/Samples/Notebooks/Demo_function_app.ipynb @@ -5,34 +5,34 @@ "id": "59a4f403", "metadata": {}, "source": [ - "# Agent Memory Toolkit \u2013 Function App (Remote Processor) Demo\n", + "# Agent Memory Toolkit – Function App (Remote Processor) Demo\n", "\n", "This notebook demonstrates the **`DurableFunctionProcessor`** hand-off pattern: the SDK only\n", "writes raw turns to Cosmos DB, and a sibling **Azure Function App** (deployed from the\n", "`function_app/` folder via `azd up`) picks them up from the Cosmos change feed and produces\n", - "summaries, facts, episodic memories, and procedural rules asynchronously \u2014 server-side.\n", + "summaries, facts, episodic memories, and procedural rules asynchronously - server-side.\n", "\n", "## When to use this pattern\n", "\n", "| Use case | Recommended processor |\n", "|---|---|\n", - "| Local dev, scripts, single agent process | **`InProcessProcessor`** (default) \u2014 no extra infra |\n", - "| Production, multi-agent fleets, server-managed processing, audit trail | **`DurableFunctionProcessor`** \u2014 durable, scalable, isolated |\n", + "| Local dev, scripts, single agent process | **`InProcessProcessor`** (default) - no extra infra |\n", + "| Production, multi-agent fleets, server-managed processing, audit trail | **`DurableFunctionProcessor`** - durable, scalable, isolated |\n", "\n", "## Prerequisites\n", "\n", - "1. **Function App deployed** \u2014 run `azd up` (or deploy `function_app/` manually) against the same Cosmos\n", + "1. **Function App deployed** - run `azd up` (or deploy `function_app/` manually) against the same Cosmos\n", " account/database/container the SDK targets.\n", "2. **Function App env vars** include `THREAD_SUMMARY_EVERY_N`, `FACT_EXTRACTION_EVERY_N`,\n", - " `USER_SUMMARY_EVERY_N` (defaults: 4 / 6 / 10). Each one \u2265 1 enables the corresponding orchestrator.\n", + " `USER_SUMMARY_EVERY_N` (defaults: 4 / 6 / 10). Each one ≥ 1 enables the corresponding orchestrator.\n", "3. **`.env`** in this repo has `COSMOS_DB_ENDPOINT`, `AI_FOUNDRY_ENDPOINT`, deployment names. The Function App\n", " uses managed identity to reach the same Cosmos + AI Foundry resources.\n", "\n", "## What the SDK does in this mode\n", "\n", - "* `add_cosmos(..., memory_type=\"turn\")` \u2192 writes the raw turn to Cosmos.\n", + "* `add_cosmos(..., 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** \u2014 the Function App owns processing.\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)." ] }, @@ -82,7 +82,7 @@ "## 2. Construct the client with `DurableFunctionProcessor`\n", "\n", "The single change vs. the in-process demo is `processor=DurableFunctionProcessor()`. After this, the\n", - "SDK **never invokes the LLM or embeddings client itself** \u2014 it only writes raw turns to Cosmos." + "SDK **never invokes the LLM or embeddings client itself** - it only writes raw turns to Cosmos." ] }, { @@ -107,7 +107,7 @@ " ai_foundry_endpoint=os.environ[\"AI_FOUNDRY_ENDPOINT\"],\n", " chat_deployment_name=os.environ[\"AI_FOUNDRY_CHAT_DEPLOYMENT_NAME\"],\n", " embedding_deployment_name=os.environ[\"AI_FOUNDRY_EMBEDDING_DEPLOYMENT_NAME\"],\n", - " # \u2500\u2500 The hand-off \u2500\u2500\n", + " # ── The hand-off ──\n", " processor=DurableFunctionProcessor(),\n", ")\n", "print(f\"Processor: {type(memory._processor).__name__}\")" @@ -145,15 +145,15 @@ "\n", "transcript = [\n", " (\"user\", \"What's the weather like in Seattle this weekend?\"),\n", - " (\"agent\", \"Around 55\u00b0F with partly cloudy skies on Saturday and light rain on Sunday.\"),\n", + " (\"agent\", \"Around 55°F with partly cloudy skies on Saturday and light rain on Sunday.\"),\n", " (\"user\", \"Can you book me a weekend trip there? Flights under $300, hotel under $200.\"),\n", " (\"agent\", \"I found a $275 Alaska Airlines round-trip and a $185/night hotel in Belltown.\"),\n", - " (\"user\", \"Whenever you book a flight for me, always book an aisle seat \u2014 never window or middle.\"),\n", - " (\"agent\", \"Got it \u2014 aisle seats only.\"),\n", + " (\"user\", \"Whenever you book a flight for me, always book an aisle seat - never window or middle.\"),\n", + " (\"agent\", \"Got it - aisle seats only.\"),\n", " (\"user\", \"For trip planning, my workflow is: first weather, then flights, then hotel.\"),\n", - " (\"agent\", \"Noted \u2014 weather \u2192 flights \u2192 hotel.\"),\n", + " (\"agent\", \"Noted - weather → flights → hotel.\"),\n", " (\"user\", \"Never book me into anything between midnight and 6am unless I explicitly approve.\"),\n", - " (\"agent\", \"Understood \u2014 no overnight bookings without your approval.\"),\n", + " (\"agent\", \"Understood - no overnight bookings without your approval.\"),\n", "]\n", "for role, content in transcript:\n", " memory.add_cosmos(\n", @@ -163,7 +163,7 @@ " content=content,\n", " memory_type=\"turn\",\n", " )\n", - " print(f\" wrote {role:>5}: {content[:60]}{'\u2026' if len(content) > 60 else ''}\")" + " print(f\" wrote {role:>5}: {content[:60]}{'…' if len(content) > 60 else ''}\")" ] }, { @@ -174,7 +174,7 @@ "## 4. Verify the SDK did NOT process locally\n", "\n", "`process_now()` is a debug-logged no-op when using `DurableFunctionProcessor`. The Function App's change-feed\n", - "trigger fires asynchronously \u2014 usually within a second or two." + "trigger fires asynchronously - usually within a second or two." ] }, { @@ -193,7 +193,7 @@ "source": [ "# No-op: the function app owns processing.\n", "memory.process_now(user_id=USER_ID, thread_id=THREAD_ID)\n", - "print(\"process_now() returned \u2014 no LLM call was made by the SDK.\")" + "print(\"process_now() returned - no LLM call was made by the SDK.\")" ] }, { @@ -204,11 +204,11 @@ "## 5. Wait for the Function App to produce a summary\n", "\n", "`process_now_and_wait` polls Cosmos for the thread's summary doc until `timeout` seconds. This is **RU-costly**\n", - "(repeated `get_thread_summary(user_id, thread_id)` queries) \u2014 only use it for demos and tests.\n", + "(repeated `get_thread_summary(user_id, thread_id)` queries) - only use it for demos and tests.\n", "\n", "> If this returns `False`, check that:\n", "> * The Function App is deployed and running.\n", - "> * `THREAD_SUMMARY_EVERY_N` is `\u2265 1` and `\u2264` the number of turns written above.\n", + "> * `THREAD_SUMMARY_EVERY_N` is `≥ 1` and `≤` the number of turns written above.\n", "> * The Function App's managed identity has Cosmos data + Cognitive Services OpenAI User roles." ] }, @@ -226,7 +226,7 @@ }, "outputs": [], "source": [ - "print(\"Polling Cosmos for the auto-generated summary\u2026\")\n", + "print(\"Polling Cosmos for the auto-generated summary…\")\n", "ok = memory.process_now_and_wait(user_id=USER_ID, thread_id=THREAD_ID, timeout=120.0)\n", "print(f\"Summary available: {ok}\")" ] @@ -282,7 +282,7 @@ " counts[mt] = len(docs)\n", " print(f\"\\n{mt.upper()}S ({len(docs)}):\")\n", " for d in docs:\n", - " print(f\" \u2022 [{d['id'][:32]}\u2026] {d['content'][:90]}\")\n", + " print(f\" • [{d['id'][:32]}…] {d['content'][:90]}\")\n", " return counts\n", "\n", "print(\"Initial state:\")\n", @@ -309,9 +309,9 @@ "## 7. Going further\n", "\n", "* **Per-turn embedding**: the SDK auto-embeds non-`turn` documents you `add_cosmos` directly. Raw `turn`\n", - " records are intentionally not embedded \u2014 the function app does that during summary/extraction.\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=\"\u2026\", memory_types=[\"fact\"], user_id=USER_ID)`.\n", + " `memory.search_cosmos(search_terms=\"…\", memory_types=[\"fact\"], user_id=USER_ID)`.\n", "* **Switch back to in-process** for ad-hoc work: instantiate the client without the `processor=` kwarg\n", " (or pass `InProcessProcessor()` explicitly) and use `generate_thread_summary` / `extract_memories`\n", " directly." @@ -339,4 +339,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +} diff --git a/Samples/Notebooks/Demo_function_app_async.ipynb b/Samples/Notebooks/Demo_function_app_async.ipynb index b6002ea..77b8167 100644 --- a/Samples/Notebooks/Demo_function_app_async.ipynb +++ b/Samples/Notebooks/Demo_function_app_async.ipynb @@ -5,7 +5,7 @@ "id": "52af19ff", "metadata": {}, "source": [ - "# Agent Memory Toolkit \u2013 Function App (Remote Processor) Demo (async)\n", + "# Agent Memory Toolkit – Function App (Remote Processor) Demo (async)\n", "\n", "Async variant of `Demo_function_app.ipynb`. Wires `AsyncDurableFunctionProcessor` to\n", "`AsyncCosmosMemoryClient` so the SDK only writes raw turns and the deployed sibling Azure Function\n", @@ -114,15 +114,15 @@ "\n", "transcript = [\n", " (\"user\", \"What's the weather like in Seattle this weekend?\"),\n", - " (\"agent\", \"Around 55\u00b0F with partly cloudy skies on Saturday and light rain on Sunday.\"),\n", + " (\"agent\", \"Around 55°F with partly cloudy skies on Saturday and light rain on Sunday.\"),\n", " (\"user\", \"Can you book me a weekend trip there? Flights under $300, hotel under $200.\"),\n", " (\"agent\", \"I found a $275 Alaska Airlines round-trip and a $185/night hotel in Belltown.\"),\n", - " (\"user\", \"Whenever you book a flight for me, always book an aisle seat \u2014 never window or middle.\"),\n", - " (\"agent\", \"Got it \u2014 aisle seats only.\"),\n", + " (\"user\", \"Whenever you book a flight for me, always book an aisle seat - never window or middle.\"),\n", + " (\"agent\", \"Got it - aisle seats only.\"),\n", " (\"user\", \"For trip planning, my workflow is: first weather, then flights, then hotel.\"),\n", - " (\"agent\", \"Noted \u2014 weather \u2192 flights \u2192 hotel.\"),\n", + " (\"agent\", \"Noted - weather → flights → hotel.\"),\n", " (\"user\", \"Never book me into anything between midnight and 6am unless I explicitly approve.\"),\n", - " (\"agent\", \"Understood \u2014 no overnight bookings without your approval.\"),\n", + " (\"agent\", \"Understood - no overnight bookings without your approval.\"),\n", "]\n", "for role, content in transcript:\n", " await memory.add_cosmos(\n", @@ -132,7 +132,7 @@ " content=content,\n", " memory_type=\"turn\",\n", " )\n", - " print(f\" wrote {role:>5}: {content[:60]}{'\u2026' if len(content) > 60 else ''}\")" + " print(f\" wrote {role:>5}: {content[:60]}{'…' if len(content) > 60 else ''}\")" ] }, { @@ -158,7 +158,7 @@ "outputs": [], "source": [ "await memory.process_now(user_id=USER_ID, thread_id=THREAD_ID)\n", - "print(\"process_now() returned \u2014 no LLM call was made by the SDK.\")" + "print(\"process_now() returned - no LLM call was made by the SDK.\")" ] }, { @@ -183,7 +183,7 @@ }, "outputs": [], "source": [ - "print(\"Polling Cosmos for the auto-generated summary\u2026\")\n", + "print(\"Polling Cosmos for the auto-generated summary…\")\n", "ok = await memory.process_now_and_wait(user_id=USER_ID, thread_id=THREAD_ID, timeout=120.0)\n", "print(f\"Summary available: {ok}\")" ] @@ -229,7 +229,7 @@ " counts[mt] = len(docs)\n", " print(f\"\\n{mt.upper()}S ({len(docs)}):\")\n", " for d in docs:\n", - " print(f\" \u2022 [{d['id'][:32]}\u2026] {d['content'][:90]}\")\n", + " print(f\" • [{d['id'][:32]}…] {d['content'][:90]}\")\n", " return counts\n", "\n", "print(\"Initial state:\")\n", @@ -297,4 +297,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +} diff --git a/Samples/Processing/processing_episodic_memory.py b/Samples/Processing/processing_episodic_memory.py new file mode 100644 index 0000000..1cad3a0 --- /dev/null +++ b/Samples/Processing/processing_episodic_memory.py @@ -0,0 +1,248 @@ +"""Demonstrate episodic memory extraction and retrieval. + +Episodic memory captures bounded user experiences as structured episodes: +what happened, who participated, when it happened, how it ended, and what +lessons were learned. Episodes are boundary-based: the toolkit segments the +turn stream into coherent experiences automatically (at an idle time-gap, a +topic shift, or a max-size cap) - the caller never signals "session end". This +sample writes a short single-topic conversation and calls +``CosmosMemoryClient.extract_episodes(..., flush=True)`` to finalize the open +segment at the end of the conversation, prints the resulting episode shape, +demonstrates blended retrieval with episodic results, and then cleans up all +created records. + +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 dotenv import load_dotenv + +from azure.cosmos.agent_memory import CosmosMemoryClient + +load_dotenv() + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +DIVIDER = "-" * 60 + + +def banner(title: str) -> None: + """Print a section banner.""" + print(f"\n{DIVIDER}") + print(f" {title}") + print(DIVIDER) + + +def print_outcome(outcome: dict | None) -> None: + """Pretty-print an episode outcome.""" + if not outcome: + print(" outcome: none") + return + status = outcome.get("status") or "unknown" + description = outcome.get("description") or "" + print(f" outcome: {status} - {description}") + + +def print_episodes(episodes: list[dict]) -> None: + """Pretty-print episodic memory records.""" + if not episodes: + print(" (none)") + return + + for episode in episodes: + print(f" title: {episode.get('title', '')}") + print(f" summary: {episode.get('content', '')}") + print(f" started_at: {episode.get('started_at')}") + print(f" ended_at: {episode.get('ended_at')}") + participants = ", ".join(episode.get("participants") or []) + print(f" participants: {participants or '(none)'}") + print(" events:") + for event in sorted(episode.get("events") or [], key=lambda item: item.get("sequence", 0)): + print(f" {event.get('sequence')}. {event.get('description', '')}") + print_outcome(episode.get("outcome")) + lessons = episode.get("lessons") or [] + print(" lessons:") + if lessons: + for lesson in lessons: + print(f" - {lesson}") + else: + print(" - none") + + +def print_search_results(results: list[dict]) -> None: + """Pretty-print blended search results.""" + if not results: + print(" (none)") + return + + for result in results: + content = str(result.get("content") or "").replace("\n", " ") + print(f" [{result.get('type', 'unknown')}] {content}") + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +CONVERSATION = [ + ( + "user", + "On Friday afternoon, I drove from Seattle to Leavenworth with Maya and Jordan for a weekend hike.", + ), + ( + "agent", + "That sounds like a fun start. What trail did you choose for Saturday?", + ), + ( + "user", + "We chose the Colchuck Lake trail for Saturday morning and planned to start by 6:30 AM.", + ), + ( + "agent", + "An early start should help with parking and give you cooler hiking weather.", + ), + ( + "user", + "Jordan forgot his trekking poles, so we stopped at a gear shop before heading to the trailhead.", + ), + ( + "user", + "The hike was harder than expected after the boulder field, but Maya paced us and kept everyone steady.", + ), + ( + "agent", + "It sounds like teamwork helped the group handle the hard section.", + ), + ( + "user", + "We reached the lake by noon, ate lunch by the water, and decided the trip was a success.", + ), +] + + +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)}") + 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"episodic-demo-{uuid.uuid4().hex[:8]}" + thread_id = f"episodic-demo-thread-{uuid.uuid4().hex[:8]}" + print(f"User ID: {user_id}") + print(f"Thread ID: {thread_id}") + + 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) + print(f" [{role:>5}] {content}") + + banner("2. Finalizing episodes (flush the open segment)") + # In production, boundary evaluation runs automatically every + # EPISODE_EVAL_EVERY_N turns and closes an episode at each detected + # boundary (idle time-gap, topic shift, or max-size cap). Here the whole + # conversation is one bounded experience with no gap, so we flush=True to + # drain the trailing open segment at the end of the chat. + stats = mem.extract_episodes(user_id, thread_id, flush=True) + print(f" stats: {json.dumps(stats, indent=2)}") + + banner("3. Retrieved episodes") + episodes = mem.get_episodes(user_id) + print_episodes(episodes) + + banner("4. Blended search (facts + episodes in one combined query)") + # With include_episodes, facts and episodes are returned by a single + # ranked query sharing one top_k budget (no separate episodic query). + results = mem.search_cosmos( + search_terms="Colchuck Lake hiking outcome", + user_id=user_id, + include_episodes=True, + ) + print_search_results(results) + + finally: + banner("5. Cleanup") + deleted = 0 + try: + memory_records = mem.get_memories(user_id=user_id, include_superseded=True) + thread_records = mem.get_thread(thread_id=thread_id, user_id=user_id, include_superseded=True) + summary_records = mem.get_thread_summary(user_id=user_id, thread_id=thread_id) + user_summary = mem.get_user_summary(user_id) + all_records = [*memory_records, *thread_records, *summary_records] + 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) + try: + mem.delete_cosmos( + memory_id=memory_id, + user_id=user_id, + thread_id=record.get("thread_id", 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/azure/cosmos/agent_memory/__init__.py b/azure/cosmos/agent_memory/__init__.py index dbe8b4f..970e07e 100644 --- a/azure/cosmos/agent_memory/__init__.py +++ b/azure/cosmos/agent_memory/__init__.py @@ -23,11 +23,21 @@ UserSummaryResult, ) from azure.cosmos.agent_memory.thresholds import ( + 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_THREAD_SUMMARY_EVERY_N, DEFAULT_USER_SUMMARY_EVERY_N, PROCESSOR_OWNER_DURABLE, PROCESSOR_OWNER_INPROCESS, + get_episode_eval_every_n, + get_episode_idle_gap_seconds, + get_episode_max_turns, + get_episode_min_turns, + get_episode_topic_drift, get_fact_extraction_every_n, get_processor_owner, get_thread_summary_every_n, @@ -56,11 +66,21 @@ "MemoryNotFoundError", "MemoryTypeMismatchError", "ValidationError", + "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_THREAD_SUMMARY_EVERY_N", "DEFAULT_USER_SUMMARY_EVERY_N", "PROCESSOR_OWNER_DURABLE", "PROCESSOR_OWNER_INPROCESS", + "get_episode_eval_every_n", + "get_episode_idle_gap_seconds", + "get_episode_max_turns", + "get_episode_min_turns", + "get_episode_topic_drift", "get_fact_extraction_every_n", "get_processor_owner", "get_thread_summary_every_n", diff --git a/azure/cosmos/agent_memory/_utils.py b/azure/cosmos/agent_memory/_utils.py index e3cafae..7cac875 100644 --- a/azure/cosmos/agent_memory/_utils.py +++ b/azure/cosmos/agent_memory/_utils.py @@ -7,6 +7,7 @@ from __future__ import annotations import hashlib +import math import os import re import uuid @@ -343,6 +344,48 @@ def _resolve_vector_index_type(val: Optional[str]) -> str: _SIMILARITY_DESCENDING_FUNCTIONS = frozenset({"cosine", "dotproduct"}) +def cosine_similarity(a: list[float], b: list[float]) -> float: + """Cosine similarity of two equal-length vectors. + + Returns 0.0 for empty, mismatched-length, or zero-magnitude inputs (treated + as "unrelated"). Pure Python so the SDK stays numpy-free. + """ + if not a or not b or len(a) != len(b): + return 0.0 + dot = 0.0 + norm_a = 0.0 + norm_b = 0.0 + for x, y in zip(a, b): + dot += x * y + norm_a += x * x + norm_b += y * y + if norm_a <= 0.0 or norm_b <= 0.0: + return 0.0 + return dot / math.sqrt(norm_a * norm_b) + + +def vector_centroid(vectors: list[list[float]]) -> list[float]: + """Component-wise mean of equal-length vectors; ``[]`` when none are usable.""" + dim = 0 + for vec in vectors: + if vec: + dim = len(vec) + break + if dim == 0: + return [] + acc = [0.0] * dim + count = 0 + for vec in vectors: + if not vec or len(vec) != dim: + continue + for i in range(dim): + acc[i] += vec[i] + count += 1 + if count == 0: + return [] + return [value / count for value in acc] + + def vector_order_direction(distance_function: str) -> str: """Return the ``ORDER BY VectorDistance(...)`` direction for most-similar-first. diff --git a/azure/cosmos/agent_memory/aio/auto_trigger.py b/azure/cosmos/agent_memory/aio/auto_trigger.py index 60df3ae..22cb6af 100644 --- a/azure/cosmos/agent_memory/aio/auto_trigger.py +++ b/azure/cosmos/agent_memory/aio/auto_trigger.py @@ -58,9 +58,10 @@ async def maybe_trigger_steps( n_facts = _threshold_int(thresholds, "get_fact_extraction_every_n", "FACT_EXTRACTION_EVERY_N") n_summary = _threshold_int(thresholds, "get_thread_summary_every_n", "THREAD_SUMMARY_EVERY_N") + n_episode = _threshold_int(thresholds, "get_episode_eval_every_n", "EPISODE_EVAL_EVERY_N") n_user = _threshold_int(thresholds, "get_user_summary_every_n", "USER_SUMMARY_EVERY_N") n_dedup = _threshold_int(thresholds, "get_dedup_every_n", "DEDUP_EVERY_N") - if n_facts == 0 and n_summary == 0 and n_user == 0: + if n_facts == 0 and n_summary == 0 and n_episode == 0 and n_user == 0: return n_dedup_turns = n_facts * n_dedup if n_facts > 0 and n_dedup > 0 else 0 @@ -70,6 +71,7 @@ async def maybe_trigger_steps( turn_counts, n_facts=n_facts, n_summary=n_summary, + n_episode=n_episode, n_dedup_turns=n_dedup_turns, thresholds=thresholds, ) @@ -83,6 +85,7 @@ async def _trigger_thread_steps( *, n_facts: int, n_summary: int, + n_episode: int, n_dedup_turns: int, thresholds: Any = None, ) -> dict[str, int]: @@ -113,6 +116,7 @@ async def _trigger_thread_steps( new_count=new_count, fire_extract=n_facts > 0 and _counters.crosses_threshold(old_count, new_count, n_facts), fire_summary=n_summary > 0 and _counters.crosses_threshold(old_count, new_count, n_summary), + fire_episode=n_episode > 0 and _counters.crosses_threshold(old_count, new_count, n_episode), fire_dedup=n_dedup_turns > 0 and _counters.crosses_threshold(old_count, new_count, n_dedup_turns), thresholds=thresholds, ) @@ -129,6 +133,7 @@ async def _fire_thread_steps( new_count: int, fire_extract: bool, fire_summary: bool, + fire_episode: bool, fire_dedup: bool, thresholds: Any = None, ) -> None: @@ -161,6 +166,12 @@ async def _fire_thread_steps( processor.synthesize_procedural, {"user_id": user_id}, ), + ( + fire_episode, + "process_extract_episodes", + processor.process_extract_episodes, + {"user_id": user_id, "thread_id": thread_id}, + ), ( fire_summary, "process_thread_summary", diff --git a/azure/cosmos/agent_memory/aio/chat.py b/azure/cosmos/agent_memory/aio/chat.py index d83d7d5..75f98c6 100644 --- a/azure/cosmos/agent_memory/aio/chat.py +++ b/azure/cosmos/agent_memory/aio/chat.py @@ -16,6 +16,7 @@ TOKEN_SCOPE, extract_content, resolve_api_version, + retry_delay, unsupported_param, ) from azure.cosmos.agent_memory.exceptions import ConfigurationError @@ -156,7 +157,7 @@ async def generate( messages: list[dict[str, str]], *, response_format: dict | None = None, - max_retries: int = 3, + max_retries: int = 6, base_delay: float = 2.0, **extra: Any, ) -> str: @@ -205,7 +206,7 @@ async def generate( return extract_content(response, self._model) except openai.RateLimitError as exc: if attempt < max_retries - 1: - delay = base_delay * (2**attempt) + delay = retry_delay(exc, attempt, base_delay) logger.warning( "LLM rate-limited (attempt %d/%d), retrying in %.1fs: %s", attempt + 1, @@ -230,7 +231,7 @@ async def generate( unsupported_strips += 1 continue if status in RETRYABLE_STATUS_CODES and attempt < max_retries - 1: - delay = base_delay * (2**attempt) + delay = retry_delay(exc, attempt, base_delay) logger.warning( "LLM API error %s (attempt %d/%d), retrying in %.1fs: %s", status, diff --git a/azure/cosmos/agent_memory/aio/cosmos_memory_client.py b/azure/cosmos/agent_memory/aio/cosmos_memory_client.py index 089c6a9..008e0e8 100644 --- a/azure/cosmos/agent_memory/aio/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/aio/cosmos_memory_client.py @@ -709,20 +709,35 @@ async def search_cosmos( min_confidence: Optional[float] = None, created_after: Optional[str | datetime] = None, created_before: Optional[str | datetime] = None, + include_episodes: bool = False, include_turns: bool = False, turn_top_k: Optional[int] = None, include_summaries: bool = False, summary_top_k: Optional[int] = None, ) -> list[dict[str, Any]]: - """Search memories using vector similarity, with optional summary / raw-turn - blending. See the sync client for details.""" + """Search memories using vector similarity, with optional retrieval blending. + + Facts + episodes in a single ranked query sharing one ``top_k`` budget + when ``include_episodes`` is True; facts only when False. Callers may + pass ``memory_types`` for other non-episodic types (``episodic`` is added + or removed by ``include_episodes``). Optional summaries/turns blend in + after the base block. See the sync client for full details.""" store = self._get_store() - results = await store.search( + # 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: + base_memory_types = [t for t in memory_types if t != "episodic"] + else: + base_memory_types = ["fact"] + if include_episodes: + base_memory_types = [*base_memory_types, "episodic"] + base_memory_types = base_memory_types or ["fact"] + base = await store.search( search_terms=search_terms, memory_id=memory_id, user_id=user_id, role=role, - memory_types=memory_types, + memory_types=base_memory_types, thread_id=thread_id, top_k=top_k, tags_all=tags_all, @@ -735,9 +750,20 @@ async def search_cosmos( created_before=created_before, ) if not user_id: - return results + return base + + results: list[dict[str, Any]] = [] + seen_content: set[str] = set() + + def _extend(docs: list[dict[str, Any]]) -> None: + for doc in docs: + content = str(doc.get("content") or "").strip() + if content and content not in seen_content: + seen_content.add(content) + results.append(doc) + + _extend(base) - seen_content = {str(r.get("content") or "").strip() for r in results} if include_summaries: try: summaries = await store.search_summaries( @@ -750,13 +776,7 @@ async def search_cosmos( except Exception as exc: # noqa: BLE001 logger.warning("search_cosmos: include_summaries search failed (%s); skipping summaries", exc) summaries = [] - # Order is facts -> summaries -> raw turns: append summaries after the - # relevance-ranked memory hits (and before turns). - for s in summaries: - content = str(s.get("content") or "").strip() - if content and content not in seen_content: - seen_content.add(content) - results.append(s) + _extend(summaries) if include_turns: try: @@ -773,11 +793,7 @@ async def search_cosmos( except Exception as exc: # noqa: BLE001 logger.warning("search_cosmos: include_turns turn search failed (%s); returning memories only", exc) turns = [] - for turn in turns: - content = str(turn.get("content") or "").strip() - if content and content not in seen_content: - seen_content.add(content) - results.append(turn) + _extend(turns) return results async def search_summaries( @@ -898,6 +914,19 @@ async def get_thread_summary( recent_k=recent_k, ) + async def get_episodes( + self, + user_id: str, + thread_id: Optional[str] = None, + recent_k: Optional[int] = None, + ) -> list[dict[str, Any]]: + """Retrieve active episodic memories for a user, newest first.""" + return await self._get_store().get_episodes( + user_id=user_id, + thread_id=thread_id, + recent_k=recent_k, + ) + async def get_user_summary(self, user_id: str) -> Optional[dict[str, Any]]: return await self._get_store().get_user_summary(user_id=user_id) @@ -989,6 +1018,30 @@ async def build_episodic_context( async def extract_memories(self, user_id: str, thread_id: str, recent_k: Optional[int] = None) -> dict[str, int]: return await self._get_pipeline().extract_memories(user_id, thread_id, recent_k) + async def extract_episodes(self, user_id: str, thread_id: str, *, flush: bool = False) -> dict[str, int]: + """Segment the thread's open turn stream into episodes at detected boundaries. + + Episodes are created at idle time-gaps (detected only once a later turn + reveals the gap), topic shifts, and a max-size cap. A focused session + shorter than the max-size cap therefore episodizes only lazily - on the + next turn after the idle gap - and a one-shot session that never resumes + is not episodized at all under the auto path. Pass ``flush=True`` at the + end of a conversation to drain the trailing open + segment immediately; integrators that know when a session ends should + call this on session close. + + Only supported when the in-process backend owns processing; when a + Durable Function app is the active processor this raises + ``NotImplementedError`` so writes are not split away from that backend. + """ + processor = self._get_processor() + if not isinstance(processor, AsyncInProcessProcessor): + raise NotImplementedError( + "Episode extraction runs in-process; manual invocation via the SDK is not " + "supported when the Durable Function app is the active processor." + ) + return await self._get_pipeline().extract_episodes(user_id, thread_id, flush=flush) + async def synthesize_procedural(self, user_id: str, *, force: bool = False) -> dict[str, Any]: processor = self._get_processor() if not isinstance(processor, AsyncInProcessProcessor): diff --git a/azure/cosmos/agent_memory/aio/processors/base.py b/azure/cosmos/agent_memory/aio/processors/base.py index aacc427..62a4173 100644 --- a/azure/cosmos/agent_memory/aio/processors/base.py +++ b/azure/cosmos/agent_memory/aio/processors/base.py @@ -36,6 +36,21 @@ async def process_extract_memories( recent_k: Optional[int] = None, ) -> dict[str, int]: ... + async def process_extract_episodes( + self, + *, + user_id: str, + thread_id: str, + ) -> dict[str, int]: + """Segment the open turn stream into episodes at detected boundaries. + + Deferred backends (e.g. the Durable Functions app) that do not yet + implement episodic segmentation may no-op (return an empty result) or + raise ``NotImplementedError``; the auto-trigger only invokes this on the + in-process backend. + """ + ... + async def process_thread_summary( self, *, diff --git a/azure/cosmos/agent_memory/aio/processors/durable.py b/azure/cosmos/agent_memory/aio/processors/durable.py index ed38d8f..52cf32a 100644 --- a/azure/cosmos/agent_memory/aio/processors/durable.py +++ b/azure/cosmos/agent_memory/aio/processors/durable.py @@ -12,6 +12,10 @@ 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`. @@ -49,6 +53,30 @@ async def process_extract_memories( ) return {} + async def process_extract_episodes( + self, + *, + 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." + ) + logger.debug( + "AsyncDurableFunctionProcessor.process_extract_episodes no-op user_id=%s thread_id=%s", + user_id, + thread_id, + ) + return {} + async def process_thread_summary( self, *, diff --git a/azure/cosmos/agent_memory/aio/processors/inprocess.py b/azure/cosmos/agent_memory/aio/processors/inprocess.py index 55cdf91..1ae3b3a 100644 --- a/azure/cosmos/agent_memory/aio/processors/inprocess.py +++ b/azure/cosmos/agent_memory/aio/processors/inprocess.py @@ -94,6 +94,15 @@ async def process_extract_memories( extracted = await self._pipeline.extract_memories(user_id, thread_id, recent_k=recent_k) return {k: v for k, v in extracted.items() if isinstance(v, int)} if isinstance(extracted, dict) else {} + async def process_extract_episodes( + self, + *, + user_id: str, + thread_id: str, + ) -> dict[str, int]: + extracted = await self._pipeline.extract_episodes(user_id, thread_id) + return {k: v for k, v in extracted.items() if isinstance(v, int)} if isinstance(extracted, dict) else {} + async def process_thread_summary( self, *, diff --git a/azure/cosmos/agent_memory/aio/services/pipeline.py b/azure/cosmos/agent_memory/aio/services/pipeline.py index 81369b9..74b0520 100644 --- a/azure/cosmos/agent_memory/aio/services/pipeline.py +++ b/azure/cosmos/agent_memory/aio/services/pipeline.py @@ -52,7 +52,6 @@ ID_SEED_SEP as _ID_SEED_SEP, ) from azure.cosmos.agent_memory.services._pipeline_helpers import ( - VALID_VALENCES, PromptyLoader, _normalize_metadata_keys, batch_turns_by_tokens, @@ -60,10 +59,16 @@ build_transcript, cap_structured_summary, chat_text, - check_extracted_fact_grounding, - coerce_valence, + clamp_unit_interval, + created_at_sort_key, + deterministic_episode_id, + extract_memories_prompt_file, + find_episode_boundary, is_retryable_llm_error, + is_valid_time_pair, parse_llm_json, + segment_time_bounds, + turn_gap_seconds, ) from azure.cosmos.agent_memory.services._pipeline_helpers import ( is_real_number as _is_real_number, @@ -75,18 +80,27 @@ 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, + get_episode_topic_drift, get_extraction_batch_max_tokens, ) logger = get_logger("azure.cosmos.agent_memory.pipeline.aio") -_coerce_valence = coerce_valence _cap_structured_summary = cap_structured_summary _ACTIVE_DOC_FILTER = "(NOT IS_DEFINED(c.superseded_by) OR IS_NULL(c.superseded_by))" _PROCEDURAL_MAX_CREATE_ATTEMPTS = 5 +# Safety cap on how many episodes one extract_episodes call may close in a +# single pass (mirror of the sync pipeline). Boundary evaluation runs on a small +# turn cadence, so this only bounds a pathological drain and never fans out into +# an unbounded burst of LLM extractions. +_EPISODE_MAX_SEGMENTS_PER_RUN = 50 + class _AsyncStoreContainerAdapter: """Expose one split ``AsyncMemoryStore`` container via Cosmos method shapes.""" @@ -464,6 +478,8 @@ def _empty_extract_counts() -> dict[str, int]: "contradicted_count": 0, "exact_dedup_skipped": 0, "dropped_episodic_count": 0, + "deferred_turn_count": 0, + "quarantined_turn_count": 0, } @staticmethod @@ -487,7 +503,7 @@ async def _mark_superseded( def _parse_llm_json(text: str | None) -> dict[str, Any]: return parse_llm_json(text) - async def extract_memories_dry( + async def extract_memories_durable( self, user_id: str, thread_id: str, @@ -501,7 +517,7 @@ async def extract_memories_dry( if not thread_id: raise ValidationError("thread_id is required") - logger.info("extract_memories_dry started user_id=%s thread_id=%s", user_id, thread_id) + logger.info("extract_memories_durable started user_id=%s thread_id=%s", user_id, thread_id) if turns is None: query = ( @@ -528,7 +544,7 @@ async def extract_memories_dry( items.reverse() if not items: - logger.warning("extract_memories_dry no memories found user_id=%s thread_id=%s", user_id, thread_id) + 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"]) @@ -544,19 +560,16 @@ async def extract_memories_dry( # *retryable* error are left un-stamped and retried on the next run. batches = batch_turns_by_tokens(items, get_extraction_batch_max_tokens()) facts: list[dict[str, Any]] = [] - episodic: list[dict[str, Any]] = [] processed_turns: list[dict[str, Any]] = [] deferred_turn_count = 0 quarantined_turn_count = 0 + extract_prompt = extract_memories_prompt_file() for batch in batches: batch_transcript = self._build_transcript(batch, include_timestamp=True) try: - response_text = await self._run_prompty( - "extract_memories.prompty", inputs={"transcript": batch_transcript} - ) + response_text = await self._run_prompty(extract_prompt, inputs={"transcript": batch_transcript}) parsed = self._parse_llm_json(response_text) facts.extend(parsed.get("facts", [])) - episodic.extend(parsed.get("episodic", [])) processed_turns.extend(batch) except Exception as exc: # noqa: BLE001 if is_retryable_llm_error(exc): @@ -584,10 +597,8 @@ async def extract_memories_dry( doc_timestamp = self._stable_source_timestamp(items) fact_docs: list[dict[str, Any]] = [] - episodic_docs: list[dict[str, Any]] = [] updates: list[dict[str, Any]] = [] exact_dedup_skipped = 0 - dropped_episodic_count = 0 for fact in facts: text = fact.get("text") @@ -627,14 +638,14 @@ async def extract_memories_dry( "type": "fact", "content": text, "content_hash": new_content_hash, - "confidence": 0.5 if confidence is None else confidence, - **self._prompt_lineage("extract_memories.prompty"), + "confidence": clamp_unit_interval(confidence, 0.5), + **self._prompt_lineage(extract_prompt), "metadata": { "category": fact.get("category") or "other", "temporal_context": fact.get("temporal_context"), "source": fact_source, }, - "salience": fact.get("salience") if fact.get("salience") is not None else 0.5, + "salience": clamp_unit_interval(fact.get("salience"), 0.5), "tags": ["sys:fact", "sys:auto-extracted"] + source_tags + topic_tags, "created_at": doc_timestamp, "updated_at": doc_timestamp, @@ -643,81 +654,8 @@ async def extract_memories_dry( fact_docs.append(self._validate_extracted_doc(doc)) existing_fact_hashes.add(new_content_hash) - for ep in episodic: - scope_type_raw = ep.get("scope_type") - scope_value_raw = ep.get("scope_value") - scope_type = scope_type_raw.strip() if isinstance(scope_type_raw, str) else None - scope_value = scope_value_raw.strip() if isinstance(scope_value_raw, str) else None - if not scope_type or not scope_value: - logger.warning( - "extract_memories: dropping malformed episodic (missing scope_type/scope_value) " - "user_id=%s thread_id=%s reason=malformed_scope payload=%r", - user_id, - thread_id, - ep, - ) - dropped_episodic_count += 1 - continue - - situation = ep.get("situation") - action_taken = ep.get("action_taken") - outcome = ep.get("outcome") - if situation and action_taken and outcome: - text = f"{situation} → {action_taken} → {outcome}" - else: - text = f"For the user's {scope_value} {scope_type}, intent recorded." - - content_hash = compute_content_hash(text) - seed = _ID_SEED_SEP.join((user_id, thread_id, content_hash)) - det_id = f"ep_{hashlib.sha256(seed.encode()).hexdigest()[:32]}" - topic_tags = build_topic_tags(ep.get("tags", [])) - confidence = ep.get("confidence") - raw_valence = ep.get("outcome_valence") - coerced_valence = _coerce_valence(raw_valence) - if raw_valence is not None and raw_valence not in VALID_VALENCES: - logger.warning( - "extract_memories: coercing unknown outcome_valence=%r → %r user_id=%s thread_id=%s", - raw_valence, - coerced_valence, - user_id, - thread_id, - ) - doc = { - "id": det_id, - "user_id": user_id, - "thread_id": thread_id, - "role": "system", - "type": "episodic", - "content": text, - "content_hash": content_hash, - "confidence": 0.5 if confidence is None else confidence, - "ttl": DEFAULT_TTL_BY_TYPE.get("episodic", 7_776_000), - **self._prompt_lineage("extract_memories.prompty"), - "metadata": { - "scope_type": scope_type, - "scope_value": scope_value, - "situation": situation, - "action_taken": action_taken, - "outcome": outcome, - "reasoning": ep.get("reasoning"), - "outcome_valence": coerced_valence, - "lesson": ep.get("lesson") - or ( - f"{situation} → {action_taken} → {outcome}" if situation and action_taken and outcome else text - ), - "domain": ep.get("domain"), - }, - "salience": ep.get("salience"), - "tags": ["sys:episodic", "sys:auto-extracted"] + topic_tags, - "created_at": doc_timestamp, - "updated_at": doc_timestamp, - } - episodic_docs.append(self._validate_extracted_doc(doc)) - if exact_dedup_skipped: updates.append({"op": "stats", "exact_dedup_skipped": exact_dedup_skipped}) - if dropped_episodic_count: - updates.append({"op": "stats", "dropped_episodic_count": dropped_episodic_count}) if deferred_turn_count or quarantined_turn_count: updates.append( { @@ -727,27 +665,18 @@ async def extract_memories_dry( } ) - check_extracted_fact_grounding( - fact_docs, - processed_turns, - existing_for_hash, - user_id=user_id, - thread_id=thread_id, - logger=logger, - ) - result = { "facts": fact_docs, - "episodic": episodic_docs, + "episodic": [], "updates": updates, "processed_turn_docs": processed_turns, } logger.info( - "extract_memories_dry completed user_id=%s thread_id=%s fact_docs=%d episodic_docs=%d updates=%d", + "extract_memories_durable completed user_id=%s thread_id=%s fact_docs=%d episodic_docs=%d updates=%d", user_id, thread_id, len(fact_docs), - len(episodic_docs), + 0, len(updates), ) return result @@ -973,10 +902,7 @@ async def persist_extracted_memories( validated = self._validate_extracted_doc(doc) doc_type = validated.get("type") try: - if doc_type == "episodic": - await self._upsert_memory(validated) - else: - await self._create_memory(validated) + await self._create_memory(validated) except CosmosResourceExistsError: logger.info("persist_extracted_memories skipped existing id=%s", validated.get("id")) continue @@ -990,16 +916,19 @@ async def persist_extracted_memories( if op.get("op") == "stats": result["exact_dedup_skipped"] += int(op.get("exact_dedup_skipped") or 0) result["dropped_episodic_count"] += int(op.get("dropped_episodic_count") or 0) - if "inplace_updated" in op: - result["inplace_updated"] = result.get("inplace_updated", 0) + int(op.get("inplace_updated") or 0) + for key in ("inplace_updated", "deferred_turn_count", "quarantined_turn_count"): + if key in op: + result[key] = result.get(key, 0) + int(op.get(key) or 0) logger.info("persist_extracted_memories completed user_id=%s counts=%s", user_id, result) return result - async def _mark_turns_extracted(self, turn_docs: list[dict[str, Any]]) -> int: - """Stamp ``extracted_at`` on each turn doc and upsert. Mirror of - the sync helper - per-turn failures are logged but never raise. + 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. """ if not turn_docs: return 0 @@ -1011,12 +940,13 @@ async def _mark_turns_extracted(self, turn_docs: list[dict[str, Any]]) -> int: continue try: doc_to_write = dict(turn) - doc_to_write["extracted_at"] = now_iso + doc_to_write[field] = now_iso await self._upsert_item(self._turns_container, body=doc_to_write) marked += 1 except Exception as exc: logger.warning( - "_mark_turns_extracted failed for turn_id=%s err=%s (turn may be re-extracted on next call)", + "_mark_turns_extracted(%s) failed for turn_id=%s err=%s (turn may be re-processed on next call)", + field, turn_id, exc, ) @@ -1031,8 +961,8 @@ async def extract_memories( turns: Optional[list[dict[str, Any]]] = None, ) -> dict[str, int]: """Extract facts and episodic memories from a thread and persist them.""" - extracted = await self.extract_memories_dry(user_id, thread_id, recent_k, turns=turns) - # Capture the processed turns from the DRY output as the single source of + 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. @@ -1060,6 +990,341 @@ async def extract_memories( ) return counts + async def _load_turn_window( + self, + user_id: str, + thread_id: str, + recent_k: int | None = None, + ) -> list[dict[str, Any]]: + """Load the newest turns for a thread, returned in chronological order.""" + query = "SELECT * FROM c WHERE c.user_id = @user_id AND c.thread_id = @thread_id AND c.type = 'turn'" + items = await self._query_items( + self._turns_container, + query=query, + parameters=[ + {"name": "@user_id", "value": user_id}, + {"name": "@thread_id", "value": thread_id}, + ], + partition_key=[user_id, thread_id], + ) + items.sort(key=lambda m: m.get("created_at", ""), reverse=True) + if recent_k is not None: + items = items[:recent_k] + items.reverse() + return items + + @staticmethod + def _ground_episode_events( + events: Any, + *, + turn_ids: list[str], + ) -> tuple[list[dict[str, Any]], list[str]]: + """Normalize event source ids to real ids from the current turn window.""" + valid_turn_ids = set(turn_ids) + label_to_id = {f"turn-{i}": turn_id for i, turn_id in enumerate(turn_ids, start=1)} + grounded_events: list[dict[str, Any]] = [] + source_turn_ids: list[str] = [] + for event in events if isinstance(events, list) else []: + if not isinstance(event, dict): + continue + grounded = dict(event) + grounded_sources: list[str] = [] + for turn_id in event.get("source_turn_ids") or []: + source = str(turn_id).strip() + mapped = source if source in valid_turn_ids else label_to_id.get(source.lower()) + if mapped and mapped not in grounded_sources: + grounded_sources.append(mapped) + if mapped and mapped not in source_turn_ids: + source_turn_ids.append(mapped) + grounded["source_turn_ids"] = grounded_sources + grounded_events.append(grounded) + return grounded_events, source_turn_ids + + async def _build_episode_docs( + self, + user_id: str, + thread_id: str, + items: list[dict[str, Any]], + *, + segment_key: str, + ) -> list[dict[str, Any]]: + """Run the episode-extraction prompt over one bounded, already-closed turn + segment and return episode docs (no embeddings, no writes). + + Episode ids are DETERMINISTIC from the segment identity (its turn range) + plus each episode's ordinal - not the summary text - so a re-run over the + same segment collides on id and is skipped rather than duplicated. Mirror + of the sync helper. + """ + if not items: + return [] + + transcript_items: list[dict[str, Any]] = [] + for index, item in enumerate(items, start=1): + turn_id = str(item.get("id") or f"turn-{index}") + copied = dict(item) + copied["content"] = f"Turn {turn_id}: {item.get('content', '')}" + transcript_items.append(copied) + transcript = self._build_transcript(transcript_items, include_timestamp=True) + response_text = await self._run_prompty("extract_episode.prompty", inputs={"transcript": transcript}) + parsed = self._parse_llm_json(response_text) + episodes = parsed.get("episodes", []) + if not isinstance(episodes, list): + logger.warning( + "_build_episode_docs dropping malformed response user_id=%s thread_id=%s payload=%r", + user_id, + thread_id, + parsed, + ) + return [] + + doc_timestamp = self._stable_source_timestamp(items) + turn_ids = [str(item.get("id")) for item in items if item.get("id")] + segment_started, segment_ended = segment_time_bounds(items) + + episode_docs: list[dict[str, Any]] = [] + for index, episode in enumerate(episodes): + if not isinstance(episode, dict): + logger.warning( + "_build_episode_docs dropping malformed episode user_id=%s thread_id=%s payload=%r", + user_id, + thread_id, + episode, + ) + continue + + summary = episode.get("summary") + if not isinstance(summary, str) or not summary.strip(): + logger.warning( + "_build_episode_docs dropping malformed episode (missing summary) " + "user_id=%s thread_id=%s payload=%r", + user_id, + thread_id, + episode, + ) + continue + + title = episode.get("title") + if not isinstance(title, str) or not title.strip(): + logger.warning( + "_build_episode_docs dropping malformed episode (missing title) user_id=%s thread_id=%s payload=%r", + user_id, + thread_id, + episode, + ) + continue + + if not isinstance(episode.get("events"), list): + logger.warning( + "_build_episode_docs dropping malformed episode (missing events) " + "user_id=%s thread_id=%s payload=%r", + user_id, + thread_id, + episode, + ) + continue + + content = summary + events, source_turn_ids = self._ground_episode_events(episode.get("events"), turn_ids=turn_ids) + content_hash = compute_content_hash(content) + llm_started, llm_ended = episode.get("started_at"), episode.get("ended_at") + # Trust model-supplied times only when they form a valid, self-consistent + # ISO pair; otherwise fall back to the grounded segment bounds rather than + # dropping the whole episode (malformed or mixed-tz strings are common). + if is_valid_time_pair(llm_started, llm_ended): + started_at, ended_at = str(llm_started).strip(), str(llm_ended).strip() + else: + started_at, ended_at = segment_started, segment_ended + try: + doc = construct_internal( + EpisodicRecord, + { + "id": self._deterministic_episode_id(segment_key, index), + "user_id": user_id, + "thread_id": thread_id, + "role": "system", + "type": "episodic", + "content": content, + "title": title, + "started_at": started_at, + "ended_at": ended_at, + "participants": episode.get("participants") or [], + "events": events, + "outcome": episode.get("outcome"), + "lessons": episode.get("lessons") or [], + "source_turn_ids": source_turn_ids, + "content_hash": content_hash, + "salience": clamp_unit_interval(episode.get("salience"), 0.5), + "confidence": clamp_unit_interval(episode.get("confidence"), 0.5), + "ttl": DEFAULT_TTL_BY_TYPE.get("episodic", 7_776_000), + "tags": ["sys:episodic", "sys:auto-extracted"], + "created_at": doc_timestamp, + "updated_at": doc_timestamp, + **self._prompt_lineage("extract_episode.prompty"), + }, + ).to_doc() + except Exception as exc: # noqa: BLE001 + logger.warning( + "_build_episode_docs dropping malformed episode user_id=%s thread_id=%s err=%s payload=%r", + user_id, + thread_id, + exc, + episode, + ) + continue + episode_docs.append(doc) + + return episode_docs + + @staticmethod + def _deterministic_episode_id(segment_key: str, index: int) -> str: + return deterministic_episode_id(segment_key, index) + + 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).""" + 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))" + ), + parameters=[ + {"name": "@user_id", "value": user_id}, + {"name": "@thread_id", "value": thread_id}, + ], + partition_key=[user_id, thread_id], + ) + items.sort(key=created_at_sort_key) + return items + + async def _episode_segment_embeddings(self, segment: list[dict[str, Any]]) -> list[list[float]]: + """One embedding per segment turn for drift detection, or ``[]`` when drift + is disabled or embeddings are unavailable (mirror of the sync helper).""" + if get_episode_topic_drift() <= 0: + return [] + embeddings: list[Optional[list[float]]] = [ + turn.get("embedding") if isinstance(turn.get("embedding"), list) else None for turn in segment + ] + missing = [i for i, emb in enumerate(embeddings) if emb is None] + if missing: + try: + fresh = await self._embed_batch([str(segment[i].get("content") or "") for i in missing]) + except Exception as exc: # noqa: BLE001 + logger.warning("episode drift embedding failed (%s); skipping drift this evaluation", exc) + return [] + for pos, i in enumerate(missing): + embeddings[i] = fresh[pos] if pos < len(fresh) else None + if any(emb is None for emb in embeddings): + return [] + return [emb for emb in embeddings if emb is not None] + + @staticmethod + def _turn_gap_seconds(prev_turn: dict[str, Any], cur_turn: dict[str, Any]) -> Optional[float]: + return turn_gap_seconds(prev_turn, cur_turn) + + def _find_episode_boundary( + self, + segment: list[dict[str, Any]], + embeddings: list[list[float]], + ) -> Optional[int]: + return find_episode_boundary( + segment, + embeddings, + max_turns=get_episode_max_turns(), + idle_gap=get_episode_idle_gap_seconds(), + drift=get_episode_topic_drift(), + min_turns=get_episode_min_turns(), + ) + + async def extract_episodes( + self, + user_id: str, + thread_id: str, + *, + flush: bool = False, + ) -> 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. + """ + if not user_id: + raise ValidationError("user_id is required") + if not thread_id: + raise ValidationError("thread_id is required") + + segment = await self._load_open_episode_segment(user_id, thread_id) + total = 0 + guard = 0 + while segment and guard < _EPISODE_MAX_SEGMENTS_PER_RUN: + guard += 1 + embeddings = await self._episode_segment_embeddings(segment) + boundary = self._find_episode_boundary(segment, embeddings) + if boundary is None: + if not flush: + break + boundary = len(segment) + closing = segment[:boundary] + if not closing: + break + first_id = str(closing[0].get("id") or "") + last_id = str(closing[-1].get("id") or "") + segment_key = _ID_SEED_SEP.join((user_id, thread_id, first_id, last_id)) + try: + docs = await self._build_episode_docs(user_id, thread_id, closing, segment_key=segment_key) + embeddings_for_docs = await self._embed_batch([str(doc["content"]) for doc in docs]) if docs else [] + except Exception as exc: # noqa: BLE001 + if is_retryable_llm_error(exc): + # Transient provider error: leave the segment un-stamped and stop + # this run so it is retried intact next time (mirror of the fact path). + logger.warning( + "extract_episodes: deferring %d turns after retryable extraction error " + "(will retry next run) user_id=%s thread_id=%s err=%s", + len(closing), + user_id, + thread_id, + exc, + ) + 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. + logger.warning( + "extract_episodes: quarantining %d turns after non-retryable extraction error " + "(marking episode_extracted_at 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") + 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) + 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") + segment = segment[boundary:] + return {"episodes": total} + async def synthesize_procedural( self, user_id: str, @@ -1131,8 +1396,8 @@ 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.metadata.lesson) " - "AND c.metadata.lesson != null " + "AND IS_DEFINED(c.lessons) " + "AND ARRAY_LENGTH(c.lessons) > 0 " "ORDER BY c.salience DESC, c.created_at ASC, c.id ASC" ), parameters=[ @@ -1143,8 +1408,8 @@ async def _read_latest_procedural() -> Optional[dict[str, Any]]: episodic_with_lessons = [ doc for doc in episodic_docs - if isinstance(doc.get("metadata", {}).get("lesson"), str) - and doc.get("metadata", {}).get("lesson", "").strip() + 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] @@ -1183,7 +1448,12 @@ def _render_bullets(values: list[str]) -> str: static_prompty_inputs = { "behavioral_facts": _render_bullets([doc.get("content", "") for doc in behavioral_fact_docs]), "episodic_lessons": _render_bullets( - [doc.get("metadata", {}).get("lesson", "") for doc in episodic_with_lessons] + [ + lesson + for doc in episodic_with_lessons + for lesson in doc.get("lessons", []) + if isinstance(lesson, str) and lesson.strip() + ] ), "user_name": user_name, } @@ -1271,7 +1541,7 @@ def _render_bullets(values: list[str]) -> str: ) return {"status": "synthesized", "procedural": written_doc} - async def generate_thread_summary_dry( + async def generate_thread_summary_durable( self, user_id: str, thread_id: str, @@ -1283,7 +1553,7 @@ async def generate_thread_summary_dry( if not thread_id: raise ValidationError("thread_id is required") - logger.info("generate_thread_summary_dry started user_id=%s thread_id=%s", user_id, thread_id) + logger.info("generate_thread_summary_durable started user_id=%s thread_id=%s", user_id, thread_id) summary_id = f"summary_{user_id}_{thread_id}" existing_summary: Optional[dict[str, Any]] = None @@ -1314,7 +1584,7 @@ async def generate_thread_summary_dry( ) if existing_summary and not items: - logger.info("generate_thread_summary_dry no new memories, returning existing") + logger.info("generate_thread_summary_durable no new memories, returning existing") summary_doc = dict(existing_summary) summary_doc.pop("embedding", None) return summary_doc @@ -1403,10 +1673,10 @@ async def generate_thread_summary( recent_k: int | None = None, ) -> dict[str, Any]: """Generate or incrementally update a thread summary and persist it.""" - summary_doc = await self.generate_thread_summary_dry(user_id, thread_id, recent_k=recent_k) + summary_doc = await self.generate_thread_summary_durable(user_id, thread_id, recent_k=recent_k) return await self.persist_thread_summary(user_id, thread_id, summary_doc) - async def generate_user_summary_dry( + async def generate_user_summary_durable( self, user_id: str, thread_ids: list[str] | None = None, @@ -1417,7 +1687,7 @@ async def generate_user_summary_dry( raise ValidationError("user_id is required") logger.info( - "generate_user_summary_dry started user_id=%s observed_thread_ids=%s", + "generate_user_summary_durable started user_id=%s observed_thread_ids=%s", user_id, len(thread_ids) if thread_ids else 0, ) @@ -1457,7 +1727,7 @@ async def generate_user_summary_dry( ) if existing_summary and not items: - logger.info("generate_user_summary_dry no new memories, returning existing") + logger.info("generate_user_summary_durable no new memories, returning existing") user_doc = dict(existing_summary) user_doc.pop("embedding", None) return user_doc @@ -1563,7 +1833,7 @@ async def generate_user_summary( recent_k: int | None = None, ) -> dict[str, Any]: """Generate or incrementally update a user summary and persist it.""" - summary_doc = await self.generate_user_summary_dry(user_id, thread_ids=thread_ids, recent_k=recent_k) + summary_doc = await self.generate_user_summary_durable(user_id, thread_ids=thread_ids, recent_k=recent_k) return await self.persist_user_summary(user_id, summary_doc) def _emit_reconcile_outcome( diff --git a/azure/cosmos/agent_memory/aio/store/memory_store.py b/azure/cosmos/agent_memory/aio/store/memory_store.py index fce795b..ceb8fc0 100644 --- a/azure/cosmos/agent_memory/aio/store/memory_store.py +++ b/azure/cosmos/agent_memory/aio/store/memory_store.py @@ -197,10 +197,11 @@ async def add( if memory_type == "fact": meta.setdefault("category", "unclassified:manual") elif memory_type == "episodic": - meta.setdefault("lesson", content) - meta.setdefault("scope_type", "manual") - meta.setdefault("scope_value", "manual") - meta.setdefault("outcome_valence", "neutral") + kwargs.setdefault("title", content[:80] or "Manual episode") + kwargs.setdefault("events", []) + kwargs.setdefault("participants", []) + kwargs.setdefault("lessons", []) + kwargs.setdefault("source_turn_ids", []) elif memory_type == "procedural": kwargs.setdefault("source_fact_ids", ["manual"]) kwargs["metadata"] = meta @@ -801,6 +802,36 @@ async def get_procedural_memories( items = [i for i in items if i.get("metadata", {}).get("category") == category] return items + async def get_episodes( + self, + user_id: str, + thread_id: Optional[str] = None, + recent_k: Optional[int] = None, + ) -> list[dict[str, Any]]: + """Retrieve active episodic memories for ``user_id``, newest first.""" + if not user_id: + raise ValidationError("user_id is required for get_episodes") + qb = _QueryBuilder() + qb.add_filter("c.type", "@type", "episodic") + qb.add_filter("c.user_id", "@user_id", user_id) + qb.add_filter("c.thread_id", "@thread_id", thread_id) + qb.add_is_null_or_undefined("c.superseded_by") + parameters = qb.get_parameters() + if recent_k is not None: + parameters.append({"name": "@recent_k", "value": recent_k}) + sql = f"SELECT TOP @recent_k * FROM c{qb.build_where()} ORDER BY c.created_at DESC" + else: + sql = f"SELECT * FROM c{qb.build_where()} ORDER BY c.created_at DESC" + + partition_key, _ = query_scope(user_id, thread_id) + logger.debug("AsyncMemoryStore.get_episodes query: %s", sql) + return await self.query( + sql, + parameters, + container_key=ContainerKey.MEMORIES, + partition_key=partition_key, + ) + async def search( self, search_terms: Optional[str] = None, @@ -1041,16 +1072,75 @@ async def search_episodic( top_k: int = 5, min_salience: Optional[float] = None, include_superseded: bool = False, + thread_id: Optional[str] = None, + tags_all: Optional[list[str]] = None, + tags_any: Optional[list[str]] = None, + exclude_tags: Optional[list[str]] = None, + created_after: Optional[str | datetime] = None, + created_before: Optional[str | datetime] = None, + started_after: Optional[str | datetime] = None, + started_before: Optional[str | datetime] = None, + ended_after: Optional[str | datetime] = None, + ended_before: Optional[str | datetime] = None, ) -> list[dict[str, Any]]: - """Semantic search across episodic memories for a user.""" - return await self.search( - search_terms=search_terms, - user_id=user_id, - memory_types=["episodic"], - top_k=top_k, - min_salience=min_salience, + """Semantic search across episodic memories for a user. + + Temporal arguments are filters only; relevance ranking is vector/FTS-only. + """ + if not user_id: + raise ValidationError("user_id is required for search_episodic") + 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", "episodic") + qb.add_filter("c.user_id", "@user_id", user_id) + qb.add_filter("c.thread_id", "@thread_id", thread_id) + add_tag_filters(qb, tags_all=tags_all, tags_any=tags_any, exclude_tags=exclude_tags) + qb.add_time_range( + "c.created_at", + after=_coerce_datetime_iso(created_after), + before=_coerce_datetime_iso(created_before), + after_param="@created_after", + before_param="@created_before", + ) + qb.add_time_range( + "c.started_at", + after=_coerce_datetime_iso(started_after), + before=_coerce_datetime_iso(started_before), + after_param="@started_after", + before_param="@started_before", + ) + qb.add_time_range( + "c.ended_at", + after=_coerce_datetime_iso(ended_after), + before=_coerce_datetime_iso(ended_before), + after_param="@ended_after", + before_param="@ended_before", + ) + add_salience_filter(qb, min_salience) + + 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, thread_id) + logger.debug("AsyncMemoryStore.search_episodic query: %s", sql) + return await self.query( + sql, + parameters, + container_key=ContainerKey.MEMORIES, + partition_key=partition_key, + ) async def build_episodic_context(self, user_id: str, query: str, top_k: int = 3) -> str: """Build formatted context of relevant past experiences.""" diff --git a/azure/cosmos/agent_memory/auto_trigger.py b/azure/cosmos/agent_memory/auto_trigger.py index ee41322..0859671 100644 --- a/azure/cosmos/agent_memory/auto_trigger.py +++ b/azure/cosmos/agent_memory/auto_trigger.py @@ -102,9 +102,10 @@ def maybe_trigger_steps( n_facts = _threshold_int(thresholds, "get_fact_extraction_every_n", "FACT_EXTRACTION_EVERY_N") n_summary = _threshold_int(thresholds, "get_thread_summary_every_n", "THREAD_SUMMARY_EVERY_N") + n_episode = _threshold_int(thresholds, "get_episode_eval_every_n", "EPISODE_EVAL_EVERY_N") n_user = _threshold_int(thresholds, "get_user_summary_every_n", "USER_SUMMARY_EVERY_N") n_dedup = _threshold_int(thresholds, "get_dedup_every_n", "DEDUP_EVERY_N") - if n_facts == 0 and n_summary == 0 and n_user == 0: + if n_facts == 0 and n_summary == 0 and n_episode == 0 and n_user == 0: return n_dedup_turns = n_facts * n_dedup if n_facts > 0 and n_dedup > 0 else 0 @@ -114,6 +115,7 @@ def maybe_trigger_steps( turn_counts, n_facts=n_facts, n_summary=n_summary, + n_episode=n_episode, n_dedup_turns=n_dedup_turns, thresholds=thresholds, ) @@ -127,6 +129,7 @@ def _trigger_thread_steps( *, n_facts: int, n_summary: int, + n_episode: int, n_dedup_turns: int, thresholds: Any = None, ) -> dict[str, int]: @@ -157,6 +160,7 @@ def _trigger_thread_steps( new_count=new_count, fire_extract=n_facts > 0 and _counters.crosses_threshold(old_count, new_count, n_facts), fire_summary=n_summary > 0 and _counters.crosses_threshold(old_count, new_count, n_summary), + fire_episode=n_episode > 0 and _counters.crosses_threshold(old_count, new_count, n_episode), fire_dedup=n_dedup_turns > 0 and _counters.crosses_threshold(old_count, new_count, n_dedup_turns), thresholds=thresholds, ) @@ -173,6 +177,7 @@ def _fire_thread_steps( new_count: int, fire_extract: bool, fire_summary: bool, + fire_episode: bool, fire_dedup: bool, thresholds: Any = None, ) -> None: @@ -203,6 +208,11 @@ def _fire_thread_steps( "synthesize_procedural", lambda: processor.synthesize_procedural(user_id=user_id), ), + ( + fire_episode, + "process_extract_episodes", + lambda: processor.process_extract_episodes(user_id=user_id, thread_id=thread_id), + ), ( fire_summary, "process_thread_summary", diff --git a/azure/cosmos/agent_memory/chat.py b/azure/cosmos/agent_memory/chat.py index 6376469..fa715c5 100644 --- a/azure/cosmos/agent_memory/chat.py +++ b/azure/cosmos/agent_memory/chat.py @@ -12,6 +12,7 @@ from __future__ import annotations import os +import random import re import time from typing import Any @@ -26,6 +27,8 @@ RETRYABLE_STATUS_CODES = (429, 500, 503) DEFAULT_AZURE_OPENAI_API_VERSION = "2024-12-01-preview" SAMPLING_PARAMS = ("temperature", "top_p", "frequency_penalty", "presence_penalty") +MAX_RETRY_AFTER_DELAY = 60.0 +RETRY_AFTER_FLOOR = 0.1 def resolve_api_version(explicit: str | None) -> str: @@ -60,6 +63,35 @@ def unsupported_param(exc: Exception) -> str | None: return None +def retry_after_delay(exc: Exception) -> float | None: + """Return a Retry-After delay from an OpenAI exception, if present.""" + response = getattr(exc, "response", None) + headers = getattr(response, "headers", None) + get_header = getattr(headers, "get", None) + if not callable(get_header): + return None + + for name, divisor in (("retry-after", 1.0), ("retry-after-ms", 1000.0)): + value = get_header(name) + if value is None: + continue + try: + delay = float(value) / divisor + except (TypeError, ValueError): + continue + if delay >= 0: + return min(max(delay, RETRY_AFTER_FLOOR), MAX_RETRY_AFTER_DELAY) + return None + + +def retry_delay(exc: Exception | None, attempt: int, base_delay: float) -> float: + """Compute retry delay with Retry-After support and jitter.""" + header_delay = retry_after_delay(exc) if exc is not None else None + if header_delay is not None: + return min(header_delay * (1.0 + 0.2 * random.random()), MAX_RETRY_AFTER_DELAY) + return base_delay * (2**attempt) * (0.8 + 0.4 * random.random()) + + def extract_content(response: Any, model: str) -> str: """Pull the assistant content out of a chat-completions response.""" if not response.choices: @@ -164,7 +196,7 @@ def generate( messages: list[dict[str, str]], *, response_format: dict | None = None, - max_retries: int = 3, + max_retries: int = 6, base_delay: float = 2.0, **extra: Any, ) -> str: @@ -215,7 +247,7 @@ def generate( return extract_content(response, self._model) except openai.RateLimitError as exc: if attempt < max_retries - 1: - delay = base_delay * (2**attempt) + delay = retry_delay(exc, attempt, base_delay) logger.warning( "LLM rate-limited (attempt %d/%d), retrying in %.1fs: %s", attempt + 1, @@ -244,7 +276,7 @@ def generate( unsupported_strips += 1 continue if status in RETRYABLE_STATUS_CODES and attempt < max_retries - 1: - delay = base_delay * (2**attempt) + delay = retry_delay(exc, attempt, base_delay) logger.warning( "LLM API error %s (attempt %d/%d), retrying in %.1fs: %s", status, diff --git a/azure/cosmos/agent_memory/cosmos_memory_client.py b/azure/cosmos/agent_memory/cosmos_memory_client.py index c1ea414..7943614 100644 --- a/azure/cosmos/agent_memory/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/cosmos_memory_client.py @@ -671,23 +671,40 @@ def search_cosmos( min_confidence: Optional[float] = None, created_after: Optional[str | datetime] = None, created_before: Optional[str | datetime] = None, + include_episodes: bool = False, include_turns: bool = False, turn_top_k: Optional[int] = None, include_summaries: bool = False, summary_top_k: Optional[int] = None, ) -> list[dict[str, Any]]: - """Search memories using vector similarity, with optional summary / raw-turn blending. - - ``include_summaries`` / ``include_turns`` prepend matching summaries / - append matching raw turns (best-effort); see Docs/concepts.md. + """Search memories using vector similarity, with optional retrieval blending. + + The base search returns facts and, when ``include_episodes`` is True, + episodes too - both in a single ranked query sharing one ``top_k`` budget, + so facts and episodes compete on relevance rather than each getting a + fixed slice. When ``include_episodes`` is False the base is facts only. + Callers may pass ``memory_types`` to search other non-episodic types; + ``episodic`` is added or removed based on ``include_episodes``. Optional + ``include_summaries`` / ``include_turns`` blend those in after the base + block, deduped by content (best-effort; a blend fetch failure never + breaks the base result). """ store = self._get_store() - results = store.search( + # 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: + base_memory_types = [t for t in memory_types if t != "episodic"] + else: + base_memory_types = ["fact"] + if include_episodes: + base_memory_types = [*base_memory_types, "episodic"] + base_memory_types = base_memory_types or ["fact"] + base = store.search( search_terms=search_terms, memory_id=memory_id, user_id=user_id, role=role, - memory_types=memory_types, + memory_types=base_memory_types, thread_id=thread_id, top_k=top_k, tags_all=tags_all, @@ -701,9 +718,20 @@ def search_cosmos( ) if not user_id: - return results + return base + + results: list[dict[str, Any]] = [] + seen_content: set[str] = set() + + def _extend(docs: list[dict[str, Any]]) -> None: + for doc in docs: + content = str(doc.get("content") or "").strip() + if content and content not in seen_content: + seen_content.add(content) + results.append(doc) + + _extend(base) - seen_content = {str(r.get("content") or "").strip() for r in results} if include_summaries: try: summaries = store.search_summaries( @@ -716,14 +744,9 @@ def search_cosmos( except Exception as exc: # noqa: BLE001 logger.warning("search_cosmos: include_summaries search failed (%s); skipping summaries", exc) summaries = [] - # Order is facts -> summaries -> raw turns: append summaries after the - # relevance-ranked memory hits (and before turns) so they neither jump - # the ranking nor evict facts under context-window truncation. - for s in summaries: - content = str(s.get("content") or "").strip() - if content and content not in seen_content: - seen_content.add(content) - results.append(s) + # Summaries come after facts/episodes (and before turns) so they + # neither jump the ranking nor evict memories under truncation. + _extend(summaries) if include_turns: try: @@ -740,11 +763,7 @@ def search_cosmos( except Exception as exc: # noqa: BLE001 logger.warning("search_cosmos: include_turns turn search failed (%s); returning memories only", exc) turns = [] - for turn in turns: - content = str(turn.get("content") or "").strip() - if content and content not in seen_content: - seen_content.add(content) - results.append(turn) + _extend(turns) return results def search_summaries( @@ -866,6 +885,19 @@ def get_thread_summary( recent_k=recent_k, ) + def get_episodes( + self, + user_id: str, + thread_id: Optional[str] = None, + recent_k: Optional[int] = None, + ) -> list[dict[str, Any]]: + """Retrieve active episodic memories for ``user_id``, newest first.""" + return self._get_store().get_episodes( + user_id=user_id, + thread_id=thread_id, + recent_k=recent_k, + ) + def get_user_summary(self, user_id: str) -> Optional[dict[str, Any]]: """Retrieve the user's summary document from Cosmos DB, or ``None`` if absent.""" return self._get_store().get_user_summary(user_id=user_id) @@ -970,6 +1002,36 @@ def extract_memories( """Extract facts and episodic memories from a thread.""" return self._get_pipeline().extract_memories(user_id, thread_id, recent_k) + def extract_episodes( + self, + user_id: str, + thread_id: str, + *, + flush: bool = False, + ) -> dict[str, int]: + """Segment the thread's open turn stream into episodes at detected boundaries. + + Episodes are created at idle time-gaps (detected only once a later turn + reveals the gap), topic shifts, and a max-size cap. A focused session + shorter than the max-size cap therefore episodizes only lazily - on the + next turn after the idle gap - and a one-shot session that never resumes + is not episodized at all under the auto path. Pass ``flush=True`` at the + end of a conversation (or benchmark run) to drain the trailing open + segment immediately; integrators that know when a session ends should + call this on session close. + + Only supported when the in-process backend owns processing; when a + Durable Function app is the active processor this raises + ``NotImplementedError`` so writes are not split away from that backend. + """ + processor = self._get_processor() + if not isinstance(processor, InProcessProcessor): + raise NotImplementedError( + "Episode extraction runs in-process; manual invocation via the SDK is not " + "supported when the Durable Function app is the active processor." + ) + return self._get_pipeline().extract_episodes(user_id, thread_id, flush=flush) + def synthesize_procedural(self, user_id: str, *, force: bool = False) -> dict[str, Any]: processor = self._get_processor() if not isinstance(processor, InProcessProcessor): diff --git a/azure/cosmos/agent_memory/models.py b/azure/cosmos/agent_memory/models.py index ac6c795..7e0683c 100644 --- a/azure/cosmos/agent_memory/models.py +++ b/azure/cosmos/agent_memory/models.py @@ -312,8 +312,6 @@ def _strip_unset_optional(data: dict[str, Any]) -> dict[str, Any]: "prompt_version", "last_used_at", "version", - "scope_type", - "scope_value", } drop_when_empty_list: set[str] = set() drop_when_zero = {"use_count"} @@ -425,46 +423,51 @@ def _require_category(self) -> "FactRecord": return self -_EPISODIC_ALLOWED_VALENCES = {"positive", "negative", "neutral", "mixed"} +class EpisodeEvent(BaseModel): + sequence: int + description: str + occurred_at: Optional[str] = None + source_turn_ids: list[str] = Field(default_factory=list) + + +class EpisodeOutcome(BaseModel): + status: Literal["successful", "partially_successful", "failed", "abandoned", "unknown"] + description: str class EpisodicRecord(MemoryRecordBase): - """A specific past experience: situation → action → outcome → lesson.""" + """A specific past experience captured as an episode timeline.""" memory_type: Literal[MemoryType.episodic] = Field( # type: ignore[assignment] alias="type", default=MemoryType.episodic ) + title: str + started_at: Optional[str] = None + ended_at: Optional[str] = None + participants: list[str] = Field(default_factory=list) + events: list[EpisodeEvent] = Field(default_factory=list) + outcome: Optional[EpisodeOutcome] = None + lessons: list[str] = Field(default_factory=list) + source_turn_ids: list[str] = Field(default_factory=list) content_hash: str - confidence: float = 0.5 - scope_type: Optional[str] = None - scope_value: Optional[str] = None prompt_id: str prompt_version: str = "v1" _ID_PREFIX: ClassVar[Optional[str]] = "ep_" @model_validator(mode="after") - def _require_episodic_metadata(self) -> "EpisodicRecord": - meta = self.metadata if isinstance(self.metadata, dict) else None - if not meta: - raise ValueError( - "EpisodicRecord requires metadata.lesson, metadata.scope_type, " - "metadata.scope_value, and metadata.outcome_valence" - ) - missing = [k for k in ("lesson", "scope_type", "scope_value", "outcome_valence") if not meta.get(k)] - if missing: - raise ValueError(f"EpisodicRecord missing required metadata field(s): {missing}") - valence = meta.get("outcome_valence") - if valence not in _EPISODIC_ALLOWED_VALENCES: - raise ValueError( - f"metadata.outcome_valence must be one of {sorted(_EPISODIC_ALLOWED_VALENCES)}, got {valence!r}" - ) - # Mirror metadata.scope_* to top-level fields so queries that filter - # on the indexed top-level columns keep working. - if not self.scope_type: - object.__setattr__(self, "scope_type", meta.get("scope_type")) - if not self.scope_value: - object.__setattr__(self, "scope_value", meta.get("scope_value")) + def _validate_time_order(self) -> "EpisodicRecord": + if self.started_at is not None and self.ended_at is not None: + started = datetime.fromisoformat(self.started_at.strip().replace("Z", "+00:00")) + ended = datetime.fromisoformat(self.ended_at.strip().replace("Z", "+00:00")) + # Normalize naive values to UTC so a mixed naive/tz-aware pair compares + # safely instead of raising TypeError. + if started.tzinfo is None: + started = started.replace(tzinfo=timezone.utc) + if ended.tzinfo is None: + ended = ended.replace(tzinfo=timezone.utc) + if ended < started: + raise ValueError("ended_at must not be before started_at") return self @@ -597,6 +600,8 @@ class OrchestrationResult(BaseModel): "ThreadSummaryRecord", "UserSummaryRecord", "FactRecord", + "EpisodeEvent", + "EpisodeOutcome", "EpisodicRecord", "ProceduralRecord", "TYPED_RECORD_CLASSES", diff --git a/azure/cosmos/agent_memory/processors/base.py b/azure/cosmos/agent_memory/processors/base.py index 164c948..c994d77 100644 --- a/azure/cosmos/agent_memory/processors/base.py +++ b/azure/cosmos/agent_memory/processors/base.py @@ -82,6 +82,21 @@ def process_extract_memories( recent_k: Optional[int] = None, ) -> dict[str, int]: ... + def process_extract_episodes( + self, + *, + user_id: str, + thread_id: str, + ) -> dict[str, int]: + """Segment the open turn stream into episodes at detected boundaries. + + Deferred backends (e.g. the Durable Functions app) that do not yet + implement episodic segmentation may no-op (return an empty result) or + raise ``NotImplementedError``; the auto-trigger only invokes this on the + in-process backend. + """ + ... + def process_thread_summary( self, *, diff --git a/azure/cosmos/agent_memory/processors/durable.py b/azure/cosmos/agent_memory/processors/durable.py index 3a2b580..e9aedc6 100644 --- a/azure/cosmos/agent_memory/processors/durable.py +++ b/azure/cosmos/agent_memory/processors/durable.py @@ -16,6 +16,10 @@ 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." @@ -55,6 +59,30 @@ def process_extract_memories( ) return {} + def process_extract_episodes( + self, + *, + 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." + ) + logger.debug( + "DurableFunctionProcessor.process_extract_episodes no-op user_id=%s thread_id=%s", + user_id, + thread_id, + ) + return {} + def process_thread_summary( self, *, diff --git a/azure/cosmos/agent_memory/processors/inprocess.py b/azure/cosmos/agent_memory/processors/inprocess.py index 734d6b0..9082d1b 100644 --- a/azure/cosmos/agent_memory/processors/inprocess.py +++ b/azure/cosmos/agent_memory/processors/inprocess.py @@ -105,6 +105,15 @@ def process_extract_memories( extracted = self._pipeline.extract_memories(user_id, thread_id, recent_k=recent_k) return {k: v for k, v in extracted.items() if isinstance(v, int)} if isinstance(extracted, dict) else {} + def process_extract_episodes( + self, + *, + user_id: str, + thread_id: str, + ) -> dict[str, int]: + extracted = self._pipeline.extract_episodes(user_id, thread_id) + return {k: v for k, v in extracted.items() if isinstance(v, int)} if isinstance(extracted, dict) else {} + def process_thread_summary( self, *, diff --git a/azure/cosmos/agent_memory/prompts/_schemas.py b/azure/cosmos/agent_memory/prompts/_schemas.py index 28fa494..a0536c4 100644 --- a/azure/cosmos/agent_memory/prompts/_schemas.py +++ b/azure/cosmos/agent_memory/prompts/_schemas.py @@ -52,7 +52,7 @@ # --------------------------------------------------------------------------- -# extract_memories.prompty - extract facts (user- or agent-sourced) + episodic +# extract_memories.prompty - extract facts (user- or agent-sourced) # --------------------------------------------------------------------------- _FACT_ITEM = { "type": "object", @@ -88,49 +88,85 @@ "additionalProperties": False, } -_EPISODIC_ITEM = { +EXTRACT_MEMORIES_SCHEMA: dict[str, Any] = { "type": "object", "properties": { - "scope_type": {"type": "string"}, - "scope_value": {"type": "string"}, - "situation": {"type": ["string", "null"]}, - "action_taken": {"type": ["string", "null"]}, - "outcome": {"type": ["string", "null"]}, - "outcome_valence": { - "type": ["string", "null"], - "enum": ["positive", "negative", "mixed", "neutral", None], + "facts": {"type": "array", "items": _FACT_ITEM}, + }, + "required": ["facts"], + "additionalProperties": False, +} + + +# --------------------------------------------------------------------------- +# extract_episode.prompty - extract bounded episodic experience records +# --------------------------------------------------------------------------- +_EPISODE_EVENT = { + "type": "object", + "properties": { + "sequence": {"type": "integer"}, + "description": {"type": "string"}, + "occurred_at": {"type": ["string", "null"]}, + "source_turn_ids": {"type": "array", "items": {"type": "string"}}, + }, + "required": ["sequence", "description", "occurred_at", "source_turn_ids"], + "additionalProperties": False, +} + +_EPISODE_OUTCOME = { + "type": ["object", "null"], + "properties": { + "status": { + "type": "string", + "enum": [ + "successful", + "partially_successful", + "failed", + "abandoned", + "unknown", + ], }, - "reasoning": {"type": ["string", "null"]}, - "lesson": {"type": ["string", "null"]}, - "domain": {"type": ["string", "null"]}, - "confidence": {"type": "number"}, + "description": {"type": "string"}, + }, + "required": ["status", "description"], + "additionalProperties": False, +} + +_EPISODE_ITEM = { + "type": "object", + "properties": { + "title": {"type": "string"}, + "summary": {"type": "string"}, + "started_at": {"type": ["string", "null"]}, + "ended_at": {"type": ["string", "null"]}, + "participants": {"type": "array", "items": {"type": "string"}}, + "events": {"type": "array", "items": _EPISODE_EVENT}, + "outcome": _EPISODE_OUTCOME, + "lessons": {"type": "array", "items": {"type": "string"}}, "salience": {"type": "number"}, - "tags": {"type": "array", "items": {"type": "string"}}, + "confidence": {"type": "number"}, }, "required": [ - "scope_type", - "scope_value", - "situation", - "action_taken", + "title", + "summary", + "started_at", + "ended_at", + "participants", + "events", "outcome", - "outcome_valence", - "reasoning", - "lesson", - "domain", - "confidence", + "lessons", "salience", - "tags", + "confidence", ], "additionalProperties": False, } -EXTRACT_MEMORIES_SCHEMA: dict[str, Any] = { +EXTRACT_EPISODE_SCHEMA: dict[str, Any] = { "type": "object", "properties": { - "facts": {"type": "array", "items": _FACT_ITEM}, - "episodic": {"type": "array", "items": _EPISODIC_ITEM}, + "episodes": {"type": "array", "items": _EPISODE_ITEM}, }, - "required": ["facts", "episodic"], + "required": ["episodes"], "additionalProperties": False, } @@ -244,7 +280,9 @@ # --------------------------------------------------------------------------- PROMPTY_SCHEMAS: dict[str, tuple[str, dict[str, Any]]] = { "dedup.prompty": ("DedupOutput", DEDUP_SCHEMA), + "extract_episode.prompty": ("ExtractEpisodesOutput", EXTRACT_EPISODE_SCHEMA), "extract_memories.prompty": ("ExtractMemoriesOutput", EXTRACT_MEMORIES_SCHEMA), + "extract_memories-v2.prompty": ("ExtractMemoriesOutput", EXTRACT_MEMORIES_SCHEMA), "summarize.prompty": ("SummarizeOutput", SUMMARIZE_SCHEMA), "summarize_update.prompty": ("SummarizeUpdateOutput", SUMMARIZE_UPDATE_SCHEMA), "user_summary.prompty": ("UserSummaryOutput", USER_SUMMARY_SCHEMA), diff --git a/azure/cosmos/agent_memory/prompts/dedup.prompty b/azure/cosmos/agent_memory/prompts/dedup.prompty index 21a80d9..85b897e 100644 --- a/azure/cosmos/agent_memory/prompts/dedup.prompty +++ b/azure/cosmos/agent_memory/prompts/dedup.prompty @@ -16,20 +16,20 @@ inputs: --- system: -You are a precision fact-reconciliation system. You receive a pool of active facts (each with an ID, content, confidence, salience, and creation timestamp) and must find **contradictions** — pairs of facts that assert opposing claims about the same subject — and pick a winner and a loser for each. +You are a precision fact-reconciliation system. You receive a pool of active facts (each with an ID, content, confidence, salience, and creation timestamp) and must find **contradictions** - pairs of facts that assert opposing claims about the same subject - and pick a winner and a loser for each. ## Your Goal Produce a clean reconciliation that: 1. Resolves opposing claims about the same subject by picking a winner and a loser. 2. Leaves everything else untouched. -Near-duplicate paraphrases are NOT your concern here — they are folded together earlier, at write time. Do not merge, rewrite, or collapse facts. Your only job is to identify genuine contradictions. +Near-duplicate paraphrases are NOT your concern here - they are folded together earlier, at write time. Do not merge, rewrite, or collapse facts. Your only job is to identify genuine contradictions. ## Two Outcomes Every input fact ends up in one of two places: -- `contradicted_pairs[*]` — the fact is either the `winner_id` or the `loser_id` of a contradiction. -- `kept_ids` — the fact is not in opposition to anything else. +- `contradicted_pairs[*]` - the fact is either the `winner_id` or the `loser_id` of a contradiction. +- `kept_ids` - the fact is not in opposition to anything else. A fact that is merely a paraphrase of, or on the same topic as, another fact (but not contradicting it) belongs in `kept_ids`. @@ -42,9 +42,9 @@ Examples: - "User's deadline is March 1" vs "User's deadline is March 15" → contradiction. Resolution: emit one entry in `contradicted_pairs` with: -- `winner_id` — the surviving fact. -- `loser_id` — the soft-deleted fact. -- `reason` — a short explanation of why the winner beats the loser. +- `winner_id` - the surviving fact. +- `loser_id` - the soft-deleted fact. +- `reason` - a short explanation of why the winner beats the loser. Pick the winner by: 1. **More recent `created_at` first.** @@ -55,7 +55,7 @@ A contradiction is a pair. If three or more facts mutually contradict, emit pair ## What is KEPT -Every fact that is not the loser or winner of a contradiction. List their IDs in `kept_ids`. Paraphrases and same-topic-but-different-claim facts go here — they are NOT contradictions. +Every fact that is not the loser or winner of a contradiction. List their IDs in `kept_ids`. Paraphrases and same-topic-but-different-claim facts go here - they are NOT contradictions. ## Rules @@ -63,9 +63,9 @@ Every fact that is not the loser or winner of a contradiction. List their IDs in 1. **Conservative bias.** If you cannot confidently classify two facts as contradictions, put them both in `kept_ids`. A false contradiction (dropping a fact the user never retracted) is worse than retaining both. -2. **Topic vs claim distinction.** Two facts on the same topic but making *different, compatible* claims are NOT a contradiction. Example: "User prefers dark mode" and "User uses VS Code" share a topic but do not oppose each other — both go in `kept_ids`. +2. **Topic vs claim distinction.** Two facts on the same topic but making *different, compatible* claims are NOT a contradiction. Example: "User prefers dark mode" and "User uses VS Code" share a topic but do not oppose each other - both go in `kept_ids`. -3. **Opposition, not refinement.** A more specific restatement of the same claim ("works at Acme" → "senior engineer at Acme") is NOT a contradiction — both go in `kept_ids`. +3. **Opposition, not refinement.** A more specific restatement of the same claim ("works at Acme" → "senior engineer at Acme") is NOT a contradiction - both go in `kept_ids`. 4. **Don't fabricate.** Never introduce facts, qualifiers, or entities that are not present in the pool. @@ -85,7 +85,7 @@ Some facts may show `Confidence: N/A`, `Salience: N/A`, or `Created: N/A`. ## Output Format -You must output ONLY valid JSON matching this exact schema. No preamble, no explanation, no markdown fences — just the JSON object. +You must output ONLY valid JSON matching this exact schema. No preamble, no explanation, no markdown fences - just the JSON object. ```json { diff --git a/azure/cosmos/agent_memory/prompts/extract_episode.prompty b/azure/cosmos/agent_memory/prompts/extract_episode.prompty new file mode 100644 index 0000000..d46c67c --- /dev/null +++ b/azure/cosmos/agent_memory/prompts/extract_episode.prompty @@ -0,0 +1,128 @@ +--- +name: extract_episode +version: v1 +description: Extract episodic experience records from a conversation window. +model: + apiType: chat + options: + seed: 42 + maxOutputTokens: 16384 + additionalProperties: + response_format: + type: json_object +inputs: + transcript: + type: string +--- + +system: +You are a precision episodic memory extraction system. Read the timestamped, speaker-tagged transcript window. Each turn is labeled `[ | user]:` or `[ | agent]:`; the speaker is the label immediately before the colon. + +## Task +Identify bounded EXPERIENCES or EVENTS that the user or agent went through: a trip, debugging session, life event, project task, booking flow, decision process, or similar coherent episode. Emit ONE episode per distinct coherent experience in the window; usually this is 0-2 episodes. + +If the window contains no coherent experience, return exactly `{"episodes": []}`. + +## Episode Fields +For every episode, produce all of these fields: +- `title`: short, specific title. +- `summary`: 2-4 self-contained, search-optimized sentences. This text is embedded, so include the core who/what/when/why/outcome without relying on transcript context. +- `started_at` and `ended_at`: ISO-8601 timestamps resolved from the turns' timestamps; use `null` if unknown. Resolve relative dates against the timestamp of the turn that stated them, never against now. +- `participants`: named people involved, using real names when the transcript gives them. Use an empty array when none are named. +- `events`: ordered timeline entries. Each event has `sequence`, `description`, `occurred_at` (ISO-8601 or `null`), and `source_turn_ids` grounded in specific turn ids from the transcript. If explicit turn ids are absent, use stable labels like `turn-1`, `turn-2` in transcript order. +- `outcome`: an object with `status` and `description` when the episode has a clear result, or `null` for life-events or open-ended experiences with no success/fail outcome. Allowed statuses are `successful`, `partially_successful`, `failed`, `abandoned`, and `unknown`. +- `lessons`: transferable takeaways, possibly empty. +- `salience` and `confidence`: numbers in `[0, 1]`. + +## Grounding Rules +- Ground everything in the transcript. Do not invent participants, timestamps, feelings, motivations, outcomes, or lessons. +- Do not turn general preferences or isolated facts into episodes unless they are part of a bounded event. +- Prefer fewer, higher-quality episodes over padded output. +- Preserve exact names, dates, products, files, commands, locations, and counts. +- The agent can be the actor in an episode when it completed a task, made a booking, changed code, or otherwise did something concrete for the user. + +## Worked Examples + +### Example 1: Task with outcome + +**Conversation:** +> [2026-03-10T09:00:00Z | user]: "Turn t-101: The checkout tests started failing after I renamed PaymentClient yesterday." +> [2026-03-10T09:08:00Z | agent]: "Turn t-102: I found the mocks still patched OldPaymentClient, updated them to PaymentClient, and reran the checkout suite. All 42 tests passed." + +**Output:** +```json +{ + "episodes": [ + { + "title": "Fixed checkout tests after PaymentClient rename", + "summary": "On 2026-03-09, the checkout tests began failing after PaymentClient was renamed. The agent found that test mocks still patched OldPaymentClient, updated them to PaymentClient, and reran the checkout suite. The checkout suite then passed with all 42 tests successful.", + "started_at": "2026-03-09T00:00:00Z", + "ended_at": "2026-03-10T09:08:00Z", + "participants": [], + "events": [ + { + "sequence": 1, + "description": "The checkout tests started failing after PaymentClient was renamed.", + "occurred_at": "2026-03-09T00:00:00Z", + "source_turn_ids": ["t-101"] + }, + { + "sequence": 2, + "description": "The agent found stale OldPaymentClient mocks, updated them to PaymentClient, and reran the checkout suite.", + "occurred_at": "2026-03-10T09:08:00Z", + "source_turn_ids": ["t-102"] + } + ], + "outcome": { + "status": "successful", + "description": "All 42 checkout tests passed after the mocks were updated." + }, + "lessons": ["After renaming a client class, update tests and mocks that patch the old class name."], + "salience": 0.78, + "confidence": 0.96 + } + ] +} +``` + +### Example 2: Life event with no success/fail outcome + +**Conversation:** +> [2025-11-02T18:20:00-05:00 | user]: "Turn t-201: My sister Elena and I drove to Burlington last weekend for Dad's 70th birthday." +> [2025-11-02T18:21:00-05:00 | user]: "Turn t-202: We had dinner at Hen of the Wood on Saturday and gave him the photo album Mom assembled." + +**Output:** +```json +{ + "episodes": [ + { + "title": "Dad's 70th birthday weekend in Burlington", + "summary": "The user and Elena drove to Burlington during the weekend before 2025-11-02 for Dad's 70th birthday. On Saturday, they had dinner at Hen of the Wood and gave Dad the photo album that Mom assembled. This was a family life event rather than a task with a success or failure outcome.", + "started_at": "2025-10-25T00:00:00-05:00", + "ended_at": "2025-10-26T23:59:59-05:00", + "participants": ["Elena", "Dad", "Mom"], + "events": [ + { + "sequence": 1, + "description": "The user and Elena drove to Burlington for Dad's 70th birthday weekend.", + "occurred_at": "2025-10-25T00:00:00-05:00", + "source_turn_ids": ["t-201"] + }, + { + "sequence": 2, + "description": "They had dinner at Hen of the Wood and gave Dad the photo album Mom assembled.", + "occurred_at": "2025-10-25T00:00:00-05:00", + "source_turn_ids": ["t-202"] + } + ], + "outcome": null, + "lessons": [], + "salience": 0.64, + "confidence": 0.92 + } + ] +} +``` + +user: +{{transcript}} diff --git a/azure/cosmos/agent_memory/prompts/extract_memories-v2.prompty b/azure/cosmos/agent_memory/prompts/extract_memories-v2.prompty new file mode 100644 index 0000000..9f7a56e --- /dev/null +++ b/azure/cosmos/agent_memory/prompts/extract_memories-v2.prompty @@ -0,0 +1,553 @@ +--- +name: extract_memories_v2 +version: v4-additive +description: Extract facts from a conversation or provided content (v2 - additive recall of assistant-provided information). +model: + apiType: chat + options: + seed: 42 + maxOutputTokens: 16384 + additionalProperties: + response_format: + type: json_object +inputs: + transcript: + type: string +--- + +system: +You are a precision memory extraction system. Your task is to read the provided content - a conversation thread, and/or documents or reference material the user has shared - and extract structured facts worth retaining for future reference. + +## What You Extract + +**Facts** - Declarative knowledge: concrete, verifiable statements about the user, their environment, decisions, experiences, or requirements - or salient factual information from the content the user is working with (documents, excerpts, or domain material they provide). This covers both what *is* true about the user and what *happened* to them, including experiences and statements scoped to a particular trip, project, or context (keep the scope inside the fact text; see below). It **also** covers substantive information the *agent* provided at the user's request - recommendations, plans, schedules, lists, tables, explanations, instructions, solutions, and researched answers - because the user will later ask about that content; capture it as an agent-sourced fact (see "Speaker Discrimination" below). + +Every fact must be explicitly grounded in the provided content - never inferred, assumed, or speculated. + +## Extraction Discipline - Read This First + +- **Extract only what was explicitly stated. Do not infer mental states.** Never manufacture a memory about what the user *thinks*, *feels*, *believes*, *wants*, *likes*, *enjoys*, *finds*, *prefers*, or *is worried/frustrated about* unless the user said so in those (or clearly equivalent) words. Turning "the API took 8 seconds" into "the user thinks the API is slow" is inference, not extraction - do not do it. +- **When the user states a concrete fact, record that fact once - and stop.** Do not also emit a second, softer "the user thinks/feels X" memory derived from it. Record each piece of information once; don't restate the same fact under a different framing. +- **Prefer high-value memories over padding, but never drop content the user requested or is likely to reference.** If a passage is pure trivia, filler, or has no future-reference value, extract nothing from it - an empty array is valid for such a passage. This is an anti-padding rule, **not** a license to skip substantive content: a list, plan, schedule, explanation, recommendation, or answer the user asked for is high-value by definition and must be captured even though it came from the agent. When a passage is borderline but the user may reference it later, keep it - a slightly redundant memory is far cheaper than losing content the user asks about. + +## Speaker Discrimination - Where Memories May Come From + +The transcript below is line-tagged by speaker. Each line is `[user]: ...` (the human's own words) or `[agent]: ...` (the agent's response); lines may also carry the turn's timestamp before the speaker, as `[ | user]: ...` / `[ | agent]: ...`. In every case the **speaker is the label immediately before the colon** (the part after `|` when a timestamp is present). Both `user` and `agent` can be a source of facts, but they carry **different** kinds of facts, and each fact you extract must declare its `source` (`"user"` or `"agent"`). + +- **User-sourced facts (`source: "user"`)** - anything the user asserts about themselves, their preferences, their world, or content they provide/share/ask to remember (including documents, excerpts, and reference material in `[user]:` lines). +- **Agent-sourced facts (`source: "agent"`)** - substantive information the agent contributes that the user may later reference. Capture BOTH kinds: + - **Actions & commitments** - concrete things the agent did on the user's behalf: bookings made, orders placed, files/PRs created, appointments scheduled, and named recommendations given (e.g. `[agent]: "I've booked you the 3 PM United flight UA327 to SFO"`; `[agent]: "I recommend Blue Bottle on Mint St for coffee"`). These answer "what did you book me?" / "what did you recommend?". + - **Information the agent provided** - the answers, explanations, lists, tables, schedules, itineraries, step-by-step instructions, solutions, stories, and descriptions the agent generated in response to the user's request. When the user asks the agent to "build a shift rotation", "list work-from-home jobs", "explain the refining process at each plant", or "write a children's story", the agent's reply **is** memorable content the user will come back to ("what was Admon's Sunday shift?", "what was the 7th job in that list?", "what color was the Plesiosaur?"). Capture the **substance** - every enumerated item **in its original order and position**, every named entity, quantity, date, and specific. If the content is detail-dense, split it into a few focused agent facts (e.g. one per section or item group) rather than compressing the specifics away. Tag all of this `source: "agent"`. + + Both kinds answer later questions like "what did you tell me about X?" and "what was the Nth item you gave me?", so they must be captured - never reduce them to a bare meta-description of the fact that a reply happened. +- **NEVER turn an agent line into a fact *about the user*.** When the agent restates, paraphrases, confirms, or assumes something on the user's behalf ("Got it, you don't eat meat", "I assume you want a luxury hotel", "So you're a vegetarian"), that is the agent's framing, **not** the user's assertion - do not extract it as a user fact, and do not extract it as an agent fact either. A user fact must be grounded in a `[user]:` line. If a claim *about the user* appears only in `[agent]:` content, drop it. (This bans importing the agent's **assumptions about the user** as user facts; it does **not** apply to information the agent *provided* at the user's request - that is genuine agent-sourced content and must still be captured under the rule above.) +- **Skip only genuinely unprompted, throwaway asides - keep information the user asked for.** A stray aside the agent volunteers that is unrelated to anything the user asked ("By the way, Python 3.13 shipped in October 2024") and that the user shows no interest in can be skipped. But when the user *requested* the information, or engages with it, the agent's answer is exactly the content they will ask about later, so it must be captured (as `source: "agent"`) with its specifics intact - do **not** discard it as "general knowledge it recited". When unsure whether the agent's information was requested or referenceable, keep it. + +## Confidence Scoring + +Every extracted fact must include a `confidence` field in `[0.0, 1.0]` indicating how strongly the conversation supports the claim: + +| Range | Meaning | +|-------|---------| +| 0.9–1.0 | Directly stated and unambiguous | +| 0.7–0.9 | Clearly implied, no contradicting evidence | +| 0.5–0.7 | Inferred from context - plausible but not explicit | +| < 0.5 | Speculative - likely should not be extracted at all | + +Retrieval may filter by `min_confidence` so under-confident extractions get suppressed automatically. + +--- + +## Facts - Declarative Knowledge + +Extract concrete, factual statements that fall into these categories: + +| Category | Description | Example | +|----------|-------------|---------| +| `preference` | Things the user likes, dislikes, or prefers | "User prefers dark mode in all IDEs" | +| `requirement` | Explicit constraints or needs | "The project must comply with SOC 2" | +| `biographical` | Stable personal or professional details | "User is a senior engineer at Contoso" | +| `other` | Any factual claim that doesn't fit the above (decisions, events, dates, relationships, tasks, domain facts, etc.) | "User chose PostgreSQL over MySQL"; "The MVP deadline is March 1st, 2025" | + +Category is a lightweight hint, not a hard decision - when a fact doesn't clearly fit `preference`, +`requirement`, or `biographical`, use `other`. **Never withhold or drop a fact because it is hard to +categorize** - extracting the fact matters far more than labeling it. + +**Capturing time:** Whenever a fact involves any time reference - an exact date, a duration, a +frequency, a deadline, or a relative expression ("yesterday", "3 weeks ago", "last June", "every +Monday") - keep that time expression verbatim in the fact `text` AND record it in `temporal_context`. +Time-bound answers ("when did X happen?", "how long?") are impossible to recover later if the time +detail is dropped, so never omit or round it. + +**Resolving relative time to absolute dates:** Each transcript line may be prefixed with the turn's +timestamp as `[ | role]: ...`. When a fact uses a **relative** expression, use that line's +timestamp as the anchor to compute the concrete calendar date, and put the resolved absolute date in +the fact `text` while keeping the original expression in `temporal_context`. For example, a turn +stamped `2024-06-20` saying "I flew to Tokyo 3 weeks ago" → `text: "The user flew to Tokyo on 2024-05-30."`, +`temporal_context: "3 weeks ago"`. If no line timestamp is present, keep the relative expression as-is +in both fields. Never invent a date when there is no anchor. + +**Two clocks - never confuse them.** A line's `[ | ...]` prefix is the fact's **event time**: +when it was actually said. The moment you are running now (processing time) is unrelated and may be +months or years later. Resolve every relative expression ("last year", "3 weeks ago", "recently", +"next month") **only** against the event time of the line it appears on - **never** against today / +the current processing date. A fact whose "last year" is computed from today instead of from when it +was said is silently wrong and stays wrong forever. + +**Grounding rules for time:** +- **Relative → absolute, never absolute → vague.** Convert "yesterday", "last week", "3 weeks ago" + into concrete dates using the event time. But never do the reverse: an exact "18 days", "June 3rd", + or "every other Tuesday" must survive verbatim - do not soften it to "some time" or "a while". +- **Preserve exact durations, frequencies, and counts** as stated: "18 days" stays "18 days" (not + "about 3 weeks"), "twice a week" stays "twice a week" (not "regularly"). +- **Keep the relationship, not just the endpoint.** "The deadline moved from March 1 to March 15" + records both the change and both dates - not merely "the deadline is March 15". + +### Fact Formatting Rules +- Each fact must be self-contained and intelligible without context - no pronouns like "it" or "they" without antecedents +- **Name people when the conversation names them.** Write in third person (never "I" / "you"), but use a person's actual name once it is given - the user's own name if stated, and always for third parties (family, colleagues, friends: "Elena", "Sara", "the user's manager Raj"). Fall back to "the user" only when no name is available. In multi-speaker transcripts this is essential: attribute each fact to the specific person who stated or is described by it (by name), never a generic "the user" - otherwise two speakers blur into one and later retrieval cannot tell who did what. +- **Preserve the exact meaning - direction and polarity matter.** Read each statement carefully before recording its meaning. "The user hasn't shipped the release yet" ≠ "shipped the release". "The user moved *off* Datadog" ≠ "uses Datadog". "The user used to run marathons" ≠ "runs marathons" (no longer). Getting the direction of a statement backwards is worse than not extracting it. +- **Preserve every specific detail stated in the source: exact dates, times, durations, numbers, quantities, amounts, proper nouns, names, brands, product/model names, and named locations or events.** NEVER generalize a specific into a category - keep "June 3rd" (not "a date"), "7 days" (not "about a week"), "Fitbit Versa 3" (not "a smartwatch"), "$4,500" (not "some money"). A fact that drops the specific detail is useless for later recall - the specific IS the memory. +- **Group by topic - one contextually rich memory per distinct topic or event, not one per claim.** Capture the whole coherent picture (the core fact plus its immediate context, qualifiers, and every specific) in a single self-contained memory instead of shattering it into fragments. Split into separate memories only when the content spans genuinely distinct topics (e.g. career vs. family vs. a trip) or when one topic is so detail-dense that a single memory would run past ~3 sentences - then split along natural sub-groupings (e.g. one memory per enemy type in a stat block), never by shaving off individual details. Do NOT merge two unrelated topics to save space, and do NOT split one topic just to make memories smaller. +- **Capture transitions as one memory - the new state AND what it replaced.** When the user changes, switches, replaces, upgrades, or stops something in favor of something else, the link between old and new is essential context: record both together. "The user switched from almond milk to oat milk after developing an almond sensitivity" - not two disconnected facts. If the change is explicitly temporary or a trial ("for a month", "trying out"), capture that too. Keeping old→new in one memory prevents the two states from later surfacing as contradictory-looking facts. +- Keep phrasing tight, but never generalize away a concrete detail (exact date, number, duration, name, brand) to save words. This preserves detail *within* the facts you choose to keep - it is not a directive to maximize how many facts you extract. +- Each memory is stored as its own document with its own vector embedding, so keep every memory to a **single coherent topic** - blending two unrelated topics into one memory produces a muddy embedding that retrieves poorly for both. The target is rich-but-focused: all the specifics of one topic, nothing borrowed from another. + +--- + +## Capturing Stories and Events as Facts + +An experience or event the user describes is captured as a fact - keep the whole arc (what happened, what they did, how it turned out) together in one rich fact rather than splitting or dropping it. Consider this exchange: + +> **User**: "I tried using Redis for session storage last week, but it kept timing out under load. We switched to DynamoDB and it's been stable since. My team lead Sarah approved the switch." + +This produces facts: +- (decision) "Last week the user's team switched session storage from Redis to DynamoDB because Redis kept timing out under load; it has been stable since." +- (relational) "Sarah is the user's team lead and approved the Redis-to-DynamoDB switch." + +The story becomes one coherent fact that keeps the problem, the change, and the outcome together - not three disconnected fragments. A statement scoped to a specific context (a trip, project, event, or session) is also a fact; keep the scope inside the text (see "Scoped statements ARE facts" below). + +--- + +## Salience Scoring Rubric + +Assign a salience score to every extracted memory using this scale: + +| Score | Level | Examples | +|-------|-------|---------| +| 0.9–1.0 | Critical / Confirmed | Name, hard deadlines, budget limits, security requirements | +| 0.7–0.8 | Strong preference / Key decision | Technology choices, communication style rules, project goals | +| 0.5–0.6 | Moderate relevance | Useful context, secondary preferences, background details | +| 0.3–0.4 | Minor detail | Incidental mentions, soft preferences, nice-to-know info | +| 0.1–0.2 | Weak / Uncertain | Offhand remarks, single-instance observations, ambiguous signals | + +--- + +## What to Exclude +Do NOT extract: +- **Inferred opinions, feelings, or mental states** - e.g. "User thinks the API might be slow", "User enjoys concise answers", "User is frustrated with the tool". Never write "the user thinks / feels / believes / wants / likes / enjoys / finds X" unless the user literally said it. Synthesizing the user's attitude from what they did or described is inference, not extraction. +- Filler or pleasantries ("User said thanks") +- Uncertain or hypothetical statements ("User mentioned they might switch tools") +- Redundant memories - if the same information appears multiple times, extract it only once +- Raw agent reasoning or intermediate steps that did not produce a confirmed outcome +- Memories that are only meaningful within the context of this conversation and have no future reference value + +**Extract the content, not the act of sharing.** When the user pastes or references a document, spec, dataset, case, stat block, or any reference material - **or when the agent produces a list, table, schedule, itinerary, story, or explanation at the user's request** - extract the concrete facts *inside* it - the parties, dates, figures, findings, named items, quantities, and every enumerated entry in order - each as its own memory. Do NOT produce a memory that merely describes the gesture ("The user shared a contract summary", "The user asked to remember a spec", **"The agent provided a shift rotation sheet with 7 agents"**). A memory of "the user pasted a case" or "the agent gave a list" is useless; the case's actual facts, or the list's actual entries, are the memory. The same applies to structured data: preserve every count and value ("4 mummies, AC 11, 45 HP"), splitting a detail-dense block into a few focused memories rather than compressing the numbers away. + +- WRONG: "The agent provided a 1-week shift rotation sheet for the social-media agents." → RIGHT: "In the shift rotation the agent produced, on Sundays Admon works the 8am–4pm Day Shift, Magdy the 4pm–12am shift, …" (keep each agent's day/shift mapping). +- WRONG: "The agent listed work-from-home jobs for seniors." → RIGHT: "The work-from-home jobs the agent listed for seniors were, in order: 1. Virtual customer service representative, 2. Telehealth professional, … 7. Transcriptionist, …" (keep the ordered enumeration). +- WRONG: "The agent explained the refining processes at CITGO's refineries." → RIGHT: "The agent said CITGO's Lake Charles refinery uses atmospheric distillation, fluid catalytic cracking (FCC), alkylation, and hydrotreating." (keep the per-item specifics). + +### Scoped statements ARE facts - keep the scope in the text +A statement tied to a particular trip, project, event, or session is still a fact - do NOT drop it. Capture it as a fact and keep the scope inside the text so it never reads as an unconditional standing claim: +- "For this Paris trip, I want luxury accommodations." → "For the user's Paris trip, the user wants luxury accommodations." +- "On the Acme project, we'll use TypeScript strict mode." → "On the Acme project, the team uses TypeScript strict mode." +- "Just for today, please give me short answers." → "For today's session, the user asked for short answers." + +Keeping the scope phrase ("for this trip", "on this project", "for today's session") in the text means a context-bound statement never collides with a standing preference on the same subject. + +--- + +## Quality Checks + +1. Could someone act on or reference this fact without reading the original thread? +2. Is this fact stated explicitly, not inferred? +3. Is each memory scoped to a single coherent topic, with all of that topic's specifics captured together - not shattered into fragments, and not blended with a second topic? +4. Does each fact retain every specific detail (dates, numbers, durations, names, brands) from the source, with nothing generalized away? + +**Before finalizing - scan the whole transcript.** Re-read the entire conversation, including the +**middle and end**, not only the opening. A common failure is capturing the first prominent topic +thoroughly and skimming the rest ("first-topic dominance"). If a single message raised several +distinct topics (e.g. a job change, a health note, and a trip), make sure each distinct topic is +represented. This is about not *missing* a stated fact - it is NOT a directive to inflate the count +or extract filler; the precision and grounding rules above still apply. + +--- + +## Few-Shot Examples + +### Example 1: Mixed conversation about project setup + +**Conversation:** +> User: "I'm Alex, a data engineer at Acme Corp. We just kicked off a new ETL pipeline project. The deadline is end of Q2." + +**Output:** +```json +{ + "facts": [ + { + "text": "The user is Alex, a data engineer at Acme Corp.", + "category": "biographical", + "source": "user", + "confidence": 1.0, + "salience": 0.9, + "temporal_context": null, + "tags": ["topic:identity", "topic:career"] + }, + { + "text": "Alex's team just kicked off a new ETL pipeline project with a deadline at end of Q2.", + "category": "other", + "source": "user", + "confidence": 0.95, + "salience": 0.9, + "temporal_context": "end of Q2", + "tags": ["topic:project", "topic:ETL"] + } + ] +} +``` + +Name and role are one coherent biographical topic → a single rich memory, not two fragments. The ETL project is a distinct topic → its own memory. + +### Example 2: Troubleshooting experience (a story becomes one rich fact) + +**Conversation:** +> User: "Last month we had a production outage because our Kubernetes pods kept OOM-killing. We bumped memory limits from 512MB to 1GB and added resource quotas per namespace. That fixed it." + +Capture the whole arc - the problem, the change, and the outcome - in a single fact. Do NOT shatter it into disconnected "uses 1GB limits" plus "had an outage" fragments. + +**Output:** +```json +{ + "facts": [ + { + "text": "Last month the user's team resolved a production outage caused by Kubernetes pods OOM-killing by raising pod memory limits from 512MB to 1GB and adding per-namespace resource quotas; that fixed it.", + "category": "other", + "source": "user", + "confidence": 0.95, + "salience": 0.7, + "temporal_context": "last month", + "tags": ["topic:kubernetes", "topic:infrastructure", "topic:outage"] + } + ] +} +``` + +### Example 3: Database and budget facts + +**Conversation:** +> User: "Our analytics DB is on BigQuery. Also, the marketing team's budget is $50,000 for this quarter." + +**Output:** +```json +{ + "facts": [ + { + "text": "The user's analytics database runs on BigQuery.", + "category": "biographical", + "source": "user", + "confidence": 0.95, + "salience": 0.7, + "temporal_context": null, + "tags": ["topic:database", "topic:BigQuery"] + }, + { + "text": "The marketing team's budget is $50,000 for the current quarter.", + "category": "requirement", + "source": "user", + "confidence": 0.95, + "salience": 0.9, + "temporal_context": "current quarter", + "tags": ["topic:budget", "topic:marketing"] + } + ] +} +``` + +### Example 4: Standing preference vs. scoped preference (both are facts) + +**Conversation:** +> User: "I usually prefer budget hotels." +> User: "For this Paris trip, I want luxury accommodations." + +Both are facts. The first is a standing preference. The second is scoped to the Paris trip - keep that scope inside the text so it never reads as a standing claim and never collides with the general preference on retrieval. + +**Output:** +```json +{ + "facts": [ + { + "text": "The user usually prefers budget hotels.", + "category": "preference", + "source": "user", + "confidence": 0.95, + "salience": 0.7, + "temporal_context": null, + "tags": ["topic:travel", "topic:hotels"] + }, + { + "text": "For the user's Paris trip, the user wants luxury accommodations.", + "category": "preference", + "source": "user", + "confidence": 0.95, + "salience": 0.7, + "temporal_context": null, + "tags": ["topic:travel", "topic:hotels", "topic:luxury"] + } + ] +} +``` + +### Example 5: Discipline - extract only what was stated, infer nothing + +**Conversation:** +> User: "The dashboard took about 8 seconds to load the report. Anyway, the weather's been nice this week." + +Extract only the literal, user-stated fact. The load time is concrete and worth keeping. "The weather's been nice" is filler with no future-reference value → excluded. Do NOT manufacture inferred opinions such as "The user thinks the dashboard is slow" or "The user enjoys nice weather" - the user asserted neither. + +**Output:** +```json +{ + "facts": [ + { + "text": "The user's dashboard took about 8 seconds to load the report.", + "category": "other", + "source": "user", + "confidence": 0.9, + "salience": 0.4, + "temporal_context": null, + "tags": ["topic:performance", "topic:dashboard"] + } + ] +} +``` + +### Example 6: Capturing a transition as a single memory + +**Conversation:** +> User: "I moved our CI from Jenkins to GitHub Actions last quarter - the Jenkins maintenance was eating too much of my time." + +The change is one coherent memory: it records the new state, what it replaced, and why. Do NOT split it into a disconnected "uses GitHub Actions" plus "used Jenkins" - that loses the relationship and later looks like two conflicting facts. + +**Output:** +```json +{ + "facts": [ + { + "text": "The user's team migrated CI from Jenkins to GitHub Actions last quarter because Jenkins maintenance was consuming too much of the user's time.", + "category": "other", + "source": "user", + "confidence": 0.95, + "salience": 0.7, + "temporal_context": "last quarter", + "tags": ["topic:CI", "topic:tooling"] + } + ] +} +``` + +### Example 7: Agent-sourced facts - what the agent did and recommended + +**Conversation:** +> User: "Book me a table for two somewhere good for our anniversary this Friday, and sort out the flights to Lisbon." +> Agent: "Done - I've reserved a table for two at Osteria Mozza at 7:30 PM this Friday, and I booked you on TAP Air Portugal flight TP204 departing SFO at 10:15 AM on June 3rd. I'd also recommend the Alfama district for your stay." + +The user's request is an intent; the concrete bookings and the named recommendation are things the **agent** did/said and are needed to answer "what did you book?" / "what did you recommend?" later. Extract them as `source: "agent"`. Keep every specific (venue, flight number, times, date, district). + +**Output:** +```json +{ + "facts": [ + { + "text": "The agent reserved a table for two at Osteria Mozza at 7:30 PM this Friday for the user's anniversary.", + "category": "other", + "source": "agent", + "confidence": 0.95, + "salience": 0.7, + "temporal_context": "this Friday", + "tags": ["topic:reservation", "topic:anniversary"] + }, + { + "text": "The agent booked the user on TAP Air Portugal flight TP204 departing SFO at 10:15 AM on June 3rd, headed to Lisbon.", + "category": "other", + "source": "agent", + "confidence": 0.95, + "salience": 0.8, + "temporal_context": "June 3rd", + "tags": ["topic:flight", "topic:travel"] + }, + { + "text": "The agent recommended the Alfama district for the user's stay in Lisbon.", + "category": "other", + "source": "agent", + "confidence": 0.9, + "salience": 0.6, + "temporal_context": null, + "tags": ["topic:travel", "topic:recommendation"] + } + ] +} +``` + +### Example 7b: Agent-provided information the user will ask about later + +**Conversation:** +> User: "Can you list some good work-from-home jobs for seniors?" +> Agent: "Sure! Here are some options: 1. Virtual customer service representative, 2. Telehealth professional, 3. Remote bookkeeper, 4. Virtual tutor or teacher, 5. Freelance writer or editor, 6. Online survey taker, 7. Transcriptionist, 8. Virtual assistant." + +The user asked for this list, so the agent's answer is memorable content they will later reference ("what was the 7th job you listed?"). Capture the **ordered enumeration** with each item in its original position, as `source: "agent"` - do NOT reduce it to "the agent listed some jobs". This is the single most common miss: an agent-provided list, table, schedule, or explanation compressed into a bare meta-description that loses the very detail the user later asks for. + +**Output:** +```json +{ + "facts": [ + { + "text": "The work-from-home jobs for seniors that the agent listed were, in order: 1. Virtual customer service representative, 2. Telehealth professional, 3. Remote bookkeeper, 4. Virtual tutor or teacher, 5. Freelance writer or editor, 6. Online survey taker, 7. Transcriptionist, 8. Virtual assistant.", + "category": "other", + "source": "agent", + "confidence": 0.95, + "salience": 0.6, + "temporal_context": null, + "tags": ["topic:work-from-home", "topic:jobs"] + } + ] +} +``` + +### Example 8: Agent assumptions about the user are NOT facts + +**Conversation:** +> User: "I need a hotel in Rome for three nights." +> Agent: "Got it - I assume you want a luxury 5-star place near the center, and since you mentioned Rome you're probably interested in historical sites. I'll put together some options." + +The user-stated need is a fact, scoped to the Rome trip - capture it with the scope in the text. The agent's "I assume you want luxury" and "you're probably interested in historical sites" are the agent's **assumptions about the user**, not the user's assertions - do NOT extract them as user facts, and they are not agent actions either, so extract nothing from the agent line. No booking was made, so there is no agent-sourced fact. + +**Output:** +```json +{ + "facts": [ + { + "text": "For the user's Rome trip, the user needs a hotel for three nights.", + "category": "requirement", + "source": "user", + "confidence": 0.9, + "salience": 0.6, + "temporal_context": null, + "tags": ["topic:travel", "topic:hotels"] + } + ] +} +``` + +### Example 9: Reference material - extract the content, not the act of sharing + +**Conversation:** +> User: "Remember this case for me. Bajimaya v Reward Homes [2021] NSWCATAP 297 - construction began in 2014, the contract was signed in 2015 with completion due by October 2015, and the owner received the keys in December 2016 with defects. The tribunal found the builder breached the contract." + +Extract the actual facts *inside* the shared material - dates, parties, findings - each as its own memory. Do NOT record "The user shared a case summary." + +**Output:** +```json +{ + "facts": [ + { + "text": "In Bajimaya v Reward Homes [2021] NSWCATAP 297, construction began in 2014, the contract was signed in 2015, and completion was due by October 2015.", + "category": "other", + "source": "user", + "confidence": 0.95, + "salience": 0.6, + "temporal_context": null, + "tags": ["topic:legal", "topic:construction"] + }, + { + "text": "In Bajimaya v Reward Homes, the owner received the keys in December 2016 and the tribunal found the builder had breached the contract.", + "category": "other", + "source": "user", + "confidence": 0.95, + "salience": 0.6, + "temporal_context": "December 2016", + "tags": ["topic:legal", "topic:construction"] + } + ] +} +``` + +### Example 10: Multi-topic message + named people - one memory per distinct topic + +**Conversation:** +> User: "Quick update - my sister Priya just moved to Portland, I got promoted to team lead last week, and my daughter Sara started aerial yoga on Tuesdays." + +Three unrelated topics in one message. Extract each separately, and use the actual names (Priya, Sara) rather than collapsing everyone into "the user". + +**Output:** +```json +{ + "facts": [ + { + "text": "The user's sister Priya recently moved to Portland.", + "category": "biographical", + "source": "user", + "confidence": 0.95, + "salience": 0.6, + "temporal_context": null, + "tags": ["topic:family"] + }, + { + "text": "The user was promoted to team lead last week.", + "category": "biographical", + "source": "user", + "confidence": 0.95, + "salience": 0.8, + "temporal_context": "last week", + "tags": ["topic:career"] + }, + { + "text": "The user's daughter Sara started aerial yoga on Tuesdays.", + "category": "biographical", + "source": "user", + "confidence": 0.95, + "salience": 0.6, + "temporal_context": null, + "tags": ["topic:family"] + } + ] +} +``` + +--- + +## Output Format + +You must output ONLY valid JSON matching the schema below. No preamble, no explanation, no closing remarks - just the JSON object. The `facts` array can be empty if no facts are found. + +```json +{ + "facts": [ + { + "text": "Self-contained fact for embedding", + "category": "preference|requirement|biographical|other", + "source": "user|agent", + "confidence": 0.95, + "salience": 0.8, + "temporal_context": "the exact time expression stated (e.g. 'June 3rd', '3 weeks ago', 'by Friday', 'every Monday') or null if none", + "tags": ["topic:x"] + } + ] +} +``` + +--- +Now extract memories from the following conversation thread: + +user: +{{transcript}} diff --git a/azure/cosmos/agent_memory/prompts/extract_memories.prompty b/azure/cosmos/agent_memory/prompts/extract_memories.prompty index d5b7cb4..f0abade 100644 --- a/azure/cosmos/agent_memory/prompts/extract_memories.prompty +++ b/azure/cosmos/agent_memory/prompts/extract_memories.prompty @@ -1,7 +1,7 @@ --- name: extract_memories -version: v3 -description: Extract facts and episodic memories from a conversation or provided content. +version: v4 +description: Extract facts from a conversation or provided content. model: apiType: chat options: @@ -16,14 +16,13 @@ inputs: --- system: -You are a precision memory extraction system. Your task is to read the provided content - a conversation thread, and/or documents or reference material the user has shared - and extract structured memories worth retaining for future reference. You extract two distinct types of memory, each serving a different purpose. +You are a precision memory extraction system. Your task is to read the provided content - a conversation thread, and/or documents or reference material the user has shared - and extract structured facts worth retaining for future reference. -## Memory Type Taxonomy +## What You Extract -1. **Facts** - Declarative knowledge: concrete, verifiable statements about the user, their environment, decisions, or requirements - or salient factual information from the content the user is working with (documents, excerpts, or domain material they provide). These are things that *are* true. -2. **Episodic** - Past experiences: specific situations the user encountered, what they tried, and what happened. These are things that *happened*. +**Facts** - Declarative knowledge: concrete, verifiable statements about the user, their environment, decisions, experiences, or requirements - or salient factual information from the content the user is working with (documents, excerpts, or domain material they provide). This covers both what *is* true about the user and what *happened* to them, including experiences and statements scoped to a particular trip, project, or context (keep the scope inside the fact text; see below). -Every memory must be explicitly grounded in the provided content - never inferred, assumed, or speculated. +Every fact must be explicitly grounded in the provided content - never inferred, assumed, or speculated. ## Extraction Discipline - Read This First @@ -35,15 +34,14 @@ Every memory must be explicitly grounded in the provided content - never inferre The transcript below is line-tagged by speaker. Each line is `[user]: ...` (the human's own words) or `[agent]: ...` (the agent's response); lines may also carry the turn's timestamp before the speaker, as `[ | user]: ...` / `[ | agent]: ...`. In every case the **speaker is the label immediately before the colon** (the part after `|` when a timestamp is present). Both `user` and `agent` can be a source of facts, but they carry **different** kinds of facts, and each fact you extract must declare its `source` (`"user"` or `"agent"`). -- **User-sourced facts (`source: "user"`)** — anything the user asserts about themselves, their preferences, their world, or content they provide/share/ask to remember (including documents, excerpts, and reference material in `[user]:` lines). -- **Agent-sourced facts (`source: "agent"`)** — the agent's **own concrete actions, commitments, and specific recommendations** carried out on the user's behalf: bookings made, orders placed, files/PRs created, appointments scheduled, and named recommendations given (e.g. `[agent]: "I've booked you the 3 PM United flight UA327 to SFO"` → agent-sourced fact; `[agent]: "I recommend Blue Bottle on Mint St for coffee"` → agent-sourced fact). These answer later questions like "what did you book me?" or "what did you recommend?", so they must be captured — tagged `source: "agent"`. -- **NEVER turn an agent line into a fact *about the user*.** When the agent restates, paraphrases, confirms, or assumes something on the user's behalf ("Got it, you don't eat meat", "I assume you want a luxury hotel", "So you're a vegetarian"), that is the agent's framing, **not** the user's assertion — do not extract it as a user fact, and do not extract it as an agent fact either. A user fact must be grounded in a `[user]:` line. If a claim about the user appears only in `[agent]:` content, drop it. -- Unsolicited world-knowledge the agent volunteers ("Python 3.13 was released in October 2024") is the agent's answer to a question, **not** a memory — exclude it. Agent-sourced facts are limited to what the agent *did* or *specifically recommended* for this user, not general knowledge it recited. -- **Episodic memories may use both speakers' content** - the user's stated intent or scope is the anchor (and must be present in `[user]:`), but the agent's content may help fill in the `action_taken` or `outcome` of a `situation → action_taken → outcome` arc when the agent carried out the action on the user's behalf. +- **User-sourced facts (`source: "user"`)** - anything the user asserts about themselves, their preferences, their world, or content they provide/share/ask to remember (including documents, excerpts, and reference material in `[user]:` lines). +- **Agent-sourced facts (`source: "agent"`)** - the agent's **own concrete actions, commitments, and specific recommendations** carried out on the user's behalf: bookings made, orders placed, files/PRs created, appointments scheduled, and named recommendations given (e.g. `[agent]: "I've booked you the 3 PM United flight UA327 to SFO"` → agent-sourced fact; `[agent]: "I recommend Blue Bottle on Mint St for coffee"` → agent-sourced fact). These answer later questions like "what did you book me?" or "what did you recommend?", so they must be captured - tagged `source: "agent"`. +- **NEVER turn an agent line into a fact *about the user*.** When the agent restates, paraphrases, confirms, or assumes something on the user's behalf ("Got it, you don't eat meat", "I assume you want a luxury hotel", "So you're a vegetarian"), that is the agent's framing, **not** the user's assertion - do not extract it as a user fact, and do not extract it as an agent fact either. A user fact must be grounded in a `[user]:` line. If a claim about the user appears only in `[agent]:` content, drop it. +- Unsolicited world-knowledge the agent volunteers ("Python 3.13 was released in October 2024") is the agent's answer to a question, **not** a memory - exclude it. Agent-sourced facts are limited to what the agent *did* or *specifically recommended* for this user, not general knowledge it recited. ## Confidence Scoring -Every extracted memory (in any bucket) must include a `confidence` field in `[0.0, 1.0]` indicating how strongly the conversation supports the claim: +Every extracted fact must include a `confidence` field in `[0.0, 1.0]` indicating how strongly the conversation supports the claim: | Range | Meaning | |-------|---------| @@ -56,7 +54,7 @@ Retrieval may filter by `min_confidence` so under-confident extractions get supp --- -## 1. Facts - Declarative Knowledge +## Facts - Declarative Knowledge Extract concrete, factual statements that fall into these categories: @@ -113,48 +111,17 @@ was said is silently wrong and stays wrong forever. --- -## 2. Episodic - Situated Memories +## Capturing Stories and Events as Facts -Extract memories that are tied to a **specific situation, scope, or context** the user is in. Episodic covers three cases: - -- **Planned / in-flight** - a stated intent or preference scoped to a particular trip, project, event, session, or other bounded context that has not yet finished. -- **Past with outcome** - a completed event the user described, following the `situation → action_taken → outcome` pattern. -- **Ongoing context** - a temporary state of affairs that will not persist as a standing fact (e.g. "right now I'm focused on X"). - -### Required Fields -- **scope_type** - short, free-form noun describing the kind of context (e.g. `trip`, `project`, `event`, `session`, `release`, `campaign`). Pick whatever vocabulary fits the user's domain. Do not invent a value if one is not implied - if no scope is present, the memory probably belongs in facts. -- **scope_value** - the specific instance of that scope (e.g. `Paris 2025`, `Acme revamp`, `Q3 launch`). - -Both must be non-empty. - -### Optional Fields (include only when applicable) -- **situation** - the context or problem faced (present for past/in-flight events) -- **action_taken** - what was specifically tried or done (present for past events) -- **outcome** - what actually happened as a result (present only when the event has concluded) -- **outcome_valence** - `positive` | `negative` | `mixed` | `neutral` (present only with `outcome`) -- **reasoning** - why it worked or failed -- **lesson** - a transferable takeaway -- **domain** - topic area - -For planned/in-flight or ongoing-context memories, leave `situation`, `action_taken`, `outcome`, `outcome_valence`, `reasoning`, and `lesson` as `null`. The scope fields alone carry the meaning. - ---- - -## Classification Guidance - -The same conversation can produce both types. Consider this exchange: +An experience or event the user describes is captured as a fact - keep the whole arc (what happened, what they did, how it turned out) together in one rich fact rather than splitting or dropping it. Consider this exchange: > **User**: "I tried using Redis for session storage last week, but it kept timing out under load. We switched to DynamoDB and it's been stable since. My team lead Sarah approved the switch." -This produces: -- **Fact** (decision): "User's team switched from Redis to DynamoDB for session storage" -- **Fact** (relational): "Sarah is the user's team lead and approved the Redis-to-DynamoDB switch" -- **Episodic**: Situation: Redis session storage timing out under load → Action: Switched to DynamoDB → Outcome: Stable performance → Lesson: DynamoDB handles session storage load better than Redis for this team's use case +This produces facts: +- (decision) "Last week the user's team switched session storage from Redis to DynamoDB because Redis kept timing out under load; it has been stable since." +- (relational) "Sarah is the user's team lead and approved the Redis-to-DynamoDB switch." -When in doubt: -- If it's a **standing state** ("X is true", with no bounded context) → Fact -- If it's a **story** ("We tried X and Y happened") → Episodic -- If it's a **state scoped to a particular context** - a trip, project, event, session, release, campaign, or any other bounded container - → Episodic, with `scope_type` and `scope_value` filled in. The classic trap is a preference that sounds general but is qualified by "for this trip / on this project / just for today" - those are episodic, not fact. +The story becomes one coherent fact that keeps the problem, the change, and the outcome together - not three disconnected fragments. A statement scoped to a specific context (a trip, project, event, or session) is also a fact; keep the scope inside the text (see "Scoped statements ARE facts" below). --- @@ -183,29 +150,23 @@ Do NOT extract: **Extract the content, not the act of sharing.** When the user pastes or references a document, spec, dataset, case, stat block, or any reference material, extract the concrete facts *inside* it - the parties, dates, figures, findings, named items, quantities - each as its own memory. Do NOT produce a memory that merely describes the gesture ("The user shared a contract summary", "The user asked to remember a spec"). A memory of "the user pasted a case" is useless; the case's actual facts are the memory. The same applies to structured data: preserve every count and value ("4 mummies, AC 11, 45 HP"), splitting a detail-dense block into a few focused memories rather than compressing the numbers away. -### Scoped intents are NOT facts -A statement that only applies inside a particular trip, project, event, session, or other bounded context is an **episodic** memory, not a fact. Examples that look like facts but are NOT: -- "For this Paris trip, I want luxury accommodations." → episodic (`scope_type=trip`, `scope_value=Paris`) -- "On the Acme project, we'll use TypeScript strict mode." → episodic (`scope_type=project`, `scope_value=Acme`) -- "Just for today, please give me short answers." → episodic (`scope_type=session`, `scope_value=today`) +### Scoped statements ARE facts - keep the scope in the text +A statement tied to a particular trip, project, event, or session is still a fact - do NOT drop it. Capture it as a fact and keep the scope inside the text so it never reads as an unconditional standing claim: +- "For this Paris trip, I want luxury accommodations." → "For the user's Paris trip, the user wants luxury accommodations." +- "On the Acme project, we'll use TypeScript strict mode." → "On the Acme project, the team uses TypeScript strict mode." +- "Just for today, please give me short answers." → "For today's session, the user asked for short answers." -A fact is a standing claim that holds outside any specific context. The test: if you cannot drop the scope qualifier ("for this trip", "on this project", "just for today") without changing the meaning, it belongs in `episodic`, not `facts`. +Keeping the scope phrase ("for this trip", "on this project", "for today's session") in the text means a context-bound statement never collides with a standing preference on the same subject. --- ## Quality Checks -**For Facts:** 1. Could someone act on or reference this fact without reading the original thread? 2. Is this fact stated explicitly, not inferred? 3. Is each memory scoped to a single coherent topic, with all of that topic's specifics captured together - not shattered into fragments, and not blended with a second topic? 4. Does each fact retain every specific detail (dates, numbers, durations, names, brands) from the source, with nothing generalized away? -**For Episodic:** -1. Did this event actually happen, or is it hypothetical? -2. Is the outcome clearly stated, not assumed? -3. Is the lesson genuinely transferable, or too specific to this one situation? - **Before finalizing - scan the whole transcript.** Re-read the entire conversation, including the **middle and end**, not only the opening. A common failure is capturing the first prominent topic thoroughly and skimming the rest ("first-topic dominance"). If a single message raised several @@ -244,46 +205,31 @@ or extract filler; the precision and grounding rules above still apply. "temporal_context": "end of Q2", "tags": ["topic:project", "topic:ETL"] } - ], - "episodic": [] + ] } ``` Name and role are one coherent biographical topic → a single rich memory, not two fragments. The ETL project is a distinct topic → its own memory. -### Example 2: Troubleshooting experience +### Example 2: Troubleshooting experience (a story becomes one rich fact) **Conversation:** > User: "Last month we had a production outage because our Kubernetes pods kept OOM-killing. We bumped memory limits from 512MB to 1GB and added resource quotas per namespace. That fixed it." +Capture the whole arc - the problem, the change, and the outcome - in a single fact. Do NOT shatter it into disconnected "uses 1GB limits" plus "had an outage" fragments. + **Output:** ```json { "facts": [ { - "text": "The user's team increased Kubernetes pod memory limits from 512MB to 1GB after OOM-killing issues.", + "text": "Last month the user's team resolved a production outage caused by Kubernetes pods OOM-killing by raising pod memory limits from 512MB to 1GB and adding per-namespace resource quotas; that fixed it.", "category": "other", "source": "user", "confidence": 0.95, "salience": 0.7, "temporal_context": "last month", - "tags": ["topic:kubernetes", "topic:infrastructure"] - } - ], - "episodic": [ - { - "scope_type": "incident", - "scope_value": "Q3 K8s OOM outage", - "situation": "Kubernetes pods in production were repeatedly OOM-killed, causing an outage.", - "action_taken": "Bumped pod memory limits from 512MB to 1GB and added resource quotas per namespace.", - "outcome": "The OOM-killing stopped and production stabilized.", - "reasoning": "Pods were exceeding the 512MB limit under normal load; quotas prevented uncontrolled resource consumption across namespaces.", - "outcome_valence": "positive", - "lesson": "Always set resource quotas per namespace and right-size memory limits to prevent OOM-related outages in Kubernetes.", - "domain": "infrastructure", - "confidence": 0.9, - "salience": 0.8, - "tags": ["topic:kubernetes", "topic:outage", "topic:resource_management"] + "tags": ["topic:kubernetes", "topic:infrastructure", "topic:outage"] } ] } @@ -316,18 +262,17 @@ Name and role are one coherent biographical topic → a single rich memory, not "temporal_context": "current quarter", "tags": ["topic:budget", "topic:marketing"] } - ], - "episodic": [] + ] } ``` -### Example 4: Standing preference vs. scoped intent +### Example 4: Standing preference vs. scoped preference (both are facts) **Conversation:** > User: "I usually prefer budget hotels." > User: "For this Paris trip, I want luxury accommodations." -The first statement is a standing preference and belongs in `facts`. The second is qualified by "for this Paris trip" - dropping the scope changes the meaning, so it belongs in `episodic` with the scope captured structurally. The two memories coexist; the trip-scoped intent never enters the fact pool and never collides with the standing preference. +Both are facts. The first is a standing preference. The second is scoped to the Paris trip - keep that scope inside the text so it never reads as a standing claim and never collides with the general preference on retrieval. **Output:** ```json @@ -341,21 +286,14 @@ The first statement is a standing preference and belongs in `facts`. The second "salience": 0.7, "temporal_context": null, "tags": ["topic:travel", "topic:hotels"] - } - ], - "episodic": [ + }, { - "scope_type": "trip", - "scope_value": "Paris", - "situation": null, - "action_taken": null, - "outcome": null, - "outcome_valence": null, - "reasoning": null, - "lesson": null, - "domain": "travel", + "text": "For the user's Paris trip, the user wants luxury accommodations.", + "category": "preference", + "source": "user", "confidence": 0.95, - "salience": 0.8, + "salience": 0.7, + "temporal_context": null, "tags": ["topic:travel", "topic:hotels", "topic:luxury"] } ] @@ -382,8 +320,7 @@ Extract only the literal, user-stated fact. The load time is concrete and worth "temporal_context": null, "tags": ["topic:performance", "topic:dashboard"] } - ], - "episodic": [] + ] } ``` @@ -407,8 +344,7 @@ The change is one coherent memory: it records the new state, what it replaced, a "temporal_context": "last quarter", "tags": ["topic:CI", "topic:tooling"] } - ], - "episodic": [] + ] } ``` @@ -451,8 +387,7 @@ The user's request is an intent; the concrete bookings and the named recommendat "temporal_context": null, "tags": ["topic:travel", "topic:recommendation"] } - ], - "episodic": [] + ] } ``` @@ -462,25 +397,19 @@ The user's request is an intent; the concrete bookings and the named recommendat > User: "I need a hotel in Rome for three nights." > Agent: "Got it - I assume you want a luxury 5-star place near the center, and since you mentioned Rome you're probably interested in historical sites. I'll put together some options." -The only user-stated fact is the request itself (a scoped intent → episodic). The agent's "I assume you want luxury" and "you're probably interested in historical sites" are the agent's **assumptions about the user**, not the user's assertions - do NOT extract them as user facts, and they are not agent actions either, so extract nothing from the agent line. No booking was made, so there is no agent-sourced fact. +The user-stated need is a fact, scoped to the Rome trip - capture it with the scope in the text. The agent's "I assume you want luxury" and "you're probably interested in historical sites" are the agent's **assumptions about the user**, not the user's assertions - do NOT extract them as user facts, and they are not agent actions either, so extract nothing from the agent line. No booking was made, so there is no agent-sourced fact. **Output:** ```json { - "facts": [], - "episodic": [ + "facts": [ { - "scope_type": "trip", - "scope_value": "Rome", - "situation": null, - "action_taken": null, - "outcome": null, - "outcome_valence": null, - "reasoning": null, - "lesson": null, - "domain": "travel", + "text": "For the user's Rome trip, the user needs a hotel for three nights.", + "category": "requirement", + "source": "user", "confidence": 0.9, "salience": 0.6, + "temporal_context": null, "tags": ["topic:travel", "topic:hotels"] } ] @@ -516,8 +445,7 @@ Extract the actual facts *inside* the shared material - dates, parties, findings "temporal_context": "December 2016", "tags": ["topic:legal", "topic:construction"] } - ], - "episodic": [] + ] } ``` @@ -559,8 +487,7 @@ Three unrelated topics in one message. Extract each separately, and use the actu "temporal_context": null, "tags": ["topic:family"] } - ], - "episodic": [] + ] } ``` @@ -568,7 +495,7 @@ Three unrelated topics in one message. Extract each separately, and use the actu ## Output Format -You must output ONLY valid JSON matching the schema below. No preamble, no explanation, no closing remarks - just the JSON object. Each array can be empty if no memories of that type are found. +You must output ONLY valid JSON matching the schema below. No preamble, no explanation, no closing remarks - just the JSON object. The `facts` array can be empty if no facts are found. ```json { @@ -582,22 +509,6 @@ You must output ONLY valid JSON matching the schema below. No preamble, no expla "temporal_context": "the exact time expression stated (e.g. 'June 3rd', '3 weeks ago', 'by Friday', 'every Monday') or null if none", "tags": ["topic:x"] } - ], - "episodic": [ - { - "scope_type": "trip|project|event|session|release|campaign|... (free-form, required, non-empty)", - "scope_value": "specific instance, e.g. Paris 2025 (required, non-empty)", - "situation": "context/problem, or null", - "action_taken": "what was tried, or null", - "outcome": "what happened, or null", - "outcome_valence": "positive|negative|mixed|neutral, or null", - "reasoning": "why it worked/failed, or null", - "lesson": "transferable takeaway, or null", - "domain": "topic area or null", - "confidence": 0.8, - "salience": 0.7, - "tags": ["topic:x"] - } ] } ``` diff --git a/azure/cosmos/agent_memory/prompts/summarize.prompty b/azure/cosmos/agent_memory/prompts/summarize.prompty index 4d1b855..31fdc03 100644 --- a/azure/cosmos/agent_memory/prompts/summarize.prompty +++ b/azure/cosmos/agent_memory/prompts/summarize.prompty @@ -17,7 +17,7 @@ inputs: --- system: -You are an expert summarization system. Your task is to read a conversation thread and produce a structured, reliable summary in JSON format that gives a future reader — whether human or AI — a complete and accurate understanding of what was discussed, concluded, and left unresolved. +You are an expert summarization system. Your task is to read a conversation thread and produce a structured, reliable summary in JSON format that gives a future reader - whether human or AI - a complete and accurate understanding of what was discussed, concluded, and left unresolved. ## Input Format You will receive a transcript of the conversation thread. Each line follows this format: @@ -31,38 +31,38 @@ You will receive a transcript of the conversation thread. Each line follows this - Messages appear in chronological order ## Your Goal -Produce a summary that fully replaces the need to re-read the original thread. A reader should be able to understand what happened, what was decided, and what comes next — purely from your summary. +Produce a summary that fully replaces the need to re-read the original thread. A reader should be able to understand what happened, what was decided, and what comes next - purely from your summary. ## What to Include Your summary must cover all of the following that are present in the thread: -- **Main subject** — What is the conversation fundamentally about? State this immediately in the overview. -- **Key points raised** — The substantive ideas, questions, problems, or information exchanged. Focus on content that shaped the conversation. -- **Decisions made** — Any conclusions reached, options selected, or agreements confirmed. -- **Open issues** — Questions left unanswered, disagreements unresolved, or topics flagged for later. -- **Action items** — Tasks committed to, next steps agreed upon, or follow-ups promised (include owner and deadline if stated). -- **Important context** — Background details necessary to make the summary intelligible on its own (e.g., who the parties are, what project this relates to). +- **Main subject** - What is the conversation fundamentally about? State this immediately in the overview. +- **Key points raised** - The substantive ideas, questions, problems, or information exchanged. Focus on content that shaped the conversation. +- **Decisions made** - Any conclusions reached, options selected, or agreements confirmed. +- **Open issues** - Questions left unanswered, disagreements unresolved, or topics flagged for later. +- **Action items** - Tasks committed to, next steps agreed upon, or follow-ups promised (include owner and deadline if stated). +- **Important context** - Background details necessary to make the summary intelligible on its own (e.g., who the parties are, what project this relates to). ## What to Exclude - Greetings, pleasantries, and filler ("Thanks!", "Sounds good", "Let me know") -- Repetition — if a point is made multiple times, mention it once +- Repetition - if a point is made multiple times, mention it once - Speculation or hypotheticals, unless they were central to the discussion - Tangents that did not influence the outcome or decisions - Verbatim quotes, unless a specific phrasing is critically important ## Tone and Style -- **Factual and neutral** — Do not editorialize, interpret intent, or add opinions -- **Third person** — Refer to participants by name or role (e.g., "the user", "the manager", "Sarah"), not as "you" or "I" -- **Past tense** — The conversation has already happened -- **Precise over vague** — Prefer "the deadline was set to April 15th" over "a deadline was mentioned" -- **Concise but complete** — Do not pad, but do not omit material details to hit an arbitrary length target +- **Factual and neutral** - Do not editorialize, interpret intent, or add opinions +- **Third person** - Refer to participants by name or role (e.g., "the user", "the manager", "Sarah"), not as "you" or "I" +- **Past tense** - The conversation has already happened +- **Precise over vague** - Prefer "the deadline was set to April 15th" over "a deadline was mentioned" +- **Concise but complete** - Do not pad, but do not omit material details to hit an arbitrary length target ## Output Format -You must output ONLY valid JSON matching the schema below. No preamble, no explanation, no closing remarks — just the JSON object. +You must output ONLY valid JSON matching the schema below. No preamble, no explanation, no closing remarks - just the JSON object. -Include only fields that have relevant content. Use empty arrays `[]` for sections with nothing to report — do not omit the field entirely. +Include only fields that have relevant content. Use empty arrays `[]` for sections with nothing to report - do not omit the field entirely. -The output will be stored as a single document and embedded as a vector for semantic search. Keep language natural and semantically rich — the summary should retrieve well when someone searches for the topics discussed. +The output will be stored as a single document and embedded as a vector for semantic search. Keep language natural and semantically rich - the summary should retrieve well when someone searches for the topics discussed. ```json { @@ -78,19 +78,19 @@ The output will be stored as a single document and embedded as a vector for sema ``` ### Field Descriptions -- **overview** — A 1–3 sentence standalone snapshot of the thread. A reader should understand the gist from this field alone. -- **key_points** — Substantive ideas, questions, problems, or information exchanged that shaped the conversation. -- **decisions** — Conclusions reached, options selected, or agreements confirmed. -- **open_issues** — Questions left unanswered, disagreements unresolved, or topics flagged for later. -- **action_items** — Tasks committed to, with owner, task description, and deadline (null if no deadline was stated). -- **topics** — Short lowercase topic labels for categorization and retrieval (e.g., "travel", "hotel booking", "kubernetes", "budget"). +- **overview** - A 1–3 sentence standalone snapshot of the thread. A reader should understand the gist from this field alone. +- **key_points** - Substantive ideas, questions, problems, or information exchanged that shaped the conversation. +- **decisions** - Conclusions reached, options selected, or agreements confirmed. +- **open_issues** - Questions left unanswered, disagreements unresolved, or topics flagged for later. +- **action_items** - Tasks committed to, with owner, task description, and deadline (null if no deadline was stated). +- **topics** - Short lowercase topic labels for categorization and retrieval (e.g., "travel", "hotel booking", "kubernetes", "budget"). ## Quality Check (apply before outputting) Before writing your final output, verify: 1. Does the overview field alone give a useful, standalone snapshot of the thread? 2. Have you omitted all filler, pleasantries, and repetition? -3. Are all decisions, open issues, and action items captured — not just the main topic? -4. Is every statement attributable to something actually said in the thread — no inferences or additions? +3. Are all decisions, open issues, and action items captured - not just the main topic? +4. Is every statement attributable to something actually said in the thread - no inferences or additions? 5. Could someone who has never seen the thread act on this summary correctly? 6. Is the output valid JSON? diff --git a/azure/cosmos/agent_memory/prompts/summarize_update.prompty b/azure/cosmos/agent_memory/prompts/summarize_update.prompty index 67ba734..da191df 100644 --- a/azure/cosmos/agent_memory/prompts/summarize_update.prompty +++ b/azure/cosmos/agent_memory/prompts/summarize_update.prompty @@ -21,23 +21,23 @@ system: You are an expert summarization system operating in update mode. You will be given an existing structured JSON summary of a conversation thread, followed by new messages from the same thread. Your task is to produce an updated JSON summary that seamlessly integrates the new information while preserving everything still valid from the original. ## Your Goal -Produce a single, authoritative JSON summary that reflects the full state of the thread as of the new messages — accurate, complete, and requiring no cross-reference with either the old summary or the new messages to understand. +Produce a single, authoritative JSON summary that reflects the full state of the thread as of the new messages - accurate, complete, and requiring no cross-reference with either the old summary or the new messages to understand. ## Inputs You Will Receive -- **Existing Summary** — A structured JSON object with the fields: `overview`, `key_points`, `decisions`, `open_issues`, `action_items`, and `topics`. This represents the conversation state before the new messages. -- **New Messages** — The latest messages added to the thread since the existing summary was written. +- **Existing Summary** - A structured JSON object with the fields: `overview`, `key_points`, `decisions`, `open_issues`, `action_items`, and `topics`. This represents the conversation state before the new messages. +- **New Messages** - The latest messages added to the thread since the existing summary was written. ## How to Handle Each Field **overview** - Rewrite to reflect the current overall state of the thread, incorporating any new direction, resolution, or development. -- Do not simply append new content — synthesize old and new into a single coherent overview. +- Do not simply append new content - synthesize old and new into a single coherent overview. **key_points** - Retain unique prior key points that remain relevant and have not been superseded; **merge near-duplicates into a single entry** rather than keeping both. - Add new key points introduced in the new messages. - Remove or rewrite any point that the new messages have contradicted, corrected, or made obsolete. -- **Self-cap at ~20-30 entries.** When the list grows beyond that, consolidate related points and drop the least informative — the summary should grow in fidelity, not in length. +- **Self-cap at ~20-30 entries.** When the list grows beyond that, consolidate related points and drop the least informative - the summary should grow in fidelity, not in length. **decisions** - Retain all prior decisions unless explicitly reversed or superseded by the new messages. @@ -60,28 +60,28 @@ Produce a single, authoritative JSON summary that reflects the full state of the - Remove topic labels only if the topic was explicitly determined to be irrelevant. ## Handling Conflicts and Corrections -- If the new messages contradict something in the existing summary, **always trust the new messages** — they represent the more current state. +- If the new messages contradict something in the existing summary, **always trust the new messages** - they represent the more current state. - If something was stated incorrectly in the existing summary (e.g., a wrong date or name), correct it silently without calling attention to the error. -- If a prior decision is reversed, do not preserve the old version — replace it entirely with the new one, unless the reversal itself is significant context worth noting. +- If a prior decision is reversed, do not preserve the old version - replace it entirely with the new one, unless the reversal itself is significant context worth noting. ## What to Exclude -- Do not include meta-commentary about what changed (e.g., "The previous summary said X, but now...") — just output the updated summary. +- Do not include meta-commentary about what changed (e.g., "The previous summary said X, but now...") - just output the updated summary. - Omit greetings, pleasantries, filler, and repetition from the new messages, exactly as you would in a fresh summary. - Do not include any content from the new messages that is speculative, hypothetical, or tangential unless it materially affects the thread's direction. ## Tone and Style -- **Factual and neutral** — no editorializing or interpretation of intent -- **Third person** — refer to participants by name or role, never as "you" or "I" -- **Past tense** — the conversation has already happened -- **Precise over vague** — specific dates, names, numbers, and decisions wherever stated -- **Concise but complete** — do not omit material content, but do not pad +- **Factual and neutral** - no editorializing or interpretation of intent +- **Third person** - refer to participants by name or role, never as "you" or "I" +- **Past tense** - the conversation has already happened +- **Precise over vague** - specific dates, names, numbers, and decisions wherever stated +- **Concise but complete** - do not omit material content, but do not pad ## Output Format -You must output ONLY valid JSON matching the schema below. No preamble, no explanation, no closing remarks — just the JSON object. +You must output ONLY valid JSON matching the schema below. No preamble, no explanation, no closing remarks - just the JSON object. -Include only fields that have relevant content. Use empty arrays `[]` for sections with nothing to report — do not omit the field entirely. +Include only fields that have relevant content. Use empty arrays `[]` for sections with nothing to report - do not omit the field entirely. -The output will be stored as a single document and embedded as a vector for semantic search. Keep language natural and semantically rich — the summary should retrieve well when someone searches for the topics discussed. +The output will be stored as a single document and embedded as a vector for semantic search. Keep language natural and semantically rich - the summary should retrieve well when someone searches for the topics discussed. ```json { @@ -98,9 +98,9 @@ The output will be stored as a single document and embedded as a vector for sema ## Quality Check (apply before outputting) Before writing your final output, verify: -1. Does every item from the existing summary still belong — is it still valid, not contradicted, not resolved? +1. Does every item from the existing summary still belong - is it still valid, not contradicted, not resolved? 2. Is every material development from the new messages reflected somewhere in the output? -3. Have all resolved open issues been moved or removed — none left as "open" if the new messages closed them? +3. Have all resolved open issues been moved or removed - none left as "open" if the new messages closed them? 4. Have all completed or cancelled action items been updated or removed? 5. Could someone who has never seen the thread, the old summary, or the new messages fully understand the current state from this output alone? 6. Have you avoided all meta-commentary about the update process itself? diff --git a/azure/cosmos/agent_memory/prompts/synthesize_procedural.prompty b/azure/cosmos/agent_memory/prompts/synthesize_procedural.prompty index aa7754c..51948f6 100644 --- a/azure/cosmos/agent_memory/prompts/synthesize_procedural.prompty +++ b/azure/cosmos/agent_memory/prompts/synthesize_procedural.prompty @@ -34,10 +34,10 @@ Produce the next version of the personalized system prompt by refining prior sta ## Inputs You will receive these inputs: -- **prior_prompt** — the previous version of this user's personalized system prompt. May be empty on first synthesis. -- **behavioral_facts** — a list of declarative preferences and constraints the user has stated, each as a self-contained imperative sentence (for example, "Always include type hints in Python code examples" or "Never suggest using var in TypeScript"). -- **episodic_lessons** — a list of transferable takeaways from past interactions (for example, "Resource quotas per namespace prevented OOM-related outages — recommend them when discussing Kubernetes deployments"). -- **user_name** — the user's preferred name, when known. +- **prior_prompt** - the previous version of this user's personalized system prompt. May be empty on first synthesis. +- **behavioral_facts** - a list of declarative preferences and constraints the user has stated, each as a self-contained imperative sentence (for example, "Always include type hints in Python code examples" or "Never suggest using var in TypeScript"). +- **episodic_lessons** - a list of transferable takeaways from past interactions (for example, "Resource quotas per namespace prevented OOM-related outages - recommend them when discussing Kubernetes deployments"). +- **user_name** - the user's preferred name, when known. Every instruction in the output must be explicitly grounded in these inputs. @@ -52,7 +52,7 @@ You must output ONLY a single valid JSON object with this exact shape: } ``` -No preamble, no explanation, no closing remarks — just the JSON object. +No preamble, no explanation, no closing remarks - just the JSON object. --- @@ -64,7 +64,7 @@ No preamble, no explanation, no closing remarks — just the JSON object. 5. If a new behavioral fact directly conflicts with something in `prior_prompt`, the new fact wins because it represents the user's most current preference. 6. Surface 1-3 of the strongest episodic lessons as concrete guidance, phrased as "When [trigger], remember [lesson]" or "In the past, [action] led to [outcome]; prefer [recommendation]". 7. Do not invent preferences, lessons, or instructions that are not grounded in the inputs. -8. Keep the result concise — typically 8-20 lines of imperative text. The output is prepended to every LLM call for this user, so brevity matters. +8. Keep the result concise - typically 8-20 lines of imperative text. The output is prepended to every LLM call for this user, so brevity matters. 9. If `user_name` is provided and is not the default placeholder, reference the user by name once at the top. 10. Write the final prompt as self-contained imperative guidance, ready to prepend to future agent calls. diff --git a/azure/cosmos/agent_memory/prompts/user_summary.prompty b/azure/cosmos/agent_memory/prompts/user_summary.prompty index d80b311..476b18e 100644 --- a/azure/cosmos/agent_memory/prompts/user_summary.prompty +++ b/azure/cosmos/agent_memory/prompts/user_summary.prompty @@ -16,23 +16,23 @@ inputs: --- system: -You are an expert user intelligence system. You will be given conversation data from multiple threads involving a single user. Your task is to synthesize this data into a structured, accurate, and actionable JSON user profile that captures everything persistently useful about this person — so that any future AI assistant or human reviewer can immediately understand who this user is, how they work, and what they need. +You are an expert user intelligence system. You will be given conversation data from multiple threads involving a single user. Your task is to synthesize this data into a structured, accurate, and actionable JSON user profile that captures everything persistently useful about this person - so that any future AI assistant or human reviewer can immediately understand who this user is, how they work, and what they need. ## Your Goal -Produce a profile that serves as a reliable, living reference document. It should be specific enough to meaningfully personalize future interactions, and disciplined enough that every entry can be traced back to something the user actually said or did — never inferred or assumed. +Produce a profile that serves as a reliable, living reference document. It should be specific enough to meaningfully personalize future interactions, and disciplined enough that every entry can be traced back to something the user actually said or did - never inferred or assumed. ## Input You Will Receive One or more conversation threads involving the same user. Threads may vary in topic, recency, and length. You should treat all threads as equally valid sources unless they contradict each other, in which case prefer the most recent data. ## Profile Sections -Each field in the JSON output corresponds to a specific category of user knowledge. Populate only fields for which you have relevant data — use empty arrays `[]` for sections with no entries rather than omitting the field. +Each field in the JSON output corresponds to a specific category of user knowledge. Populate only fields for which you have relevant data - use empty arrays `[]` for sections with no entries rather than omitting the field. The profile will be stored as a single document and embedded as a vector for semantic search. Keep content semantically rich so it retrieves well when searching for the user's interests, preferences, or context. ### Section Descriptions -**key_facts** — Concrete, identifying information about the user: +**key_facts** - Concrete, identifying information about the user: - Full name, preferred name, or username (if stated) - Role, title, or function - Organization, company, or team @@ -40,7 +40,7 @@ The profile will be stored as a single document and embedded as a vector for sem - Technical environment (OS, tools, languages, platforms regularly used) - Any other stable identifying details explicitly mentioned -**personal_preferences** — How the user likes to communicate and work: +**personal_preferences** - How the user likes to communicate and work: - Communication style (e.g., direct, detail-oriented, prefers examples) - Preferred response format (e.g., bullet points, prose, code blocks, step-by-step) - Preferred language or terminology @@ -48,48 +48,48 @@ The profile will be stored as a single document and embedded as a vector for sem - Things the user has explicitly said they dislike or want avoided - Accessibility needs or display preferences, if stated -**account_environment** — Technical and account-level context relevant to future interactions: +**account_environment** - Technical and account-level context relevant to future interactions: - Subscription tier, plan, or licensing details - Active features, integrations, or tools in use - Known limitations, restrictions, or feature flags - Usage patterns (e.g., heavy API user, primarily uses the web UI) - Any account issues, past errors, or open support items -**goals_current_work** — What the user is trying to accomplish, near-term and longer-term: +**goals_current_work** - What the user is trying to accomplish, near-term and longer-term: - Active projects or initiatives mentioned across threads - Stated objectives or success criteria - Known constraints (budget, timeline, team size, technical debt) - Problems the user is actively trying to solve -**behavioral_patterns** — Observable tendencies derived from how the user behaves across threads: +**behavioral_patterns** - Observable tendencies derived from how the user behaves across threads: - Recurring question types or topics they return to repeatedly - Common workflows or task sequences - Frequent friction points or things that consistently confuse or frustrate them - How they typically approach problems (e.g., asks for options first, prefers to try before asking) - Patterns in how they give feedback or express satisfaction/dissatisfaction -**compliance_requirements** — Constraints or obligations the user has mentioned that must be respected: +**compliance_requirements** - Constraints or obligations the user has mentioned that must be respected: - Regulatory or legal requirements (e.g., HIPAA, GDPR, SOC 2) - Data handling restrictions (e.g., no PII in logs, no third-party data sharing) - Organizational policies or approval processes - Accessibility requirements - Any stated hard limits on what solutions are acceptable -**open_items** — Things that were raised but not yet resolved, and may need follow-up: +**open_items** - Things that were raised but not yet resolved, and may need follow-up: - Questions the user asked that were not fully answered - Issues or bugs reported but not confirmed resolved - Commitments made to the user (by an assistant or support agent) that should be honored - Topics the user said they would return to -**topics** — Short lowercase topic labels for categorization and retrieval (e.g., "programming", "python", "data science", "kubernetes", "travel") +**topics** - Short lowercase topic labels for categorization and retrieval (e.g., "programming", "python", "data science", "kubernetes", "travel") ## Formatting Rules for Array Entries -- Use concise, self-contained strings — one fact per entry +- Use concise, self-contained strings - one fact per entry - Each entry must be intelligible without reading the threads - Write in third person (e.g., "The user prefers...", "The user is working on...") - Be specific: prefer "The user works in Python 3.11 on macOS" over "The user codes" - Where recency matters, note it: "As of [approximate date or thread], the user was..." -- Do not use vague qualifiers like "seems to" or "might" — if you're not certain, omit it +- Do not use vague qualifiers like "seems to" or "might" - if you're not certain, omit it ## Source Conflicts - If two threads contradict each other, prefer the more recent thread @@ -101,7 +101,7 @@ The profile will be stored as a single document and embedded as a vector for sem - One-off remarks that do not reflect a persistent pattern or stable fact - Pleasantries, filler, and conversational noise - Anything the user said hypothetically or about someone else -- Duplicate facts — if the same detail appears in multiple threads, list it once +- Duplicate facts - if the same detail appears in multiple threads, list it once ## Quality Check (apply before outputting) Before writing your final output, verify: @@ -109,12 +109,12 @@ Before writing your final output, verify: 2. Have you populated all sections with relevant data and used empty arrays for sections with none? 3. Is every entry written so it makes sense without the source threads? 4. Have you flagged conflicting data rather than silently resolving it? -5. Are there any vague entries that could be made more specific — and if so, have you done that? +5. Are there any vague entries that could be made more specific - and if so, have you done that? 6. Is the goals_current_work section populated if the user mentioned any active project, even briefly? 7. Is the output valid JSON? ## Output Format -You must output ONLY valid JSON matching the schema below. No preamble, no explanation, no closing remarks — just the JSON object. +You must output ONLY valid JSON matching the schema below. No preamble, no explanation, no closing remarks - just the JSON object. ```json { @@ -158,7 +158,7 @@ You must output ONLY valid JSON matching the schema below. No preamble, no expla "The user frequently returns to questions about data pipeline optimization." ], "compliance_requirements": [ - "All data processing must comply with GDPR — no PII in logs.", + "All data processing must comply with GDPR - no PII in logs.", "The user's organization requires SOC 2 compliance for all third-party tools." ], "open_items": [ diff --git a/azure/cosmos/agent_memory/prompts/user_summary_update.prompty b/azure/cosmos/agent_memory/prompts/user_summary_update.prompty index a4b96f5..8bc08f1 100644 --- a/azure/cosmos/agent_memory/prompts/user_summary_update.prompty +++ b/azure/cosmos/agent_memory/prompts/user_summary_update.prompty @@ -18,21 +18,21 @@ inputs: --- system: -You are an expert user intelligence system operating in update mode. You will be given an existing structured JSON user profile and new conversation data from recent threads involving the same user. Your task is to produce a single, authoritative updated JSON profile that reflects everything known about this user as of the new conversations — accurate, complete, and requiring no cross-reference with either the old profile or the new threads to understand. +You are an expert user intelligence system operating in update mode. You will be given an existing structured JSON user profile and new conversation data from recent threads involving the same user. Your task is to produce a single, authoritative updated JSON profile that reflects everything known about this user as of the new conversations - accurate, complete, and requiring no cross-reference with either the old profile or the new threads to understand. ## Your Goal -The output should be indistinguishable from a freshly built profile — not a patched or annotated version of the old one. Every field should read as a clean, current snapshot, not a changelog. +The output should be indistinguishable from a freshly built profile - not a patched or annotated version of the old one. Every field should read as a clean, current snapshot, not a changelog. ## Inputs You Will Receive -- **Existing Profile** — A structured JSON user profile built from prior conversation threads, with fields: `key_facts`, `personal_preferences`, `account_environment`, `goals_current_work`, `behavioral_patterns`, `compliance_requirements`, `open_items`, and `topics`. -- **New Conversation Data** — One or more recent threads involving the same user, not yet reflected in the existing profile. +- **Existing Profile** - A structured JSON user profile built from prior conversation threads, with fields: `key_facts`, `personal_preferences`, `account_environment`, `goals_current_work`, `behavioral_patterns`, `compliance_requirements`, `open_items`, and `topics`. +- **New Conversation Data** - One or more recent threads involving the same user, not yet reflected in the existing profile. ## General Update Principles -- **New data wins** — If new conversations contradict the existing profile, always trust the newer data. The existing profile reflects a past state; the new threads reflect the current one. -- **Preserve & consolidate** — Retain unique entries that the new conversations do not contradict, supersede, or invalidate. **Merge near-duplicates into a single entry** rather than keeping both. When a section already contains many entries, **prefer consolidation to accumulation** — combine related items, drop minor or stale entries not surfaced in recent threads, and keep the section focused on the most informative facts. The profile should grow in fidelity, not in length. -- **Target length per section** — Aim for roughly 20-40 entries per array section. If a section already exceeds that, treat it as overdue for consolidation: merge similar items and drop the least informative. Hard cap is enforced downstream, but you should self-cap by quality. -- **Correct silently** — Fix outdated or incorrect entries without meta-commentary. Do not write "previously X, now Y" unless the transition itself is useful context. -- **No empty sections** — Use empty arrays `[]` for sections with no relevant data. Do not omit fields from the JSON. +- **New data wins** - If new conversations contradict the existing profile, always trust the newer data. The existing profile reflects a past state; the new threads reflect the current one. +- **Preserve & consolidate** - Retain unique entries that the new conversations do not contradict, supersede, or invalidate. **Merge near-duplicates into a single entry** rather than keeping both. When a section already contains many entries, **prefer consolidation to accumulation** - combine related items, drop minor or stale entries not surfaced in recent threads, and keep the section focused on the most informative facts. The profile should grow in fidelity, not in length. +- **Target length per section** - Aim for roughly 20-40 entries per array section. If a section already exceeds that, treat it as overdue for consolidation: merge similar items and drop the least informative. Hard cap is enforced downstream, but you should self-cap by quality. +- **Correct silently** - Fix outdated or incorrect entries without meta-commentary. Do not write "previously X, now Y" unless the transition itself is useful context. +- **No empty sections** - Use empty arrays `[]` for sections with no relevant data. Do not omit fields from the JSON. ## How to Handle Each Field @@ -45,7 +45,7 @@ The output should be indistinguishable from a freshly built profile — not a pa ### personal_preferences - Retain all prior preferences not contradicted or withdrawn. - Add newly expressed preferences, format requests, or stated dislikes. -- If the user has reversed a prior preference (e.g., now prefers prose over bullet points), replace the old entry — do not keep both. +- If the user has reversed a prior preference (e.g., now prefers prose over bullet points), replace the old entry - do not keep both. - Note shift in tone or communication style if consistently different across the new threads. ### account_environment @@ -55,15 +55,15 @@ The output should be indistinguishable from a freshly built profile — not a pa - Add any new known limitations, errors, or open support items. ### goals_current_work -- Remove or archive completed projects — do not retain goals the user has explicitly finished or abandoned. +- Remove or archive completed projects - do not retain goals the user has explicitly finished or abandoned. - Update scope, timeline, or constraints if revised in the new threads. - Add newly mentioned projects, initiatives, or problems being actively worked on. -- If a goal from the existing profile is not mentioned in the new threads, retain it — absence of mention is not evidence of completion. +- If a goal from the existing profile is not mentioned in the new threads, retain it - absence of mention is not evidence of completion. ### behavioral_patterns - Retain established patterns not contradicted by new data. - Strengthen a pattern if the new threads provide additional confirming instances. -- If a new thread shows behavior inconsistent with an established pattern, note the exception only if it appears more than once — a single deviation is not sufficient to revise a pattern. +- If a new thread shows behavior inconsistent with an established pattern, note the exception only if it appears more than once - a single deviation is not sufficient to revise a pattern. - Add new patterns only if they appear at least twice across the combined thread history, not from a single instance. ### compliance_requirements @@ -85,7 +85,7 @@ The output should be indistinguishable from a freshly built profile — not a pa ## Handling Conflicts Between Old Profile and New Data - If new threads directly contradict the existing profile, replace the old entry with the new one. - If the conflict involves something time-sensitive (e.g., a project deadline or account tier), always use the new data without preserving the old. -- If the conflict is ambiguous — the new thread is unclear, not the user's own words, or possibly a one-off — retain the existing entry and add a note: "Recent thread suggests this may have changed; not yet confirmed." +- If the conflict is ambiguous - the new thread is unclear, not the user's own words, or possibly a one-off - retain the existing entry and add a note: "Recent thread suggests this may have changed; not yet confirmed." - Never silently blend conflicting data into a false consensus. ## What to Exclude @@ -93,30 +93,30 @@ The output should be indistinguishable from a freshly built profile — not a pa - Speculation, inference, or extrapolation beyond what is explicitly stated or observed - One-off remarks that do not reflect a stable fact or persistent pattern - Pleasantries, filler, and conversational noise from the new threads -- Duplicate facts — if a detail is confirmed again in the new threads, do not list it twice +- Duplicate facts - if a detail is confirmed again in the new threads, do not list it twice ## Formatting Rules for Array Entries -- Use concise, self-contained strings — one fact per entry +- Use concise, self-contained strings - one fact per entry - Each entry must be intelligible without reading either the old profile or the new threads - Write in third person ("The user prefers...", "The user is currently working on...") - Be specific: prefer "The user upgraded to the Team plan in March" over "The user's plan changed" - Where recency matters, annotate it: "As of [approximate date or thread], the user..." -- Do not use vague qualifiers like "seems to" or "appears to" — if uncertain, omit +- Do not use vague qualifiers like "seems to" or "appears to" - if uncertain, omit ## Quality Check (apply before outputting) Before writing your final output, verify: -1. Does every retained entry still reflect the current known state — not a past state superseded by new data? +1. Does every retained entry still reflect the current known state - not a past state superseded by new data? 2. Is every new entry traceable to something explicitly stated or demonstrated in the new threads? 3. Have all resolved open items been removed or closed? 4. Have all completed or abandoned goals been removed? 5. Have you avoided all meta-commentary about the update process? 6. Are conflicts flagged with a note rather than silently resolved into a false consensus? -7. Does the final profile read as a clean, standalone document — not an annotated or patched version of the old one? +7. Does the final profile read as a clean, standalone document - not an annotated or patched version of the old one? 8. Have you applied the two-instance rule before adding any new behavioral pattern? 9. Is the output valid JSON? ## Output Format -You must output ONLY valid JSON matching the schema below. No preamble, no explanation, no closing remarks — just the JSON object. +You must output ONLY valid JSON matching the schema below. No preamble, no explanation, no closing remarks - just the JSON object. ```json { diff --git a/azure/cosmos/agent_memory/services/_pipeline_helpers.py b/azure/cosmos/agent_memory/services/_pipeline_helpers.py index 3f9953a..e2f9ef3 100644 --- a/azure/cosmos/agent_memory/services/_pipeline_helpers.py +++ b/azure/cosmos/agent_memory/services/_pipeline_helpers.py @@ -9,19 +9,179 @@ from __future__ import annotations +import hashlib import json +import math import os import re from collections import defaultdict +from datetime import datetime, timezone from pathlib import Path from typing import Any, Iterable, Mapping, Optional from azure.cosmos.agent_memory._embedding_tokens import count_tokens +from azure.cosmos.agent_memory._utils import cosine_similarity, vector_centroid from azure.cosmos.agent_memory.exceptions import LLMError from azure.cosmos.agent_memory.logging import get_logger logger = get_logger(__name__) + +def parse_iso_datetime(value: Any) -> Optional[datetime]: + """Parse an ISO-8601 string to a tz-aware datetime (UTC assumed for naive + values), or return ``None`` if it is missing or unparseable. + + Pure: used to validate/normalize model-supplied episode timestamps without + raising, so callers can fall back to grounded segment bounds. + """ + if not isinstance(value, str) or not value.strip(): + return None + try: + dt = datetime.fromisoformat(value.strip().replace("Z", "+00:00")) + except ValueError: + return None + if dt.tzinfo is None: + dt = dt.replace(tzinfo=timezone.utc) + return dt + + +def is_valid_time_pair(started: Any, ended: Any) -> bool: + """True iff both values parse as ISO-8601 datetimes with ``started <= ended``. + + Naive and tz-aware values are normalized to UTC before comparison, so a + mixed pair (e.g. a naive date plus a tz-aware datetime) compares safely + instead of raising ``TypeError``. + """ + start = parse_iso_datetime(started) + end = parse_iso_datetime(ended) + return start is not None and end is not None and start <= end + + +def clamp_unit_interval(value: Any, default: float) -> float: + """Clamp a number into ``[0.0, 1.0]``; return ``default`` if it is not a + finite number. + + The strict LLM ``json_schema`` constrains the field type (number) but not + its range, so salience / confidence can arrive out of range (e.g. 1.4 or + -0.2). Normalizing rather than raising keeps the episode instead of dropping + it over a slightly-off score. + """ + if isinstance(value, bool) or not isinstance(value, (int, float)): + return default + v = float(value) + if not math.isfinite(v): + return default + return max(0.0, min(1.0, v)) + + +# --------------------------------------------------------------------------- +# Episode boundary segmentation - pure helpers shared by the sync and aio +# pipelines. IO-free: functions of their inputs (plus threshold values passed +# in), so the two pipelines share one implementation instead of hand-mirroring. +# --------------------------------------------------------------------------- + + +def turn_gap_seconds(prev_turn: dict[str, Any], cur_turn: dict[str, Any]) -> Optional[float]: + """Seconds between two consecutive turns' ``created_at``, or None if either + timestamp is missing/unparseable (the idle-gap check then abstains).""" + prev_ts = parse_iso_datetime(prev_turn.get("created_at")) + cur_ts = parse_iso_datetime(cur_turn.get("created_at")) + if prev_ts is None or cur_ts is None: + return None + return (cur_ts - prev_ts).total_seconds() + + +def segment_time_bounds(items: list[dict[str, Any]]) -> tuple[Optional[str], Optional[str]]: + """Return (earliest, latest) turn ``created_at`` in a segment, or (None, None). + + Used to ground an episode's started_at/ended_at in its actual turn window + when the model omits them - for the conversational benchmarks each turn's + created_at carries the session date, so this yields correct temporal spans. + + Ordering is by PARSED datetime, not lexical string order: mixed UTC offsets + (e.g. ``+05:00`` vs ``Z``) sort correctly instead of producing an inverted + ``(started, ended)`` pair that would later fail ``_validate_time_order``. The + original ISO strings are returned unchanged. + """ + parsed: list[tuple[datetime, str]] = [] + for item in items: + raw = item.get("created_at") + if not isinstance(raw, str) or not raw: + continue + dt = parse_iso_datetime(raw) + if dt is not None: + parsed.append((dt, raw)) + if not parsed: + return None, None + parsed.sort(key=lambda pair: pair[0]) + return parsed[0][1], parsed[-1][1] + + +def created_at_sort_key(item: dict[str, Any]) -> tuple[int, float, str]: + """Total, chronological, stable sort key for a turn by ``created_at``. + + Orders by PARSED instant (so mixed UTC offsets sort by true time, not lexical + string) and breaks ties on ``id`` so turns sharing one timestamp - common + when many turns carry the same session date - keep a deterministic order. A + stable order matters because the open segment's first/last turn ids seed the + ``segment_key`` that grounds the deterministic episode id: an unstable order + could shift that id between runs and admit a duplicate episode. Missing or + unparseable timestamps sort last, deterministically by id. + """ + dt = parse_iso_datetime(item.get("created_at")) + if dt is None: + return (1, 0.0, str(item.get("id") or "")) + return (0, dt.timestamp(), str(item.get("id") or "")) + + +def deterministic_episode_id(segment_key: str, index: int) -> str: + """Identity from the STABLE segment key plus the episode's ordinal within the + segment - never the LLM summary text - so re-running the same un-stamped + segment yields the same id and the duplicate write is skipped (409).""" + seed = ID_SEED_SEP.join((segment_key, str(index))) + return f"ep_{hashlib.sha256(seed.encode()).hexdigest()[:32]}" + + +def find_episode_boundary( + segment: list[dict[str, Any]], + embeddings: list[list[float]], + *, + max_turns: int, + idle_gap: int, + drift: float, + min_turns: int, +) -> Optional[int]: + """Return the exclusive end index of the earliest episode boundary in the + open segment, or ``None`` if the segment is still open. + + Signals, earliest wins: (1) an idle time-gap between two consecutive turns, + (2) a topic-drift shift of a new turn away from the segment centroid, and + (3) a max-size cap that force-closes an over-long segment. Idle-gap and drift + boundaries below ``min_turns`` are suppressed so a lone turn is not emitted as + a trivial episode; the max-size cap is a hard ceiling and is not floored. + Boundaries are temporal/semantic only; prior episodes are never mutated. + """ + n = len(segment) + if n == 0: + return None + min_turns = max(1, min_turns) + has_embeddings = bool(embeddings) and len(embeddings) == n + for i in range(1, n): + if max_turns > 0 and i >= max_turns: + return i + if idle_gap > 0 and i >= min_turns: + gap = turn_gap_seconds(segment[i - 1], segment[i]) + if gap is not None and gap > idle_gap: + return i + if drift > 0 and i >= min_turns and has_embeddings: + centroid = vector_centroid(embeddings[:i]) + if centroid and (1.0 - cosine_similarity(embeddings[i], centroid)) > drift: + return i + if max_turns > 0 and n >= max_turns: + return max_turns + return None + + _NON_RETRYABLE_LLM_MARKERS = ( "content_filter", "content management policy", @@ -29,9 +189,21 @@ "maximum context length", ) +# Programming errors, not provider failures: a bug in our own extraction/parse +# code (e.g. calling ``.get`` on a non-dict, indexing past the end) raises one of +# these. They are deterministic, so retrying re-fails identically and would wedge +# the segment/batch forever - classify them non-retryable so the poison input is +# quarantined and surfaced rather than deferred indefinitely. Note this is a +# narrow allow-list of true code-bug types, NOT "any non-LLMError": genuine +# transient provider failures (rate-limit / timeout / connection) are their own +# SDK exception types and must stay retryable. +_NON_RETRYABLE_EXC_TYPES = (AttributeError, KeyError, TypeError, IndexError, NameError) + def is_retryable_llm_error(exc: BaseException) -> bool: """Classify an extraction LLM failure as retryable (transient) or not.""" + if isinstance(exc, _NON_RETRYABLE_EXC_TYPES): + return False text = str(exc).lower() return not any(marker in text for marker in _NON_RETRYABLE_LLM_MARKERS) @@ -396,214 +568,6 @@ def _canonical_speaker(role: Any) -> str: return _SPEAKER_ALIASES.get(normalized, str(role or "unknown")) -# Stopwords stripped from grounding checks. Keep this list short and focused -# on tokens that carry no factual content; any word a memory might legitimately -# differ on (e.g. "not", "no") must NOT be added here. -_GROUNDING_STOPWORDS = frozenset( - { - "the", - "a", - "an", - "is", - "are", - "was", - "were", - "be", - "been", - "being", - "and", - "or", - "but", - "to", - "of", - "for", - "on", - "in", - "at", - "by", - "with", - "from", - "as", - "that", - "this", - "these", - "those", - "it", - "its", - "user", - "they", - "them", - "their", - "he", - "she", - "his", - "her", - "him", - "has", - "have", - "had", - "do", - "does", - "did", - "will", - "would", - "should", - "can", - "could", - "may", - "might", - "must", - "say", - "says", - "said", - "saying", - "tell", - "tells", - "told", - "ask", - "asks", - "asked", - "mention", - "mentions", - "mentioned", - "stated", - "noted", - "added", - "replied", - "want", - "wants", - "wanted", - "decide", - "decides", - "decided", - "propose", - "proposes", - "proposed", - "suggest", - "suggests", - "suggested", - "planned", - "choose", - "chooses", - "chose", - "like", - "likes", - "liked", - "later", - "then", - "also", - "again", - } -) - -_GROUNDING_TOKEN_RE = re.compile(r"[a-zA-Z]{3,}") - - -def _grounding_tokens(text: str) -> set[str]: - """Tokenize text into lowercased content words (>=3 chars, stopwords removed).""" - if not text: - return set() - return {t for t in _GROUNDING_TOKEN_RE.findall(text.lower()) if t not in _GROUNDING_STOPWORDS} - - -def check_extracted_fact_grounding( - fact_docs: list[dict[str, Any]], - turn_items: list[dict[str, Any]], - existing_facts: list[dict[str, Any]], - *, - user_id: str, - thread_id: str, - logger: Any, -) -> None: - """Warn when an extracted fact's content is not grounded in the new user turns. - - Catches two known LLM failure modes that previously corrupted the fact store: - - 1. **Synthesis from existing facts** - the LLM emits an ADD whose content - paraphrase-merges two or more existing facts (e.g. existing - "user eats meat" + "user loves steak" → emitted "user loves steak, - indicating they eat meat") even though the new user turn says nothing - on the topic. Reconciliation later catches the resulting duplicates - but the visible artefact is a chain of "duplicate" supersedes that the - user never triggered. - - 2. **Phantom explicit-negation** - the LLM emits a second CONTRADICT fact - alongside the literal user statement (e.g. user says "I love steak and - seafood"; LLM emits both "user loves steak and seafood" and an invented - "user eats meat" CONTRADICT) when the supersedes_id on the literal fact - would have sufficed. Pollutes the store with claims the user didn't make. - - Heuristic: tokenize each emitted fact's content into lowercased content - words; subtract tokens present in the new user-turn transcript; the - remainder is "ungrounded". If ungrounded tokens come from 2+ existing - facts → strong synthesis signal. If they come from a single existing - fact with >=50%% overlap → weaker phantom-negation signal. - - Logs a WARNING for each suspected fact. Does NOT drop facts - downstream - reconciliation remains the dedup authority - but the WARNING is the - deterministic test signal that catches regressions. - """ - if not fact_docs or not turn_items: - return - - user_turn_text = " ".join( - str(m.get("content") or "") for m in turn_items if (m.get("role") or "").lower() == "user" - ) - user_tokens = _grounding_tokens(user_turn_text) - - existing_with_tokens: list[tuple[str, set[str]]] = [] - for mem in existing_facts: - toks = _grounding_tokens(str(mem.get("content") or "")) - if toks: - existing_with_tokens.append((str(mem.get("id") or ""), toks)) - - for doc in fact_docs: - content = str(doc.get("content") or "") - fact_tokens = _grounding_tokens(content) - if not fact_tokens: - continue - - ungrounded = fact_tokens - user_tokens - if not ungrounded: - continue - - contributors: list[tuple[str, set[str]]] = [ - (eid, ungrounded & toks) for eid, toks in existing_with_tokens if ungrounded & toks - ] - - if len(contributors) >= 2: - logger.warning( - "extract_memories: emitted fact appears synthesized from %d existing facts " - "(ungrounded in user turns) - extract should ground only in this turn's [user] lines. " - "doc_id=%s content=%r ungrounded_tokens=%s contributor_ids=%s " - "user_id=%s thread_id=%s", - len(contributors), - doc.get("id"), - content, - sorted(ungrounded), - [eid for eid, _ in contributors], - user_id, - thread_id, - ) - elif len(contributors) == 1 and len(ungrounded) >= 2: - eid, overlap = contributors[0] - overlap_ratio = len(overlap) / len(ungrounded) - if overlap_ratio >= 0.5: - logger.warning( - "extract_memories: emitted fact has ungrounded tokens overlapping a single existing fact " - "(possible phantom-negation/restatement) - extract should ground only in this turn's " - "[user] lines. doc_id=%s content=%r ungrounded_tokens=%s overlap_existing_id=%s " - "overlap_ratio=%.2f user_id=%s thread_id=%s", - doc.get("id"), - content, - sorted(ungrounded), - eid, - overlap_ratio, - user_id, - thread_id, - ) - - def parse_llm_json(text: str | None) -> dict[str, Any]: """Parse JSON from an LLM response, stripping markdown fences.""" if text is None: @@ -618,8 +582,9 @@ def parse_llm_json(text: str | None) -> dict[str, Any]: if cleaned.endswith("```"): cleaned = cleaned[:-3] cleaned = cleaned.strip() + decoder = json.JSONDecoder() try: - obj, end = json.JSONDecoder().raw_decode(cleaned) + obj, end = decoder.raw_decode(cleaned) except json.JSONDecodeError as exc: preview = (text or "")[:200].replace("\n", " ") if _looks_truncated(cleaned, exc): @@ -631,15 +596,85 @@ def parse_llm_json(text: str | None) -> dict[str, Any]: f"recent_k, or split oversized turns). Decode error: {exc}. preview={preview!r}" ) from exc raise LLMError(f"LLM returned invalid JSON (preview={preview!r}): {exc}") from exc - trailing = cleaned[end:].strip() - if trailing: - logger.warning( - "LLM response had %d chars of extra data after the first JSON object; using the " - "first object and ignoring the remainder (trailing_preview=%r)", - len(trailing), - trailing[:120].replace("\n", " "), + remainder = cleaned[end:] + if not isinstance(obj, dict): + # Type contract: this returns a JSON object. Some deployments emit a bare + # array or scalar root (e.g. ``[{...}]``); every caller then does + # ``parsed.get("facts"/"episodes")``, so returning a non-dict would + # surface downstream as an AttributeError that the extraction + # error-handlers misclassify as a transient (retryable) failure and defer + # forever. Raise a typed LLMError so a malformed root is handled like any + # other bad output instead of wedging the segment/batch. + preview = (text or "")[:200].replace("\n", " ") + raise LLMError( + f"LLM returned a non-object JSON root ({type(obj).__name__}); expected a JSON " + f"object such as {{'facts': [...]}}. preview={preview!r}" + ) + if not remainder.strip(): + # Fast path: exactly one JSON object, no trailing content (the vast + # majority of responses). + return obj + + # Some models (notably under strict json_schema on certain deployments) emit + # MULTIPLE back-to-back top-level JSON objects for a single call - e.g. + # ``{"facts":[...]}{"facts":[...]}``. ``raw_decode`` only returns the first, + # so keeping just it would silently DROP every item in the trailing + # object(s) - and the batch's turns get stamped ``extracted_at``, so those + # facts are never recovered. Instead, decode every top-level object and merge + # them: list-valued keys (``facts``, ``events``, ...) are concatenated and any + # other/scalar key keeps its first-seen value. Downstream exact-dup hashing + # removes any repeats, so merging is always safe. + merged: dict[str, Any] = obj + object_count = 1 + pos = end + length = len(cleaned) + while pos < length: + # Skip inter-object whitespace and stray separators between objects. + while pos < length and cleaned[pos] in " \t\r\n,": + pos += 1 + if pos >= length: + break + try: + nxt, pos = decoder.raw_decode(cleaned, pos) + except json.JSONDecodeError: + leftover = cleaned[pos:].strip() + logger.warning( + "LLM response had %d chars of non-JSON trailing data after %d concatenated " + "JSON object(s); merged the objects and ignored the remainder " + "(trailing_preview=%r)", + len(leftover), + object_count, + leftover[:120].replace("\n", " "), + ) + break + object_count += 1 + if isinstance(nxt, dict): + _merge_json_objects(merged, nxt) + # A non-dict follow-on cannot be merged into a dict result; skip it. + + if object_count > 1: + logger.info( + "LLM response contained %d concatenated JSON objects; merged their list fields into " + "one result so no items were dropped.", + object_count, ) - return obj + return merged + + +def _merge_json_objects(base: dict[str, Any], extra: dict[str, Any]) -> None: + """Merge ``extra`` into ``base`` in place. + + List-valued keys present in both are concatenated (this is what recovers the + facts/events an LLM splits across multiple concatenated JSON objects); keys + only in ``extra`` are added; any other conflict keeps ``base``'s (first-seen) + value. Exact-duplicate items are pruned later by the pipeline's content-hash + dedup, so concatenating without de-duping here is safe. + """ + for key, value in extra.items(): + if key not in base: + base[key] = value + elif isinstance(base[key], list) and isinstance(value, list): + base[key].extend(value) def _looks_truncated(cleaned: str, exc: json.JSONDecodeError) -> bool: @@ -657,6 +692,26 @@ def default_prompts_dir() -> str: return os.path.join(pkg_dir, "prompts") +_EXTRACT_MEMORIES_PROMPT_DEFAULT = "extract_memories-v2.prompty" +_EXTRACT_MEMORIES_PROMPT_ALLOWED = frozenset({"extract_memories.prompty", "extract_memories-v2.prompty"}) + + +def extract_memories_prompt_file() -> str: + """Return the fact-extraction prompt filename, env-selectable. + + Defaults to the v2 ``extract_memories-v2.prompty`` extractor (higher recall; + also captures assistant-provided information: lists, tables, instructions, + and researched answers the user may later reference). Set + ``AMT_EXTRACT_MEMORIES_PROMPT=extract_memories.prompty`` to fall back to the + v1 extractor without touching the default. An unknown value falls back to + the v2 default so a typo can never point extraction at an arbitrary or + missing prompt file. The chosen filename must be registered in + ``PROMPTY_SCHEMAS`` so it still gets the structured-output response format. + """ + name = os.environ.get("AMT_EXTRACT_MEMORIES_PROMPT", _EXTRACT_MEMORIES_PROMPT_DEFAULT) + return name if name in _EXTRACT_MEMORIES_PROMPT_ALLOWED else _EXTRACT_MEMORIES_PROMPT_DEFAULT + + def _read_prompty_version(path: str | Path) -> str: """Read the ``version:`` key from a prompty file's YAML front-matter.""" text = Path(path).read_text(encoding="utf-8") @@ -695,9 +750,31 @@ def load(self, filename: str) -> Any: import prompty loaded = prompty.load(self._path_for(filename)) + self._disable_strict_format(prompty, loaded) self._cache[filename] = loaded return loaded + def _disable_strict_format(self, prompty: Any, loaded: Any) -> None: + """Disable prompty strict nonce parsing for SDK-owned internal prompts.""" + template = getattr(loaded, "template", None) + if template is None: + template_type = getattr(prompty, "Template", None) + if template_type is None: + return + template = template_type() + loaded.template = template + + format_config = getattr(template, "format", None) + if format_config is None: + format_type = getattr(prompty, "FormatConfig", None) + if format_type is None: + return + format_config = format_type() + template.format = format_config + + if hasattr(format_config, "strict"): + format_config.strict = False + def prompt_version(self, filename: str) -> str: """Return the ``version:`` declared in the prompty front-matter.""" cached = self._version_cache.get(filename) @@ -717,24 +794,6 @@ def prepare(self, filename: str, inputs: dict[str, Any]) -> tuple[list[dict[str, return messages, params -# Allowed values for the EpisodicRecord ``outcome_valence`` field - mirrors -# ``azure.cosmos.agent_memory.models._EPISODIC_ALLOWED_VALENCES`` but kept inline -# to avoid an import cycle (helpers must not import models). -VALID_VALENCES = frozenset({"positive", "negative", "neutral", "mixed"}) - - -def coerce_valence(value: Any) -> str: - """Map an LLM-emitted ``outcome_valence`` to a record-safe value. - - The strict response schema permits ``positive | negative | mixed | neutral - | null``; null and any unknown value fall through to ``"neutral"`` so a - single drifted episode never aborts the whole extract batch. - """ - if isinstance(value, str) and value in VALID_VALENCES: - return value - return "neutral" - - # Per-section caps on the persisted ``structured_summary``. Strict-mode JSON # output does not enforce ``maxItems``, so the LLM grows lists unboundedly # across incremental updates. Capping at persist time keeps both Cosmos diff --git a/azure/cosmos/agent_memory/services/pipeline.py b/azure/cosmos/agent_memory/services/pipeline.py index 3ca1c96..2584fbc 100644 --- a/azure/cosmos/agent_memory/services/pipeline.py +++ b/azure/cosmos/agent_memory/services/pipeline.py @@ -51,7 +51,6 @@ ID_SEED_SEP as _ID_SEED_SEP, ) from azure.cosmos.agent_memory.services._pipeline_helpers import ( - VALID_VALENCES, PromptyLoader, _normalize_metadata_keys, batch_turns_by_tokens, @@ -59,10 +58,16 @@ build_transcript, cap_structured_summary, chat_text, - check_extracted_fact_grounding, - coerce_valence, + clamp_unit_interval, + created_at_sort_key, + deterministic_episode_id, + extract_memories_prompt_file, + find_episode_boundary, is_retryable_llm_error, + is_valid_time_pair, parse_llm_json, + segment_time_bounds, + turn_gap_seconds, ) from azure.cosmos.agent_memory.services._pipeline_helpers import ( is_real_number as _is_real_number, @@ -75,7 +80,6 @@ logger = get_logger("azure.cosmos.agent_memory.pipeline") -_coerce_valence = coerce_valence _cap_structured_summary = cap_structured_summary # Standard SQL predicate that selects "active" (non-superseded) docs. @@ -90,6 +94,13 @@ # threads per user, occasional triple-fanout) without inflating the cost cap. _PROCEDURAL_MAX_CREATE_ATTEMPTS = 5 +# Safety cap on how many episodes one extract_episodes call may close in a +# single pass. Boundary evaluation runs on a small turn cadence so a backlog +# rarely exceeds one segment; the cap only bounds a pathological drain (e.g. +# first evaluation after a huge unprocessed backlog) so one call cannot fan out +# into an unbounded burst of LLM extractions. +_EPISODE_MAX_SEGMENTS_PER_RUN = 50 + class _StoreContainerAdapter: """Expose one split ``MemoryStore`` container via Cosmos method shapes.""" @@ -487,6 +498,8 @@ def _empty_extract_counts() -> dict[str, int]: "exact_dedup_skipped": 0, "dropped_episodic_count": 0, "inplace_updated": 0, + "deferred_turn_count": 0, + "quarantined_turn_count": 0, } @staticmethod @@ -549,7 +562,7 @@ def _mark_superseded_via_container( def _parse_llm_json(text: str | None) -> dict[str, Any]: return parse_llm_json(text) - def extract_memories_dry( + def extract_memories_durable( self, user_id: str, thread_id: str, @@ -563,7 +576,7 @@ def extract_memories_dry( if not thread_id: raise ValidationError("thread_id is required") - logger.info("extract_memories_dry started user_id=%s thread_id=%s", user_id, thread_id) + logger.info("extract_memories_durable started user_id=%s thread_id=%s", user_id, thread_id) if turns is None: query = ( @@ -591,7 +604,7 @@ def extract_memories_dry( items.reverse() if not items: - logger.warning("extract_memories_dry no memories found user_id=%s thread_id=%s", user_id, thread_id) + 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"]) @@ -608,17 +621,16 @@ def extract_memories_dry( # *retryable* error are left un-stamped and retried on the next run. batches = batch_turns_by_tokens(items, threshold_config.get_extraction_batch_max_tokens()) facts: list[dict[str, Any]] = [] - episodic: list[dict[str, Any]] = [] processed_turns: list[dict[str, Any]] = [] deferred_turn_count = 0 quarantined_turn_count = 0 + extract_prompt = extract_memories_prompt_file() for batch in batches: batch_transcript = self._build_transcript(batch, include_timestamp=True) try: - response_text = self._run_prompty("extract_memories.prompty", inputs={"transcript": batch_transcript}) + response_text = self._run_prompty(extract_prompt, inputs={"transcript": batch_transcript}) parsed = self._parse_llm_json(response_text) facts.extend(parsed.get("facts", [])) - episodic.extend(parsed.get("episodic", [])) processed_turns.extend(batch) except Exception as exc: # noqa: BLE001 if is_retryable_llm_error(exc): @@ -646,10 +658,8 @@ def extract_memories_dry( doc_timestamp = self._stable_source_timestamp(items) fact_docs: list[dict[str, Any]] = [] - episodic_docs: list[dict[str, Any]] = [] updates: list[dict[str, Any]] = [] exact_dedup_skipped = 0 - dropped_episodic_count = 0 for fact in facts: text = fact.get("text") @@ -689,14 +699,14 @@ def extract_memories_dry( "type": "fact", "content": text, "content_hash": new_content_hash, - "confidence": 0.5 if confidence is None else confidence, - **self._prompt_lineage("extract_memories.prompty"), + "confidence": clamp_unit_interval(confidence, 0.5), + **self._prompt_lineage(extract_prompt), "metadata": { "category": fact.get("category") or "other", "temporal_context": fact.get("temporal_context"), "source": fact_source, }, - "salience": fact.get("salience") if fact.get("salience") is not None else 0.5, + "salience": clamp_unit_interval(fact.get("salience"), 0.5), "tags": ["sys:fact", "sys:auto-extracted"] + source_tags + topic_tags, "created_at": doc_timestamp, "updated_at": doc_timestamp, @@ -705,81 +715,8 @@ def extract_memories_dry( fact_docs.append(self._validate_extracted_doc(doc)) existing_fact_hashes.add(new_content_hash) - for ep in episodic: - scope_type_raw = ep.get("scope_type") - scope_value_raw = ep.get("scope_value") - scope_type = scope_type_raw.strip() if isinstance(scope_type_raw, str) else None - scope_value = scope_value_raw.strip() if isinstance(scope_value_raw, str) else None - if not scope_type or not scope_value: - logger.warning( - "extract_memories: dropping malformed episodic (missing scope_type/scope_value) " - "user_id=%s thread_id=%s reason=malformed_scope payload=%r", - user_id, - thread_id, - ep, - ) - dropped_episodic_count += 1 - continue - - situation = ep.get("situation") - action_taken = ep.get("action_taken") - outcome = ep.get("outcome") - if situation and action_taken and outcome: - text = f"{situation} → {action_taken} → {outcome}" - else: - text = f"For the user's {scope_value} {scope_type}, intent recorded." - - content_hash = compute_content_hash(text) - seed = _ID_SEED_SEP.join((user_id, thread_id, content_hash)) - det_id = f"ep_{hashlib.sha256(seed.encode()).hexdigest()[:32]}" - topic_tags = build_topic_tags(ep.get("tags", [])) - confidence = ep.get("confidence") - raw_valence = ep.get("outcome_valence") - coerced_valence = _coerce_valence(raw_valence) - if raw_valence is not None and raw_valence not in VALID_VALENCES: - logger.warning( - "extract_memories: coercing unknown outcome_valence=%r → %r user_id=%s thread_id=%s", - raw_valence, - coerced_valence, - user_id, - thread_id, - ) - doc = { - "id": det_id, - "user_id": user_id, - "thread_id": thread_id, - "role": "system", - "type": "episodic", - "content": text, - "content_hash": content_hash, - "confidence": 0.5 if confidence is None else confidence, - "ttl": DEFAULT_TTL_BY_TYPE.get("episodic", 7_776_000), - **self._prompt_lineage("extract_memories.prompty"), - "metadata": { - "scope_type": scope_type, - "scope_value": scope_value, - "situation": situation, - "action_taken": action_taken, - "outcome": outcome, - "reasoning": ep.get("reasoning"), - "outcome_valence": coerced_valence, - "lesson": ep.get("lesson") - or ( - f"{situation} → {action_taken} → {outcome}" if situation and action_taken and outcome else text - ), - "domain": ep.get("domain"), - }, - "salience": ep.get("salience"), - "tags": ["sys:episodic", "sys:auto-extracted"] + topic_tags, - "created_at": doc_timestamp, - "updated_at": doc_timestamp, - } - episodic_docs.append(self._validate_extracted_doc(doc)) - if exact_dedup_skipped: updates.append({"op": "stats", "exact_dedup_skipped": exact_dedup_skipped}) - if dropped_episodic_count: - updates.append({"op": "stats", "dropped_episodic_count": dropped_episodic_count}) if deferred_turn_count or quarantined_turn_count: updates.append( { @@ -789,31 +726,358 @@ def extract_memories_dry( } ) - check_extracted_fact_grounding( - fact_docs, - processed_turns, - existing_for_hashes, - user_id=user_id, - thread_id=thread_id, - logger=logger, - ) - result = { "facts": fact_docs, - "episodic": episodic_docs, + "episodic": [], "updates": updates, "processed_turn_docs": processed_turns, } logger.info( - "extract_memories_dry completed user_id=%s thread_id=%s fact_docs=%d episodic_docs=%d updates=%d", + "extract_memories_durable completed user_id=%s thread_id=%s fact_docs=%d updates=%d", user_id, thread_id, len(fact_docs), - len(episodic_docs), len(updates), ) return result + def _build_episode_transcript(self, items: list[dict[str, Any]]) -> str: + """Build a timestamped transcript that exposes real turn ids to the LLM.""" + transcript_items: list[dict[str, Any]] = [] + for index, item in enumerate(items, start=1): + turn_id = str(item.get("id") or f"turn-{index}") + copied = dict(item) + copied["content"] = f"Turn {turn_id}: {item.get('content', '')}" + transcript_items.append(copied) + return self._build_transcript(transcript_items, include_timestamp=True) + + @staticmethod + def _ground_episode_events( + events: Any, + *, + turn_ids: list[str], + ) -> tuple[list[dict[str, Any]], list[str]]: + """Normalize event source ids to real ids from the current turn window.""" + valid_turn_ids = set(turn_ids) + label_to_id = {f"turn-{i}": turn_id for i, turn_id in enumerate(turn_ids, start=1)} + grounded_events: list[dict[str, Any]] = [] + source_turn_ids: list[str] = [] + for event in events if isinstance(events, list) else []: + if not isinstance(event, dict): + continue + grounded = dict(event) + grounded_sources: list[str] = [] + for raw_source in event.get("source_turn_ids") or []: + source = str(raw_source).strip() + mapped = source if source in valid_turn_ids else label_to_id.get(source.lower()) + if mapped and mapped not in grounded_sources: + grounded_sources.append(mapped) + if mapped and mapped not in source_turn_ids: + source_turn_ids.append(mapped) + grounded["source_turn_ids"] = grounded_sources + grounded_events.append(grounded) + return grounded_events, source_turn_ids + + def _build_episode_docs( + self, + user_id: str, + thread_id: str, + items: list[dict[str, Any]], + *, + segment_key: str, + ) -> list[dict[str, Any]]: + """Run the episode-extraction prompt over one bounded, already-closed turn + segment and return episode docs (no embeddings, no writes). + + Episode ids are DETERMINISTIC from the segment identity (its turn range) + plus each episode's ordinal - not the summary text - so a re-run over the + same segment (e.g. after a crash between persist and turn-stamping) + collides on id and is skipped rather than duplicated. + """ + if not items: + return [] + + transcript = self._build_episode_transcript(items) + response_text = self._run_prompty("extract_episode.prompty", inputs={"transcript": transcript}) + parsed = self._parse_llm_json(response_text) + episodes = parsed.get("episodes", []) + if not isinstance(episodes, list): + logger.warning( + "_build_episode_docs dropping malformed response user_id=%s thread_id=%s payload=%r", + user_id, + thread_id, + parsed, + ) + return [] + + doc_timestamp = self._stable_source_timestamp(items) + turn_ids = [str(item.get("id")) for item in items if item.get("id")] + # An episode's temporal span is grounded in its turn window: for the + # conversational benchmarks each turn's created_at carries the session + # date, so segment bounds give correct started_at/ended_at even when the + # model omits them. We only trust model-supplied times when it provides a + # complete, self-consistent pair. + segment_started, segment_ended = segment_time_bounds(items) + docs: list[dict[str, Any]] = [] + for index, episode in enumerate(episodes): + if not isinstance(episode, dict): + logger.warning( + "_build_episode_docs dropping malformed episode user_id=%s thread_id=%s payload=%r", + user_id, + thread_id, + episode, + ) + continue + + summary = episode.get("summary") + if not isinstance(summary, str) or not summary.strip(): + logger.warning( + "_build_episode_docs dropping malformed episode (missing summary) " + "user_id=%s thread_id=%s payload=%r", + user_id, + thread_id, + episode, + ) + continue + title = episode.get("title") + if not isinstance(title, str) or not title.strip(): + logger.warning( + "_build_episode_docs dropping malformed episode (missing title) user_id=%s thread_id=%s payload=%r", + user_id, + thread_id, + episode, + ) + continue + if not isinstance(episode.get("events"), list): + logger.warning( + "_build_episode_docs dropping malformed episode (missing events) " + "user_id=%s thread_id=%s payload=%r", + user_id, + thread_id, + episode, + ) + continue + + events, source_turn_ids = self._ground_episode_events(episode.get("events"), turn_ids=turn_ids) + content_hash = compute_content_hash(str(summary)) + llm_started, llm_ended = episode.get("started_at"), episode.get("ended_at") + # Trust model-supplied times only when they form a valid, self-consistent + # ISO pair; otherwise fall back to the grounded segment bounds rather than + # dropping the whole episode (malformed or mixed-tz strings are common). + if is_valid_time_pair(llm_started, llm_ended): + started_at, ended_at = str(llm_started).strip(), str(llm_ended).strip() + else: + started_at, ended_at = segment_started, segment_ended + try: + doc = construct_internal( + EpisodicRecord, + { + "id": self._deterministic_episode_id(segment_key, index), + "user_id": user_id, + "thread_id": thread_id, + "role": "system", + "type": "episodic", + "content": summary, + "title": title, + "started_at": started_at, + "ended_at": ended_at, + "participants": episode.get("participants") or [], + "events": events, + "outcome": episode.get("outcome"), + "lessons": episode.get("lessons") or [], + "source_turn_ids": source_turn_ids, + "content_hash": content_hash, + "salience": clamp_unit_interval(episode.get("salience"), 0.5), + "confidence": clamp_unit_interval(episode.get("confidence"), 0.5), + "ttl": DEFAULT_TTL_BY_TYPE.get("episodic", 7_776_000), + "tags": ["sys:episodic", "sys:auto-extracted"], + "created_at": doc_timestamp, + "updated_at": doc_timestamp, + **self._prompt_lineage("extract_episode.prompty"), + }, + ).to_doc() + except Exception as exc: # noqa: BLE001 + logger.warning( + "_build_episode_docs dropping malformed episode user_id=%s thread_id=%s err=%s payload=%r", + user_id, + thread_id, + exc, + episode, + ) + continue + docs.append(doc) + return docs + + @staticmethod + def _deterministic_episode_id(segment_key: str, index: int) -> str: + return deterministic_episode_id(segment_key, index) + + 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. + + Mirrors the fact ``extracted_at`` watermark with an independent + ``episode_extracted_at`` cursor, so episodic segmentation neither blocks + nor is blocked by fact extraction. + """ + 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))" + ) + parameters: list[dict[str, Any]] = [ + {"name": "@user_id", "value": user_id}, + {"name": "@thread_id", "value": thread_id}, + ] + items = list( + self._turns_container.query_items( + query=query, + parameters=parameters, + partition_key=[user_id, thread_id], + ) + ) + items.sort(key=created_at_sort_key) + return items + + def _episode_segment_embeddings(self, segment: list[dict[str, Any]]) -> list[list[float]]: + """Return one embedding per segment turn for drift detection, or ``[]`` when + drift is disabled or embeddings are unavailable. + + Reuses a stored turn ``embedding`` when present (deployments with + ``enable_turn_embeddings``); otherwise embeds turn text in one batch. The + segment is bounded by ``EPISODE_MAX_TURNS`` so this stays cheap, and + embeddings cost far less than the boundary-gated LLM extraction it guards. + """ + if threshold_config.get_episode_topic_drift() <= 0: + return [] + embeddings: list[Optional[list[float]]] = [ + turn.get("embedding") if isinstance(turn.get("embedding"), list) else None for turn in segment + ] + missing = [i for i, emb in enumerate(embeddings) if emb is None] + if missing: + try: + fresh = self._embed_batch([str(segment[i].get("content") or "") for i in missing]) + except Exception as exc: # noqa: BLE001 + logger.warning("episode drift embedding failed (%s); skipping drift this evaluation", exc) + return [] + for pos, i in enumerate(missing): + embeddings[i] = fresh[pos] if pos < len(fresh) else None + if any(emb is None for emb in embeddings): + return [] + return [emb for emb in embeddings if emb is not None] + + @staticmethod + def _turn_gap_seconds(prev_turn: dict[str, Any], cur_turn: dict[str, Any]) -> Optional[float]: + return turn_gap_seconds(prev_turn, cur_turn) + + def _find_episode_boundary( + self, + segment: list[dict[str, Any]], + embeddings: list[list[float]], + ) -> Optional[int]: + return find_episode_boundary( + segment, + embeddings, + max_turns=threshold_config.get_episode_max_turns(), + idle_gap=threshold_config.get_episode_idle_gap_seconds(), + drift=threshold_config.get_episode_topic_drift(), + min_turns=threshold_config.get_episode_min_turns(), + ) + + def extract_episodes( + self, + user_id: str, + thread_id: str, + *, + flush: bool = False, + ) -> 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 + 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. + """ + if not user_id: + raise ValidationError("user_id is required") + if not thread_id: + raise ValidationError("thread_id is required") + + segment = self._load_open_episode_segment(user_id, thread_id) + total = 0 + guard = 0 + while segment and guard < _EPISODE_MAX_SEGMENTS_PER_RUN: + guard += 1 + embeddings = self._episode_segment_embeddings(segment) + boundary = self._find_episode_boundary(segment, embeddings) + if boundary is None: + if not flush: + break + boundary = len(segment) + closing = segment[:boundary] + if not closing: + break + first_id = str(closing[0].get("id") or "") + last_id = str(closing[-1].get("id") or "") + segment_key = _ID_SEED_SEP.join((user_id, thread_id, first_id, last_id)) + try: + docs = self._build_episode_docs(user_id, thread_id, closing, segment_key=segment_key) + embeddings_for_docs = self._embed_batch([str(doc["content"]) for doc in docs]) if docs else [] + except Exception as exc: # noqa: BLE001 + if is_retryable_llm_error(exc): + # Transient provider error: leave the segment un-stamped and stop + # this run so it is retried intact next time (mirror of the fact path). + logger.warning( + "extract_episodes: deferring %d turns after retryable extraction error " + "(will retry next run) user_id=%s thread_id=%s err=%s", + len(closing), + user_id, + thread_id, + exc, + ) + 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. + logger.warning( + "extract_episodes: quarantining %d turns after non-retryable extraction error " + "(marking episode_extracted_at 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") + segment = segment[boundary:] + continue + for doc, embedding in zip(docs, embeddings_for_docs): + doc["embedding"] = embedding + try: + self._create_memory(doc) + total += 1 + except CosmosResourceExistsError: + logger.info("extract_episodes idempotent skip duplicate episode id=%s", doc.get("id")) + # 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") + segment = segment[boundary:] + return {"episodes": total} + def dedup_extracted_memories( self, user_id: str, @@ -1084,7 +1348,7 @@ def persist_extracted_memories( doc["embedding"] = embedding for doc in docs_to_create: - # Intentionally re-validates even though extract_memories_dry did: + # Intentionally re-validates even though extract_memories_durable did: # persist_extracted_memories is a public surface (recovery scripts, # custom processors, third-party callers) so we don't trust input # shape. Cost is microseconds per doc; the safety boundary is the @@ -1092,10 +1356,7 @@ def persist_extracted_memories( validated = self._validate_extracted_doc(doc) doc_type = validated.get("type") try: - if doc_type == "episodic": - self._upsert_memory(validated) - else: - self._create_memory(validated) + self._create_memory(validated) except CosmosResourceExistsError: logger.info("persist_extracted_memories skipped existing id=%s", validated.get("id")) continue @@ -1109,7 +1370,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",): + for key in ("inplace_updated", "deferred_turn_count", "quarantined_turn_count"): if key in op: result[key] = result.get(key, 0) + int(op.get(key) or 0) @@ -1117,12 +1378,17 @@ def persist_extracted_memories( return result - def _mark_turns_extracted(self, turn_docs: list[dict[str, Any]]) -> int: - """Stamp ``extracted_at`` on each turn doc and upsert. + 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. + + ``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). We upsert the full doc (rather than patch) because the container adapter only exposes upsert. Per-turn failures are logged but do - not raise - the worst case is one turn gets re-extracted on the + not raise - the worst case is one turn gets re-processed on the next call, which is bounded and recoverable. """ if not turn_docs: @@ -1135,12 +1401,13 @@ def _mark_turns_extracted(self, turn_docs: list[dict[str, Any]]) -> int: continue try: doc_to_write = dict(turn) - doc_to_write["extracted_at"] = now_iso + doc_to_write[field] = now_iso self._turns_container.upsert_item(body=doc_to_write) marked += 1 except Exception as exc: logger.warning( - "_mark_turns_extracted failed for turn_id=%s err=%s (turn may be re-extracted on next call)", + "_mark_turns_extracted(%s) failed for turn_id=%s err=%s (turn may be re-processed on next call)", + field, turn_id, exc, ) @@ -1155,8 +1422,8 @@ def extract_memories( turns: Optional[list[dict[str, Any]]] = None, ) -> dict[str, int]: """Extract facts and episodic memories from a thread and persist them.""" - extracted = self.extract_memories_dry(user_id, thread_id, recent_k, turns=turns) - # Capture the processed turns from the DRY output as the single source of + 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. @@ -1259,8 +1526,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.metadata.lesson) " - "AND c.metadata.lesson != null " + "AND IS_DEFINED(c.lessons) " + "AND ARRAY_LENGTH(c.lessons) > 0 " "ORDER BY c.salience DESC, c.created_at ASC, c.id ASC" ), parameters=[ @@ -1273,8 +1540,8 @@ def _read_latest_procedural() -> Optional[dict[str, Any]]: episodic_with_lessons = [ doc for doc in episodic_docs - if isinstance(doc.get("metadata", {}).get("lesson"), str) - and doc.get("metadata", {}).get("lesson", "").strip() + 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] @@ -1313,7 +1580,12 @@ def _render_bullets(values: list[str]) -> str: static_prompty_inputs = { "behavioral_facts": _render_bullets([doc.get("content", "") for doc in behavioral_fact_docs]), "episodic_lessons": _render_bullets( - [doc.get("metadata", {}).get("lesson", "") for doc in episodic_with_lessons] + [ + lesson + for doc in episodic_with_lessons + for lesson in doc.get("lessons", []) + if isinstance(lesson, str) and lesson.strip() + ] ), "user_name": user_name, } @@ -1401,7 +1673,7 @@ def _render_bullets(values: list[str]) -> str: ) return {"status": "synthesized", "procedural": written_doc} - def generate_thread_summary_dry( + def generate_thread_summary_durable( self, user_id: str, thread_id: str, @@ -1413,7 +1685,7 @@ def generate_thread_summary_dry( if not thread_id: raise ValidationError("thread_id is required") - logger.info("generate_thread_summary_dry started user_id=%s thread_id=%s", user_id, thread_id) + logger.info("generate_thread_summary_durable started user_id=%s thread_id=%s", user_id, thread_id) summary_id = f"summary_{user_id}_{thread_id}" existing_summary: Optional[dict[str, Any]] = None @@ -1441,7 +1713,7 @@ def generate_thread_summary_dry( ) if existing_summary and not items: - logger.info("generate_thread_summary_dry no new memories, returning existing") + logger.info("generate_thread_summary_durable no new memories, returning existing") summary_doc = dict(existing_summary) summary_doc.pop("embedding", None) return summary_doc @@ -1530,10 +1802,10 @@ def generate_thread_summary( recent_k: int | None = None, ) -> dict[str, Any]: """Generate or incrementally update a thread summary and persist it.""" - summary_doc = self.generate_thread_summary_dry(user_id, thread_id, recent_k=recent_k) + summary_doc = self.generate_thread_summary_durable(user_id, thread_id, recent_k=recent_k) return self.persist_thread_summary(user_id, thread_id, summary_doc) - def generate_user_summary_dry( + def generate_user_summary_durable( self, user_id: str, thread_ids: list[str] | None = None, @@ -1544,7 +1816,7 @@ def generate_user_summary_dry( raise ValidationError("user_id is required") logger.info( - "generate_user_summary_dry started user_id=%s observed_thread_ids=%s", + "generate_user_summary_durable started user_id=%s observed_thread_ids=%s", user_id, len(thread_ids) if thread_ids else 0, ) @@ -1585,7 +1857,7 @@ def generate_user_summary_dry( ) if existing_summary and not items: - logger.info("generate_user_summary_dry no new memories, returning existing") + logger.info("generate_user_summary_durable no new memories, returning existing") user_doc = dict(existing_summary) user_doc.pop("embedding", None) return user_doc @@ -1691,7 +1963,7 @@ def generate_user_summary( recent_k: int | None = None, ) -> dict[str, Any]: """Generate or incrementally update a user summary and persist it.""" - summary_doc = self.generate_user_summary_dry(user_id, thread_ids=thread_ids, recent_k=recent_k) + summary_doc = self.generate_user_summary_durable(user_id, thread_ids=thread_ids, recent_k=recent_k) return self.persist_user_summary(user_id, summary_doc) def _emit_reconcile_outcome( diff --git a/azure/cosmos/agent_memory/store/_search_helpers.py b/azure/cosmos/agent_memory/store/_search_helpers.py index 5e58269..84a4efb 100644 --- a/azure/cosmos/agent_memory/store/_search_helpers.py +++ b/azure/cosmos/agent_memory/store/_search_helpers.py @@ -16,6 +16,8 @@ MEMORY_PROJECTION = ( "c.id, c.user_id, c.thread_id, c.role, c.type, c.content, " "c.metadata, c.created_at, c.tags, c.salience, c.confidence, " + "c.title, c.started_at, c.ended_at, c.participants, c.events, " + "c.outcome, c.lessons, c.source_turn_ids, " "c.superseded_by, c.superseded_at, c.supersede_reason" ) @@ -84,10 +86,10 @@ def format_episodic_context(memories: Iterable[dict[str, Any]]) -> str: return "" lines = ["## Relevant Past Experiences"] for i, memory in enumerate(memories_list, 1): - metadata = memory.get("metadata") or {} - domain = metadata.get("domain", "general") - valence = metadata.get("outcome_valence", "neutral") - lines.append(f"{i}. [{domain}] {memory['content']} ({valence})") + title = memory.get("title") or "Episode" + outcome = memory.get("outcome") or {} + status = outcome.get("status", "unknown") if isinstance(outcome, dict) else "unknown" + lines.append(f"{i}. [{status}] {title}: {memory.get('content', '')}") 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 7bf2072..f54833e 100644 --- a/azure/cosmos/agent_memory/store/memory_store.py +++ b/azure/cosmos/agent_memory/store/memory_store.py @@ -231,10 +231,11 @@ def add( if memory_type == "fact": meta.setdefault("category", "unclassified:manual") elif memory_type == "episodic": - meta.setdefault("lesson", content) - meta.setdefault("scope_type", "manual") - meta.setdefault("scope_value", "manual") - meta.setdefault("outcome_valence", "neutral") + kwargs.setdefault("title", content[:80] or "Manual episode") + kwargs.setdefault("events", []) + kwargs.setdefault("participants", []) + kwargs.setdefault("lessons", []) + kwargs.setdefault("source_turn_ids", []) elif memory_type == "procedural": kwargs.setdefault("source_fact_ids", ["manual"]) kwargs["metadata"] = meta @@ -570,6 +571,35 @@ def get_thread_summary( container=self._summaries_container, ) + def get_episodes( + self, + user_id: str, + thread_id: Optional[str] = None, + recent_k: Optional[int] = None, + ) -> list[dict[str, Any]]: + """Retrieve active episodic memories for ``user_id``, newest first.""" + if not user_id: + raise ValidationError("user_id is required for get_episodes") + qb = _QueryBuilder() + qb.add_filter("c.type", "@type", "episodic") + qb.add_filter("c.user_id", "@user_id", user_id) + qb.add_filter("c.thread_id", "@thread_id", thread_id) + qb.add_is_null_or_undefined("c.superseded_by") + parameters = qb.get_parameters() + if recent_k is not None: + parameters.append({"name": "@recent_k", "value": recent_k}) + sql = f"SELECT TOP @recent_k * FROM c{qb.build_where()} ORDER BY c.created_at DESC" + else: + sql = f"SELECT * FROM c{qb.build_where()} ORDER BY c.created_at DESC" + partition_key, cross_partition = query_scope(user_id, thread_id) + return self.query( + sql, + parameters, + container_key=ContainerKey.MEMORIES, + partition_key=partition_key, + cross_partition=cross_partition, + ) + def get_user_summary(self, user_id: str) -> Optional[dict[str, Any]]: """Retrieve the user's summary document from Cosmos DB, or ``None`` if absent.""" from azure.cosmos.exceptions import CosmosResourceNotFoundError @@ -1083,16 +1113,76 @@ def search_episodic( top_k: int = 5, min_salience: Optional[float] = None, include_superseded: bool = False, + thread_id: Optional[str] = None, + tags_all: Optional[list[str]] = None, + tags_any: Optional[list[str]] = None, + exclude_tags: Optional[list[str]] = None, + created_after: Optional[str | datetime] = None, + created_before: Optional[str | datetime] = None, + started_after: Optional[str | datetime] = None, + started_before: Optional[str | datetime] = None, + ended_after: Optional[str | datetime] = None, + ended_before: Optional[str | datetime] = None, ) -> list[dict[str, Any]]: - """Semantic search across episodic memories for a user.""" - return self.search( - search_terms=search_terms, - user_id=user_id, - memory_types=["episodic"], - top_k=top_k, - min_salience=min_salience, + """Semantic search across episodic memories for a user. + + Temporal arguments are filters only; relevance ranking is vector/FTS-only. + """ + if not user_id: + raise ValidationError("user_id is required for search_episodic") + 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", "episodic") + qb.add_filter("c.user_id", "@user_id", user_id) + qb.add_filter("c.thread_id", "@thread_id", thread_id) + add_tag_filters(qb, tags_all=tags_all, tags_any=tags_any, exclude_tags=exclude_tags) + qb.add_time_range( + "c.created_at", + after=_coerce_datetime_iso(created_after), + before=_coerce_datetime_iso(created_before), + after_param="@created_after", + before_param="@created_before", + ) + qb.add_time_range( + "c.started_at", + after=_coerce_datetime_iso(started_after), + before=_coerce_datetime_iso(started_before), + after_param="@started_after", + before_param="@started_before", + ) + qb.add_time_range( + "c.ended_at", + after=_coerce_datetime_iso(ended_after), + before=_coerce_datetime_iso(ended_before), + after_param="@ended_after", + before_param="@ended_before", + ) + add_salience_filter(qb, min_salience) + + 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, thread_id) + logger.debug("MemoryStore.search_episodic query: %s", sql) + return self.query( + sql, + parameters, + container_key=ContainerKey.MEMORIES, + partition_key=partition_key, + cross_partition=cross_partition, + ) def build_episodic_context(self, user_id: str, query: str, top_k: int = 3) -> str: """Build formatted context of relevant past experiences.""" diff --git a/azure/cosmos/agent_memory/thresholds.py b/azure/cosmos/agent_memory/thresholds.py index 6d24dce..1f29056 100644 --- a/azure/cosmos/agent_memory/thresholds.py +++ b/azure/cosmos/agent_memory/thresholds.py @@ -4,10 +4,15 @@ InProcess and Durable backends fire on the same turn boundaries by default. 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 +import math import os from typing import Optional @@ -17,6 +22,37 @@ DEFAULT_FACT_EXTRACTION_EVERY_N = 1 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 +# immutable episode. EPISODE_EVAL_EVERY_N is the cheap boundary-evaluation +# cadence (how often we CHECK for a boundary), NOT how often we create an +# episode. 0 disables episodic memory entirely. When enabled, boundaries are +# detected automatically from an idle time-gap, a topic-drift shift, or a +# max-size safety cap - the caller never has to signal "session end". Default 4 +# checks the open segment every 4 turns: responsive enough to close a boundary +# promptly while keeping the per-check overhead (a turn query + drift embeddings) +# low. Set to 0 to disable episodic memory. +DEFAULT_EPISODE_EVAL_EVERY_N = 4 +# A gap larger than this many seconds between two consecutive turns closes the +# open episode (natural session / idle boundary). +DEFAULT_EPISODE_IDLE_GAP_SECONDS = 1800 +# Cosine distance from the open segment's centroid past which a new turn counts +# as a topic/goal shift and closes the prior episode. Requires embeddings. +# Default 0 (OFF): episodes are segmented purely by the idle time-gap and the +# max-size cap, so a dated multi-session conversation yields one episode per +# session. Set to a positive value (e.g. 0.35) to additionally split a long +# single-session run into per-topic episodes; this is heuristic and adds one +# embedding pass over the open segment per boundary evaluation. +DEFAULT_EPISODE_TOPIC_DRIFT = 0.0 +# Hard cap on an open segment: force a boundary so neither an episode nor its +# extraction prompt grows unbounded during a long single-topic session. +DEFAULT_EPISODE_MAX_TURNS = 40 +# Minimum turns before a drift signal may close an episode, and a floor on all +# natural boundaries: idle-gap and drift boundaries below this many turns are +# suppressed so a lone turn is not emitted as a trivial episode (an explicit +# flush still drains a sub-min trailing segment). The max-size cap is a hard +# ceiling and is not floored. +DEFAULT_EPISODE_MIN_TURNS = 2 DEFAULT_USER_SUMMARY_EVERY_N = 20 # Dedup runs on its own cadence - every Nth extract (NOT every Nth turn), # because dedup is O(N²) over all active facts and dominates per-push cost @@ -106,6 +142,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 default_ttl_for(memory_type: str) -> Optional[int]: """Return the per-type default TTL, or None for 'use container default'. @@ -127,6 +191,32 @@ def get_thread_summary_every_n() -> int: return _parse_threshold("THREAD_SUMMARY_EVERY_N", DEFAULT_THREAD_SUMMARY_EVERY_N) +def get_episode_eval_every_n() -> int: + """Boundary-evaluation cadence in turns. 0 disables episodic memory entirely. + + This is how often we cheaply CHECK the open segment for a boundary, not how + often an episode is created; an episode is created only when a boundary is + actually detected (idle gap, topic drift, or max-size cap). + """ + 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: return _parse_threshold("USER_SUMMARY_EVERY_N", DEFAULT_USER_SUMMARY_EVERY_N) @@ -237,6 +327,11 @@ def get_processor_owner() -> Optional[str]: __all__ = [ "DEFAULT_FACT_EXTRACTION_EVERY_N", "DEFAULT_THREAD_SUMMARY_EVERY_N", + "DEFAULT_EPISODE_EVAL_EVERY_N", + "DEFAULT_EPISODE_IDLE_GAP_SECONDS", + "DEFAULT_EPISODE_TOPIC_DRIFT", + "DEFAULT_EPISODE_MAX_TURNS", + "DEFAULT_EPISODE_MIN_TURNS", "DEFAULT_USER_SUMMARY_EVERY_N", "DEFAULT_DEDUP_EVERY_N", "DEFAULT_DEDUP_POOL_SIZE", @@ -249,6 +344,11 @@ def get_processor_owner() -> Optional[str]: "default_ttl_for", "get_fact_extraction_every_n", "get_thread_summary_every_n", + "get_episode_eval_every_n", + "get_episode_idle_gap_seconds", + "get_episode_topic_drift", + "get_episode_max_turns", + "get_episode_min_turns", "get_user_summary_every_n", "get_dedup_every_n", "get_dedup_pool_size", diff --git a/function_app/orchestrators/extract_memories.py b/function_app/orchestrators/extract_memories.py index 4952b11..dc30ef6 100644 --- a/function_app/orchestrators/extract_memories.py +++ b/function_app/orchestrators/extract_memories.py @@ -107,7 +107,7 @@ def em_Extract(payload: dict) -> dict: recent_k = payload.get("recent_k") if recent_k is None: recent_k = config.get_max_batch_size() - extracted = get_pipeline().extract_memories_dry( + extracted = get_pipeline().extract_memories_durable( user_id=user_id, thread_id=thread_id, recent_k=recent_k, diff --git a/function_app/orchestrators/thread_summary.py b/function_app/orchestrators/thread_summary.py index b76cf09..a149a83 100644 --- a/function_app/orchestrators/thread_summary.py +++ b/function_app/orchestrators/thread_summary.py @@ -49,7 +49,7 @@ def ts_Extract(payload: dict) -> dict: """Generate (or incrementally update) the thread summary body only.""" user_id = payload["user_id"] thread_id = payload["thread_id"] - summary = get_pipeline().generate_thread_summary_dry( + summary = get_pipeline().generate_thread_summary_durable( user_id=user_id, thread_id=thread_id, recent_k=payload.get("limit"), diff --git a/function_app/orchestrators/user_summary.py b/function_app/orchestrators/user_summary.py index abe23da..e4024b5 100644 --- a/function_app/orchestrators/user_summary.py +++ b/function_app/orchestrators/user_summary.py @@ -48,7 +48,7 @@ def UserSummaryOrchestrator(context: df.DurableOrchestrationContext): def us_Extract(payload: dict) -> dict: """Generate a cross-thread user summary body only.""" user_id = payload["user_id"] - summary = get_pipeline().generate_user_summary_dry( + summary = get_pipeline().generate_user_summary_durable( user_id=user_id, recent_k=payload.get("limit"), thread_ids=payload.get("thread_ids") or None, diff --git a/infra/README.md b/infra/README.md index 76bf353..0b17a6b 100644 --- a/infra/README.md +++ b/infra/README.md @@ -6,7 +6,7 @@ This folder provisions everything the Agent Memory Toolkit needs in a single Azu `azd up` creates **all** of the following: -- **Cosmos DB for NoSQL** — serverless account with the `ai_memory` database and this container topology: +- **Cosmos DB for NoSQL** - serverless account with the `ai_memory` database and this container topology: | Container | Purpose | | --- | --- | @@ -15,16 +15,16 @@ This folder provisions everything the Agent Memory Toolkit needs in a single Azu | `memories_summaries` | Thread + user summaries; minimal index, point-read access pattern | | `leases` | Change-feed checkpoints | | `counter` | Atomic counters | -- **AI Foundry** (`Microsoft.CognitiveServices/accounts` with `kind: AIServices`) — with `gpt-4o-mini` and `text-embedding-3-large` deployments -- **User-assigned managed identity (UAMI)** — used by the Function app -- **RBAC role assignments** — Cosmos DB Built-in Data Reader + Contributor, Cognitive Services OpenAI User, Storage Blob/Queue/Table data roles, granted to both the UAMI and the deploying user (full table in [Identity & RBAC](#identity--rbac)) -- **Function app** — Flex Consumption (Python 3.11), Storage account, App Insights, Log Analytics +- **AI Foundry** (`Microsoft.CognitiveServices/accounts` with `kind: AIServices`) - with `gpt-4o-mini` and `text-embedding-3-large` deployments +- **User-assigned managed identity (UAMI)** - used by the Function app +- **RBAC role assignments** - Cosmos DB Built-in Data Reader + Contributor, Cognitive Services OpenAI User, Storage Blob/Queue/Table data roles, granted to both the UAMI and the deploying user (full table in [Identity & RBAC](#identity--rbac)) +- **Function app** - Flex Consumption (Python 3.11), Storage account, App Insights, Log Analytics -The Function app is **always provisioned**, even if you plan to use `InProcessProcessor` only. Flex Consumption is pay-per-execution — at zero traffic the Function app is essentially free (idle cost is the Storage account, ~$0.05/month). The Function app sits idle and unused for in-process workloads. +The Function app is **always provisioned**, even if you plan to use `InProcessProcessor` only. Flex Consumption is pay-per-execution - at zero traffic the Function app is essentially free (idle cost is the Storage account, ~$0.05/month). The Function app sits idle and unused for in-process workloads. > Advanced escape hatch: set `azd env set DEPLOY_FUNCTION_APP false` and run `azd provision` to skip the Function app + its supporting resources entirely. Not recommended unless you have a strong reason. See [SDK-only mode](#sdk-only-mode-skip-the-function-app) below for the full procedure. -> **Bring-your-own-resources is not supported.** If you already have a Cosmos account or AI Foundry account you want to reuse, point the SDK and Function app at them via the standard `COSMOS_DB_ENDPOINT` / `AI_FOUNDRY_ENDPOINT` environment variables and skip `azd up` entirely — you only need the Bicep when you want the toolkit to manage the accounts for you. Wiring BYO accounts into this template proved fragile (cross-RG scoping, account-already-exists races, partial RBAC) for low real-world value. +> **Bring-your-own-resources is not supported.** If you already have a Cosmos account or AI Foundry account you want to reuse, point the SDK and Function app at them via the standard `COSMOS_DB_ENDPOINT` / `AI_FOUNDRY_ENDPOINT` environment variables and skip `azd up` entirely - you only need the Bicep when you want the toolkit to manage the accounts for you. Wiring BYO accounts into this template proved fragile (cross-RG scoping, account-already-exists races, partial RBAC) for low real-world value. ## Prereqs @@ -38,7 +38,7 @@ az login azd auth login azd env new memorytoolkit-dev -azd env set AZURE_LOCATION eastus2 # required — subscription-scoped Bicep needs it +azd env set AZURE_LOCATION eastus2 # required - subscription-scoped Bicep needs it # Optional: pin a different region # azd env set AZURE_LOCATION swedencentral @@ -67,7 +67,7 @@ set -a && . ./.azure/memorytoolkit-dev/.env && set +a ## SDK-only mode (skip the Function app) -If you only ever plan to use the in-process `MemoryProcessor` and want to keep the Bicep footprint minimal — no Function app, no Storage account, no App Insights, no Log Analytics: +If you only ever plan to use the in-process `MemoryProcessor` and want to keep the Bicep footprint minimal - no Function app, no Storage account, no App Insights, no Log Analytics: ```bash azd env new memorytoolkit-sdkonly @@ -108,7 +108,7 @@ The `*_DEPLOYMENT_NAME` value is what the SDK and Function app pass as the `mode ## Counter-based trigger configuration (Function app only) -The Function app uses a counter document per `(user_id, thread_id)` to decide when to fire each orchestrator. Every knob is a Bicep parameter bound to an `azd env` variable — `azd up` re-renders them on every deploy, so changing a value is a one-line `azd env set ...` followed by `azd up`. +The Function app uses a counter document per `(user_id, thread_id)` to decide when to fire each orchestrator. Every knob is a Bicep parameter bound to an `azd env` variable - `azd up` re-renders them on every deploy, so changing a value is a one-line `azd env set ...` followed by `azd up`. | `azd env` variable | Bicep param | Default | Effect | |---|---|---|---| @@ -132,16 +132,16 @@ azd up # re-runs provisioning and pushes new App Settings ## Identity & RBAC -Every data-plane role is granted at the resource (account) scope — never at subscription or RG level. `principalType` is set explicitly on every standard `Microsoft.Authorization/roleAssignments` so first-deploy from a freshly-created service principal succeeds without the usual "PrincipalNotFound" RBAC race. +Every data-plane role is granted at the resource (account) scope - never at subscription or RG level. `principalType` is set explicitly on every standard `Microsoft.Authorization/roleAssignments` so first-deploy from a freshly-created service principal succeeds without the usual "PrincipalNotFound" RBAC race. | Resource | Built-in role | Granted to | Why | |---|---|---|---| -| Cosmos DB account | `00000000-0000-0000-0000-000000000001` — Cosmos DB Built-in Data Reader | UAMI + deploying user | Explicit read-only scope. Granted alongside Data Contributor so downstream consumers (audit dashboards, analytics jobs) can run as the same identity but be validated by security review as needing only the Reader scope. Cosmos uses its own `sqlRoleAssignments` resource type which does not accept `principalType` (Cosmos enforces it internally via `principalId`). | -| Cosmos DB account | `00000000-0000-0000-0000-000000000002` — Cosmos DB Built-in Data Contributor | UAMI + deploying user | Data-plane reads/writes from Function app + local samples. | -| AI Foundry account | `5e0bd9bd-7b93-4f28-af87-19fc36ad61bd` — Cognitive Services OpenAI User | UAMI + deploying user | Inference (chat + embeddings) from Function app + local samples. | -| Storage account | `b7e6dc6d-f1e8-4753-8033-0f276bb0955b` — Storage Blob Data Owner | UAMI + deploying user | Function-app deployment-from-blob, Durable history blobs, local sample blob inspection. Owner (not Contributor) keeps `azd deploy` symmetric with manual ops scripts that may need lease/ACL operations. | -| Storage account | `974c5e8b-45b9-4653-ba55-5f855dd0fb88` — Storage Queue Data Contributor | UAMI only | Durable Functions task hub (default Azure Storage provider) uses Queues for orchestration messages. Without this, the very first orchestration start returns 403 even though Blob is fine. | -| Storage account | `0a9a7e1f-b9d0-4cc4-a60d-0319b160aaa3` — Storage Table Data Contributor | UAMI only | Durable Functions history Tables. Same 403 symptom as queues if missing. | +| Cosmos DB account | `00000000-0000-0000-0000-000000000001` - Cosmos DB Built-in Data Reader | UAMI + deploying user | Explicit read-only scope. Granted alongside Data Contributor so downstream consumers (audit dashboards, analytics jobs) can run as the same identity but be validated by security review as needing only the Reader scope. Cosmos uses its own `sqlRoleAssignments` resource type which does not accept `principalType` (Cosmos enforces it internally via `principalId`). | +| Cosmos DB account | `00000000-0000-0000-0000-000000000002` - Cosmos DB Built-in Data Contributor | UAMI + deploying user | Data-plane reads/writes from Function app + local samples. | +| AI Foundry account | `5e0bd9bd-7b93-4f28-af87-19fc36ad61bd` - Cognitive Services OpenAI User | UAMI + deploying user | Inference (chat + embeddings) from Function app + local samples. | +| Storage account | `b7e6dc6d-f1e8-4753-8033-0f276bb0955b` - Storage Blob Data Owner | UAMI + deploying user | Function-app deployment-from-blob, Durable history blobs, local sample blob inspection. Owner (not Contributor) keeps `azd deploy` symmetric with manual ops scripts that may need lease/ACL operations. | +| Storage account | `974c5e8b-45b9-4653-ba55-5f855dd0fb88` - Storage Queue Data Contributor | UAMI only | Durable Functions task hub (default Azure Storage provider) uses Queues for orchestration messages. Without this, the very first orchestration start returns 403 even though Blob is fine. | +| Storage account | `0a9a7e1f-b9d0-4cc4-a60d-0319b160aaa3` - Storage Table Data Contributor | UAMI only | Durable Functions history Tables. Same 403 symptom as queues if missing. | All Storage roles are skipped when `DEPLOY_FUNCTION_APP=false` (no Storage account = nothing to grant on). @@ -166,14 +166,14 @@ azd pipeline config | Gotcha | Mitigation | | --- | --- | | First-time provisioning is slow (8–15 min for Cosmos + AI Foundry + Function app) | `azd up` shows progress; just wait | -| `location` property missing — subscription-scoped Bicep requires `AZURE_LOCATION` | Always run `azd env set AZURE_LOCATION eastus2` after `azd env new` | -| AI Foundry region constraints — many regions don't have all features / models | Default `AZURE_LOCATION=eastus2`; supported: `eastus2`, `swedencentral`, `westus3`, `eastus` | -| Model deployment quota — fails if the subscription has zero quota for the model in the chosen region | Request quota or change region; error from Azure points to the right doc | -| Cosmos free-tier limit (one per subscription) | Default is **serverless** — no idle cost, no free-tier conflict | -| AAD propagation — RBAC takes 30–90s; the Function app may briefly 403 on its first invocation after deploy | Retry after a minute. `dependsOn` chains in Bicep ensure roles exist before the Function app starts | -| Resource naming rules — Storage ≤24 chars lowercase, AI Foundry has its own | Naming uses `take(uniqueString(...), 13)` and `toLower()` to satisfy all rules | - -## Architecture choice — AI Foundry +| `location` property missing - subscription-scoped Bicep requires `AZURE_LOCATION` | Always run `azd env set AZURE_LOCATION eastus2` after `azd env new` | +| AI Foundry region constraints - many regions don't have all features / models | Default `AZURE_LOCATION=eastus2`; supported: `eastus2`, `swedencentral`, `westus3`, `eastus` | +| Model deployment quota - fails if the subscription has zero quota for the model in the chosen region | Request quota or change region; error from Azure points to the right doc | +| Cosmos free-tier limit (one per subscription) | Default is **serverless** - no idle cost, no free-tier conflict | +| AAD propagation - RBAC takes 30–90s; the Function app may briefly 403 on its first invocation after deploy | Retry after a minute. `dependsOn` chains in Bicep ensure roles exist before the Function app starts | +| Resource naming rules - Storage ≤24 chars lowercase, AI Foundry has its own | Naming uses `take(uniqueString(...), 13)` and `toLower()` to satisfy all rules | + +## Architecture choice - AI Foundry The Bicep uses a single `Microsoft.CognitiveServices/accounts` resource with `kind: AIServices` (named `aif-`) instead of the full AI Foundry hub + project + ML workspace. The AIServices account exposes the same Azure OpenAI endpoint and supports the same `Cognitive Services OpenAI User` RBAC role, which is everything the toolkit needs for embeddings and chat completions. This avoids the extra Storage / Key Vault / App Insights / ML workspace resources a hub-style deployment would create. This is the pattern used by most current `azd`-based Microsoft samples (e.g. `azure-search-openai-demo`, `openai-chat-app-quickstart`). diff --git a/infra/main.bicep b/infra/main.bicep index 4c77f8a..eaf3573 100644 --- a/infra/main.bicep +++ b/infra/main.bicep @@ -1,4 +1,4 @@ -// Agent Memory Toolkit — main entry point. +// Agent Memory Toolkit - main entry point. // See infra/README.md for architecture and operational knobs. targetScope = 'subscription' @@ -30,7 +30,7 @@ param principalId string = '' ]) param principalType string = 'User' -@description('Whether to deploy the Function app. Defaults to true. Set false only if you have a strong reason to skip it (Flex Consumption is pay-per-execution — idle cost is ~$0).') +@description('Whether to deploy the Function app. Defaults to true. Set false only if you have a strong reason to skip it (Flex Consumption is pay-per-execution - idle cost is ~$0).') param deployFunctionApp bool = true @description('Cosmos database name.') diff --git a/infra/modules/ai-foundry-rbac.bicep b/infra/modules/ai-foundry-rbac.bicep index 267294d..902f9b6 100644 --- a/infra/modules/ai-foundry-rbac.bicep +++ b/infra/modules/ai-foundry-rbac.bicep @@ -2,10 +2,10 @@ // principalId) on a single AI Foundry / Cognitive Services account. // // This module must be scoped to the resource group that contains the AI -// Foundry account — even when that's not the toolkit's own RG. +// Foundry account - even when that's not the toolkit's own RG. // // Built-in role: -// - 5e0bd9bd-7b93-4f28-af87-19fc36ad61bd — Cognitive Services OpenAI User. +// - 5e0bd9bd-7b93-4f28-af87-19fc36ad61bd - Cognitive Services OpenAI User. @description('Name of the AI Foundry account (must already exist in this module\'s scope).') param aiFoundryAccountName string diff --git a/infra/modules/ai-foundry.bicep b/infra/modules/ai-foundry.bicep index 26050e0..eba3323 100644 --- a/infra/modules/ai-foundry.bicep +++ b/infra/modules/ai-foundry.bicep @@ -1,5 +1,5 @@ // AI Foundry (Cognitive Services kind=AIServices) account + chat + embedding -// model deployments. Single-file module — no `existing` keyword tricks because +// model deployments. Single-file module - no `existing` keyword tricks because // the account is always created fresh by this template. // // We use a single Microsoft.CognitiveServices/accounts resource with @@ -73,7 +73,7 @@ resource account 'Microsoft.CognitiveServices/accounts@2024-10-01' = { // --- Model deployments ---------------------------------------------------- // -// Deployments are serialized via dependsOn — Cognitive Services rejects +// Deployments are serialized via dependsOn - Cognitive Services rejects // concurrent deployment writes on the same account. resource llmDeployment 'Microsoft.CognitiveServices/accounts/deployments@2024-10-01' = { diff --git a/infra/modules/cosmos-rbac.bicep b/infra/modules/cosmos-rbac.bicep index 8373252..79d0359 100644 --- a/infra/modules/cosmos-rbac.bicep +++ b/infra/modules/cosmos-rbac.bicep @@ -3,11 +3,11 @@ // // Cosmos data-plane access is granted via sqlRoleAssignments (children of // the account), so this module must be scoped to the resource group that -// contains the Cosmos account — even when that's not the toolkit's own RG. +// contains the Cosmos account - even when that's not the toolkit's own RG. // // Built-in roles: -// - 00000000-0000-0000-0000-000000000001 — Cosmos DB Built-in Data Reader. -// - 00000000-0000-0000-0000-000000000002 — Cosmos DB Built-in Data Contributor. +// - 00000000-0000-0000-0000-000000000001 - Cosmos DB Built-in Data Reader. +// - 00000000-0000-0000-0000-000000000002 - Cosmos DB Built-in Data Contributor. // // Both roles are granted so the principal has explicit read-only access in // addition to read/write. Useful for downstream consumers (audit dashboards, diff --git a/infra/modules/cosmos.bicep b/infra/modules/cosmos.bicep index f11098e..74ef0e7 100644 --- a/infra/modules/cosmos.bicep +++ b/infra/modules/cosmos.bicep @@ -1,5 +1,5 @@ // Cosmos DB NoSQL serverless account + database + containers for the Agent -// Memory Toolkit. Single-file module — no `existing` keyword tricks because +// Memory Toolkit. Single-file module - no `existing` keyword tricks because // the account is always created fresh by this template. @description('Name of the Cosmos account to create.') diff --git a/infra/modules/functions.bicep b/infra/modules/functions.bicep index aad4f63..55e50ab 100644 --- a/infra/modules/functions.bicep +++ b/infra/modules/functions.bicep @@ -59,7 +59,7 @@ param embeddingDimensions int = 1536 @description('LLM model deployment name.') param chatDeploymentName string = 'gpt-4o-mini' -@description('Azure OpenAI REST API version pinned for both chat and embedding clients. Always supplied by main.bicep — declared here without a default so the wiring stays explicit.') +@description('Azure OpenAI REST API version pinned for both chat and embedding clients. Always supplied by main.bicep - declared here without a default so the wiring stays explicit.') param azureOpenAiApiVersion string // --- Function-app threshold / batching knobs ------------------------------ @@ -67,7 +67,7 @@ param azureOpenAiApiVersion string // All knobs are surfaced as Bicep params in main.bicep (bound to // `${THREAD_SUMMARY_EVERY_N=10}` etc. in main.parameters.json) so customers // can override them via `azd env set ...` before `azd up`. The defaults live -// in main.bicep — these module params are declared without defaults so +// in main.bicep - these module params are declared without defaults so // main.bicep stays the single source of truth. @description('Run thread-summary orchestration every N turns within a (user_id, thread_id). 0 = disabled.') diff --git a/infra/modules/storage-rbac.bicep b/infra/modules/storage-rbac.bicep index f91b30f..4fd0162 100644 --- a/infra/modules/storage-rbac.bicep +++ b/infra/modules/storage-rbac.bicep @@ -6,9 +6,9 @@ // scoped locally. // // Built-in roles: -// - b7e6dc6d-f1e8-4753-8033-0f276bb0955b — Storage Blob Data Owner. -// - 974c5e8b-45b9-4653-ba55-5f855dd0fb88 — Storage Queue Data Contributor. -// - 0a9a7e1f-b9d0-4cc4-a60d-0319b160aaa3 — Storage Table Data Contributor. +// - b7e6dc6d-f1e8-4753-8033-0f276bb0955b - Storage Blob Data Owner. +// - 974c5e8b-45b9-4653-ba55-5f855dd0fb88 - Storage Queue Data Contributor. +// - 0a9a7e1f-b9d0-4cc4-a60d-0319b160aaa3 - Storage Table Data Contributor. // // Durable Functions (default Azure Storage provider) talks to Storage Queues // + Tables under the function app's identity. Without those two roles, the diff --git a/tests/integration/test_episodic_pipeline.py b/tests/integration/test_episodic_pipeline.py new file mode 100644 index 0000000..a8feebb --- /dev/null +++ b/tests/integration/test_episodic_pipeline.py @@ -0,0 +1,231 @@ +"""Live episodic-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 time + +import pytest + +from azure.cosmos.agent_memory import CosmosMemoryClient +from tests.conftest import INTEGRATION_ENABLED + +pytestmark = [ + pytest.mark.integration, + pytest.mark.skipif( + not INTEGRATION_ENABLED, + reason="Set AGENT_MEMORY_RUN_INTEGRATION=true", + ), +] + +VALID_OUTCOMES = {"successful", "partially_successful", "failed", "abandoned", "unknown"} + + +@pytest.fixture(scope="module") +def episodic_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_hiking_thread(mem: CosmosMemoryClient, user_id: str, thread_id: str) -> set[str]: + turns = [ + ( + "user", + "On Friday 2026-07-17, Priya, Marco, and I planned a weekend hiking trip to " + "Mount Rainier's Skyline Trail for Saturday morning.", + ), + ( + "agent", + "That sounds like a clear plan: Saturday morning on Skyline Trail with Priya and Marco.", + ), + ( + "user", + "We left Seattle at 6:30 AM on Saturday, stopped in Ashford for coffee, and reached " + "Paradise before the main parking lot filled.", + ), + ( + "agent", + "Getting to Paradise early likely helped the group start the hike before the crowds.", + ), + ( + "user", + "Near Panorama Point, Marco slipped on a wet rock and scraped his knee, so Priya used " + "the small first-aid kit I had packed.", + ), + ( + "user", + "After a short rest, Marco felt okay, and we continued slowly to see the wildflowers " + "and the Nisqually Glacier views.", + ), + ( + "agent", + "The first-aid kit turned the mishap into a manageable pause rather than ending the hike.", + ), + ( + "user", + "We got back to Seattle by 7 PM Saturday, tired but happy, and decided next time we " + "would bring trekking poles for the steeper wet sections.", + ), + ] + turn_ids = set() + for role, content in turns: + turn_ids.add( + mem.add_cosmos( + user_id=user_id, + role=role, + content=content, + memory_type="turn", + thread_id=thread_id, + ) + ) + return turn_ids + + +def test_live_episodic_extraction_and_blended_search( + episodic_memory, + unique_user_id, + unique_thread_id, +): + try: + turn_ids = _write_hiking_thread(episodic_memory, unique_user_id, unique_thread_id) + time.sleep(1) + + stats = episodic_memory.extract_episodes(unique_user_id, unique_thread_id, flush=True) + assert stats.get("episodes", 0) >= 1, f"Expected at least one extracted episode, got {stats}" + + episodes = episodic_memory.get_episodes(unique_user_id) + assert len(episodes) >= 1 + episode = episodes[0] + assert episode.get("content") + assert episode.get("title") + assert episode.get("started_at") + + events = episode.get("events") or [] + assert events, f"Expected at least one grounded event, got {episode}" + assert any(set(event.get("source_turn_ids") or []) & turn_ids for event in events), ( + f"Expected event source_turn_ids to reference written turns {turn_ids}, got {events}" + ) + + outcome = episode.get("outcome") + assert outcome is None or outcome.get("status") in VALID_OUTCOMES + + results = episodic_memory.search_cosmos( + search_terms="the hiking trip", + user_id=unique_user_id, + include_episodes=True, + ) + assert any(result.get("type") == "episodic" for result in results), results + finally: + _delete_user_records(episodic_memory, unique_user_id) + + +def test_live_episodic_lessons_feed_procedural_synthesis( + episodic_memory, + unique_user_id, + unique_thread_id, +): + """Episodes carry first-class ``lessons`` that must flow into procedural synthesis. + + Uses a fresh user with no extracted facts, so the only possible source for the + synthesized prompt is episodic lessons - this isolates the episodic->procedural + seam end-to-end against live Cosmos (real ``IS_DEFINED(c.lessons)`` filtering). + """ + try: + _write_hiking_thread(episodic_memory, unique_user_id, unique_thread_id) + time.sleep(1) + + stats = episodic_memory.extract_episodes(unique_user_id, unique_thread_id, flush=True) + assert stats.get("episodes", 0) >= 1, f"Expected at least one extracted episode, got {stats}" + + episodes = episodic_memory.get_episodes(unique_user_id) + lesson_bearing = { + ep["id"] + for ep in episodes + if isinstance(ep.get("lessons"), list) + and any(isinstance(lesson, str) and lesson.strip() for lesson in ep.get("lessons", [])) + } + assert lesson_bearing, ( + "Expected extract_episodes to write at least one episode with first-class lessons, " + f"got {[ep.get('lessons') for ep in episodes]}" + ) + + result = episodic_memory.synthesize_procedural(unique_user_id, force=True) + assert result.get("status") == "synthesized", result + proc = result.get("procedural") or {} + assert isinstance(proc.get("content"), str) and proc["content"].strip(), proc + assert set(proc.get("source_episodic_ids") or []) == lesson_bearing, ( + "Expected every lesson-bearing episode to feed procedural synthesis; " + f"lesson_bearing={lesson_bearing} source_episodic_ids={proc.get('source_episodic_ids')}" + ) + finally: + _delete_user_records(episodic_memory, unique_user_id) diff --git a/tests/unit/aio/processors/test_inprocess.py b/tests/unit/aio/processors/test_inprocess.py index c933e7d..03cb0ed 100644 --- a/tests/unit/aio/processors/test_inprocess.py +++ b/tests/unit/aio/processors/test_inprocess.py @@ -91,6 +91,21 @@ async def test_process_extract_memories_invokes_pipeline_and_filters_to_ints(): assert result == {"fact_count": 3} +@pytest.mark.asyncio +async def test_process_extract_episodes_invokes_pipeline_and_filters_to_ints(): + pipeline = AsyncMock() + pipeline.extract_episodes.return_value = { + "episodes": 2, + "non_int_field": "skip me", + } + + proc = AsyncInProcessProcessor(pipeline=pipeline) + result = await proc.process_extract_episodes(user_id="u", thread_id="t") + + pipeline.extract_episodes.assert_called_once_with("u", "t") + assert result == {"episodes": 2} + + @pytest.mark.asyncio async def test_process_thread_summary_invokes_pipeline_and_returns_dict(): pipeline = AsyncMock() diff --git a/tests/unit/aio/processors/test_protocol_satisfaction.py b/tests/unit/aio/processors/test_protocol_satisfaction.py index 2a23515..8ae2f76 100644 --- a/tests/unit/aio/processors/test_protocol_satisfaction.py +++ b/tests/unit/aio/processors/test_protocol_satisfaction.py @@ -36,6 +36,14 @@ async def process_extract_memories( ) -> dict[str, int]: return {} + async def process_extract_episodes( + self, + *, + user_id: str, + thread_id: str, + ) -> dict[str, int]: + return {} + async def process_thread_summary( self, *, diff --git a/tests/unit/aio/services/test_dedup_vector_async.py b/tests/unit/aio/services/test_dedup_vector_async.py index 122126c..ed679ed 100644 --- a/tests/unit/aio/services/test_dedup_vector_async.py +++ b/tests/unit/aio/services/test_dedup_vector_async.py @@ -57,12 +57,15 @@ def _episode(eid: str, content: str) -> dict: "confidence": 0.8, "salience": 0.7, "tags": ["sys:episodic", "sys:dup-candidate"], - "metadata": { - "scope_type": "project", - "scope_value": "CI", - "lesson": content, - "outcome_valence": "positive", - }, + "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], } diff --git a/tests/unit/aio/services/test_episode_boundary_async.py b/tests/unit/aio/services/test_episode_boundary_async.py new file mode 100644 index 0000000..5961637 --- /dev/null +++ b/tests/unit/aio/services/test_episode_boundary_async.py @@ -0,0 +1,233 @@ +"""Boundary-based episodic segmentation (async mirror of test_episode_boundary).""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from azure.cosmos.agent_memory.aio.services.pipeline import AsyncPipelineService +from tests.unit.aio.services.test_extract_episodes_async import _AsyncTrackingStore +from tests.unit.services.test_extract_dry import ( + _async_containers_for_store, + _AsyncChat, + _AsyncEmbeddings, + _AsyncStore, +) +from tests.unit.services.test_extract_episodes import _episode + + +def _turn_at(i: int, minute: int, *, content: str | None = None) -> dict[str, Any]: + return { + "id": f"turn-{i}", + "user_id": "u1", + "thread_id": "t1", + "role": "user", + "type": "turn", + "content": content if content is not None else f"Turn {i}: routine content", + "created_at": f"2025-01-01T00:{minute:02d}:00+00:00", + } + + +class _AsyncDriftEmbeddings(_AsyncEmbeddings): + async def generate_batch(self, texts: list[str]) -> list[list[float]]: + self.calls.append(list(texts)) + return [[0.0, 1.0] if "B:" in text else [1.0, 0.0] for text in texts] + + +def _service( + turns: list[dict[str, Any]], + responses: list[dict[str, Any]] | None = None, + *, + embeddings: _AsyncEmbeddings | None = None, +) -> tuple[AsyncPipelineService, _AsyncTrackingStore, _AsyncStore, _AsyncChat]: + memories = _AsyncTrackingStore([]) + turns_store = _AsyncStore(turns) + chat = _AsyncChat(responses or [{"episodes": [_episode()]} for _ in range(20)]) + service = AsyncPipelineService( + memories, + chat, + embeddings or _AsyncEmbeddings(), + containers=_async_containers_for_store(memories, turns_store=turns_store), + ) + + # Stub _run_prompty so boundary tests return canned episode JSON without + # invoking the real prompty template renderer (fast and hermetic). + async def _fake_run_prompty(*a: Any, **k: Any) -> str: + return await chat.generate([]) + + service._run_prompty = _fake_run_prompty # type: ignore[assignment] + return service, memories, turns_store, chat + + +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")) + + +@pytest.mark.asyncio +async def test_time_gap_closes_prior_episode_and_leaves_tail_open(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) + + result = await service.extract_episodes("u1", "t1") + + assert result == {"episodes": 1} + assert len(_episodes(memories)) == 1 + assert _stamped(turns_store) == ["turn-1", "turn-2"] + + +@pytest.mark.asyncio +async def test_reevaluation_is_idempotent_via_watermark(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, _, _ = _service(turns) + + first = await service.extract_episodes("u1", "t1") + second = await service.extract_episodes("u1", "t1") + + assert first == {"episodes": 1} + assert second == {"episodes": 0} + assert len(_episodes(memories)) == 1 + + +@pytest.mark.asyncio +async def test_flush_drains_open_tail(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") + flushed = await service.extract_episodes("u1", "t1", flush=True) + + assert flushed == {"episodes": 1} + assert len(_episodes(memories)) == 2 + assert _stamped(turns_store) == ["turn-1", "turn-2", "turn-3", "turn-4"] + assert await service.extract_episodes("u1", "t1", flush=True) == {"episodes": 0} + + +@pytest.mark.asyncio +async def test_max_turns_forces_a_boundary(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_IDLE_GAP_SECONDS", "0") + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + monkeypatch.setenv("EPISODE_MAX_TURNS", "2") + turns = [_turn_at(i, i) for i in range(1, 5)] + service, memories, turns_store, _ = _service(turns) + + result = await service.extract_episodes("u1", "t1") + + assert result == {"episodes": 2} + assert _stamped(turns_store) == ["turn-1", "turn-2", "turn-3", "turn-4"] + + +@pytest.mark.asyncio +async def test_topic_drift_closes_episode(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_IDLE_GAP_SECONDS", "0") + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0.5") + monkeypatch.setenv("EPISODE_MIN_TURNS", "2") + turns = [ + _turn_at(1, 1, content="A: apples and orchards"), + _turn_at(2, 2, content="A: more about apples"), + _turn_at(3, 3, content="B: rockets and orbits"), + _turn_at(4, 4, content="B: more about rockets"), + ] + service, memories, turns_store, _ = _service(turns, embeddings=_AsyncDriftEmbeddings()) + + result = await service.extract_episodes("u1", "t1") + + assert result == {"episodes": 1} + assert _stamped(turns_store) == ["turn-1", "turn-2"] + + +@pytest.mark.asyncio +async def test_no_boundary_keeps_segment_open_without_calling_the_llm(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_IDLE_GAP_SECONDS", "1800") + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + monkeypatch.setenv("EPISODE_MAX_TURNS", "40") + turns = [_turn_at(1, 1), _turn_at(2, 2), _turn_at(3, 3)] + service, memories, turns_store, chat = _service(turns) + + result = await service.extract_episodes("u1", "t1") + + assert result == {"episodes": 0} + assert _episodes(memories) == [] + assert _stamped(turns_store) == [] + assert chat.calls == 0 + + +@pytest.mark.asyncio +async def test_idle_gap_below_min_turns_does_not_close_episode(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_IDLE_GAP_SECONDS", "120") + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + monkeypatch.setenv("EPISODE_MAX_TURNS", "40") + monkeypatch.setenv("EPISODE_MIN_TURNS", "2") + turns = [_turn_at(1, 1), _turn_at(2, 30), _turn_at(3, 31)] + service, memories, turns_store, _ = _service(turns) + + result = await service.extract_episodes("u1", "t1") + + assert result == {"episodes": 0} + assert _episodes(memories) == [] + assert _stamped(turns_store) == [] + + +@pytest.mark.asyncio +async def test_idle_gap_below_min_turns_still_flushes_as_one_episode(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_IDLE_GAP_SECONDS", "120") + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + monkeypatch.setenv("EPISODE_MAX_TURNS", "40") + monkeypatch.setenv("EPISODE_MIN_TURNS", "2") + turns = [_turn_at(1, 1), _turn_at(2, 30), _turn_at(3, 31)] + service, memories, turns_store, _ = _service(turns) + + result = await service.extract_episodes("u1", "t1", flush=True) + + assert result == {"episodes": 1} + assert _stamped(turns_store) == ["turn-1", "turn-2", "turn-3"] + + +@pytest.mark.asyncio +async def test_extract_episodes_defers_segment_on_retryable_error(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_IDLE_GAP_SECONDS", "120") + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + monkeypatch.setenv("EPISODE_MAX_TURNS", "40") + turns = [_turn_at(1, 1), _turn_at(2, 2), _turn_at(3, 30)] # gap closes [1,2] + service, memories, turns_store, _ = _service(turns) + + async def _boom(*a: Any, **k: Any) -> str: + raise RuntimeError("transient rate limit 429") + + service._run_prompty = _boom # type: ignore[assignment] + + result = await service.extract_episodes("u1", "t1") + + assert result == {"episodes": 0} + assert _episodes(memories) == [] + assert _stamped(turns_store) == [] # un-stamped -> retried next run + + +@pytest.mark.asyncio +async def test_extract_episodes_quarantines_segment_on_non_retryable_error(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_IDLE_GAP_SECONDS", "120") + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + monkeypatch.setenv("EPISODE_MAX_TURNS", "40") + turns = [_turn_at(1, 1), _turn_at(2, 2), _turn_at(3, 30)] + service, memories, turns_store, _ = _service(turns) + + async def _boom(*a: Any, **k: Any) -> str: + raise RuntimeError("content_filter triggered") + + service._run_prompty = _boom # type: ignore[assignment] + + result = await service.extract_episodes("u1", "t1") + + assert result == {"episodes": 0} + assert _episodes(memories) == [] + assert _stamped(turns_store) == ["turn-1", "turn-2"] # quarantined + advanced diff --git a/tests/unit/aio/services/test_episodic_retrieval_async.py b/tests/unit/aio/services/test_episodic_retrieval_async.py new file mode 100644 index 0000000..4b0973b --- /dev/null +++ b/tests/unit/aio/services/test_episodic_retrieval_async.py @@ -0,0 +1,195 @@ +from __future__ import annotations + +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock + +from azure.cosmos.agent_memory._container_routing import ContainerKey +from azure.cosmos.agent_memory.aio.cosmos_memory_client import AsyncCosmosMemoryClient +from azure.cosmos.agent_memory.aio.store import AsyncMemoryStore + + +class AsyncIterator: + def __init__(self, items): + self._items = iter(items) + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self._items) + except StopIteration: + raise StopAsyncIteration + + +def _connected_client() -> tuple[AsyncCosmosMemoryClient, MagicMock]: + client = AsyncCosmosMemoryClient(use_default_credential=False) + memories = MagicMock() + turns = MagicMock() + summaries = MagicMock() + for container in (memories, turns, summaries): + container.query_items = MagicMock(return_value=AsyncIterator([])) + container.upsert_item = AsyncMock() + client._memories_container_client = memories + client._turns_container_client = turns + client._summaries_container_client = summaries + return client, memories + + +def _containers(*, memories=None): + return { + ContainerKey.TURNS: MagicMock(), + ContainerKey.MEMORIES: memories if memories is not None else MagicMock(), + ContainerKey.SUMMARIES: MagicMock(), + } + + +def _params_by_name(call_kwargs): + return {p["name"]: p["value"] for p in call_kwargs["parameters"]} + + +async def test_async_get_episodes_user_wide_newest_first(): + mem, memories = _connected_client() + docs = [ + {"id": "ep-new", "type": "episodic", "content": "new"}, + {"id": "ep-old", "type": "episodic", "content": "old"}, + ] + memories.query_items = MagicMock(return_value=AsyncIterator(docs)) + + results = await mem.get_episodes(user_id="u1", recent_k=2) + + assert [doc["id"] for doc in results] == ["ep-new", "ep-old"] + call_kwargs = memories.query_items.call_args.kwargs + assert "SELECT TOP @recent_k * FROM c" in call_kwargs["query"] + assert "c.user_id = @user_id" in call_kwargs["query"] + assert "c.type = @type" in call_kwargs["query"] + assert "ORDER BY c.created_at DESC" in call_kwargs["query"] + assert "partition_key" not in call_kwargs + params = _params_by_name(call_kwargs) + assert params["@user_id"] == "u1" + assert params["@type"] == "episodic" + assert params["@recent_k"] == 2 + + +async def test_async_search_cosmos_base_is_facts_only_no_episodes_without_optin(): + # Base search is facts-only; without include_episodes no episodic query runs. + mem, _ = _connected_client() + store = MagicMock() + store.search = AsyncMock(return_value=[{"content": "fact A", "type": "fact"}]) + store.search_summaries = AsyncMock(return_value=[{"content": "summary B", "type": "thread_summary"}]) + store.search_episodic = AsyncMock(return_value=[]) + store.search_turns = AsyncMock(return_value=[{"content": "turn D", "type": "turn"}]) + mem._get_store = MagicMock(return_value=store) + + results = await mem.search_cosmos( + "weather", + user_id="u1", + thread_id="t1", + top_k=4, + include_summaries=True, + include_turns=True, + ) + + assert [doc["content"] for doc in results] == ["fact A", "summary B", "turn D"] + assert store.search.call_args.kwargs["memory_types"] == ["fact"] + store.search_episodic.assert_not_awaited() + + +async def test_async_search_cosmos_include_episodes_combines_facts_and_episodes_in_base_query(): + mem, _ = _connected_client() + store = MagicMock() + store.search = AsyncMock( + return_value=[ + {"content": "fact A", "type": "fact"}, + {"content": "episode C", "type": "episodic"}, + ] + ) + store.search_episodic = AsyncMock() + store.search_summaries = AsyncMock(return_value=[{"content": "summary B", "type": "thread_summary"}]) + store.search_turns = AsyncMock(return_value=[{"content": "turn D", "type": "turn"}]) + mem._get_store = MagicMock(return_value=store) + + results = await mem.search_cosmos( + "weather", + user_id="u1", + thread_id="t1", + top_k=100, + include_episodes=True, + include_summaries=True, + include_turns=True, + ) + # Combined base (facts + episodes) -> summaries -> turns; one shared budget. + assert [doc["content"] for doc in results] == ["fact A", "episode C", "summary B", "turn D"] + assert store.search.call_args.kwargs["top_k"] == 100 + assert store.search.call_args.kwargs["memory_types"] == ["fact", "episodic"] + store.search_episodic.assert_not_awaited() + + +async def test_async_search_cosmos_include_episodes_combined_query_hits_real_store(): + # Drive the REAL AsyncMemoryStore.search (not a mock) through search_cosmos: + # include_episodes folds "episodic" into a single combined base query that + # applies the caller's tag/salience filters uniformly - no separate episodic + # query, so facts and episodes share one top_k and the same filters. + memories = MagicMock() + memories.query_items = MagicMock(return_value=AsyncIterator([])) + embeddings = MagicMock() + embeddings.generate = AsyncMock(return_value=[0.1, 0.2]) + store = AsyncMemoryStore(containers=_containers(memories=memories), embeddings_client=embeddings) + + mem, _ = _connected_client() + mem._get_store = MagicMock(return_value=store) + + results = await mem.search_cosmos( + "weather", + user_id="u1", + thread_id="t1", + top_k=5, + include_episodes=True, + tags_all=["trip"], + min_salience=0.5, + ) + + assert results == [] + # Exactly one combined base query ran, scoped to fact + episodic, with the + # caller's filters applied uniformly. + assert memories.query_items.call_count == 1 + call_kwargs = memories.query_items.call_args.kwargs + params = _params_by_name(call_kwargs) + type_values = {v for k, v in params.items() if k.startswith("@memory_type")} + assert type_values == {"fact", "episodic"} + assert params["@min_salience"] == 0.5 + assert params["@tag_0"] == "trip" + + +async def test_async_search_episodic_temporal_filters_do_not_rank_by_time(): + memories = MagicMock() + memories.query_items.return_value = AsyncIterator([]) + embeddings = MagicMock() + embeddings.generate = AsyncMock(return_value=[0.1, 0.2]) + store = AsyncMemoryStore(containers=_containers(memories=memories), embeddings_client=embeddings) + created_after = datetime(2026, 1, 1, tzinfo=timezone.utc) + + await store.search_episodic( + user_id="u1", + search_terms="checkout hotel", + top_k=3, + created_after=created_after, + created_before="2026-02-01T00:00:00+00:00", + started_after="2026-01-10T00:00:00+00:00", + ended_before="2026-01-20T00:00:00+00:00", + ) + + call_kwargs = memories.query_items.call_args.kwargs + query = call_kwargs["query"] + assert "c.created_at >= @created_after" in query + assert "c.created_at <= @created_before" in query + assert "c.started_at >= @started_after" in query + assert "c.ended_at <= @ended_before" in query + assert "ORDER BY RANK RRF(VectorDistance(c.embedding, @embedding), FullTextScore(c.content, @kw0, @kw1))" in query + assert "ORDER BY c.created_at" not in query + assert "ORDER BY c.started_at" not in query + params = _params_by_name(call_kwargs) + assert params["@created_after"] == created_after.isoformat() + assert params["@created_before"] == "2026-02-01T00:00:00+00:00" + assert params["@started_after"] == "2026-01-10T00:00:00+00:00" + assert params["@ended_before"] == "2026-01-20T00:00:00+00:00" diff --git a/tests/unit/aio/services/test_extract_episodes_async.py b/tests/unit/aio/services/test_extract_episodes_async.py new file mode 100644 index 0000000..2815a83 --- /dev/null +++ b/tests/unit/aio/services/test_extract_episodes_async.py @@ -0,0 +1,255 @@ +from __future__ import annotations + +import logging +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, + _response, + _turn, +) +from tests.unit.services.test_extract_episodes import _episode + + +class _AsyncTrackingStore(_AsyncStore): + def __init__(self, docs: list[dict[str, Any]]): + super().__init__(docs) + self.supersede_calls: list[dict[str, Any]] = [] + + async def mark_superseded(self, old_doc: dict[str, Any], superseder_id: str, *, reason: str) -> bool: + self.supersede_calls.append({"old_doc": old_doc, "superseder_id": superseder_id, "reason": reason}) + return True + + +def _service( + responses: list[dict[str, Any]], + *, + memories_store: _AsyncTrackingStore | None = None, + embeddings: _AsyncEmbeddings | None = None, +) -> tuple[AsyncPipelineService, _AsyncTrackingStore, _AsyncEmbeddings]: + store = memories_store or _AsyncTrackingStore([]) + embedding_client = embeddings or _AsyncEmbeddings() + turns_store = _AsyncStore([_turn(1), _turn(2)]) + service = AsyncPipelineService( + store, + _AsyncChat(responses), + embedding_client, + containers=_async_containers_for_store(store, turns_store=turns_store), + ) + return service, store, embedding_client + + +@pytest.mark.asyncio +async def test_build_episode_docs_returns_multiple_docs_without_embeddings() -> None: + outcome = {"status": "successful", "description": "The tests passed."} + service, _, embeddings = _service( + [ + { + "episodes": [ + _episode(outcome=outcome), + _episode(title="Planned vacation", summary="The user planned a vacation.", outcome=None), + ] + } + ] + ) + + docs = await service._build_episode_docs("u1", "t1", [_turn(1), _turn(2)], segment_key="seg-1") + + assert [doc["content"] for doc in docs] == [ + "The user fixed flaky CI retries and the tests passed.", + "The user planned a vacation.", + ] + assert all(doc["id"].startswith("ep_") for doc in docs) + assert all("embedding" not in doc for doc in docs) + assert docs[0]["source_turn_ids"] == ["turn-1", "turn-2"] + assert docs[0]["events"][0]["source_turn_ids"] == ["turn-1"] + assert docs[0]["outcome"] == outcome + assert docs[1]["outcome"] is None + assert embeddings.calls == [] + + +@pytest.mark.asyncio +async def test_extract_episodes_embeds_content_persists_append_only(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + service, store, embeddings = _service( + [ + { + "episodes": [ + _episode(), + _episode(title="Planned vacation", summary="The user planned a vacation.", outcome=None), + ] + } + ] + ) + + result = await service.extract_episodes("u1", "t1", flush=True) + + assert result == {"episodes": 2} + assert embeddings.calls == [ + [ + "The user fixed flaky CI retries and the tests passed.", + "The user planned a vacation.", + ] + ] + assert [doc["content"] for doc in store.docs] == [ + "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 store.supersede_calls == [] + assert store.search_calls == [] + + +@pytest.mark.asyncio +async def test_extract_episodes_empty_window_persists_nothing(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + service, store, embeddings = _service([{"episodes": []}]) + + result = await service.extract_episodes("u1", "t1", flush=True) + + assert result == {"episodes": 0} + assert store.docs == [] + assert embeddings.calls == [] + + +@pytest.mark.asyncio +async def test_extract_episodes_skips_malformed_episode_with_warning(caplog, monkeypatch) -> None: + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + service, store, _ = _service( + [ + { + "episodes": [ + {**_episode(), "title": ""}, + _episode(title="Valid episode", summary="The valid episode persisted."), + ] + } + ] + ) + + with caplog.at_level(logging.WARNING): + result = await service.extract_episodes("u1", "t1", flush=True) + + assert result == {"episodes": 1} + assert [doc["title"] for doc in store.docs] == ["Valid episode"] + assert "dropping malformed episode" in caplog.text + + +@pytest.mark.asyncio +async def test_extract_memories_durable_keeps_episodic_empty_regression_guard() -> None: + store = _AsyncTrackingStore([]) + service = AsyncPipelineService( + store, + _AsyncChat([_response()]), + _AsyncEmbeddings(), + containers=_async_containers_for_store(store, turns_store=_AsyncStore([_turn(1)])), + ) + + output = await service.extract_memories_durable("u1", "t1") + + assert output["facts"] + assert output["episodic"] == [] + + +class _AsyncIdUniqueStore(_AsyncTrackingStore): + """Async store enforcing id-uniqueness on create, like Cosmos (409 on dup).""" + + 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) + + +@pytest.mark.asyncio +async def test_build_episode_docs_id_stable_across_summary_text() -> None: + service, _, _ = _service( + [ + {"episodes": [_episode(summary="One phrasing of the CI-retry episode.")]}, + {"episodes": [_episode(summary="A completely different phrasing entirely.")]}, + ] + ) + first = await service._build_episode_docs("u1", "t1", [_turn(1), _turn(2)], segment_key="seg-1") + second = await service._build_episode_docs("u1", "t1", [_turn(1), _turn(2)], segment_key="seg-1") + + assert first[0]["content"] != second[0]["content"] + assert first[0]["id"] == second[0]["id"] + assert first[0]["content_hash"] != second[0]["content_hash"] + + +@pytest.mark.asyncio +async def test_extract_episodes_skips_duplicate_when_segment_reprocessed(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + store = _AsyncIdUniqueStore([]) + turns = _AsyncStore([_turn(1), _turn(2)]) + service = AsyncPipelineService( + store, + _AsyncChat( + [ + {"episodes": [_episode(summary="First run prose.")]}, + {"episodes": [_episode(summary="Second run, different prose.")]}, + ] + ), + _AsyncEmbeddings(), + containers=_async_containers_for_store(store, turns_store=turns), + ) + + assert await service.extract_episodes("u1", "t1", flush=True) == {"episodes": 1} + for turn in turns.docs: + turn.pop("episode_extracted_at", None) + 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 + + +@pytest.mark.asyncio +async def test_build_episode_docs_falls_back_to_segment_times_on_unparseable_llm_times() -> None: + bad = _episode() + bad["started_at"] = "March 9th" + bad["ended_at"] = "2025-01-01T00:02:00+00:00" + service, _, _ = _service([{"episodes": [bad]}]) + + docs = await service._build_episode_docs("u1", "t1", [_turn(1), _turn(2)], segment_key="seg-1") + + assert len(docs) == 1 + assert docs[0]["started_at"] == "2025-01-01T00:01:00+00:00" + assert docs[0]["ended_at"] == "2025-01-01T00:02:00+00:00" + + +@pytest.mark.asyncio +async def test_build_episode_docs_keeps_mixed_tz_llm_times_after_normalization() -> None: + mixed = _episode() + mixed["started_at"] = "2026-03-09" + mixed["ended_at"] = "2026-03-10T09:08:00+00:00" + service, _, _ = _service([{"episodes": [mixed]}]) + + docs = await service._build_episode_docs("u1", "t1", [_turn(1), _turn(2)], segment_key="seg-1") + + assert len(docs) == 1 + assert docs[0]["started_at"] == "2026-03-09" + assert docs[0]["ended_at"] == "2026-03-10T09:08:00+00:00" + + +@pytest.mark.asyncio +async def test_build_episode_docs_keeps_padded_timestamps_and_clamps_scores() -> None: + ep = _episode() + ep["started_at"] = " 2025-01-01T00:01:00+00:00" + ep["ended_at"] = "2025-01-01T00:02:00+00:00 " + ep["salience"] = 1.4 + ep["confidence"] = -0.2 + service, _, _ = _service([{"episodes": [ep]}]) + + docs = await service._build_episode_docs("u1", "t1", [_turn(1), _turn(2)], segment_key="seg-1") + + assert len(docs) == 1 + assert docs[0]["started_at"] == "2025-01-01T00:01:00+00:00" + assert docs[0]["ended_at"] == "2025-01-01T00:02:00+00:00" + assert docs[0]["salience"] == 1.0 + assert docs[0]["confidence"] == 0.0 diff --git a/tests/unit/aio/test_auto_trigger.py b/tests/unit/aio/test_auto_trigger.py index 668b134..f51eb92 100644 --- a/tests/unit/aio/test_auto_trigger.py +++ b/tests/unit/aio/test_auto_trigger.py @@ -13,6 +13,7 @@ import pytest from azure.cosmos.exceptions import CosmosResourceNotFoundError +from azure.cosmos.agent_memory.aio.auto_trigger import maybe_trigger_steps from azure.cosmos.agent_memory.aio.cosmos_memory_client import AsyncCosmosMemoryClient from azure.cosmos.agent_memory.aio.processors import AsyncInProcessProcessor @@ -56,6 +57,7 @@ class TestAsyncAutoTriggerNonBlocking: async def test_push_to_cosmos_does_not_await_auto_trigger(self, monkeypatch): monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") + monkeypatch.setenv("EPISODE_EVAL_EVERY_N", "0") monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") processor = AsyncInProcessProcessor(pipeline=MagicMock()) @@ -99,12 +101,83 @@ async def fake_upsert(body): class TestAsyncExtractRecentK: + @pytest.mark.asyncio + async def test_episode_zero_does_not_fire(self): + processor = AsyncInProcessProcessor(pipeline=MagicMock()) + processor.process_extract_episodes = AsyncMock() + counter_container = _AsyncFakeCounterContainer() + + await maybe_trigger_steps( + processor, + counter_container, + {("u1", "t1"): 1}, + thresholds={ + "FACT_EXTRACTION_EVERY_N": 0, + "THREAD_SUMMARY_EVERY_N": 0, + "EPISODE_EVAL_EVERY_N": 0, + "USER_SUMMARY_EVERY_N": 0, + "MEMORY_PROCESSOR_OWNER": "inprocess", + }, + ) + + processor.process_extract_episodes.assert_not_awaited() + assert counter_container.store == {} + + @pytest.mark.asyncio + async def test_episode_fires_when_threshold_crossed(self): + processor = AsyncInProcessProcessor(pipeline=MagicMock()) + processor.process_extract_episodes = AsyncMock(return_value={}) + counter_container = _AsyncFakeCounterContainer() + thresholds = { + "FACT_EXTRACTION_EVERY_N": 0, + "THREAD_SUMMARY_EVERY_N": 0, + "EPISODE_EVAL_EVERY_N": 3, + "USER_SUMMARY_EVERY_N": 0, + "MEMORY_PROCESSOR_OWNER": "inprocess", + } + + await maybe_trigger_steps(processor, counter_container, {("u1", "t1"): 2}, thresholds=thresholds) + processor.process_extract_episodes.assert_not_awaited() + + await maybe_trigger_steps(processor, counter_container, {("u1", "t1"): 1}, thresholds=thresholds) + + processor.process_extract_episodes.assert_awaited_once_with(user_id="u1", thread_id="t1") + + @pytest.mark.asyncio + async def test_episode_failure_is_caught_and_other_steps_continue(self): + processor = AsyncInProcessProcessor(pipeline=MagicMock()) + processor.process_extract_episodes = AsyncMock(side_effect=RuntimeError("episode boom")) + processor.process_thread_summary = AsyncMock(return_value={}) + counter_container = _AsyncFakeCounterContainer() + + with patch( + "azure.cosmos.agent_memory._counters.stamp_failure_async", + new=AsyncMock(), + ) as stamp: + await maybe_trigger_steps( + processor, + counter_container, + {("u1", "t1"): 1}, + thresholds={ + "FACT_EXTRACTION_EVERY_N": 0, + "THREAD_SUMMARY_EVERY_N": 1, + "EPISODE_EVAL_EVERY_N": 1, + "USER_SUMMARY_EVERY_N": 0, + "MEMORY_PROCESSOR_OWNER": "inprocess", + }, + ) + + processor.process_extract_episodes.assert_awaited_once_with(user_id="u1", thread_id="t1") + processor.process_thread_summary.assert_awaited_once_with(user_id="u1", thread_id="t1") + stamp.assert_awaited_once() + @pytest.mark.asyncio async def test_extract_fires_without_recent_k_or_watermark(self, monkeypatch): """Async: extraction covers all un-extracted turns (extracted_at gated) and batches internally, so it fires with NO recent_k and NO success watermark.""" monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") + monkeypatch.setenv("EPISODE_EVAL_EVERY_N", "0") monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") processor = AsyncInProcessProcessor(pipeline=MagicMock()) @@ -132,6 +205,7 @@ async def test_extract_failure_stamps_failure(self, monkeypatch): """A total async extract failure is recorded via stamp_failure_async.""" monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") + monkeypatch.setenv("EPISODE_EVAL_EVERY_N", "0") monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") processor = AsyncInProcessProcessor(pipeline=MagicMock()) @@ -169,6 +243,7 @@ class TestPushToCosmosUnflushedDelta: @pytest.mark.asyncio async def test_repeat_push_does_not_re_increment(self, monkeypatch): monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") + monkeypatch.setenv("EPISODE_EVAL_EVERY_N", "0") client = AsyncCosmosMemoryClient(use_default_credential=False) @@ -212,6 +287,7 @@ async def capture(turn_counts): @pytest.mark.asyncio async def test_only_new_adds_count_after_partial_push(self, monkeypatch): monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") + monkeypatch.setenv("EPISODE_EVAL_EVERY_N", "0") client = AsyncCosmosMemoryClient(use_default_credential=False) diff --git a/tests/unit/aio/test_chat.py b/tests/unit/aio/test_chat.py index 1d79132..ab0f846 100644 --- a/tests/unit/aio/test_chat.py +++ b/tests/unit/aio/test_chat.py @@ -145,3 +145,62 @@ async def test_generate_returns_content(): result = await client.generate([{"role": "user", "content": "hi"}]) assert result == "hello world" + + +def _api_status_error(status_code: int, headers: dict[str, str] | None = None): + import httpx + import openai + + request = httpx.Request("POST", "https://test.openai.azure.com/openai/deployments/test/chat/completions") + response = httpx.Response(status_code, headers=headers or {}, request=request) + return openai.APIStatusError(message=f"status {status_code}", response=response, body=None) + + +@pytest.mark.asyncio +async def test_generate_retryable_api_error_honors_retry_after(monkeypatch): + client = AsyncChatClient(endpoint="https://test.openai.azure.com", api_key="key") + fake = MagicMock() + fake.chat.completions.create = AsyncMock( + side_effect=[ + _api_status_error(503, {"retry-after": "30"}), + MagicMock(choices=[MagicMock(message=MagicMock(content="recovered"))], usage=None), + ] + ) + client._client = fake + sleeps: list[float] = [] + + async def capture_sleep(delay: float) -> None: + sleeps.append(delay) + + monkeypatch.setattr("azure.cosmos.agent_memory.chat.random.random", lambda: 0.0) + monkeypatch.setattr("azure.cosmos.agent_memory.aio.chat.asyncio.sleep", capture_sleep) + + result = await client.generate([{"role": "user", "content": "hi"}], max_retries=2, base_delay=2.0) + + assert result == "recovered" + assert sleeps == [30.0] + + +@pytest.mark.asyncio +async def test_generate_retryable_api_error_falls_back_without_retry_after(monkeypatch): + client = AsyncChatClient(endpoint="https://test.openai.azure.com", api_key="key") + fake = MagicMock() + fake.chat.completions.create = AsyncMock( + side_effect=[ + _api_status_error(503), + MagicMock(choices=[MagicMock(message=MagicMock(content="recovered"))], usage=None), + ] + ) + client._client = fake + sleeps: list[float] = [] + + async def capture_sleep(delay: float) -> None: + sleeps.append(delay) + + monkeypatch.setattr("azure.cosmos.agent_memory.chat.random.random", lambda: 0.0) + monkeypatch.setattr("azure.cosmos.agent_memory.aio.chat.asyncio.sleep", capture_sleep) + + result = await client.generate([{"role": "user", "content": "hi"}], max_retries=2, base_delay=2.0) + + assert result == "recovered" + assert sleeps == [1.6] diff --git a/tests/unit/aio/test_cosmos_memory_client.py b/tests/unit/aio/test_cosmos_memory_client.py index 098fcae..ae58496 100644 --- a/tests/unit/aio/test_cosmos_memory_client.py +++ b/tests/unit/aio/test_cosmos_memory_client.py @@ -888,7 +888,8 @@ async def test_search_whitespace_only_terms(self): async def test_search_episodic_forwards_search_options(self): containers = {key: MagicMock() for key in ContainerKey} store = AsyncMemoryStore(containers=containers) - store.search = AsyncMock(return_value=[]) + store.query = AsyncMock(return_value=[]) + store._embed = AsyncMock(return_value=[0.1]) await store.search_episodic( user_id="u1", @@ -898,14 +899,10 @@ async def test_search_episodic_forwards_search_options(self): include_superseded=True, ) - store.search.assert_awaited_once_with( - search_terms="weather", - user_id="u1", - memory_types=["episodic"], - top_k=2, - min_salience=0.4, - include_superseded=True, - ) + query, parameters = store.query.await_args.args[:2] + assert "c.type = @type" in query + assert "VectorDistance(c.embedding, @embedding)" in query + assert {"name": "@type", "value": "episodic"} in parameters async def test_build_episodic_context_forwards_search_options(self): containers = {key: MagicMock() for key in ContainerKey} diff --git a/tests/unit/function_app/test_orchestrators.py b/tests/unit/function_app/test_orchestrators.py index 634514b..14b149e 100644 --- a/tests/unit/function_app/test_orchestrators.py +++ b/tests/unit/function_app/test_orchestrators.py @@ -376,22 +376,22 @@ def test_missing_thread_id_raises(self): class TestExtractMemoryActivities: def test_em_extract_uses_payload_recent_k(self): pipeline = MagicMock() - pipeline.extract_memories_dry.return_value = {"facts": [], "episodic": [], "updates": []} + pipeline.extract_memories_durable.return_value = {"facts": [], "episodic": [], "updates": []} with patch.object(em_mod, "get_pipeline", return_value=pipeline): result = em_mod.em_Extract({"user_id": "u1", "thread_id": "t1", "recent_k": 3}) - pipeline.extract_memories_dry.assert_called_once_with(user_id="u1", thread_id="t1", recent_k=3) + pipeline.extract_memories_durable.assert_called_once_with(user_id="u1", thread_id="t1", recent_k=3) assert result == {"facts": [], "episodic": [], "updates": []} def test_em_extract_falls_back_to_max_batch_size_when_recent_k_absent(self): pipeline = MagicMock() - pipeline.extract_memories_dry.return_value = {"facts": [], "episodic": [], "updates": []} + pipeline.extract_memories_durable.return_value = {"facts": [], "episodic": [], "updates": []} with patch.object(em_mod, "get_pipeline", return_value=pipeline): em_mod.em_Extract({"user_id": "u1", "thread_id": "t1"}) - pipeline.extract_memories_dry.assert_called_once_with(user_id="u1", thread_id="t1", recent_k=20) + 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": []} diff --git a/tests/unit/processors/test_protocol_satisfaction.py b/tests/unit/processors/test_protocol_satisfaction.py index 6869f00..f08a016 100644 --- a/tests/unit/processors/test_protocol_satisfaction.py +++ b/tests/unit/processors/test_protocol_satisfaction.py @@ -36,6 +36,14 @@ def process_extract_memories( ) -> dict[str, int]: return {} + def process_extract_episodes( + self, + *, + user_id: str, + thread_id: str, + ) -> dict[str, int]: + return {} + def process_thread_summary( self, *, diff --git a/tests/unit/services/test_chaos_extract_persist.py b/tests/unit/services/test_chaos_extract_persist.py index 33afa5e..9a40976 100644 --- a/tests/unit/services/test_chaos_extract_persist.py +++ b/tests/unit/services/test_chaos_extract_persist.py @@ -178,7 +178,7 @@ def test_persist_retry_reuses_extract_output_without_second_llm_call() -> None: containers=_containers_for_store(store, turns_store=turns_store), ) - extracted = service.extract_memories_dry("u1", "t1") + extracted = service.extract_memories_durable("u1", "t1") with pytest.raises(RuntimeError, match="transient"): service.persist_extracted_memories("u1", extracted) result = service.persist_extracted_memories("u1", extracted) @@ -201,7 +201,7 @@ async def test_async_persist_retry_reuses_extract_output_without_second_llm_call containers=_async_containers_for_store(store, turns_store=turns_store), ) - extracted = await service.extract_memories_dry("u1", "t1") + extracted = await service.extract_memories_durable("u1", "t1") with pytest.raises(RuntimeError, match="transient"): await service.persist_extracted_memories("u1", extracted) result = await service.persist_extracted_memories("u1", extracted) diff --git a/tests/unit/services/test_dedup_vector.py b/tests/unit/services/test_dedup_vector.py index 02beaff..e2d58f6 100644 --- a/tests/unit/services/test_dedup_vector.py +++ b/tests/unit/services/test_dedup_vector.py @@ -36,16 +36,9 @@ def _doc(mid: str, content: str, memory_type: str = "fact", **extra: Any) -> dic tags = extra.pop("tags", [f"sys:{memory_type}"]) metadata = extra.pop( "metadata", - {"category": "preference"} - if memory_type == "fact" - else { - "scope_type": "project", - "scope_value": "demo", - "lesson": content, - "outcome_valence": "neutral", - }, + {"category": "preference"} if memory_type == "fact" else {}, ) - return { + doc = { "id": mid, "user_id": "u1", "thread_id": "t1", @@ -63,6 +56,13 @@ def _doc(mid: str, content: str, memory_type: str = "fact", **extra: Any) -> dic "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: diff --git a/tests/unit/services/test_episode_boundary.py b/tests/unit/services/test_episode_boundary.py new file mode 100644 index 0000000..97b9a4e --- /dev/null +++ b/tests/unit/services/test_episode_boundary.py @@ -0,0 +1,265 @@ +"""Boundary-based episodic segmentation (sync). + +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. +""" + +from __future__ import annotations + +from typing import Any + +from azure.cosmos.agent_memory.services.pipeline import PipelineService +from tests.unit.services.test_extract_dry import ( + _containers_for_store, + _Store, + _SyncChat, + _SyncEmbeddings, +) +from tests.unit.services.test_extract_episodes import _episode, _TrackingStore + + +def _turn_at(i: int, minute: int, *, content: str | None = None) -> dict[str, Any]: + return { + "id": f"turn-{i}", + "user_id": "u1", + "thread_id": "t1", + "role": "user", + "type": "turn", + "content": content if content is not None else f"Turn {i}: routine content", + "created_at": f"2025-01-01T00:{minute:02d}:00+00:00", + } + + +class _DriftEmbeddings(_SyncEmbeddings): + """Content-keyed embeddings: turns tagged 'A:' vs 'B:' land on orthogonal axes.""" + + def generate_batch(self, texts: list[str]) -> list[list[float]]: + self.calls.append(list(texts)) + return [[0.0, 1.0] if "B:" in text else [1.0, 0.0] for text in texts] + + +def _service( + turns: list[dict[str, Any]], + responses: list[dict[str, Any]] | None = None, + *, + embeddings: _SyncEmbeddings | None = None, +) -> tuple[PipelineService, _TrackingStore, _Store, _SyncChat]: + memories = _TrackingStore([]) + turns_store = _Store(turns) + chat = _SyncChat(responses or [{"episodes": [_episode()]} for _ in range(20)]) + service = PipelineService( + memories, + chat, + embeddings or _SyncEmbeddings(), + containers=_containers_for_store(memories, turns_store=turns_store), + ) + # Boundary tests exercise segmentation logic, not prompt rendering: stub + # _run_prompty so we return canned episode JSON without invoking the real + # prompty template renderer (keeps these tests fast and hermetic). + service._run_prompty = lambda *a, **k: chat.generate([]) # type: ignore[assignment] + return service, memories, turns_store, chat + + +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 test_time_gap_closes_prior_episode_and_leaves_tail_open(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_IDLE_GAP_SECONDS", "120") + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + monkeypatch.setenv("EPISODE_MAX_TURNS", "40") + # 00:01, 00:02 then a 28-minute gap to 00:30, 00:31. + turns = [_turn_at(1, 1), _turn_at(2, 2), _turn_at(3, 30), _turn_at(4, 31)] + service, memories, turns_store, _ = _service(turns) + + result = service.extract_episodes("u1", "t1") + + 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"] + + +def test_reevaluation_is_idempotent_via_watermark(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) + + first = service.extract_episodes("u1", "t1") + second = service.extract_episodes("u1", "t1") + + assert first == {"episodes": 1} + # The open tail has no further boundary: no new episode, no duplicate. + assert second == {"episodes": 0} + assert len(_episodes(memories)) == 1 + + +def test_flush_drains_open_tail(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") # closes [turn-1, turn-2] + flushed = service.extract_episodes("u1", "t1", flush=True) # drains [turn-3, turn-4] + + assert flushed == {"episodes": 1} + assert len(_episodes(memories)) == 2 + assert _stamped(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} + + +def test_max_turns_forces_a_boundary(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_IDLE_GAP_SECONDS", "0") # disable gap + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") # disable drift + monkeypatch.setenv("EPISODE_MAX_TURNS", "2") # force a boundary every 2 turns + turns = [_turn_at(i, i) for i in range(1, 5)] # 4 turns, no large gaps + service, memories, turns_store, _ = _service(turns) + + result = service.extract_episodes("u1", "t1") + + # 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"] + + +def test_topic_drift_closes_episode(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_IDLE_GAP_SECONDS", "0") # isolate drift + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0.5") + monkeypatch.setenv("EPISODE_MIN_TURNS", "2") + monkeypatch.setenv("EPISODE_MAX_TURNS", "40") + turns = [ + _turn_at(1, 1, content="A: apples and orchards"), + _turn_at(2, 2, content="A: more about apples"), + _turn_at(3, 3, content="B: rockets and orbits"), + _turn_at(4, 4, content="B: more about rockets"), + ] + service, memories, turns_store, _ = _service(turns, embeddings=_DriftEmbeddings()) + + result = service.extract_episodes("u1", "t1") + + # 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"] + + +def test_no_boundary_keeps_segment_open_without_calling_the_llm(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_IDLE_GAP_SECONDS", "1800") + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + monkeypatch.setenv("EPISODE_MAX_TURNS", "40") + turns = [_turn_at(1, 1), _turn_at(2, 2), _turn_at(3, 3)] # close together, small + service, memories, turns_store, chat = _service(turns) + + result = service.extract_episodes("u1", "t1") + + assert result == {"episodes": 0} + assert _episodes(memories) == [] + assert _stamped(turns_store) == [] # nothing closed + assert chat.calls == 0 # extraction LLM only runs at a boundary + + +def test_deterministic_episode_id_is_stable_and_ordinal_scoped() -> None: + service, *_ = _service([]) + key = "u1\x00t1\x00turn-1\x00turn-2" + id_0 = service._deterministic_episode_id(key, 0) + id_0_again = service._deterministic_episode_id(key, 0) + id_1 = service._deterministic_episode_id(key, 1) + id_other_segment = service._deterministic_episode_id("u1\x00t1\x00turn-3\x00turn-4", 0) + + assert id_0.startswith("ep_") + assert id_0 == id_0_again # same segment + same ordinal -> same id (idempotent) + assert id_0 != id_1 # different ordinal -> different id + assert id_0 != id_other_segment # different segment -> different id + + +def test_idle_gap_below_min_turns_does_not_close_episode(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_IDLE_GAP_SECONDS", "120") + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + monkeypatch.setenv("EPISODE_MAX_TURNS", "40") + monkeypatch.setenv("EPISODE_MIN_TURNS", "2") + # Gap between turn-1 (i=0) and turn-2 (i=1); i=1 < min_turns=2 -> not closed. + turns = [_turn_at(1, 1), _turn_at(2, 30), _turn_at(3, 31)] + service, memories, turns_store, _ = _service(turns) + + result = service.extract_episodes("u1", "t1") + + assert result == {"episodes": 0} + assert _episodes(memories) == [] + assert _stamped(turns_store) == [] + + +def test_idle_gap_below_min_turns_still_flushes_as_one_episode(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_IDLE_GAP_SECONDS", "120") + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + monkeypatch.setenv("EPISODE_MAX_TURNS", "40") + monkeypatch.setenv("EPISODE_MIN_TURNS", "2") + turns = [_turn_at(1, 1), _turn_at(2, 30), _turn_at(3, 31)] + service, memories, turns_store, _ = _service(turns) + + result = service.extract_episodes("u1", "t1", flush=True) + + assert result == {"episodes": 1} + assert _stamped(turns_store) == ["turn-1", "turn-2", "turn-3"] + + +def test_closed_segment_with_no_episode_still_stamps_turns(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_IDLE_GAP_SECONDS", "120") + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + monkeypatch.setenv("EPISODE_MAX_TURNS", "40") + # 00:01, 00:02 then a 28-minute gap to 00:30 -> the pre-gap segment closes. + turns = [_turn_at(1, 1), _turn_at(2, 2), _turn_at(3, 30)] + service, memories, turns_store, _ = _service(turns, responses=[{"episodes": []}]) + + result = service.extract_episodes("u1", "t1") + + 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"] + + +def test_extract_episodes_defers_segment_on_retryable_error(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_IDLE_GAP_SECONDS", "120") + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + monkeypatch.setenv("EPISODE_MAX_TURNS", "40") + turns = [_turn_at(1, 1), _turn_at(2, 2), _turn_at(3, 30)] # gap closes [1,2] + service, memories, turns_store, _ = _service(turns) + + def _boom(*a, **k): + raise RuntimeError("transient rate limit 429") + + service._run_prompty = _boom # type: ignore[assignment] + + result = service.extract_episodes("u1", "t1") + + assert result == {"episodes": 0} + assert _episodes(memories) == [] + assert _stamped(turns_store) == [] # un-stamped -> retried next run + + +def test_extract_episodes_quarantines_segment_on_non_retryable_error(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_IDLE_GAP_SECONDS", "120") + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + monkeypatch.setenv("EPISODE_MAX_TURNS", "40") + turns = [_turn_at(1, 1), _turn_at(2, 2), _turn_at(3, 30)] + service, memories, turns_store, _ = _service(turns) + + def _boom(*a, **k): + raise RuntimeError("content_filter triggered") + + service._run_prompty = _boom # type: ignore[assignment] + + result = service.extract_episodes("u1", "t1") + + assert result == {"episodes": 0} + assert _episodes(memories) == [] + assert _stamped(turns_store) == ["turn-1", "turn-2"] # quarantined + advanced diff --git a/tests/unit/services/test_episodic_retrieval.py b/tests/unit/services/test_episodic_retrieval.py new file mode 100644 index 0000000..d72eaf1 --- /dev/null +++ b/tests/unit/services/test_episodic_retrieval.py @@ -0,0 +1,156 @@ +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.cosmos_memory_client import CosmosMemoryClient +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 _client_with_store(store: Any) -> CosmosMemoryClient: + client = CosmosMemoryClient( + use_default_credential=False, + embeddings_client=MagicMock(), + chat_client=MagicMock(), + ) + client._get_store = MagicMock(return_value=store) # type: ignore[method-assign] + return client + + +def test_get_episodes_returns_user_episodes_newest_first() -> None: + episodes = [ + {"id": "ep-new", "type": "episodic", "user_id": "u1", "thread_id": "t2", "created_at": "2026-01-02"}, + {"id": "ep-old", "type": "episodic", "user_id": "u1", "thread_id": "t1", "created_at": "2026-01-01"}, + ] + memories = MagicMock() + memories.query_items.return_value = episodes + store = MemoryStore(containers=_containers(memories=memories)) + client = _client_with_store(store) + + result = client.get_episodes("u1", recent_k=2) + + assert result == episodes + kwargs = memories.query_items.call_args.kwargs + assert "c.type = @type" in kwargs["query"] + assert "c.user_id = @user_id" in kwargs["query"] + assert "ORDER BY c.created_at DESC" in kwargs["query"] + assert kwargs["enable_cross_partition_query"] is True + + +def test_extract_episodes_routes_to_pipeline_with_flush() -> None: + client = CosmosMemoryClient( + use_default_credential=False, + embeddings_client=MagicMock(), + chat_client=MagicMock(), + ) + pipeline = MagicMock() + pipeline.extract_episodes.return_value = {"episodes": 2} + client._get_pipeline = MagicMock(return_value=pipeline) # type: ignore[method-assign] + + assert client.extract_episodes("u1", "t1", flush=True) == {"episodes": 2} + pipeline.extract_episodes.assert_called_once_with("u1", "t1", flush=True) + + +def test_search_cosmos_base_is_facts_only_no_episodes_without_optin() -> None: + # Episodes never enter via the base query: base search is scoped to facts, + # and without include_episodes no episodic query runs at all. + store = MagicMock() + store.search.return_value = [{"id": "fact", "content": "fact", "type": "fact"}] + store.search_summaries.return_value = [{"id": "summary", "content": "summary", "type": "thread_summary"}] + store.search_turns.return_value = [{"id": "turn", "content": "turn", "type": "turn"}] + client = _client_with_store(store) + + result = client.search_cosmos( + "ci retries", + user_id="u1", + thread_id="t1", + top_k=3, + include_summaries=True, + include_turns=True, + ) + + assert [doc["id"] for doc in result] == ["fact", "summary", "turn"] + assert store.search.call_args.kwargs["memory_types"] == ["fact"] + store.search_episodic.assert_not_called() + + +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. + store.search.return_value = [ + {"id": "fact", "content": "fact", "type": "fact"}, + {"id": "episode", "content": "episode", "type": "episodic"}, + ] + store.search_summaries.return_value = [{"id": "summary", "content": "summary", "type": "thread_summary"}] + store.search_turns.return_value = [{"id": "turn", "content": "turn", "type": "turn"}] + client = _client_with_store(store) + + result = client.search_cosmos( + "ci retries", + user_id="u1", + thread_id="t1", + top_k=100, + include_episodes=True, + include_summaries=True, + include_turns=True, + ) + + # Combined base (facts + episodes) -> summaries -> turns. + assert [doc["id"] for doc in result] == ["fact", "episode", "summary", "turn"] + # One query, one shared top_k budget; episodic is folded into the base types. + assert store.search.call_args.kwargs["top_k"] == 100 + assert store.search.call_args.kwargs["memory_types"] == ["fact", "episodic"] + # No separate episodic query is issued. + store.search_episodic.assert_not_called() + + +class _RankedEpisodeContainer: + def __init__(self) -> None: + self.query: str | None = None + self.parameters: list[dict[str, Any]] | None = None + self.docs = [ + {"id": "best-old", "created_at": "2026-01-01T00:00:00+00:00"}, + {"id": "best-new", "created_at": "2026-01-03T00:00:00+00:00"}, + {"id": "third-new", "created_at": "2026-01-04T00:00:00+00:00"}, + ] + + def query_items(self, **kwargs: Any) -> list[dict[str, Any]]: + self.query = kwargs["query"] + self.parameters = kwargs["parameters"] + params = {p["name"]: p["value"] for p in self.parameters or []} + after = params.get("@created_after") + before = params.get("@created_before") + docs = list(self.docs) + if after is not None: + docs = [doc for doc in docs if doc["created_at"] >= after] + if before is not None: + docs = [doc for doc in docs if doc["created_at"] <= before] + return docs + + +def test_search_episodic_temporal_filter_narrows_without_time_ranking() -> None: + memories = _RankedEpisodeContainer() + embeddings = MagicMock() + embeddings.generate.return_value = [0.1, 0.2] + store = MemoryStore(containers=_containers(memories=memories), embeddings_client=embeddings) + + result = store.search_episodic( + user_id="u1", + search_terms="ci retries", + created_after="2026-01-02T00:00:00+00:00", + ) + + assert [doc["id"] for doc in result] == ["best-new", "third-new"] + assert memories.query is not None + assert "c.created_at >= @created_after" in memories.query + order_by = memories.query.split("ORDER BY", 1)[1] + assert "created_at" not in order_by diff --git a/tests/unit/services/test_extract_dry.py b/tests/unit/services/test_extract_dry.py index 2ad225a..6e86486 100644 --- a/tests/unit/services/test_extract_dry.py +++ b/tests/unit/services/test_extract_dry.py @@ -77,7 +77,9 @@ 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 "extracted_at" in sql: + if "episode_extracted_at" in sql: + docs = [doc for doc in docs if not doc.get("episode_extracted_at")] + elif "extracted_at" in sql: docs = [doc for doc in docs if not doc.get("extracted_at")] return docs @@ -195,19 +197,11 @@ def _response() -> dict[str, Any]: "tags": ["ui"], } ], - "episodic": [ - { - "scope_type": "project", - "scope_value": "CI", - "text": "CI retries resolved flaky tests.", - "lesson": "Use retries for flaky CI tests.", - "confidence": 0.8, - } - ], + "episodic": [], } -def test_extract_memories_dry_shape_is_small_and_has_no_embeddings() -> None: +def test_extract_memories_durable_shape_is_small_and_has_no_embeddings() -> None: chat = _SyncChat([_response()]) embeddings = _SyncEmbeddings() memories_store = _Store([]) @@ -219,16 +213,120 @@ def test_extract_memories_dry_shape_is_small_and_has_no_embeddings() -> None: containers=_containers_for_store(memories_store, turns_store=turns_store), ) - output = service.extract_memories_dry("u1", "t1") + output = service.extract_memories_durable("u1", "t1") assert set(output) == {"facts", "episodic", "updates", "processed_turn_docs"} assert len(json.dumps(output)) < 32 * 1024 - assert output["facts"] and output["episodic"] + assert output["facts"] + assert output["episodic"] == [] assert all("embedding" not in doc for docs in (output["facts"], output["episodic"]) for doc in docs) assert embeddings.calls == [] -def test_extract_memories_dry_is_byte_deterministic_for_same_llm_response() -> None: +def test_build_episode_docs_builds_new_episode_shape_without_embeddings() -> None: + chat = _SyncChat( + [ + { + "episodes": [ + { + "title": "Fixed CI retries", + "summary": "The user fixed flaky CI retries and the tests passed.", + "started_at": "2025-01-01T00:01:00+00:00", + "ended_at": "2025-01-01T00:02:00+00:00", + "participants": [], + "events": [ + { + "sequence": 1, + "description": "CI retries were added.", + "occurred_at": "2025-01-01T00:01:00+00:00", + "source_turn_ids": ["turn-1", "missing"], + }, + { + "sequence": 2, + "description": "The tests passed.", + "occurred_at": "2025-01-01T00:02:00+00:00", + "source_turn_ids": ["turn-2"], + }, + ], + "outcome": {"status": "successful", "description": "The tests passed."}, + "lessons": ["Use retries for flaky CI."], + "salience": 0.8, + "confidence": 0.9, + } + ] + } + ] + ) + embeddings = _SyncEmbeddings() + memories_store = _Store([]) + turns_store = _Store([_turn(1), _turn(2)]) + service = PipelineService( + memories_store, + chat, + embeddings, + containers=_containers_for_store(memories_store, turns_store=turns_store), + ) + + [doc] = service._build_episode_docs("u1", "t1", [_turn(1), _turn(2)], segment_key="seg-1") + + assert doc["id"].startswith("ep_") + assert doc["type"] == "episodic" + assert doc["content"] == "The user fixed flaky CI retries and the tests passed." + assert len(doc["content_hash"]) == 32 + assert doc["title"] == "Fixed CI retries" + assert doc["source_turn_ids"] == ["turn-1", "turn-2"] + assert doc["events"][0]["source_turn_ids"] == ["turn-1"] + assert doc["outcome"]["status"] == "successful" + assert doc["prompt_id"] == "extract_episode.prompty" + assert "embedding" not in doc + assert embeddings.calls == [] + + +def test_build_episode_docs_allows_null_outcome_and_maps_stable_turn_labels() -> None: + chat = _SyncChat( + [ + { + "episodes": [ + { + "title": "Family dinner", + "summary": "The user had a family dinner with Elena.", + "started_at": None, + "ended_at": None, + "participants": ["Elena"], + "events": [ + { + "sequence": 1, + "description": "The user had dinner with Elena.", + "occurred_at": None, + "source_turn_ids": ["turn-1"], + } + ], + "outcome": None, + "lessons": [], + "salience": 0.6, + "confidence": 0.95, + } + ] + } + ] + ) + memories_store = _Store([]) + turns_store = _Store([{**_turn(1), "id": "actual-turn-id"}]) + service = PipelineService( + memories_store, + chat, + _SyncEmbeddings(), + containers=_containers_for_store(memories_store, turns_store=turns_store), + ) + + [doc] = service._build_episode_docs("u1", "t1", [{**_turn(1), "id": "actual-turn-id"}], segment_key="seg-2") + + assert doc["outcome"] is None + assert doc["source_turn_ids"] == ["actual-turn-id"] + assert doc["events"][0]["source_turn_ids"] == ["actual-turn-id"] + + +def test_extract_memories_durable_is_byte_deterministic_for_same_llm_response() -> None: store = _Store([]) turns_store = _Store([_turn(1)]) service = PipelineService( @@ -238,15 +336,15 @@ def test_extract_memories_dry_is_byte_deterministic_for_same_llm_response() -> N containers=_containers_for_store(store, turns_store=turns_store), ) - first = service.extract_memories_dry("u1", "t1") - second = service.extract_memories_dry("u1", "t1") + first = service.extract_memories_durable("u1", "t1") + second = service.extract_memories_durable("u1", "t1") assert json.dumps(first, sort_keys=True, separators=(",", ":")) == json.dumps( second, sort_keys=True, separators=(",", ":") ) -def test_extract_memories_dry_does_not_call_store_search() -> None: +def test_extract_memories_durable_does_not_call_store_search() -> None: """Extraction is single-pass and existing-memory-free: it must not issue a dedup-context vector search (dedup is handled by hash + reconciliation).""" memories_store = _Store([]) @@ -258,13 +356,13 @@ def test_extract_memories_dry_does_not_call_store_search() -> None: containers=_containers_for_store(memories_store, turns_store=_Store([_turn(1)])), ) - service.extract_memories_dry("u1", "t1") + service.extract_memories_durable("u1", "t1") assert memories_store.search_calls == [] @pytest.mark.asyncio -async def test_async_extract_memories_dry_shape_is_small_and_has_no_embeddings() -> None: +async def test_async_extract_memories_durable_shape_is_small_and_has_no_embeddings() -> None: chat = _AsyncChat([_response()]) embeddings = _AsyncEmbeddings() memories_store = _AsyncStore([]) @@ -276,7 +374,7 @@ async def test_async_extract_memories_dry_shape_is_small_and_has_no_embeddings() containers=_async_containers_for_store(memories_store, turns_store=turns_store), ) - output = await service.extract_memories_dry("u1", "t1") + output = await service.extract_memories_durable("u1", "t1") assert set(output) == {"facts", "episodic", "updates", "processed_turn_docs"} assert len(json.dumps(output)) < 32 * 1024 @@ -285,7 +383,7 @@ async def test_async_extract_memories_dry_shape_is_small_and_has_no_embeddings() @pytest.mark.asyncio -async def test_async_extract_memories_dry_is_byte_deterministic_for_same_llm_response() -> None: +async def test_async_extract_memories_durable_is_byte_deterministic_for_same_llm_response() -> None: store = _AsyncStore([]) turns_store = _AsyncStore([_turn(1)]) service = AsyncPipelineService( @@ -295,8 +393,8 @@ async def test_async_extract_memories_dry_is_byte_deterministic_for_same_llm_res containers=_async_containers_for_store(store, turns_store=turns_store), ) - first = await service.extract_memories_dry("u1", "t1") - second = await service.extract_memories_dry("u1", "t1") + first = await service.extract_memories_durable("u1", "t1") + second = await service.extract_memories_durable("u1", "t1") assert json.dumps(first, sort_keys=True, separators=(",", ":")) == json.dumps( second, sort_keys=True, separators=(",", ":") @@ -304,7 +402,7 @@ async def test_async_extract_memories_dry_is_byte_deterministic_for_same_llm_res @pytest.mark.asyncio -async def test_async_extract_memories_dry_does_not_call_store_search() -> None: +async def test_async_extract_memories_durable_does_not_call_store_search() -> None: store = _AsyncStore([]) store.search = AsyncMock(return_value=[]) @@ -315,7 +413,7 @@ async def test_async_extract_memories_dry_does_not_call_store_search() -> None: containers=_async_containers_for_store(store, turns_store=_AsyncStore([_turn(1)])), ) - await service.extract_memories_dry("u1", "t1") + await service.extract_memories_durable("u1", "t1") store.search.assert_not_awaited() @@ -350,6 +448,11 @@ def generate(self, messages, **opts): ) +class _AsyncBatchChat(_BatchChat): + async def generate(self, messages, **opts): + return super().generate(messages, **opts) + + def _one_turn_per_batch(monkeypatch): # Force each small turn into its own extraction batch. monkeypatch.setattr("azure.cosmos.agent_memory.thresholds.get_extraction_batch_max_tokens", lambda: 5) @@ -367,7 +470,7 @@ def test_extract_batches_run_independently_one_call_per_batch(monkeypatch) -> No containers=_containers_for_store(memories_store, turns_store=turns_store), ) - out = service.extract_memories_dry("u1", "t1") + out = service.extract_memories_durable("u1", "t1") assert chat.calls == 3 # one LLM call per batch assert len(out["facts"]) == 3 @@ -386,7 +489,7 @@ def test_extract_quarantines_non_retryable_batch_but_keeps_others(monkeypatch) - containers=_containers_for_store(memories_store, turns_store=turns_store), ) - out = service.extract_memories_dry("u1", "t1") + out = service.extract_memories_durable("u1", "t1") # Batches 1 and 3 produced facts; batch 2 was quarantined (no fact) ... assert len(out["facts"]) == 2 @@ -408,7 +511,7 @@ def test_extract_defers_retryable_batch_leaving_turns_unstamped(monkeypatch) -> containers=_containers_for_store(memories_store, turns_store=turns_store), ) - out = service.extract_memories_dry("u1", "t1") + out = service.extract_memories_durable("u1", "t1") # Batches 1 and 3 produced facts; batch 2 deferred (retryable) ... assert len(out["facts"]) == 2 @@ -418,6 +521,84 @@ def test_extract_defers_retryable_batch_leaving_turns_unstamped(monkeypatch) -> assert stats and stats[0]["deferred_turn_count"] == 1 +def test_extract_memories_returns_deferred_turn_count_for_retryable_batch() -> None: + chat = _BatchChat(fail_on_call=1, error=Exception("Error code: 429 rate limit")) + memories_store = _Store([]) + turns = [_turn(i) for i in range(3)] + turns_store = _Store(turns) + service = PipelineService( + memories_store, + chat, + _SyncEmbeddings(), + containers=_containers_for_store(memories_store, turns_store=turns_store), + ) + + counts = service.extract_memories("u1", "t1") + + assert counts["deferred_turn_count"] == len(turns) + assert counts["quarantined_turn_count"] == 0 + assert all("extracted_at" not in turn for turn in turns_store.docs) + + +def test_extract_memories_returns_quarantined_turn_count_for_non_retryable_batch() -> None: + chat = _BatchChat(fail_on_call=1, error=Exception("Error code: 400 content_filter")) + memories_store = _Store([]) + turns = [_turn(i) for i in range(3)] + turns_store = _Store(turns) + service = PipelineService( + memories_store, + chat, + _SyncEmbeddings(), + containers=_containers_for_store(memories_store, turns_store=turns_store), + ) + + counts = service.extract_memories("u1", "t1") + + assert counts["deferred_turn_count"] == 0 + assert counts["quarantined_turn_count"] == len(turns) + assert all("extracted_at" in turn for turn in turns_store.docs) + + +@pytest.mark.asyncio +async def test_async_extract_memories_returns_deferred_turn_count_for_retryable_batch() -> None: + chat = _AsyncBatchChat(fail_on_call=1, error=Exception("Error code: 429 rate limit")) + memories_store = _AsyncStore([]) + turns = [_turn(i) for i in range(3)] + turns_store = _AsyncStore(turns) + service = AsyncPipelineService( + memories_store, + chat, + _AsyncEmbeddings(), + containers=_async_containers_for_store(memories_store, turns_store=turns_store), + ) + + counts = await service.extract_memories("u1", "t1") + + assert counts["deferred_turn_count"] == len(turns) + assert counts["quarantined_turn_count"] == 0 + assert all("extracted_at" not in turn for turn in turns_store.docs) + + +@pytest.mark.asyncio +async def test_async_extract_memories_returns_quarantined_turn_count_for_non_retryable_batch() -> None: + chat = _AsyncBatchChat(fail_on_call=1, error=Exception("Error code: 400 content_filter")) + memories_store = _AsyncStore([]) + turns = [_turn(i) for i in range(3)] + turns_store = _AsyncStore(turns) + service = AsyncPipelineService( + memories_store, + chat, + _AsyncEmbeddings(), + containers=_async_containers_for_store(memories_store, turns_store=turns_store), + ) + + counts = await service.extract_memories("u1", "t1") + + assert counts["deferred_turn_count"] == 0 + assert counts["quarantined_turn_count"] == len(turns) + assert all("extracted_at" in turn for turn in turns_store.docs) + + def _agent_source_response() -> dict[str, Any]: # One agent-sourced fact, one user-sourced fact, one with source omitted # (must default to "user"). Mirrors the extract_memories.prompty schema. @@ -470,7 +651,7 @@ def test_agent_sourced_fact_is_tagged_and_stamped() -> None: containers=_containers_for_store(memories_store, turns_store=_Store([_turn(1)])), ) - out = service.extract_memories_dry("u1", "t1") + out = service.extract_memories_durable("u1", "t1") agent_fact = _fact_by_text(out["facts"], "booked the user on flight") assert agent_fact["metadata"]["source"] == "agent" @@ -496,7 +677,7 @@ async def test_async_agent_sourced_fact_is_tagged_and_stamped() -> None: containers=_async_containers_for_store(memories_store, turns_store=_AsyncStore([_turn(1)])), ) - out = await service.extract_memories_dry("u1", "t1") + out = await service.extract_memories_durable("u1", "t1") agent_fact = _fact_by_text(out["facts"], "booked the user on flight") assert agent_fact["metadata"]["source"] == "agent" @@ -520,7 +701,7 @@ def test_extraction_transcript_includes_turn_timestamps() -> None: containers=_containers_for_store(memories_store, turns_store=turns_store), ) - service.extract_memories_dry("u1", "t1") + service.extract_memories_durable("u1", "t1") prompt_text = json.dumps(chat.messages) assert "2025-01-01T00:01:00+00:00 | user" in prompt_text @@ -541,8 +722,38 @@ def test_extraction_transcript_canonicalizes_speaker_role() -> None: containers=_containers_for_store(memories_store, turns_store=turns_store), ) - service.extract_memories_dry("u1", "t1") + service.extract_memories_durable("u1", "t1") prompt_text = json.dumps(chat.messages) assert "| agent]" in prompt_text assert "| assistant]" not in prompt_text + + +def test_extract_memories_durable_clamps_out_of_range_fact_scores() -> None: + resp = { + "facts": [ + { + "text": "The user loves hiking.", + "action": "ADD", + "category": "preference", + "confidence": 1.4, + "salience": -0.2, + "tags": [], + } + ], + "episodic": [], + } + memories_store = _Store([]) + turns_store = _Store([_turn(1), _turn(2)]) + service = PipelineService( + memories_store, + _SyncChat([resp]), + _SyncEmbeddings(), + containers=_containers_for_store(memories_store, turns_store=turns_store), + ) + + output = service.extract_memories_durable("u1", "t1") + + assert len(output["facts"]) == 1 + assert output["facts"][0]["confidence"] == 1.0 + assert output["facts"][0]["salience"] == 0.0 diff --git a/tests/unit/services/test_extract_episodes.py b/tests/unit/services/test_extract_episodes.py new file mode 100644 index 0000000..c4b9abd --- /dev/null +++ b/tests/unit/services/test_extract_episodes.py @@ -0,0 +1,304 @@ +from __future__ import annotations + +import logging +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, + _response, + _Store, + _SyncChat, + _SyncEmbeddings, + _turn, +) + + +class _TrackingStore(_Store): + def __init__(self, docs: list[dict[str, Any]]): + super().__init__(docs) + self.supersede_calls: list[dict[str, Any]] = [] + + def mark_superseded(self, old_doc: dict[str, Any], superseder_id: str, *, reason: str) -> bool: + self.supersede_calls.append({"old_doc": old_doc, "superseder_id": superseder_id, "reason": reason}) + return True + + +def _episode( + *, + title: str = "Fixed CI retries", + summary: str = "The user fixed flaky CI retries and the tests passed.", + outcome: dict[str, str] | None = None, +) -> dict[str, Any]: + return { + "title": title, + "summary": summary, + "started_at": "2025-01-01T00:01:00+00:00", + "ended_at": "2025-01-01T00:02:00+00:00", + "participants": [], + "events": [ + { + "sequence": 1, + "description": "CI retries were added.", + "occurred_at": "2025-01-01T00:01:00+00:00", + "source_turn_ids": ["turn-1", "missing"], + }, + { + "sequence": 2, + "description": "The tests passed.", + "occurred_at": "2025-01-01T00:02:00+00:00", + "source_turn_ids": ["turn-2", "turn-1"], + }, + ], + "outcome": outcome, + "lessons": ["Use retries for flaky CI."], + "salience": 0.8, + "confidence": 0.9, + } + + +def _service( + responses: list[dict[str, Any]], + *, + memories_store: _TrackingStore | None = None, + embeddings: _SyncEmbeddings | None = None, +) -> tuple[PipelineService, _TrackingStore, _SyncEmbeddings]: + store = memories_store or _TrackingStore([]) + embedding_client = embeddings or _SyncEmbeddings() + turns_store = _Store([_turn(1), _turn(2)]) + service = PipelineService( + store, + _SyncChat(responses), + embedding_client, + containers=_containers_for_store(store, turns_store=turns_store), + ) + return service, store, embedding_client + + +def test_build_episode_docs_returns_multiple_docs_without_embeddings() -> None: + outcome = {"status": "successful", "description": "The tests passed."} + service, _, embeddings = _service( + [ + { + "episodes": [ + _episode(outcome=outcome), + _episode(title="Planned vacation", summary="The user planned a vacation.", outcome=None), + ] + } + ] + ) + + docs = service._build_episode_docs("u1", "t1", [_turn(1), _turn(2)], segment_key="seg-1") + + assert [doc["content"] for doc in docs] == [ + "The user fixed flaky CI retries and the tests passed.", + "The user planned a vacation.", + ] + assert all(doc["id"].startswith("ep_") for doc in docs) + assert all("embedding" not in doc for doc in docs) + assert docs[0]["source_turn_ids"] == ["turn-1", "turn-2"] + assert docs[0]["events"][0]["source_turn_ids"] == ["turn-1"] + assert docs[0]["outcome"] == outcome + assert docs[1]["outcome"] is None + assert embeddings.calls == [] + + +def test_extract_episodes_embeds_content_persists_append_only(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + service, store, embeddings = _service( + [ + { + "episodes": [ + _episode(), + _episode(title="Planned vacation", summary="The user planned a vacation.", outcome=None), + ] + } + ] + ) + + result = service.extract_episodes("u1", "t1", flush=True) + + assert result == {"episodes": 2} + assert embeddings.calls == [ + [ + "The user fixed flaky CI retries and the tests passed.", + "The user planned a vacation.", + ] + ] + assert [doc["content"] for doc in store.docs] == [ + "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 store.supersede_calls == [] + assert store.search_calls == [] + + +def test_extract_episodes_empty_window_persists_nothing(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + service, store, embeddings = _service([{"episodes": []}]) + + result = service.extract_episodes("u1", "t1", flush=True) + + assert result == {"episodes": 0} + assert store.docs == [] + assert embeddings.calls == [] + + +def test_extract_episodes_skips_malformed_episode_with_warning(caplog, monkeypatch) -> None: + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + service, store, _ = _service( + [ + { + "episodes": [ + {**_episode(), "title": ""}, + _episode(title="Valid episode", summary="The valid episode persisted."), + ] + } + ] + ) + + with caplog.at_level(logging.WARNING): + result = service.extract_episodes("u1", "t1", flush=True) + + assert result == {"episodes": 1} + assert [doc["title"] for doc in store.docs] == ["Valid episode"] + assert "dropping malformed episode" in caplog.text + + +def test_extract_memories_durable_keeps_episodic_empty_regression_guard() -> None: + store = _TrackingStore([]) + service = PipelineService( + store, + _SyncChat([_response()]), + _SyncEmbeddings(), + containers=_containers_for_store(store, turns_store=_Store([_turn(1)])), + ) + + output = service.extract_memories_durable("u1", "t1") + + assert output["facts"] + assert output["episodic"] == [] + + +class _IdUniqueStore(_TrackingStore): + """Store that enforces id-uniqueness on create, like Cosmos (409 on dup).""" + + 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 test_build_episode_docs_id_stable_across_summary_text() -> None: + service, _, _ = _service( + [ + {"episodes": [_episode(summary="One phrasing of the CI-retry episode.")]}, + {"episodes": [_episode(summary="A completely different phrasing entirely.")]}, + ] + ) + first = service._build_episode_docs("u1", "t1", [_turn(1), _turn(2)], segment_key="seg-1") + second = service._build_episode_docs("u1", "t1", [_turn(1), _turn(2)], segment_key="seg-1") + + assert first[0]["content"] != second[0]["content"] + assert first[0]["id"] == second[0]["id"] + assert first[0]["content_hash"] != second[0]["content_hash"] + + +def test_build_episode_docs_multiple_episodes_get_distinct_ordinal_ids() -> None: + service, _, _ = _service( + [{"episodes": [_episode(), _episode(title="Vacation", summary="The user planned a vacation.")]}] + ) + docs = service._build_episode_docs("u1", "t1", [_turn(1), _turn(2)], segment_key="seg-1") + assert len({doc["id"] for doc in docs}) == 2 + + +def test_extract_episodes_skips_duplicate_when_segment_reprocessed(monkeypatch) -> None: + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0") + store = _IdUniqueStore([]) + turns = _Store([_turn(1), _turn(2)]) + service = PipelineService( + store, + _SyncChat( + [ + {"episodes": [_episode(summary="First run prose.")]}, + {"episodes": [_episode(summary="Second run, different prose.")]}, + ] + ), + _SyncEmbeddings(), + containers=_containers_for_store(store, turns_store=turns), + ) + + 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) + + assert 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 + + +def test_build_episode_docs_falls_back_to_segment_times_on_unparseable_llm_times() -> None: + bad = _episode() + bad["started_at"] = "March 9th" + bad["ended_at"] = "2025-01-01T00:02:00+00:00" + service, _, _ = _service([{"episodes": [bad]}]) + + docs = service._build_episode_docs("u1", "t1", [_turn(1), _turn(2)], segment_key="seg-1") + + assert len(docs) == 1 + assert docs[0]["started_at"] == "2025-01-01T00:01:00+00:00" + assert docs[0]["ended_at"] == "2025-01-01T00:02:00+00:00" + + +def test_build_episode_docs_keeps_mixed_tz_llm_times_after_normalization() -> None: + mixed = _episode() + mixed["started_at"] = "2026-03-09" + mixed["ended_at"] = "2026-03-10T09:08:00+00:00" + service, _, _ = _service([{"episodes": [mixed]}]) + + docs = service._build_episode_docs("u1", "t1", [_turn(1), _turn(2)], segment_key="seg-1") + + assert len(docs) == 1 + assert docs[0]["started_at"] == "2026-03-09" + assert docs[0]["ended_at"] == "2026-03-10T09:08:00+00:00" + + +def test_build_episode_docs_keeps_episode_with_padded_timestamps() -> None: + ep = _episode() + ep["started_at"] = " 2025-01-01T00:01:00+00:00" + ep["ended_at"] = "2025-01-01T00:02:00+00:00 " + service, _, _ = _service([{"episodes": [ep]}]) + + docs = service._build_episode_docs("u1", "t1", [_turn(1), _turn(2)], segment_key="seg-1") + + assert len(docs) == 1 + assert docs[0]["started_at"] == "2025-01-01T00:01:00+00:00" + assert docs[0]["ended_at"] == "2025-01-01T00:02:00+00:00" + + +def test_build_episode_docs_clamps_out_of_range_salience_confidence() -> None: + ep = _episode() + ep["salience"] = 1.4 + ep["confidence"] = -0.2 + service, _, _ = _service([{"episodes": [ep]}]) + + docs = service._build_episode_docs("u1", "t1", [_turn(1), _turn(2)], segment_key="seg-1") + + assert len(docs) == 1 + assert docs[0]["salience"] == 1.0 + assert docs[0]["confidence"] == 0.0 + + +def test_build_episode_docs_drops_whitespace_only_summary() -> None: + service, _, _ = _service([{"episodes": [_episode(summary=" ")]}]) + + docs = service._build_episode_docs("u1", "t1", [_turn(1), _turn(2)], segment_key="seg-1") + + assert docs == [] diff --git a/tests/unit/services/test_extract_prompt_selection.py b/tests/unit/services/test_extract_prompt_selection.py new file mode 100644 index 0000000..24a97c3 --- /dev/null +++ b/tests/unit/services/test_extract_prompt_selection.py @@ -0,0 +1,26 @@ +"""Tests for the env-selectable fact-extraction prompt. + +The v2 extractor is the shipped default; an env override can select v1, and an +unknown value falls back to the v2 default (safe allowlist). +""" + +from __future__ import annotations + +import pytest + +from azure.cosmos.agent_memory.services._pipeline_helpers import extract_memories_prompt_file + + +def test_default_is_v2(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("AMT_EXTRACT_MEMORIES_PROMPT", raising=False) + assert extract_memories_prompt_file() == "extract_memories-v2.prompty" + + +def test_env_override_selects_v1(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("AMT_EXTRACT_MEMORIES_PROMPT", "extract_memories.prompty") + assert extract_memories_prompt_file() == "extract_memories.prompty" + + +def test_unknown_value_falls_back_to_v2(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("AMT_EXTRACT_MEMORIES_PROMPT", "bogus.prompty") + assert extract_memories_prompt_file() == "extract_memories-v2.prompty" diff --git a/tests/unit/services/test_extraction_batching.py b/tests/unit/services/test_extraction_batching.py index 5793b17..a3ba107 100644 --- a/tests/unit/services/test_extraction_batching.py +++ b/tests/unit/services/test_extraction_batching.py @@ -29,6 +29,27 @@ def test_unknown_error_defaults_retryable(self) -> None: # Conservative: never quarantine (drop) turns on an unclassified error. assert is_retryable_llm_error(Exception("something weird happened")) is True + def test_programming_errors_are_non_retryable(self) -> None: + # A bug in our own parse/extract code (e.g. .get on a non-dict) is + # deterministic - retrying re-fails identically and would wedge the + # segment/batch forever. Classify by type as non-retryable so it is + # quarantined and surfaced instead of deferred indefinitely. + assert is_retryable_llm_error(AttributeError("'list' object has no attribute 'get'")) is False + assert is_retryable_llm_error(KeyError("episodes")) is False + assert is_retryable_llm_error(TypeError("unhashable type")) is False + assert is_retryable_llm_error(IndexError("list index out of range")) is False + + def test_programming_error_type_beats_retryable_looking_message(self) -> None: + # Type wins over text: a programming error stays non-retryable even when + # its message would otherwise look transient. + assert is_retryable_llm_error(KeyError("429 rate limit")) is False + + def test_transient_provider_error_stays_retryable(self) -> None: + # Genuine transient provider failures are their own SDK exception types + # (not the programming-error allow-list) and must remain retryable - the + # allow-list is deliberately narrow, NOT "any non-LLMError". + assert is_retryable_llm_error(RuntimeError("Error code: 429 - rate limit")) is True + class TestBatchTurnsByTokens: def _turns(self, n, content="word " * 10): diff --git a/tests/unit/services/test_parse_llm_json.py b/tests/unit/services/test_parse_llm_json.py index 2f7ddd2..02b2e06 100644 --- a/tests/unit/services/test_parse_llm_json.py +++ b/tests/unit/services/test_parse_llm_json.py @@ -38,10 +38,26 @@ def test_object_then_garbage_is_salvaged(self) -> None: def test_fenced_object_with_trailing_duplicate(self) -> None: assert parse_llm_json('```json\n{"a": 1}\n``` {"a": 1}') == {"a": 1} - def test_trailing_data_emits_warning(self, caplog: pytest.LogCaptureFixture) -> None: - with caplog.at_level(logging.WARNING, logger=_HELPER_LOGGER): + def test_doubled_object_emits_merge_info(self, caplog: pytest.LogCaptureFixture) -> None: + # Two clean concatenated objects are merged (so no items are dropped); + # this is a recoverable case logged at INFO, not a warning. + with caplog.at_level(logging.INFO, logger=_HELPER_LOGGER): parse_llm_json(_DOUBLED) - assert any("extra data after the first JSON object" in r.message for r in caplog.records) + assert any("concatenated JSON objects" in r.message for r in caplog.records) + + def test_doubled_object_with_nonempty_lists_concatenates(self) -> None: + # Discriminating merge test: list fields from BOTH objects are combined, + # so a broken (no-op) merge that kept only the first object would fail here. + doubled = '{"facts":[{"t":"a"}],"episodic":[]} {"facts":[{"t":"b"}],"episodic":[{"e":"x"}]}' + merged = parse_llm_json(doubled) + assert merged["facts"] == [{"t": "a"}, {"t": "b"}] + assert merged["episodic"] == [{"e": "x"}] + + def test_trailing_garbage_emits_warning(self, caplog: pytest.LogCaptureFixture) -> None: + # A valid object followed by non-JSON garbage is salvaged but warned. + with caplog.at_level(logging.WARNING, logger=_HELPER_LOGGER): + parse_llm_json('{"facts": [{"text": "x"}]} ') + assert any("non-JSON trailing data" in r.message for r in caplog.records) def test_clean_object_emits_no_warning(self, caplog: pytest.LogCaptureFixture) -> None: with caplog.at_level(logging.WARNING, logger=_HELPER_LOGGER): @@ -78,3 +94,30 @@ def test_non_json_raises_invalid_not_truncated(self) -> None: parse_llm_json("I could not find any memories.") assert "invalid JSON" in str(exc.value) assert "TRUNCATED" not in str(exc.value) + + +class TestNonObjectRoot: + """A non-object JSON root must raise LLMError, not leak a non-dict. + + Callers immediately do ``parsed.get("facts"/"episodes")``; a bare list/scalar + would raise AttributeError downstream, which the extraction error-handlers + then misclassify as a transient (retryable) failure and defer forever. + """ + + def test_bare_array_root_raises(self) -> None: + with pytest.raises(LLMError) as exc: + parse_llm_json('[{"text": "x"}]') + assert "non-object JSON root" in str(exc.value) + + def test_bare_array_root_with_trailing_data_raises(self) -> None: + with pytest.raises(LLMError) as exc: + parse_llm_json('[1, 2] {"a": 1}') + assert "non-object JSON root" in str(exc.value) + + def test_scalar_number_root_raises(self) -> None: + with pytest.raises(LLMError): + parse_llm_json("42") + + def test_scalar_string_root_raises(self) -> None: + with pytest.raises(LLMError): + parse_llm_json('"just a sentence"') diff --git a/tests/unit/services/test_persist_extracted.py b/tests/unit/services/test_persist_extracted.py index 01907f7..366537b 100644 --- a/tests/unit/services/test_persist_extracted.py +++ b/tests/unit/services/test_persist_extracted.py @@ -131,6 +131,30 @@ def _fact_doc(content: str = "The user prefers dark mode.") -> dict[str, Any]: } +def _episodic_doc(content: str = "The user debugged a flaky CI test together with the agent.") -> dict[str, Any]: + content_hash = compute_content_hash(content) + seed = ID_SEED_SEP.join(("u1", "t1", content_hash)) + return { + "id": f"ep_{hashlib.sha256(seed.encode()).hexdigest()[:32]}", + "user_id": "u1", + "thread_id": "t1", + "type": "episodic", + "title": "CI flake debugging", + "content": content, + "content_hash": content_hash, + "started_at": "2025-01-01T00:00:00+00:00", + "ended_at": "2025-01-01T00:05:00+00:00", + "confidence": 0.8, + "salience": 0.7, + "tags": ["sys:episodic"], + "prompt_id": "extract_episode.prompty", + "prompt_version": "v1", + "metadata": {}, + "created_at": "2025-01-01T00:00:00+00:00", + "updated_at": "2025-01-01T00:00:00+00:00", + } + + def test_persist_extracted_memories_uses_deterministic_ids_and_skips_replay() -> None: container = _Container() store = _Store(container) @@ -196,6 +220,29 @@ def test_persist_extracted_memories_409_skip_continues_to_next_doc() -> None: assert container.created_ids == [second["id"]] +def test_persist_extracted_memories_episodic_creates_not_upserts_and_skips_replay() -> None: + # Episodic docs persist via create + 409 (first-write-wins), matching the + # extract_episodes path - never a last-write-wins upsert. A replayed + # identical episode is skipped, not overwritten. + container = _Container() + store = _Store(container) + service = PipelineService( + store, + chat_client=object(), + embeddings_client=_Embeddings(), + containers=_containers_for_store(store), + ) + doc = _episodic_doc() + + first = service.persist_extracted_memories("u1", {"facts": [], "episodic": [doc], "updates": []}) + second = service.persist_extracted_memories("u1", {"facts": [], "episodic": [doc], "updates": []}) + + assert first["episodic_count"] == 1 + assert second["episodic_count"] == 0 + assert container.created_ids == [doc["id"]] + assert store.upserts == [] + + @pytest.mark.asyncio async def test_async_persist_extracted_memories_uses_deterministic_ids_and_skips_replay() -> None: container = _AsyncContainer() @@ -234,3 +281,26 @@ async def test_async_persist_extracted_memories_409_skip_continues_to_next_doc() assert result["fact_count"] == 1 assert container.created_ids == [second["id"]] + + +@pytest.mark.asyncio +async def test_async_persist_extracted_memories_episodic_creates_not_upserts_and_skips_replay() -> None: + # Episodic docs persist via create + 409 (first-write-wins), matching the + # extract_episodes path - never a last-write-wins upsert. + container = _AsyncContainer() + store = _AsyncStore(container) + service = AsyncPipelineService( + store, + chat_client=object(), + embeddings_client=_AsyncEmbeddings(), + containers=_async_containers_for_store(store), + ) + doc = _episodic_doc() + + first = await service.persist_extracted_memories("u1", {"facts": [], "episodic": [doc], "updates": []}) + second = await service.persist_extracted_memories("u1", {"facts": [], "episodic": [doc], "updates": []}) + + assert first["episodic_count"] == 1 + assert second["episodic_count"] == 0 + assert container.created_ids == [doc["id"]] + assert store.upserts == [] diff --git a/tests/unit/services/test_pipeline_helpers_episode.py b/tests/unit/services/test_pipeline_helpers_episode.py new file mode 100644 index 0000000..4349e4c --- /dev/null +++ b/tests/unit/services/test_pipeline_helpers_episode.py @@ -0,0 +1,137 @@ +"""Unit tests for the shared, IO-free episode segmentation helpers. + +These pure functions live in ``services/_pipeline_helpers.py`` so the sync and +aio pipelines share one implementation. A single test file covers both. +""" + +from __future__ import annotations + +from azure.cosmos.agent_memory.services._pipeline_helpers import ( + created_at_sort_key, + deterministic_episode_id, + find_episode_boundary, + is_valid_time_pair, + parse_iso_datetime, + segment_time_bounds, + turn_gap_seconds, +) + + +def _turn(minute: int) -> dict[str, object]: + return {"id": f"turn-{minute}", "created_at": f"2025-01-01T00:{minute:02d}:00+00:00"} + + +class TestDeterministicEpisodeId: + def test_stable_for_same_segment_and_index(self) -> None: + assert deterministic_episode_id("seg", 0) == deterministic_episode_id("seg", 0) + + def test_differs_by_index(self) -> None: + assert deterministic_episode_id("seg", 0) != deterministic_episode_id("seg", 1) + + def test_differs_by_segment_key(self) -> None: + assert deterministic_episode_id("seg-a", 0) != deterministic_episode_id("seg-b", 0) + + def test_ep_prefix(self) -> None: + assert deterministic_episode_id("seg", 0).startswith("ep_") + + +class TestTimeHelpers: + def test_parse_naive_gets_utc(self) -> None: + dt = parse_iso_datetime("2025-01-01T00:00:00") + assert dt is not None and dt.tzinfo is not None + + def test_parse_unparseable_is_none(self) -> None: + assert parse_iso_datetime("not a date") is None + + def test_is_valid_time_pair_mixed_tz(self) -> None: + assert is_valid_time_pair("2026-03-09", "2026-03-10T09:08:00+00:00") is True + + def test_is_valid_time_pair_reversed(self) -> None: + assert is_valid_time_pair("2026-03-10T00:00:00+00:00", "2026-03-09") is False + + def test_turn_gap_seconds(self) -> None: + assert turn_gap_seconds(_turn(1), _turn(3)) == 120.0 + + def test_turn_gap_seconds_missing_is_none(self) -> None: + assert turn_gap_seconds({"id": "x"}, _turn(3)) is None + + def test_segment_time_bounds(self) -> None: + assert segment_time_bounds([_turn(3), _turn(1), _turn(2)]) == ( + "2025-01-01T00:01:00+00:00", + "2025-01-01T00:03:00+00:00", + ) + + def test_segment_time_bounds_empty(self) -> None: + assert segment_time_bounds([]) == (None, None) + + def test_segment_time_bounds_mixed_offsets_sort_chronologically(self) -> None: + items = [ + {"created_at": "2025-01-01T05:00:00Z"}, + {"created_at": "2025-01-01T09:00:00+05:00"}, + ] + assert segment_time_bounds(items) == ( + "2025-01-01T09:00:00+05:00", + "2025-01-01T05:00:00Z", + ) + + +class TestCreatedAtSortKey: + def test_tied_timestamps_break_deterministically_on_id(self) -> None: + # Turns sharing one timestamp (common when many turns carry the same + # session date) must keep a stable, id-ordered sequence so the segment's + # first/last ids - and thus the deterministic episode id - do not drift. + items = [ + {"id": "c", "created_at": "2025-01-01T00:00:00+00:00"}, + {"id": "a", "created_at": "2025-01-01T00:00:00+00:00"}, + {"id": "b", "created_at": "2025-01-01T00:00:00+00:00"}, + ] + assert [i["id"] for i in sorted(items, key=created_at_sort_key)] == ["a", "b", "c"] + + def test_mixed_offsets_sort_by_true_instant_not_lexical(self) -> None: + # 09:00+05:00 (=04:00Z) precedes 05:00Z chronologically, though the raw + # string "09..." sorts after "05..." lexically. + items = [ + {"id": "later", "created_at": "2025-01-01T05:00:00Z"}, + {"id": "earlier", "created_at": "2025-01-01T09:00:00+05:00"}, + ] + assert [i["id"] for i in sorted(items, key=created_at_sort_key)] == ["earlier", "later"] + + def test_missing_or_unparseable_sorts_last_by_id(self) -> None: + items = [ + {"id": "no-ts"}, + {"id": "bad-ts", "created_at": "not-a-date"}, + {"id": "has-ts", "created_at": "2025-01-01T00:00:00Z"}, + ] + assert [i["id"] for i in sorted(items, key=created_at_sort_key)] == ["has-ts", "bad-ts", "no-ts"] + + +class TestFindEpisodeBoundary: + _KN = {"max_turns": 40, "idle_gap": 120, "drift": 0.0, "min_turns": 2} + + def test_no_boundary_returns_none(self) -> None: + seg = [_turn(1), _turn(2), _turn(3)] + assert find_episode_boundary(seg, [], **self._KN) is None + + def test_idle_gap_closes_at_index(self) -> None: + # 00:01, 00:02, then a 28-minute gap to 00:30 -> boundary at i=2. + seg = [_turn(1), _turn(2), _turn(30)] + assert find_episode_boundary(seg, [], **self._KN) == 2 + + def test_idle_gap_below_min_turns_is_suppressed(self) -> None: + # Gap between i=0 and i=1 is below min_turns=2 -> not closed. + seg = [_turn(1), _turn(30), _turn(31)] + assert find_episode_boundary(seg, [], **self._KN) is None + + def test_max_turns_cap_closes(self) -> None: + seg = [_turn(i) for i in range(1, 6)] + assert find_episode_boundary(seg, [], max_turns=3, idle_gap=0, drift=0.0, min_turns=2) == 3 + + def test_max_turns_not_floored_when_below_min_turns(self) -> None: + # The max-size cap is a hard ceiling and is NOT floored by min_turns. + seg = [_turn(i) for i in range(1, 5)] + assert find_episode_boundary(seg, [], max_turns=2, idle_gap=0, drift=0.0, min_turns=3) == 2 + + def test_earliest_boundary_wins_across_signals(self) -> None: + # Idle gap at i=2 vs a max cap at i=3: the earliest boundary (i=2) wins. + seg = [_turn(1), _turn(2), _turn(30), _turn(31)] + assert find_episode_boundary(seg, [], max_turns=3, idle_gap=120, drift=0.0, min_turns=2) == 2 diff --git a/tests/unit/services/test_pipeline_service.py b/tests/unit/services/test_pipeline_service.py index 0800ab1..bbec77e 100644 --- a/tests/unit/services/test_pipeline_service.py +++ b/tests/unit/services/test_pipeline_service.py @@ -94,8 +94,13 @@ def query( docs = [doc for doc in docs if doc.get("metadata", {}).get("predicate") == params["@predicate"]] if "superseded_by" in sql: docs = [doc for doc in docs if not doc.get("superseded_by")] - if "IS_DEFINED(c.metadata.lesson)" in sql: - docs = [doc for doc in docs if doc.get("metadata", {}).get("lesson")] + if "IS_DEFINED(c.lessons)" in sql: + docs = [ + doc + for doc in docs + if isinstance(doc.get("lessons"), list) + and any(isinstance(lesson, str) and lesson.strip() for lesson in doc.get("lessons", [])) + ] if "source_memory_ids" not in sql and "ORDER BY c.created_at DESC" in sql: docs.sort(key=lambda doc: doc.get("created_at", ""), reverse=True) elif "ORDER BY c.version DESC" in sql: @@ -191,7 +196,7 @@ def _fact(fid: str, content: str, **extra: Any) -> dict[str, Any]: } -def test_extract_memories_happy_path_writes_fact_and_episodic() -> None: +def test_extract_memories_happy_path_writes_fact_only() -> None: store = FakeStore() turns_store = FakeStore([_turn("I prefer dark mode and learned CI needs retries.")]) llm = FakeLLMService( @@ -207,17 +212,7 @@ def test_extract_memories_happy_path_writes_fact_and_episodic() -> None: "tags": ["ui"], } ], - "episodic": [ - { - "scope_type": "project", - "scope_value": "CI", - "situation": "CI tests flaked intermittently", - "action_taken": "Added retries", - "outcome": "Tests stabilized", - "lesson": "Use retries for flaky CI tests.", - "confidence": 0.8, - } - ], + "episodic": [], } ] ) @@ -225,17 +220,12 @@ def test_extract_memories_happy_path_writes_fact_and_episodic() -> None: result = _pipeline(store, llm, turns_store=turns_store).extract_memories("u1", "t1") assert result["fact_count"] == 1 - assert result["episodic_count"] == 1 + assert result["episodic_count"] == 0 assert result["updated_count"] == 0 - assert [doc["type"] for doc in store.upserts] == ["fact", "episodic"] + assert [doc["type"] for doc in store.upserts] == ["fact"] assert set(store.upserts[0]["tags"]) == {"sys:fact", "sys:auto-extracted", "topic:ui"} assert llm.chat_calls - assert llm.embed_calls == [ - [ - "The user prefers dark mode.", - "CI tests flaked intermittently → Added retries → Tests stabilized", - ] - ] + assert llm.embed_calls == [["The user prefers dark mode."]] def test_extract_memories_creates_new_fact_without_superseding() -> None: @@ -277,7 +267,7 @@ def test_synthesize_procedural_produces_procedural_memory() -> None: "role": "system", "type": "episodic", "content": "Past project", - "metadata": {"lesson": "Keep examples small."}, + "lessons": ["Keep examples small."], "salience": 0.8, "created_at": "2025-01-02T00:00:00+00:00", }, diff --git a/tests/unit/services/test_prompty_loader.py b/tests/unit/services/test_prompty_loader.py index 76c46da..29b75ed 100644 --- a/tests/unit/services/test_prompty_loader.py +++ b/tests/unit/services/test_prompty_loader.py @@ -56,10 +56,12 @@ def test_loader_prompt_version_is_cached(tmp_path: Path) -> None: def test_all_shipped_prompts_declare_version() -> None: loader = PromptyLoader() - # extract_memories bumped to v2 when agent-sourced fact extraction landed; - # the rest remain v1. Every shipped prompt must declare *some* version. + # extract_memories bumped to v4 when episodic extraction was removed + # (facts-only); the rest remain v1. Every shipped prompt must declare + # *some* version. expected = { - "extract_memories.prompty": "v3", + "extract_memories.prompty": "v4", + "extract_episode.prompty": "v1", "dedup.prompty": "v1", "summarize.prompty": "v1", "summarize_update.prompty": "v1", diff --git a/tests/unit/store/test_memory_store.py b/tests/unit/store/test_memory_store.py index aaf9ff5..a54bf0e 100644 --- a/tests/unit/store/test_memory_store.py +++ b/tests/unit/store/test_memory_store.py @@ -344,19 +344,21 @@ def test_search_all_stopwords_falls_back_to_vector_only(): def test_search_episodic_forwards_search_options(): - store = MemoryStore(containers=_containers()) - store.search = MagicMock(return_value=[]) + memories = MagicMock() + memories.query_items.return_value = [] + embeddings = MagicMock() + embeddings.generate.return_value = [0.1, 0.2] + store = MemoryStore(containers=_containers(memories=memories), embeddings_client=embeddings) - store.search_episodic("u1", "weather") + store.search_episodic("u1", "weather", top_k=3, min_salience=0.5) - store.search.assert_called_once_with( - search_terms="weather", - user_id="u1", - memory_types=["episodic"], - top_k=5, - min_salience=None, - include_superseded=False, - ) + kwargs = memories.query_items.call_args.kwargs + assert "TOP 3" in kwargs["query"] + assert "c.type = @type" in kwargs["query"] + params = _params_by_name(kwargs) + assert params["@type"] == "episodic" + assert params["@user_id"] == "u1" + assert params["@min_salience"] == 0.5 def test_build_episodic_context_forwards_search_options(): diff --git a/tests/unit/test_auto_trigger.py b/tests/unit/test_auto_trigger.py index a2891eb..5ed82aa 100644 --- a/tests/unit/test_auto_trigger.py +++ b/tests/unit/test_auto_trigger.py @@ -12,6 +12,7 @@ from azure.cosmos.exceptions import CosmosResourceNotFoundError +from azure.cosmos.agent_memory.auto_trigger import maybe_trigger_steps from azure.cosmos.agent_memory.cosmos_memory_client import CosmosMemoryClient from azure.cosmos.agent_memory.processors import DurableFunctionProcessor, InProcessProcessor @@ -65,6 +66,7 @@ def _connected(processor=None) -> CosmosMemoryClient: def test_push_to_cosmos_fires_inprocess_trigger_per_turn(monkeypatch): monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") + monkeypatch.setenv("EPISODE_EVAL_EVERY_N", "0") client = _connected(processor=InProcessProcessor(pipeline=MagicMock())) counter_container = MagicMock() @@ -88,6 +90,7 @@ def test_push_to_cosmos_fires_inprocess_trigger_per_turn(monkeypatch): def test_push_to_cosmos_durable_does_not_fire_trigger(monkeypatch): monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") + monkeypatch.setenv("EPISODE_EVAL_EVERY_N", "0") client = _connected(processor=DurableFunctionProcessor()) client._counter_container_client = MagicMock() @@ -104,6 +107,7 @@ def test_push_to_cosmos_durable_does_not_fire_trigger(monkeypatch): def test_push_to_cosmos_skips_trigger_when_thresholds_zero(monkeypatch): monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "0") monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") + monkeypatch.setenv("EPISODE_EVAL_EVERY_N", "0") monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") client = _connected(processor=InProcessProcessor(pipeline=MagicMock())) @@ -122,6 +126,7 @@ def test_push_to_cosmos_skips_trigger_when_thresholds_zero(monkeypatch): def test_push_to_cosmos_swallows_trigger_failures(monkeypatch): """Auto-trigger errors must never propagate from push_to_cosmos.""" monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") + monkeypatch.setenv("EPISODE_EVAL_EVERY_N", "0") pipeline = MagicMock() pipeline.generate_thread_summary.side_effect = RuntimeError("boom") @@ -138,6 +143,7 @@ def test_push_to_cosmos_swallows_trigger_failures(monkeypatch): def test_push_to_cosmos_skips_when_counter_container_unavailable(monkeypatch): monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") + monkeypatch.setenv("EPISODE_EVAL_EVERY_N", "0") client = _connected(processor=InProcessProcessor(pipeline=MagicMock())) # Counter container handle stays None; lazy getter would normally try to @@ -162,10 +168,75 @@ def test_push_to_cosmos_skips_when_counter_container_unavailable(monkeypatch): class TestPerStepAutoTrigger: + def test_episode_zero_does_not_fire(self): + processor = InProcessProcessor(pipeline=MagicMock()) + processor.process_extract_episodes = MagicMock() + counter_container = _FakeCounterContainer() + + maybe_trigger_steps( + processor, + counter_container, + {("u1", "t1"): 1}, + thresholds={ + "FACT_EXTRACTION_EVERY_N": 0, + "THREAD_SUMMARY_EVERY_N": 0, + "EPISODE_EVAL_EVERY_N": 0, + "USER_SUMMARY_EVERY_N": 0, + "MEMORY_PROCESSOR_OWNER": "inprocess", + }, + ) + + processor.process_extract_episodes.assert_not_called() + assert counter_container.store == {} + + def test_episode_fires_when_threshold_crossed(self): + processor = InProcessProcessor(pipeline=MagicMock()) + processor.process_extract_episodes = MagicMock(return_value={}) + counter_container = _FakeCounterContainer() + thresholds = { + "FACT_EXTRACTION_EVERY_N": 0, + "THREAD_SUMMARY_EVERY_N": 0, + "EPISODE_EVAL_EVERY_N": 3, + "USER_SUMMARY_EVERY_N": 0, + "MEMORY_PROCESSOR_OWNER": "inprocess", + } + + maybe_trigger_steps(processor, counter_container, {("u1", "t1"): 2}, thresholds=thresholds) + processor.process_extract_episodes.assert_not_called() + + maybe_trigger_steps(processor, counter_container, {("u1", "t1"): 1}, thresholds=thresholds) + + processor.process_extract_episodes.assert_called_once_with(user_id="u1", thread_id="t1") + + def test_episode_failure_is_caught_and_other_steps_continue(self): + processor = InProcessProcessor(pipeline=MagicMock()) + processor.process_extract_episodes = MagicMock(side_effect=RuntimeError("episode boom")) + processor.process_thread_summary = MagicMock(return_value={}) + counter_container = _FakeCounterContainer() + + with patch("azure.cosmos.agent_memory._counters.stamp_failure_sync") as stamp: + maybe_trigger_steps( + processor, + counter_container, + {("u1", "t1"): 1}, + thresholds={ + "FACT_EXTRACTION_EVERY_N": 0, + "THREAD_SUMMARY_EVERY_N": 1, + "EPISODE_EVAL_EVERY_N": 1, + "USER_SUMMARY_EVERY_N": 0, + "MEMORY_PROCESSOR_OWNER": "inprocess", + }, + ) + + processor.process_extract_episodes.assert_called_once_with(user_id="u1", thread_id="t1") + processor.process_thread_summary.assert_called_once_with(user_id="u1", thread_id="t1") + stamp.assert_called_once() + def test_extract_fires_independently_of_summary(self, monkeypatch): """N_facts=1 alone fires extract; summary/user-summary stay quiet.""" monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "10") + monkeypatch.setenv("EPISODE_EVAL_EVERY_N", "0") monkeypatch.setenv("USER_SUMMARY_EVERY_N", "20") processor = InProcessProcessor(pipeline=MagicMock()) @@ -193,6 +264,7 @@ def test_extract_fires_without_recent_k_or_watermark(self, monkeypatch): NO recent_k and tracks NO success-gated watermark.""" monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") + monkeypatch.setenv("EPISODE_EVAL_EVERY_N", "0") monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") processor = InProcessProcessor(pipeline=MagicMock()) @@ -216,6 +288,7 @@ def test_extract_failure_stamps_failure(self, monkeypatch): the pipeline, so this outer path only sees unexpected total failures.""" monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") + monkeypatch.setenv("EPISODE_EVAL_EVERY_N", "0") monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") processor = InProcessProcessor(pipeline=MagicMock()) @@ -242,6 +315,7 @@ def test_summary_fires_independently_when_threshold_crossed(self, monkeypatch): """N_summary=10 boundary fires summary; N_facts=0 prevents extract.""" monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "0") monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "10") + monkeypatch.setenv("EPISODE_EVAL_EVERY_N", "0") monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") processor = InProcessProcessor(pipeline=MagicMock()) @@ -265,6 +339,7 @@ def test_user_summary_fires_at_user_threshold(self, monkeypatch): """The user-scoped counter is incremented separately from the thread counter.""" monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "0") monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") + monkeypatch.setenv("EPISODE_EVAL_EVERY_N", "0") monkeypatch.setenv("USER_SUMMARY_EVERY_N", "2") processor = InProcessProcessor(pipeline=MagicMock()) @@ -295,6 +370,7 @@ def test_user_summary_fires_at_user_threshold(self, monkeypatch): class TestProcessorOwner: def test_durable_owner_suppresses_sdk_trigger(self, monkeypatch): monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") + monkeypatch.setenv("EPISODE_EVAL_EVERY_N", "0") monkeypatch.setenv("MEMORY_PROCESSOR_OWNER", "durable") processor = InProcessProcessor(pipeline=MagicMock()) @@ -316,6 +392,7 @@ def test_durable_owner_suppresses_sdk_trigger(self, monkeypatch): def test_inprocess_owner_allows_sdk_trigger(self, monkeypatch): monkeypatch.setenv("FACT_EXTRACTION_EVERY_N", "1") monkeypatch.setenv("THREAD_SUMMARY_EVERY_N", "0") + monkeypatch.setenv("EPISODE_EVAL_EVERY_N", "0") monkeypatch.setenv("USER_SUMMARY_EVERY_N", "0") monkeypatch.setenv("MEMORY_PROCESSOR_OWNER", "inprocess") diff --git a/tests/unit/test_chat.py b/tests/unit/test_chat.py index a5ad5d9..9cf6bac 100644 --- a/tests/unit/test_chat.py +++ b/tests/unit/test_chat.py @@ -185,6 +185,59 @@ def test_generate_exhausts_retries_on_rate_limit(): ) +def _api_status_error(status_code: int, headers: dict[str, str] | None = None): + import httpx + import openai + + request = httpx.Request("POST", "https://test.openai.azure.com/openai/deployments/test/chat/completions") + response = httpx.Response(status_code, headers=headers or {}, request=request) + return openai.APIStatusError(message=f"status {status_code}", response=response, body=None) + + +def test_generate_retryable_api_error_honors_retry_after(monkeypatch): + client = ChatClient(endpoint="https://test.openai.azure.com", api_key="test-key") + + mock_choice = MagicMock() + mock_choice.message.content = "recovered" + mock_response = MagicMock(choices=[mock_choice], usage=None) + mock_openai_client = MagicMock() + mock_openai_client.chat.completions.create.side_effect = [ + _api_status_error(503, {"retry-after": "30"}), + mock_response, + ] + client._client = mock_openai_client + sleeps: list[float] = [] + monkeypatch.setattr("azure.cosmos.agent_memory.chat.random.random", lambda: 0.0) + monkeypatch.setattr("azure.cosmos.agent_memory.chat.time.sleep", sleeps.append) + + result = client.generate([{"role": "user", "content": "test"}], max_retries=2, base_delay=2.0) + + assert result == "recovered" + assert sleeps == [30.0] + + +def test_generate_retryable_api_error_falls_back_without_retry_after(monkeypatch): + client = ChatClient(endpoint="https://test.openai.azure.com", api_key="test-key") + + mock_choice = MagicMock() + mock_choice.message.content = "recovered" + mock_response = MagicMock(choices=[mock_choice], usage=None) + mock_openai_client = MagicMock() + mock_openai_client.chat.completions.create.side_effect = [ + _api_status_error(503), + mock_response, + ] + client._client = mock_openai_client + sleeps: list[float] = [] + monkeypatch.setattr("azure.cosmos.agent_memory.chat.random.random", lambda: 0.0) + monkeypatch.setattr("azure.cosmos.agent_memory.chat.time.sleep", sleeps.append) + + result = client.generate([{"role": "user", "content": "test"}], max_retries=2, base_delay=2.0) + + assert result == "recovered" + assert sleeps == [1.6] + + # --------------------------------------------------------------------------- # generate() – non-retryable errors # --------------------------------------------------------------------------- diff --git a/tests/unit/test_cosmos_memory_client.py b/tests/unit/test_cosmos_memory_client.py index eade659..61b0142 100644 --- a/tests/unit/test_cosmos_memory_client.py +++ b/tests/unit/test_cosmos_memory_client.py @@ -1049,7 +1049,7 @@ def test_search_cosmos_forwards_search_options_to_store(self): memory_id=None, user_id=None, role=None, - memory_types=None, + memory_types=["fact"], thread_id=None, top_k=5, tags_all=None, diff --git a/tests/unit/test_memory_type_multi.py b/tests/unit/test_memory_type_multi.py index f76f357..df4528f 100644 --- a/tests/unit/test_memory_type_multi.py +++ b/tests/unit/test_memory_type_multi.py @@ -4,8 +4,11 @@ forward to it: ``search_cosmos``, ``get_memories``, ``get_thread``. A non-empty list emits ``c.type IN (@memory_type_0, @memory_type_1, ...)``. -``None`` (default) or an empty list disables the type filter so the call -returns every memory type. +``None`` (default) or an empty list disables the type filter at the +``_build_memory_query_builder`` level. Note ``search_cosmos`` overrides this: +its base search defaults to facts, and ``episodic`` is folded into the same +base query only when ``include_episodes=True`` (no separate episodic query), +so None/empty there resolves to ``["fact"]``. """ from __future__ import annotations @@ -183,8 +186,10 @@ def test_get_thread_does_not_accept_memory_types(): client.get_thread(thread_id="t1", memory_types=["turn", "thread_summary"]) -def test_search_cosmos_accepts_list(): - """search_cosmos must thread a list of memory types through to the WHERE.""" +def test_search_cosmos_threads_non_episodic_types_and_excludes_episodic_without_optin(): + """search_cosmos threads non-episodic memory types through to the WHERE, and + excludes episodic unless include_episodes is set - so episodes join the base + query only when explicitly requested.""" client, container = _connected_client() client._embeddings_client = MagicMock() client._embeddings_client.generate.return_value = [0.0] * 8 @@ -194,19 +199,49 @@ def test_search_cosmos_accepts_list(): memory_types=["fact", "procedural", "episodic"], ) query = _captured_query(container) + # episodic dropped (no include_episodes) -> only fact + procedural remain. + assert "c.type IN (@memory_type_0, @memory_type_1)" in query + params = {p["name"]: p["value"] for p in _captured_params(container)} + type_values = {v for k, v in params.items() if k.startswith("@memory_type")} + assert params.get("@memory_type_0") == "fact" + assert params.get("@memory_type_1") == "procedural" + assert "episodic" not in type_values + + +def test_search_cosmos_include_episodes_adds_episodic_to_base_query(): + """include_episodes folds ``episodic`` into the single base query alongside + the caller's other non-episodic types (no separate episodic query).""" + client, container = _connected_client() + client._embeddings_client = MagicMock() + client._embeddings_client.generate.return_value = [0.0] * 8 + client.search_cosmos( + search_terms="user preferences", + user_id="u1", + memory_types=["fact", "procedural"], + include_episodes=True, + ) + query = _captured_query(container) assert "c.type IN (@memory_type_0, @memory_type_1, @memory_type_2)" in query + params = {p["name"]: p["value"] for p in _captured_params(container)} + type_values = {v for k, v in params.items() if k.startswith("@memory_type")} + assert type_values == {"fact", "procedural", "episodic"} -def test_search_cosmos_empty_list_disables_type_filter(): +def test_search_cosmos_empty_or_none_types_default_to_facts_only(): + """With include_episodes off, the base search is facts-only: empty/None + memory_types resolve to a fact type filter, never 'all types' (which would + leak episodes into the base result).""" client, container = _connected_client() client._embeddings_client = MagicMock() client._embeddings_client.generate.return_value = [0.0] * 8 client.search_cosmos(search_terms="x", user_id="u1", memory_types=[]) query = _captured_query(container) + params = {p["name"]: p["value"] for p in _captured_params(container)} + type_values = {v for k, v in params.items() if k.startswith("@memory_type")} where_clause = query.split("FROM c", 1)[1] - assert "c.type =" not in where_clause - assert "c.type IN" not in where_clause - assert all(not p["name"].startswith("@memory_type") for p in _captured_params(container)) + assert "c.type" in where_clause + assert params.get("@memory_type_0") == "fact" + assert "episodic" not in type_values def test_get_memories_default_uses_all_memories_types(): diff --git a/tests/unit/test_models.py b/tests/unit/test_models.py index d3a91b9..550c1d2 100644 --- a/tests/unit/test_models.py +++ b/tests/unit/test_models.py @@ -17,6 +17,8 @@ import pytest from azure.cosmos.agent_memory.models import ( + EpisodeEvent, + EpisodeOutcome, EpisodicRecord, FactRecord, MemoryRecord, @@ -56,14 +58,24 @@ def _episodic_kwargs(**overrides: Any) -> dict[str, Any]: base = { "id": "ep_" + _HEX32, "user_id": "u1", - "content": "Trip planning worked.", + "content": "Paris trip planning worked because the team planned early.", + "title": "Paris trip planning", + "started_at": "2026-01-01T09:00:00+00:00", + "ended_at": "2026-01-02T10:00:00+00:00", + "participants": ["user", "agent"], + "events": [ + { + "sequence": 1, + "description": "User asked for Paris trip planning help.", + "occurred_at": "2026-01-01T09:00:00+00:00", + "source_turn_ids": ["turn-1"], + }, + EpisodeEvent(sequence=2, description="Agent recommended booking early.", source_turn_ids=["turn-2"]), + ], + "outcome": {"status": "successful", "description": "The trip plan was completed."}, + "lessons": ["Plan early."], + "source_turn_ids": ["turn-1", "turn-2"], "content_hash": _HEX32, - "metadata": { - "lesson": "Plan early.", - "scope_type": "trip", - "scope_value": "Paris", - "outcome_valence": "positive", - }, "prompt_id": "extract_memories.prompty", } base.update(overrides) @@ -281,42 +293,61 @@ class TestEpisodicRecord: def test_minimal_valid(self): rec = EpisodicRecord(**_episodic_kwargs()) assert rec.memory_type == "episodic" - assert rec.scope_type == "trip" - assert rec.scope_value == "Paris" - - def test_requires_lesson_in_metadata(self): - meta = {"scope_type": "trip", "scope_value": "Paris", "outcome_valence": "positive"} - with pytest.raises(pydantic.ValidationError, match="lesson"): - EpisodicRecord(**_episodic_kwargs(metadata=meta)) - - def test_requires_scope_in_metadata(self): - meta = {"lesson": "x", "outcome_valence": "positive"} - with pytest.raises(pydantic.ValidationError, match="scope_type"): - EpisodicRecord(**_episodic_kwargs(metadata=meta)) - - def test_outcome_valence_enum(self): - meta = { - "lesson": "x", - "scope_type": "t", - "scope_value": "v", - "outcome_valence": "bogus", - } - with pytest.raises(pydantic.ValidationError, match="outcome_valence"): - EpisodicRecord(**_episodic_kwargs(metadata=meta)) - - @pytest.mark.parametrize("valence", ["positive", "negative", "neutral", "mixed"]) - def test_outcome_valence_accepts_all_schema_permitted_values(self, valence): - """Round-trip regression: every value the strict schema permits must - also be accepted by ``EpisodicRecord``. Previously ``"mixed"`` slipped - through schema validation but crashed the whole extract batch.""" - meta = { - "lesson": "x", - "scope_type": "t", - "scope_value": "v", - "outcome_valence": valence, - } - rec = EpisodicRecord(**_episodic_kwargs(metadata=meta)) - assert rec.metadata["outcome_valence"] == valence + assert rec.title == "Paris trip planning" + assert rec.started_at == "2026-01-01T09:00:00+00:00" + assert rec.ended_at == "2026-01-02T10:00:00+00:00" + assert rec.participants == ["user", "agent"] + assert rec.events[0] == EpisodeEvent( + sequence=1, + description="User asked for Paris trip planning help.", + occurred_at="2026-01-01T09:00:00+00:00", + source_turn_ids=["turn-1"], + ) + assert rec.events[1].description == "Agent recommended booking early." + assert rec.outcome == EpisodeOutcome(status="successful", description="The trip plan was completed.") + assert rec.lessons == ["Plan early."] + assert rec.source_turn_ids == ["turn-1", "turn-2"] + + def test_outcome_is_optional(self): + rec = EpisodicRecord(**_episodic_kwargs(outcome=None)) + restored = MemoryRecordBase.from_doc(rec.to_doc()) + assert rec.outcome is None + assert isinstance(restored, EpisodicRecord) + assert restored.outcome is None + + def test_outcome_status_literal(self): + with pytest.raises(pydantic.ValidationError, match="outcome"): + EpisodicRecord(**_episodic_kwargs(outcome={"status": "bogus", "description": "Nope."})) + + def test_ended_at_must_not_be_before_started_at(self): + with pytest.raises(pydantic.ValidationError, match="ended_at must not be before started_at"): + EpisodicRecord( + **_episodic_kwargs( + started_at="2026-01-02T10:00:00+00:00", + ended_at="2026-01-01T09:00:00+00:00", + ) + ) + + def test_mixed_naive_and_aware_times_compare_without_typeerror(self): + # A naive date + a tz-aware datetime must not raise TypeError; the naive + # value is normalized to UTC and the (valid) pair is accepted. + rec = EpisodicRecord( + **_episodic_kwargs( + started_at="2026-03-09", + ended_at="2026-03-10T09:08:00+00:00", + ) + ) + assert rec.started_at == "2026-03-09" + assert rec.ended_at == "2026-03-10T09:08:00+00:00" + + def test_mixed_tz_reversed_order_still_rejected(self): + with pytest.raises(pydantic.ValidationError, match="ended_at must not be before started_at"): + EpisodicRecord( + **_episodic_kwargs( + started_at="2026-03-11T00:00:00+00:00", + ended_at="2026-03-10", + ) + ) def test_id_must_start_with_ep_prefix(self): with pytest.raises(pydantic.ValidationError, match="id must start with 'ep_'"): @@ -556,10 +587,25 @@ def test_fact_round_trip(self, sample_embedding): def test_episodic_round_trip(self): original = EpisodicRecord(**_episodic_kwargs()) doc = original.to_doc() + assert doc["events"][0] == { + "sequence": 1, + "description": "User asked for Paris trip planning help.", + "occurred_at": "2026-01-01T09:00:00+00:00", + "source_turn_ids": ["turn-1"], + } + assert doc["outcome"] == {"status": "successful", "description": "The trip plan was completed."} restored = MemoryRecordBase.from_doc(doc) assert isinstance(restored, EpisodicRecord) - assert restored.scope_type == "trip" - assert restored.scope_value == "Paris" + assert restored.title == original.title + assert restored.participants == ["user", "agent"] + assert restored.events[1] == EpisodeEvent( + sequence=2, + description="Agent recommended booking early.", + source_turn_ids=["turn-2"], + ) + assert restored.outcome == EpisodeOutcome(status="successful", description="The trip plan was completed.") + assert restored.lessons == ["Plan early."] + assert restored.source_turn_ids == ["turn-1", "turn-2"] def test_from_doc_strips_cosmos_system_fields(self, sample_embedding): original = FactRecord(**_fact_kwargs(embedding=sample_embedding)) diff --git a/tests/unit/test_pipeline_confidence.py b/tests/unit/test_pipeline_confidence.py index d31c3a4..61dbe21 100644 --- a/tests/unit/test_pipeline_confidence.py +++ b/tests/unit/test_pipeline_confidence.py @@ -93,15 +93,7 @@ def test_extract_defaults_confidence_to_half_when_missing(): pipeline, upserted = _make_pipeline( { "facts": [{"text": "User likes coffee", "action": "ADD"}], - "episodic": [ - { - "scope_type": "project", - "scope_value": "X rollout", - "situation": "Trying X", - "action_taken": "Did Y", - "outcome": "Worked", - } - ], + "episodic": [], } ) @@ -111,27 +103,6 @@ def test_extract_defaults_confidence_to_half_when_missing(): assert doc["confidence"] == 0.5, f"missing default for {doc['type']} {doc['id']}" -def test_extract_episodic_carries_confidence(): - pipeline, upserted = _make_pipeline( - { - "episodic": [ - { - "scope_type": "project", - "scope_value": "CI revamp", - "situation": "Setup CI", - "action_taken": "Added Ruff", - "outcome": "Faster lint", - "confidence": 0.8, - "salience": 0.7, - } - ] - } - ) - pipeline.extract_memories("u1", "t1") - [ep] = [d for d in upserted if d["type"] == "episodic"] - assert ep["confidence"] == pytest.approx(0.8) - - class TestMarkSupersededDoesNotMutate: """``_mark_superseded`` must not mutate its input dict before the write. @@ -260,16 +231,11 @@ def test_thread_ids_does_not_appear_in_query_or_parameters(self): # --------------------------------------------------------------------------- -# Scoped episodic memories (scope_type / scope_value) +# Legacy episodic fields in extract_memories.prompty responses # --------------------------------------------------------------------------- -def test_extract_scoped_intent_without_outcome_stores_correctly(caplog): - """An episodic with only scope fields (no situation/action/outcome) is kept. - - The doc must use the deterministic fallback content string, expose the - scope fields at the top level, and not emit a "dropping malformed" warning. - """ +def test_extract_memories_ignores_legacy_episodic_payloads(caplog): pipeline, upserted = _make_pipeline( { "episodic": [ @@ -284,183 +250,11 @@ def test_extract_scoped_intent_without_outcome_stores_correctly(caplog): ) with caplog.at_level("WARNING", logger="azure.cosmos.agent_memory.pipeline"): - pipeline.extract_memories("u1", "t1") - - eps = [d for d in upserted if d["type"] == "episodic"] - assert len(eps) == 1 - ep = eps[0] - assert ep["scope_type"] == "trip" - assert ep["scope_value"] == "Paris" - assert ep["metadata"]["scope_type"] == "trip" - assert ep["metadata"]["scope_value"] == "Paris" - assert ep["metadata"]["situation"] is None - assert ep["metadata"]["action_taken"] is None - assert ep["metadata"]["outcome"] is None - assert ep["content"] == "For the user's Paris trip, intent recorded." - assert ep["confidence"] == pytest.approx(0.95) - assert not any("dropping malformed episodic" in rec.getMessage() for rec in caplog.records) - - -def test_extract_past_event_episodic_uses_arrow_form_and_keeps_scope(): - pipeline, upserted = _make_pipeline( - { - "episodic": [ - { - "scope_type": "project", - "scope_value": "Acme revamp", - "situation": "Migrated DB", - "action_taken": "Ran the script", - "outcome": "All rows migrated", - "outcome_valence": "positive", - "reasoning": "Schema was simple", - "lesson": "Test on staging first", - "domain": "engineering", - "confidence": 0.88, - "salience": 0.6, - "tags": ["db"], - } - ] - } - ) - - pipeline.extract_memories("u1", "t1") - - [ep] = [d for d in upserted if d["type"] == "episodic"] - assert ep["content"] == "Migrated DB → Ran the script → All rows migrated" - assert ep["scope_type"] == "project" - assert ep["scope_value"] == "Acme revamp" - md = ep["metadata"] - assert md["situation"] == "Migrated DB" - assert md["action_taken"] == "Ran the script" - assert md["outcome"] == "All rows migrated" - assert md["outcome_valence"] == "positive" - assert md["reasoning"] == "Schema was simple" - assert md["lesson"] == "Test on staging first" - assert md["domain"] == "engineering" - assert "topic:db" in ep["tags"] - - -def test_extract_episodic_falls_back_to_arrow_form_when_summary_field_present(): - """The schema dropped ``summary``; pipeline now always uses arrow form. - - Even if a non-strict LLM smuggles a ``summary`` field through, the - pipeline ignores it and builds content from - ``situation → action_taken → outcome``. - """ - pipeline, upserted = _make_pipeline( - { - "episodic": [ - { - "scope_type": "trip", - "scope_value": "Paris", - "summary": "User wants luxury hotels for the Paris trip.", - "situation": "Planning Paris trip", - "action_taken": "Said luxury", - "outcome": "Pending", - } - ] - } - ) - - pipeline.extract_memories("u1", "t1") - - [ep] = [d for d in upserted if d["type"] == "episodic"] - assert ep["content"] == "Planning Paris trip → Said luxury → Pending" - - -def test_extract_drops_episodic_missing_scope_type(caplog): - pipeline, upserted = _make_pipeline( - { - "episodic": [ - { - "scope_value": "Paris", - "situation": "Planning", - "action_taken": "Booked", - "outcome": "Confirmed", - } - ] - } - ) - - with caplog.at_level("WARNING", logger="azure.cosmos.agent_memory.pipeline"): - pipeline.extract_memories("u1", "t1") + result = pipeline.extract_memories("u1", "t1") + assert result["episodic_count"] == 0 assert not any(d["type"] == "episodic" for d in upserted) - assert any("dropping malformed episodic" in rec.getMessage() for rec in caplog.records) - - -def test_extract_drops_episodic_missing_scope_value(caplog): - pipeline, upserted = _make_pipeline( - { - "episodic": [ - { - "scope_type": "trip", - "situation": "Planning", - "action_taken": "Booked", - "outcome": "Confirmed", - } - ] - } - ) - - with caplog.at_level("WARNING", logger="azure.cosmos.agent_memory.pipeline"): - pipeline.extract_memories("u1", "t1") - - assert not any(d["type"] == "episodic" for d in upserted) - assert any("dropping malformed episodic" in rec.getMessage() for rec in caplog.records) - - -@pytest.mark.parametrize( - "scope_type,scope_value", - [ - ("", "Paris"), - (" ", "Paris"), - ("trip", ""), - ("trip", " "), - (None, "Paris"), - ("trip", None), - (123, "Paris"), - ], -) -def test_extract_drops_episodic_with_blank_or_invalid_scope(scope_type, scope_value, caplog): - pipeline, upserted = _make_pipeline( - { - "episodic": [ - { - "scope_type": scope_type, - "scope_value": scope_value, - "confidence": 0.9, - } - ] - } - ) - - with caplog.at_level("WARNING", logger="azure.cosmos.agent_memory.pipeline"): - pipeline.extract_memories("u1", "t1") - - assert not any(d["type"] == "episodic" for d in upserted) - assert any("dropping malformed episodic" in rec.getMessage() for rec in caplog.records) - - -def test_extract_strips_whitespace_from_scope_fields(): - pipeline, upserted = _make_pipeline( - { - "episodic": [ - { - "scope_type": " trip ", - "scope_value": " Paris ", - "confidence": 0.9, - } - ] - } - ) - - pipeline.extract_memories("u1", "t1") - - [ep] = [d for d in upserted if d["type"] == "episodic"] - assert ep["scope_type"] == "trip" - assert ep["scope_value"] == "Paris" - assert ep["content"] == "For the user's Paris trip, intent recorded." + assert not any("dropping malformed episodic" in rec.getMessage() for rec in caplog.records) def test_extract_compound_statement_yields_facts_across_categories(): diff --git a/tests/unit/test_procedural_synthesis.py b/tests/unit/test_procedural_synthesis.py index e96ace4..2a47dcb 100644 --- a/tests/unit/test_procedural_synthesis.py +++ b/tests/unit/test_procedural_synthesis.py @@ -110,10 +110,13 @@ def _fact_doc( def _episodic_doc( doc_id: str, *, - lesson: 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 [] return { "id": doc_id, "user_id": "u1", @@ -121,7 +124,7 @@ def _episodic_doc( "role": "system", "type": "episodic", "content": f"Episode {doc_id}", - "metadata": {"lesson": lesson}, + "lessons": lessons, "salience": salience, "created_at": created_at, } @@ -213,7 +216,7 @@ def _make_client(*, processor=None) -> CosmosMemoryClient: return client -def test_extract_memories_without_procedural_bucket_returns_new_count_shape(): +def test_extract_memories_returns_count_shape_and_ignores_legacy_episodic_payload(): pipeline, _, upserted = _make_extract_pipeline( { "facts": [ @@ -241,7 +244,7 @@ def test_extract_memories_without_procedural_bucket_returns_new_count_shape(): legacy_proc_key = "_".join(("procedural", "count")) assert result["fact_count"] == 1 - assert result["episodic_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) @@ -305,6 +308,38 @@ def test_synthesize_procedural_first_synthesis_from_empty_prior(): 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() diff --git a/tests/unit/test_thresholds.py b/tests/unit/test_thresholds.py index eb350d0..43a5872 100644 --- a/tests/unit/test_thresholds.py +++ b/tests/unit/test_thresholds.py @@ -58,6 +58,7 @@ def test_enable_turn_embeddings_falsy_values(monkeypatch, raw) -> None: [ ("FACT_EXTRACTION_EVERY_N", "get_fact_extraction_every_n", 1), ("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), ("DEDUP_EVERY_N", "get_dedup_every_n", 5), ("DEDUP_POOL_SIZE", "get_dedup_pool_size", 50), @@ -80,6 +81,7 @@ def test_env_config_getters_defaults( [ ("FACT_EXTRACTION_EVERY_N", "get_fact_extraction_every_n", "2", 2), ("THREAD_SUMMARY_EVERY_N", "get_thread_summary_every_n", "11", 11), + ("EPISODE_EVAL_EVERY_N", "get_episode_eval_every_n", "3", 3), ("USER_SUMMARY_EVERY_N", "get_user_summary_every_n", "21", 21), ("DEDUP_EVERY_N", "get_dedup_every_n", "3", 3), ("DEDUP_POOL_SIZE", "get_dedup_pool_size", "75", 75), @@ -103,6 +105,7 @@ def test_env_config_getters_parse_env( [ ("FACT_EXTRACTION_EVERY_N", "get_fact_extraction_every_n", 1), ("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), ("DEDUP_EVERY_N", "get_dedup_every_n", 5), ("DEDUP_POOL_SIZE", "get_dedup_pool_size", 50), @@ -124,6 +127,7 @@ def test_int_getters_reject_negative( [ ("FACT_EXTRACTION_EVERY_N", "get_fact_extraction_every_n", 1), ("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), ("DEDUP_EVERY_N", "get_dedup_every_n", 5), ("DEDUP_POOL_SIZE", "get_dedup_pool_size", 50), @@ -152,6 +156,19 @@ def test_dedup_pool_size_rejects_zero(monkeypatch: pytest.MonkeyPatch) -> None: assert thresholds.get_dedup_pool_size() == 50 +@pytest.mark.parametrize("bad_value", ["nan", "inf", "-inf", "-1"]) +def test_episode_topic_drift_rejects_non_finite_and_negative(monkeypatch: pytest.MonkeyPatch, bad_value: str) -> None: + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", bad_value) + + assert thresholds.get_episode_topic_drift() == thresholds.DEFAULT_EPISODE_TOPIC_DRIFT + + +def test_episode_topic_drift_accepts_valid_float(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("EPISODE_TOPIC_DRIFT", "0.35") + + assert thresholds.get_episode_topic_drift() == 0.35 + + def test_procedural_synthesis_auto_invalid_uses_default(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("PROCEDURAL_SYNTHESIS_AUTO", "bogus")