Skip to content

Improving the extraction prompt with Agent-sourced facts, event-time capture#31

Merged
aayush3011 merged 3 commits into
AzureCosmosDB:mainfrom
aayush3011:users/akataria/improvements
Jul 15, 2026
Merged

Improving the extraction prompt with Agent-sourced facts, event-time capture#31
aayush3011 merged 3 commits into
AzureCosmosDB:mainfrom
aayush3011:users/akataria/improvements

Conversation

@aayush3011

@aayush3011 aayush3011 commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

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 everything
the assistant said. It can now also capture the assistant's own concrete
actions and recommendations
(bookings made, items recommended), tagged
source=agent on the fact and sys:agent-fact in tags. Assistant assumptions
about 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 required source field to the
fact schema; extraction prompt bumped to v2.

Event time on write

add_cosmos(...) (sync + async) gains an optional created_at (str or
datetime) that sets a memory's event time instead of defaulting to ingestion
wall-clock. Extracted facts inherit their source turns' timestamps, so backdating
turns propagates real event time to facts — making the existing
created_after / created_before filters and temporal reasoning actually usable
for 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 a
memory's created_at (the backdated source-turn / session time), not the
resolved per-fact date embedded in the text.

Unified retrieval

search_cosmos(include_turns=True, turn_top_k=...) (sync + async) blends raw
conversation 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_ENABLED is now an environment knob (default false = 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_at
passthrough (str/datetime), timestamped transcript + extraction wiring,
include_turns blend/dedup/failure paths, and DEDUP_VECTOR_ENABLED env
parsing. Full unit suite green; ruff clean.

Benchmark Results

LongMemEval

lme-main-vs-current

LoCoMo

locomo-main-vs-current

Beam 1M

beam-1m-current

Copilot AI review requested due to automatic review settings July 14, 2026 14:10

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.prompty to v2, add source to the fact schema, and tag agent-sourced facts (sys:agent-fact).
  • Prefix extraction transcripts with created_at timestamps and plumb created_at through manual add APIs/stores.
  • Make DEDUP_VECTOR_ENABLED env-configurable (default false) and add optional include_turns retrieval to search_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.

Comment thread azure/cosmos/agent_memory/prompts/extract_memories.prompty Outdated
Comment thread azure/cosmos/agent_memory/cosmos_memory_client.py Outdated
Comment thread azure/cosmos/agent_memory/cosmos_memory_client.py Outdated
Comment thread azure/cosmos/agent_memory/aio/cosmos_memory_client.py Outdated
Comment thread azure/cosmos/agent_memory/aio/cosmos_memory_client.py Outdated
@aayush3011 aayush3011 changed the title Improving the extraction prompt Improving the extraction prompt with Agent-sourced facts, event-time capture Jul 14, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 25 out of 25 changed files in this pull request and generated 3 comments.

Comment thread azure/cosmos/agent_memory/prompts/extract_memories.prompty Outdated
Comment thread Docs/concepts.md
Comment thread Docs/concepts.md

@mrrahman1517 mrrahman1517 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.0 telemetry asserts, and they reproduce on origin/main too — pre-existing timer-resolution flakiness (mocked reconcile finishes sub-microsecond), not caused by this PR. Every PR-specific test passes. (test_embedding_tokens.py errored only because tiktoken isn't installed locally — unrelated.)
  • ruff check on the changed source + tests: clean.
  • Confirmed the event-time path: add_cosmos(created_at=)store.addnormalize_created_at_isoMemoryRecord (Pydantic default_factory correctly overridden). add_cosmosstore.add positional args stay aligned. Agent facts reuse the existing exclude_tags plumbing via sys:agent-fact.

Broader notes (don't map to one line):

  • include_turns result shape: appended turn docs lack metadata.category and 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.

Comment thread azure/cosmos/agent_memory/thresholds.py
Comment thread azure/cosmos/agent_memory/services/_pipeline_helpers.py
Comment thread azure/cosmos/agent_memory/cosmos_memory_client.py
Comment thread azure/cosmos/agent_memory/_utils.py
@aayush3011 aayush3011 merged commit ddd8fc7 into AzureCosmosDB:main Jul 15, 2026
7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants