Episodic memory: boundary-based segmentation, facts/episode retrieval separation, and v2 extraction default

[Aayush Kataria (aayush3011)]

Summary

Adds episodic memory as a first-class memory type and reworks how episodes are extracted and retrieved. Episodes self-segment a continuous, timestamped turn stream (no caller "session end" signal required) and are stored append-only. Retrieval returns facts by default and, when include_episodes is set, facts and episodes together in a single ranked query. Also promotes the higher-recall v2 fact extractor to the default and fixes the episodic -> procedural "lessons" seam.

What's new

Episodic memory

• New EpisodicRecord model (with lessons, EpisodeEvent, EpisodeOutcome) and a dedicated extract_episode.prompty.
• Boundary-based self-segmentation of the open turn stream. A boundary is the earliest of: idle time-gap, topic drift, or a max-size cap, floored by EPISODE_MIN_TURNS so a lone turn is not emitted as a trivial episode. Segmentation grounds started_at / ended_at in turn timestamps.
• Append-only writes with best-effort idempotency: an episode_extracted_at watermark marks consumed turns, and each episode 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; episodic reconciliation is a no-op and does not fold duplicates. flush=True drains the trailing open segment (useful for batch/benchmark ingestion).
• Store support: search_episodic and get_episodes.
• Auto-trigger evaluates episode boundaries on cadence during ingestion (in-process backend only).

Retrieval (facts + optional episodes)

• search_cosmos returns facts by default. With include_episodes=True, episodic is folded into the SAME ranked base query, so facts and episodes share one top_k budget and compete on relevance (a highly relevant episode can outrank facts). Callers may pass memory_types to search other non-episodic types.
• One query means tag / salience / time filters apply uniformly to facts and episodes. search_episodic_memories remains available for episode-only retrieval.

Extraction v2 is now the default

• extract_memories-v2.prompty becomes the default extractor: higher recall, and it also captures assistant-provided information the user may later reference (lists, tables, instructions, researched answers).
• Override with AMT_EXTRACT_MEMORIES_PROMPT=extract_memories.prompty to fall back to v1; an unknown value safely falls back to the default. The v1 prompt is slimmed (episode emission removed; episodes now have a dedicated extractor).

Episodic -> procedural lessons seam fix

• Procedural synthesis now reads first-class lessons from EpisodicRecords (previously it looked for metadata.lesson, so lessons from new-style episodes never reached synthesis).

Internal / robustness

• Renamed the pipeline's LLM-compute-only _dry methods to _durable across sync + aio pipelines and the Durable Functions orchestrators, to make the checkpoint seam (compute vs embed+persist) explicit.
• parse_llm_json now merges multiple concatenated top-level JSON objects that some deployments emit for a single call, so list items (facts/events) are no longer silently dropped.
• Malformed / mixed-timezone model timestamps fall back to grounded segment bounds instead of dropping the whole episode.

New configuration knobs (env-selectable, with defaults)

• EPISODE_EVAL_EVERY_N = 4 (how often to evaluate boundaries)
• EPISODE_IDLE_GAP_SECONDS = 1800
• EPISODE_TOPIC_DRIFT = 0.0 (drift splitting OFF by default)
• EPISODE_MAX_TURNS = 40
• EPISODE_MIN_TURNS = 2 (floors all natural boundaries)
• AMT_EXTRACT_MEMORIES_PROMPT (extractor selection)

Compatibility

• Facts extraction, reconciliation, and summaries are unchanged in shape; reconciliation remains the dedup authority for facts.
• Episodic wiring is in-process only; the Durable Functions backend has no episodic path yet (deferred). The EPISODE_* knobs are in-process only and have no function-app mirror.

[Copilot]

Pull request overview

This PR introduces a new episodic-memory pipeline that segments turn streams into bounded “episodes” (with first-class timeline, outcome, and lessons), while making extract_memories facts-only and updating retrieval to keep episodes opt-in (so they don’t dilute fact search budgets).

Changes:

• Add boundary-based episodic extraction (extract_episode.prompty, segmentation + watermarking, new episode model shape).
• Update retrieval/search to be facts-only by default, with opt-in episodic blending (include_episodes, separate episode_top_k budget) plus new get_episodes API.
• Refresh tests/samples to reflect the new episodic schema, extraction behavior, and durable extraction naming; add Retry-After–aware retry delay handling.

Reviewed changes

Copilot reviewed 53 out of 53 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
tests/unit/test_thresholds.py	Adds coverage for the new EPISODE_EVAL_EVERY_N threshold getter.
tests/unit/test_procedural_synthesis.py	Updates episodic fixtures to use lessons and adds test for flattening multiple lessons.
tests/unit/test_pipeline_confidence.py	Removes legacy episodic confidence tests; adds regression guard that legacy episodic payloads are ignored.
tests/unit/test_models.py	Updates episodic record tests to new episode-first schema (events/outcome/lessons).
tests/unit/test_memory_type_multi.py	Adjusts search_cosmos type-filter semantics (facts-only base; episodic stripped).
tests/unit/test_cosmos_memory_client.py	Updates expectations to facts-only memory_types default.
tests/unit/test_chat.py	Adds Retry-After–aware retry tests for sync chat client.
tests/unit/test_auto_trigger.py	Adds episodic cadence to auto-trigger thresholds and per-step triggering tests.
tests/unit/store/test_memory_store.py	Updates episodic search tests to assert query construction instead of delegating to generic search.
tests/unit/services/test_prompty_loader.py	Updates shipped prompt version expectations and adds extract_episode prompt.
tests/unit/services/test_pipeline_service.py	Updates pipeline happy-path expectations to facts-only extraction and lessons field.
tests/unit/services/test_extract_episodes.py	Adds unit tests for episode doc construction and episodic extraction persistence behavior.
tests/unit/services/test_extract_dry.py	Renames “dry” to “durable” path tests and adds episode-doc building tests.
tests/unit/services/test_episodic_retrieval.py	Adds tests for get_episodes and opt-in episodic blending in search.
tests/unit/services/test_episode_boundary.py	Adds sync boundary-based segmentation tests (idle gap, drift, max turns, flush, idempotency).
tests/unit/services/test_dedup_vector.py	Updates episodic doc fixtures for new schema shape.
tests/unit/services/test_chaos_extract_persist.py	Updates chaos tests to use durable extraction path.
tests/unit/function_app/test_orchestrators.py	Updates orchestrator tests to call durable extraction endpoints.
tests/unit/aio/test_cosmos_memory_client.py	Updates async episodic search tests to assert query construction for episodic retrieval.
tests/unit/aio/test_chat.py	Adds Retry-After–aware retry tests for async chat client.
tests/unit/aio/test_auto_trigger.py	Adds episodic cadence to async auto-trigger behavior and tests.
tests/unit/aio/services/test_extract_episodes_async.py	Adds async unit tests for episode doc building and episodic extraction.
tests/unit/aio/services/test_episodic_retrieval_async.py	Adds async tests for get_episodes and blended search ordering/budgeting.
tests/unit/aio/services/test_episode_boundary_async.py	Adds async mirror tests for boundary segmentation behavior.
tests/unit/aio/services/test_dedup_vector_async.py	Updates async episodic fixtures to new schema shape.
tests/unit/aio/processors/test_inprocess.py	Adds async processor test for process_extract_episodes.
tests/integration/test_episodic_pipeline.py	Adds live integration tests for episodic extraction + blended retrieval + lessons→procedural synthesis.
Samples/Processing/processing_episodic_memory.py	Adds a runnable sample demonstrating episodic extraction, retrieval, blended search, and cleanup.
function_app/orchestrators/user_summary.py	Switches to durable user summary generation.
function_app/orchestrators/thread_summary.py	Switches to durable thread summary generation.
function_app/orchestrators/extract_memories.py	Switches to durable extract_memories.
azure/cosmos/agent_memory/thresholds.py	Adds episodic thresholds (cadence + boundary tuning) and float parsing helper.
azure/cosmos/agent_memory/store/memory_store.py	Adds get_episodes; implements episodic-specific semantic search with temporal filters.
azure/cosmos/agent_memory/store/_search_helpers.py	Extends projection and episodic context formatting for new episode fields.
azure/cosmos/agent_memory/services/_pipeline_helpers.py	Enhances JSON parsing to merge concatenated objects; adds prompt selection; adjusts prompty loader strictness.
azure/cosmos/agent_memory/prompts/extract_memories.prompty	Updates extract_memories to facts-only schema (v4).
azure/cosmos/agent_memory/prompts/extract_memories-v2.prompty	Adds an alternate facts-only prompt variant (v4-additive, agent-additive recall).
azure/cosmos/agent_memory/prompts/extract_episode.prompty	Adds new episodic extraction prompt and schema expectations.
azure/cosmos/agent_memory/prompts/_schemas.py	Splits schemas: facts-only extract_memories + new extract_episode output schema.
azure/cosmos/agent_memory/processors/inprocess.py	Adds process_extract_episodes processor entrypoint.
azure/cosmos/agent_memory/models.py	Introduces EpisodeEvent/EpisodeOutcome and updates EpisodicRecord schema + validations.
azure/cosmos/agent_memory/cosmos_memory_client.py	Adds get_episodes/extract_episodes and facts-only base search with opt-in episodic blending.
azure/cosmos/agent_memory/chat.py	Adds Retry-After parsing and unified retry delay function; increases retry count.
azure/cosmos/agent_memory/auto_trigger.py	Adds episodic cadence to auto-trigger and wires processor call.
azure/cosmos/agent_memory/aio/store/memory_store.py	Adds async get_episodes and episodic search implementation (vector + FTS + temporal filters).
azure/cosmos/agent_memory/aio/services/pipeline.py	Renames dry→durable, adds episode segmentation/extraction pipeline and new watermarks, updates procedural lesson handling.
azure/cosmos/agent_memory/aio/processors/inprocess.py	Adds async process_extract_episodes.
azure/cosmos/agent_memory/aio/cosmos_memory_client.py	Adds async get_episodes/extract_episodes and facts-only base search with opt-in episodes.
azure/cosmos/agent_memory/aio/chat.py	Uses shared retry delay logic; increases retry count.
azure/cosmos/agent_memory/aio/auto_trigger.py	Adds episodic cadence to async auto-trigger and wires processor call.
azure/cosmos/agent_memory/_utils.py	Adds cosine similarity + centroid helpers used for topic drift segmentation.
azure/cosmos/agent_memory/init.py	Exports episodic threshold defaults and accessors via package API.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 54 out of 54 changed files in this pull request and generated no new comments.

Suppressed comments (4)

azure/cosmos/agent_memory/cosmos_memory_client.py:695

• The docstring (and PR description) state that the base search returns facts only and that top_k is a guaranteed fact budget, but the implementation uses caller-provided memory_types (minus only "episodic"). If a caller passes e.g. ["procedural"], the base query will include procedural docs and can displace facts, violating the documented/advertised behavior.

        """Search memories using vector similarity, with optional retrieval blending.

        The base search returns FACTS ONLY - episodes never enter the result set
        through the base query, so ``top_k`` is a guaranteed fact budget and
        episodes cannot dilute it. Episodes are opt-in via ``include_episodes``
        and get their own ``episode_top_k`` budget. Order: facts -> episodes ->
        summaries -> raw turns, all deduped by content (best-effort; a blend
        fetch failure never breaks the base result).
        """
        store = self._get_store()
        # Hard-scope the base search to facts: episodes are served exclusively by
        # the include_episodes pass (with its own budget), never through the base.
        base_memory_types = memory_types if memory_types is not None else ["fact"]
        base_memory_types = [t for t in base_memory_types if t != "episodic"] or ["fact"]
        facts = store.search(

azure/cosmos/agent_memory/aio/store/memory_store.py:1107

• AsyncMemoryStore.search_episodic exposes started_at/ended_at parameters with one-sided range semantics (started_at => >=, ended_at => <=), while the sync MemoryStore.search_episodic exposes started_after/started_before and ended_after/ended_before (full range on both). This breaks the stated goal that sync and aio paths stay mirrored and makes the async API strictly less expressive.

    async def search_episodic(
        self,
        user_id: str,
        search_terms: str,
        top_k: int = 5,
        min_salience: Optional[float] = None,
        include_superseded: bool = False,
        thread_id: Optional[str] = None,
        created_after: Optional[str | datetime] = None,
        created_before: Optional[str | datetime] = None,
        started_at: Optional[str | datetime] = None,
        ended_at: Optional[str | datetime] = None,
    ) -> list[dict[str, Any]]:

azure/cosmos/agent_memory/aio/store/memory_store.py:206

• AsyncMemoryStore.add() applies different defaults for manual episodic records than the sync MemoryStore.add() (e.g., lessons defaults to [content] here but [] in the sync store). This makes manual episodic writes inconsistent across sync vs aio clients and can change whether an episode contributes lessons into procedural synthesis.

            if memory_type == "fact":
                meta.setdefault("category", "unclassified:manual")
            elif memory_type == "episodic":
                kwargs.setdefault("title", meta.get("title") or content[:80] or "Manual episode")
                kwargs.setdefault("participants", [])
                kwargs.setdefault("events", [])
                kwargs.setdefault("outcome", None)
                kwargs.setdefault("lessons", [content] if content else [])
                kwargs.setdefault("source_turn_ids", [])
            elif memory_type == "procedural":

azure/cosmos/agent_memory/store/memory_store.py:239

• The defaults applied when manually adding an episodic memory differ between sync and async stores: sync defaults lessons to an empty list, while AsyncMemoryStore.add defaults lessons to [content] (and sets outcome=None). This leads to inconsistent episodic document shapes/behavior depending on client type (and may affect procedural synthesis which now reads first-class lessons).

            if memory_type == "fact":
                meta.setdefault("category", "unclassified:manual")
            elif memory_type == "episodic":
                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":