Improving the extraction prompt with Agent-sourced facts, event-time capture#31
Conversation
There was a problem hiding this comment.
Pull request overview
This PR updates the memory-extraction prompt and pipeline to (1) capture agent-sourced facts explicitly, (2) include per-turn timestamps in extraction transcripts to resolve relative dates, and (3) make write-time vector dedup folding operator-tunable via an environment variable. It also extends Cosmos search to optionally include raw turns alongside derived memories.
Changes:
- Bump
extract_memories.promptyto v2, addsourceto the fact schema, and tag agent-sourced facts (sys:agent-fact). - Prefix extraction transcripts with
created_attimestamps and plumbcreated_atthrough manual add APIs/stores. - Make
DEDUP_VECTOR_ENABLEDenv-configurable (defaultfalse) and add optionalinclude_turnsretrieval tosearch_cosmos()(sync + async).
Reviewed changes
Copilot reviewed 21 out of 21 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit/test_thresholds.py | Updates/extends tests for DEDUP_VECTOR_ENABLED default/env behavior. |
| tests/unit/test_cosmos_memory_client.py | Adds coverage for add_cosmos(created_at=...) and unified retrieval (include_turns). |
| tests/unit/store/test_memory_store.py | Adds created_at defaulting/normalization tests for sync store add. |
| tests/unit/services/test_transcript_metadata.py | Adds tests for timestamp-prefixed transcript formatting. |
| tests/unit/services/test_prompty_loader.py | Updates expected prompt versions (extract prompt now v2). |
| tests/unit/services/test_extract_dry.py | Verifies agent-sourced tagging + timestamp propagation into extraction prompt. |
| tests/unit/services/test_dedup_vector.py | Enables vector folding via env for this suite now that default is off. |
| tests/unit/aio/test_cosmos_memory_client.py | Adds async coverage for unified retrieval (include_turns). |
| tests/unit/aio/store/test_memory_store.py | Adds created_at normalization tests for async store add. |
| tests/unit/aio/services/test_dedup_vector_async.py | Enables vector folding via env for async suite now that default is off. |
| Docs/concepts.md | Documents DEDUP_VECTOR_ENABLED as an operator-facing knob (default off). |
| azure/cosmos/agent_memory/thresholds.py | Introduces DEFAULT_DEDUP_VECTOR_ENABLED and reads DEDUP_VECTOR_ENABLED from env. |
| azure/cosmos/agent_memory/store/memory_store.py | Adds created_at plumbing to MemoryStore.add(). |
| azure/cosmos/agent_memory/services/pipeline.py | Adds timestamped extraction transcript + agent-sourced fact metadata/tags. |
| azure/cosmos/agent_memory/services/_pipeline_helpers.py | Adds include_timestamp option to build_transcript(). |
| azure/cosmos/agent_memory/prompts/extract_memories.prompty | Prompt v2: adds source, agent-sourced rules, relative-time anchoring guidance. |
| azure/cosmos/agent_memory/prompts/_schemas.py | Requires source on extracted fact items. |
| azure/cosmos/agent_memory/cosmos_memory_client.py | Adds created_at param to add; extends search_cosmos() with include_turns. |
| azure/cosmos/agent_memory/aio/store/memory_store.py | Adds created_at plumbing to AsyncMemoryStore.add(). |
| azure/cosmos/agent_memory/aio/services/pipeline.py | Async parity: timestamped transcript + agent-sourced fact metadata/tags. |
| azure/cosmos/agent_memory/aio/cosmos_memory_client.py | Async parity: created_at param + search_cosmos(include_turns=...). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
mrrahman1517
left a comment
There was a problem hiding this comment.
Reviewed the full diff with local context, traced the new mechanisms end-to-end, and ran the suite against the PR branch. Overall this is a strong, well-tested, backward-compatible change — async/sync parity is maintained throughout and there are good defensive guards at every new decision point. No blocking issues. A few non-blocking notes below (inline + here).
Verification on my machine (Windows, py3.12):
- Unit suite on the PR branch: 1049 passed, 5 failed. All 5 failures are the reconcile
duration_ms > 0.0telemetry asserts, and they reproduce onorigin/maintoo — pre-existing timer-resolution flakiness (mocked reconcile finishes sub-microsecond), not caused by this PR. Every PR-specific test passes. (test_embedding_tokens.pyerrored only becausetiktokenisn't installed locally — unrelated.) ruff checkon the changed source + tests: clean.- Confirmed the event-time path:
add_cosmos(created_at=)→store.add→normalize_created_at_iso→MemoryRecord(Pydantic default_factory correctly overridden).add_cosmos↔store.addpositional args stay aligned. Agent facts reuse the existingexclude_tagsplumbing viasys:agent-fact.
Broader notes (don't map to one line):
include_turnsresult shape: appended turn docs lackmetadata.categoryand have different salience/score semantics from extracted memories, so they're heterogeneous within the returned list. It's opt-in and documented, but a consumer iterating results expecting the memory shape could trip — worth calling out prominently or tagging turn-sourced results.- Reconcile excludes agent facts entirely: two contradictory agent actions ("booked X" then "rebooked Y") won't reconcile against each other. Almost certainly intended (both are true historical events) — just flagging the consequence.
- Benchmarks: the LongMemEval/LoCoMo/Beam-1M gains are embedded as images; I can't validate the runs from the diff. If those numbers are load-bearing for the merge decision, linking the eval config/commit used would help reviewers reproduce.
Summary
Four independent, additive retrieval/extraction improvements on top of the
dedup pipeline from #26. Each targets a specific memory-quality gap; all are
opt-in or backward-compatible, and no existing behavior changes by default.
Changes
Agent-sourced fact extraction
Extraction previously only kept facts from
[user]lines and dropped everythingthe assistant said. It can now also capture the assistant's own concrete
actions and recommendations (bookings made, items recommended), tagged
source=agenton the fact andsys:agent-factin tags. Assistant assumptionsabout the user ("I assume you want luxury") are still excluded, so user-fact
quality is unchanged. Retrieval can include agent facts by default or filter them
out via
exclude_tags=["sys:agent-fact"]. Adds a requiredsourcefield to thefact schema; extraction prompt bumped to v2.
Event time on write
add_cosmos(...)(sync + async) gains an optionalcreated_at(strordatetime) that sets a memory's event time instead of defaulting to ingestionwall-clock. Extracted facts inherit their source turns' timestamps, so backdating
turns propagates real event time to facts — making the existing
created_after/created_beforefilters and temporal reasoning actually usablefor backfilled conversations.
Relative-date resolution in extraction
The extraction transcript now includes each turn's timestamp, and the prompt
resolves relative expressions ("3 weeks ago") to absolute dates in the fact
text using that anchor, keeping the original expression in
temporal_context.Note: time-range filtering (
created_after/created_before) matches on amemory's
created_at(the backdated source-turn / session time), not theresolved per-fact date embedded in the text.
Unified retrieval
search_cosmos(include_turns=True, turn_top_k=...)(sync + async) blends rawconversation turns into results after the extracted-memory hits — recovering
detail lost during extraction. Turns are appended (not score-fused) so memory
hits keep priority, duplicates are skipped, and the turn search is best-effort
(never fails the primary search). Requires
enable_turn_embeddings.Configurable write-time dedup
DEDUP_VECTOR_ENABLEDis now an environment knob (defaultfalse= add-only)instead of a fixed internal constant, so in-place folding can be toggled without
editing source.
Tests
Unit coverage added for each: agent-fact tagging/defaulting,
created_atpassthrough (str/datetime), timestamped transcript + extraction wiring,
include_turnsblend/dedup/failure paths, andDEDUP_VECTOR_ENABLEDenvparsing. Full unit suite green; ruff clean.
Benchmark Results
LongMemEval
LoCoMo
Beam 1M