From bcad810f0ff05057801db9cccff6efe8e9cbd551 Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Tue, 14 Jul 2026 07:10:11 -0700 Subject: [PATCH 1/3] Improving the extraction prompt --- Docs/concepts.md | 5 +- .../agent_memory/aio/cosmos_memory_client.py | 31 ++++- .../agent_memory/aio/services/pipeline.py | 9 +- .../agent_memory/aio/store/memory_store.py | 3 + .../agent_memory/cosmos_memory_client.py | 36 +++++- azure/cosmos/agent_memory/prompts/_schemas.py | 7 +- .../prompts/extract_memories.prompty | 102 ++++++++++++++++- .../services/_pipeline_helpers.py | 29 +++-- .../cosmos/agent_memory/services/pipeline.py | 9 +- .../cosmos/agent_memory/store/memory_store.py | 3 + azure/cosmos/agent_memory/thresholds.py | 14 ++- .../aio/services/test_dedup_vector_async.py | 8 ++ tests/unit/aio/store/test_memory_store.py | 29 +++++ tests/unit/aio/test_cosmos_memory_client.py | 47 ++++++++ tests/unit/services/test_dedup_vector.py | 10 ++ tests/unit/services/test_extract_dry.py | 108 ++++++++++++++++++ tests/unit/services/test_prompty_loader.py | 23 ++-- .../unit/services/test_transcript_metadata.py | 31 +++++ tests/unit/store/test_memory_store.py | 38 ++++++ tests/unit/test_cosmos_memory_client.py | 70 ++++++++++++ tests/unit/test_thresholds.py | 13 ++- 21 files changed, 580 insertions(+), 45 deletions(-) diff --git a/Docs/concepts.md b/Docs/concepts.md index 5ac1942..ee19084 100644 --- a/Docs/concepts.md +++ b/Docs/concepts.md @@ -159,12 +159,13 @@ Extraction only ever reads turns not yet stamped `extracted_at`; `persist` stamp ### Tunable -Only two reconcile knobs are operator-configurable: +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. + -The similarity threshold (`DEDUP_SIM_HIGH`, `0.97`) and the vector-fold on/off switch (`DEDUP_VECTOR_ENABLED`) ship as **fixed internal constants** in `azure.cosmos.agent_memory.thresholds` — they have no env plumbing and ignore any environment variable, so the write-time fold behavior is not operator-tunable today. > **Indexing note.** The reconcile pool query orders by `created_at` (matching the prompt's "more recent first" tiebreaker). Cosmos's default indexing policy includes every property, so this works out of the box. If you customize the indexing policy to reduce write RU, ensure `/created_at/?` remains indexed or the query will fail with a 400 (`Order-by over a non-indexed path`). diff --git a/azure/cosmos/agent_memory/aio/cosmos_memory_client.py b/azure/cosmos/agent_memory/aio/cosmos_memory_client.py index d922221..89676fa 100644 --- a/azure/cosmos/agent_memory/aio/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/aio/cosmos_memory_client.py @@ -585,6 +585,7 @@ async def add_cosmos( salience: Optional[float] = None, embedding: Optional[list[float]] = None, embed: Optional[bool] = None, + created_at: Optional[str | datetime] = None, ) -> str: """Add a memory directly to Cosmos DB, bypassing the local buffer. @@ -612,6 +613,7 @@ async def add_cosmos( salience, embedding, embed, + created_at, ) if memory_type == "turn" and thread_id: task = asyncio.create_task(self._maybe_auto_trigger({(user_id, thread_id): 1})) @@ -717,8 +719,11 @@ async def search_cosmos( min_confidence: Optional[float] = None, created_after: Optional[str | datetime] = None, created_before: Optional[str | datetime] = None, + include_turns: bool = False, + turn_top_k: Optional[int] = None, ) -> list[dict[str, Any]]: - return await self._get_store().search( + """Search memories in Cosmos DB using vector similarity.""" + results = await self._get_store().search( search_terms=search_terms, memory_id=memory_id, user_id=user_id, @@ -735,6 +740,30 @@ async def search_cosmos( created_after=created_after, created_before=created_before, ) + if not include_turns or not user_id: + return results + + seen_content = {str(r.get("content") or "").strip() for r in results} + try: + turns = await self._get_store().search_turns( + search_terms=search_terms, + user_id=user_id, + thread_id=thread_id, + role=role, + top_k=turn_top_k if turn_top_k is not None else top_k, + exclude_tags=exclude_tags, + created_after=created_after, + created_before=created_before, + ) + except Exception as exc: # noqa: BLE001 + logger.warning("search_cosmos: include_turns turn search failed (%s); returning memories only", exc) + return results + 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) + return results async def search_turns( self, diff --git a/azure/cosmos/agent_memory/aio/services/pipeline.py b/azure/cosmos/agent_memory/aio/services/pipeline.py index f28ad5d..fc03034 100644 --- a/azure/cosmos/agent_memory/aio/services/pipeline.py +++ b/azure/cosmos/agent_memory/aio/services/pipeline.py @@ -398,6 +398,7 @@ def _build_transcript( items: list[dict[str, Any]], *, group_by_thread: bool = False, + include_timestamp: bool = False, ) -> str: # getattr fallback covers unit tests that build AsyncPipelineService # via __new__ to bypass __init__ (and therefore the metadata-keys stash). @@ -405,6 +406,7 @@ def _build_transcript( items, group_by_thread=group_by_thread, metadata_keys=getattr(self, "_transcript_metadata_keys", None), + include_timestamp=include_timestamp, ) async def _load_existing_memories( @@ -547,7 +549,7 @@ async def extract_memories_dry( deferred_turn_count = 0 quarantined_turn_count = 0 for batch in batches: - batch_transcript = self._build_transcript(batch) + batch_transcript = self._build_transcript(batch, include_timestamp=True) try: response_text = await self._run_prompty( "extract_memories.prompty", inputs={"transcript": batch_transcript} @@ -607,6 +609,8 @@ async def extract_memories_dry( seed = _ID_SEED_SEP.join((user_id, thread_id, new_content_hash)) det_id = f"fact_{hashlib.sha256(seed.encode()).hexdigest()[:32]}" topic_tags = build_topic_tags(fact.get("tags", [])) + fact_source = fact.get("source") if fact.get("source") in ("user", "agent") else "user" + source_tags = ["sys:agent-fact"] if fact_source == "agent" else [] confidence = fact.get("confidence") doc: dict[str, Any] = { "id": det_id, @@ -621,9 +625,10 @@ async def extract_memories_dry( "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, - "tags": ["sys:fact", "sys:auto-extracted"] + topic_tags, + "tags": ["sys:fact", "sys:auto-extracted"] + source_tags + topic_tags, "created_at": doc_timestamp, "updated_at": doc_timestamp, } diff --git a/azure/cosmos/agent_memory/aio/store/memory_store.py b/azure/cosmos/agent_memory/aio/store/memory_store.py index 4d5942b..986d486 100644 --- a/azure/cosmos/agent_memory/aio/store/memory_store.py +++ b/azure/cosmos/agent_memory/aio/store/memory_store.py @@ -168,6 +168,7 @@ async def add( salience: Optional[float] = None, embedding: Optional[list[float]] = None, embed: Optional[bool] = None, + created_at: Optional[str | datetime] = None, ) -> str: """Add a memory document to Cosmos DB and return its id.""" kwargs: dict[str, Any] = { @@ -185,6 +186,8 @@ async def add( kwargs["ttl"] = ttl if salience is not None: kwargs["salience"] = salience + if created_at is not None: + kwargs["created_at"] = created_at.isoformat() if isinstance(created_at, datetime) else created_at if memory_type != "turn": kwargs.setdefault("content_hash", compute_content_hash(content)) kwargs.setdefault("prompt_id", "manual:add") diff --git a/azure/cosmos/agent_memory/cosmos_memory_client.py b/azure/cosmos/agent_memory/cosmos_memory_client.py index 5fec042..0e26ca7 100644 --- a/azure/cosmos/agent_memory/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/cosmos_memory_client.py @@ -537,6 +537,7 @@ def add_cosmos( salience: Optional[float] = None, embedding: Optional[list[float]] = None, embed: Optional[bool] = None, + created_at: Optional[str | datetime] = None, ) -> str: """Add a memory directly to Cosmos DB, bypassing the local buffer. @@ -565,6 +566,7 @@ def add_cosmos( salience, embedding, embed, + created_at, ) if memory_type == "turn" and thread_id: try: @@ -672,13 +674,11 @@ def search_cosmos( min_confidence: Optional[float] = None, created_after: Optional[str | datetime] = None, created_before: Optional[str | datetime] = None, + include_turns: bool = False, + turn_top_k: Optional[int] = None, ) -> list[dict[str, Any]]: - """Search memories in Cosmos DB using vector similarity. - - Searches the derived memories container (facts/episodic/procedural). Use - :meth:`search_turns` to vector-search the raw conversation log instead. - """ - return self._get_store().search( + """Search memories in Cosmos DB using vector similarity.""" + results = self._get_store().search( search_terms=search_terms, memory_id=memory_id, user_id=user_id, @@ -695,6 +695,30 @@ def search_cosmos( created_after=created_after, created_before=created_before, ) + if not include_turns or not user_id: + return results + + seen_content = {str(r.get("content") or "").strip() for r in results} + try: + turns = self._get_store().search_turns( + search_terms=search_terms, + user_id=user_id, + thread_id=thread_id, + role=role, + top_k=turn_top_k if turn_top_k is not None else top_k, + exclude_tags=exclude_tags, + created_after=created_after, + created_before=created_before, + ) + except Exception as exc: # noqa: BLE001 + logger.warning("search_cosmos: include_turns turn search failed (%s); returning memories only", exc) + return results + 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) + return results def search_turns( self, diff --git a/azure/cosmos/agent_memory/prompts/_schemas.py b/azure/cosmos/agent_memory/prompts/_schemas.py index 2f43a12..28fa494 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 + episodic + unclassified +# extract_memories.prompty - extract facts (user- or agent-sourced) + episodic # --------------------------------------------------------------------------- _FACT_ITEM = { "type": "object", @@ -67,6 +67,10 @@ "other", ], }, + "source": { + "type": "string", + "enum": ["user", "agent"], + }, "confidence": {"type": "number"}, "salience": {"type": "number"}, "temporal_context": {"type": ["string", "null"]}, @@ -75,6 +79,7 @@ "required": [ "text", "category", + "source", "confidence", "salience", "temporal_context", diff --git a/azure/cosmos/agent_memory/prompts/extract_memories.prompty b/azure/cosmos/agent_memory/prompts/extract_memories.prompty index 1ce30a9..20caa51 100644 --- a/azure/cosmos/agent_memory/prompts/extract_memories.prompty +++ b/azure/cosmos/agent_memory/prompts/extract_memories.prompty @@ -1,6 +1,6 @@ --- name: extract_memories -version: v1 +version: v2 description: Extract facts and episodic memories from a conversation or provided content. model: apiType: chat @@ -33,11 +33,13 @@ Every memory must be explicitly grounded in the provided content - never inferre ## Speaker Discrimination - Where Memories May Come From -The transcript below is line-tagged: `[user]:` lines are the human's own words, `[agent]:` lines are the agent's response. These two sources are NOT interchangeable. +The transcript below is line-tagged: `[user]:` lines are the human's own words, `[agent]:` lines are the agent's response. Both 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"`). -- **Facts may ONLY come from `[user]:` lines.** The agent may restate, paraphrase, confirm, or make assumptions on the user's behalf ("Got it, you don't eat meat", "I assume you want a luxury hotel") - those are the agent's response, NOT the user's assertion. Never treat agent text as a new source of facts. If a fact appears only in `[agent]:` content and is not present in the user's lines, do not extract it. +- **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. -- Unsolicited world-knowledge the agent volunteers in `[agent]:` lines (e.g. "Python 3.13 was released in October 2024") is the agent's answer, not a memory. However, factual and domain content the user provides, shares, or asks to remember - including documents, excerpts, and reference material in `[user]:` lines - IS extractable. ## Confidence Scoring @@ -75,6 +77,14 @@ Monday") - keep that time expression verbatim in the fact `text` AND record it i 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. + ### Fact Formatting Rules - Each fact must be self-contained and intelligible without context - no pronouns like "it" or "they" without antecedents - Write in third person ("The user...", "The project...") @@ -193,6 +203,7 @@ A fact is a standing claim that holds outside any specific context. The test: if { "text": "The user is Alex, a data engineer at Acme Corp.", "category": "biographical", + "source": "user", "confidence": 1.0, "salience": 0.9, "temporal_context": null, @@ -201,6 +212,7 @@ A fact is a standing claim that holds outside any specific context. The test: if { "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", @@ -225,6 +237,7 @@ Name and role are one coherent biographical topic → a single rich memory, not { "text": "The user's team increased Kubernetes pod memory limits from 512MB to 1GB after OOM-killing issues.", "category": "other", + "source": "user", "confidence": 0.95, "salience": 0.7, "temporal_context": "last month", @@ -262,6 +275,7 @@ Name and role are one coherent biographical topic → a single rich memory, not { "text": "The user's analytics database runs on BigQuery.", "category": "biographical", + "source": "user", "confidence": 0.95, "salience": 0.7, "temporal_context": null, @@ -270,6 +284,7 @@ Name and role are one coherent biographical topic → a single rich memory, not { "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", @@ -295,6 +310,7 @@ The first statement is a standing preference and belongs in `facts`. The second { "text": "The user usually prefers budget hotels.", "category": "preference", + "source": "user", "confidence": 0.95, "salience": 0.7, "temporal_context": null, @@ -334,6 +350,7 @@ Extract only the literal, user-stated fact. The load time is concrete and worth { "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, @@ -358,6 +375,7 @@ The change is one coherent memory: it records the new state, what it replaced, a { "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", @@ -368,6 +386,81 @@ The change is one coherent memory: it records the new state, what it replaced, a } ``` +### 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"] + } + ], + "episodic": [] +} +``` + +### 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 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. + +**Output:** +```json +{ + "facts": [], + "episodic": [ + { + "scope_type": "trip", + "scope_value": "Rome", + "situation": null, + "action_taken": null, + "outcome": null, + "outcome_valence": null, + "reasoning": null, + "lesson": null, + "domain": "travel", + "confidence": 0.9, + "salience": 0.6, + "tags": ["topic:travel", "topic:hotels"] + } + ] +} +``` + --- ## Output Format @@ -380,6 +473,7 @@ You must output ONLY valid JSON matching the schema below. No preamble, no expla { "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", diff --git a/azure/cosmos/agent_memory/services/_pipeline_helpers.py b/azure/cosmos/agent_memory/services/_pipeline_helpers.py index a44314a..0a88e05 100644 --- a/azure/cosmos/agent_memory/services/_pipeline_helpers.py +++ b/azure/cosmos/agent_memory/services/_pipeline_helpers.py @@ -307,6 +307,7 @@ def build_transcript( *, group_by_thread: bool = False, metadata_keys: Optional[Iterable[str]] = None, + include_timestamp: bool = False, ) -> str: """Build a formatted transcript from memory documents. @@ -333,16 +334,25 @@ def build_transcript( Accepts any iterable of strings except ``str`` itself (which would be interpreted char-by-char). Generators are coerced to a tuple so the allow-list is reusable across turns. + include_timestamp: + If *True*, prefix each line with the turn's top-level ``created_at`` + (its event time) as ``[ | role]: content``. This lets the + extraction LLM anchor relative time expressions ("3 weeks ago", "last + June") to absolute dates instead of leaving them unresolved. Turns + without a ``created_at`` fall back to the plain ``[role]:`` form. """ keys = _normalize_metadata_keys(metadata_keys) + + def _line(m: dict[str, Any]) -> str: + role = m.get("role", "unknown") + content = m.get("content", "") + meta_str = _format_metadata_segment(m.get("metadata", {}), keys) + created_at = m.get("created_at") if include_timestamp else None + prefix = f"[{created_at} | {role}]" if created_at else f"[{role}]" + return f"{prefix}: {content}{meta_str}" + if not group_by_thread: - lines: list[str] = [] - for m in items: - role = m.get("role", "unknown") - content = m.get("content", "") - meta_str = _format_metadata_segment(m.get("metadata", {}), keys) - lines.append(f"[{role}]: {content}{meta_str}") - return "\n".join(lines) + return "\n".join(_line(m) for m in items) threads: dict[str, list[dict[str, Any]]] = defaultdict(list) for m in items: @@ -352,10 +362,7 @@ def build_transcript( for tid, thread_items in threads.items(): parts.append(f"=== Thread {tid} ===") for m in thread_items: - role = m.get("role", "unknown") - content = m.get("content", "") - meta_str = _format_metadata_segment(m.get("metadata", {}), keys) - parts.append(f"[{role}]: {content}{meta_str}") + parts.append(_line(m)) parts.append("") return "\n".join(parts) diff --git a/azure/cosmos/agent_memory/services/pipeline.py b/azure/cosmos/agent_memory/services/pipeline.py index c99a2c2..e340c30 100644 --- a/azure/cosmos/agent_memory/services/pipeline.py +++ b/azure/cosmos/agent_memory/services/pipeline.py @@ -236,6 +236,7 @@ def _build_transcript( items: list[dict[str, Any]], *, group_by_thread: bool = False, + include_timestamp: bool = False, ) -> str: # getattr fallback covers unit tests that build PipelineService via # __new__ to bypass __init__ (and therefore the metadata-keys stash). @@ -243,6 +244,7 @@ def _build_transcript( items, group_by_thread=group_by_thread, metadata_keys=getattr(self, "_transcript_metadata_keys", None), + include_timestamp=include_timestamp, ) def _load_existing_memories( @@ -611,7 +613,7 @@ def extract_memories_dry( deferred_turn_count = 0 quarantined_turn_count = 0 for batch in batches: - batch_transcript = self._build_transcript(batch) + batch_transcript = self._build_transcript(batch, include_timestamp=True) try: response_text = self._run_prompty("extract_memories.prompty", inputs={"transcript": batch_transcript}) parsed = self._parse_llm_json(response_text) @@ -669,6 +671,8 @@ def extract_memories_dry( seed = _ID_SEED_SEP.join((user_id, thread_id, new_content_hash)) det_id = f"fact_{hashlib.sha256(seed.encode()).hexdigest()[:32]}" topic_tags = build_topic_tags(fact.get("tags", [])) + fact_source = fact.get("source") if fact.get("source") in ("user", "agent") else "user" + source_tags = ["sys:agent-fact"] if fact_source == "agent" else [] confidence = fact.get("confidence") doc: dict[str, Any] = { "id": det_id, @@ -683,9 +687,10 @@ def extract_memories_dry( "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, - "tags": ["sys:fact", "sys:auto-extracted"] + topic_tags, + "tags": ["sys:fact", "sys:auto-extracted"] + source_tags + topic_tags, "created_at": doc_timestamp, "updated_at": doc_timestamp, } diff --git a/azure/cosmos/agent_memory/store/memory_store.py b/azure/cosmos/agent_memory/store/memory_store.py index f67b532..2ccd8e9 100644 --- a/azure/cosmos/agent_memory/store/memory_store.py +++ b/azure/cosmos/agent_memory/store/memory_store.py @@ -202,6 +202,7 @@ def add( salience: Optional[float] = None, embedding: Optional[list[float]] = None, embed: Optional[bool] = None, + created_at: Optional[str | datetime] = None, ) -> str: """Add a memory document to Cosmos DB and return its id.""" kwargs: dict[str, Any] = { @@ -219,6 +220,8 @@ def add( kwargs["ttl"] = ttl if salience is not None: kwargs["salience"] = salience + if created_at is not None: + kwargs["created_at"] = created_at.isoformat() if isinstance(created_at, datetime) else created_at if memory_type != "turn": kwargs.setdefault("content_hash", compute_content_hash(content)) kwargs.setdefault("prompt_id", "manual:add") diff --git a/azure/cosmos/agent_memory/thresholds.py b/azure/cosmos/agent_memory/thresholds.py index a1edc6a..6d24dce 100644 --- a/azure/cosmos/agent_memory/thresholds.py +++ b/azure/cosmos/agent_memory/thresholds.py @@ -27,16 +27,21 @@ # parameter of :py:meth:`ProcessingPipeline.reconcile_memories`. Hard cap # of 500 (enforced by the pipeline) bounds prompt size and LLM cost. DEFAULT_DEDUP_POOL_SIZE = 50 +# Write-time in-place near-duplicate folding. When enabled, a freshly +# extracted memory that is >= DEDUP_SIM_HIGH similar to an existing active +# memory is folded into that record in place instead of persisting as a new +# doc. Default OFF (add-only): keeping every extracted memory preserves the +# retrieval surface, which benchmarked better than folding. Operators set +# ``DEDUP_VECTOR_ENABLED=true`` to turn folding back on. +DEFAULT_DEDUP_VECTOR_ENABLED = False # --------------------------------------------------------------------------- # INTERNAL dedup/search tuning - NOT customer-configurable. # These ship as fixed feature constants (no env vars, not in any settings # template). They are maintainer-tunable here in code only; if a knob ever # needs to become operator-facing we add the env plumbing back deliberately. -# The dedup + hybrid-search features ship ON via these values. # --------------------------------------------------------------------------- EXTRACTION_BATCH_MAX_TOKENS = 7000 -DEDUP_VECTOR_ENABLED = True # write-time in-place near-dup folding DEDUP_SIM_HIGH = 0.97 # >= -> fold new memory into existing canonical in place DEFAULT_TTL_BY_TYPE: dict[str, int] = { @@ -160,8 +165,8 @@ def get_extraction_batch_max_tokens() -> int: def get_dedup_vector_enabled() -> bool: - """Whether Stage-3 vector deduplication is enabled (internal; on).""" - return DEDUP_VECTOR_ENABLED + """Whether write-time vector deduplication (in-place folding) is enabled.""" + return _parse_bool("DEDUP_VECTOR_ENABLED", DEFAULT_DEDUP_VECTOR_ENABLED) def get_dedup_sim_high() -> float: @@ -235,6 +240,7 @@ def get_processor_owner() -> Optional[str]: "DEFAULT_USER_SUMMARY_EVERY_N", "DEFAULT_DEDUP_EVERY_N", "DEFAULT_DEDUP_POOL_SIZE", + "DEFAULT_DEDUP_VECTOR_ENABLED", "DEFAULT_TTL_BY_TYPE", "DEFAULT_PROCEDURAL_SYNTHESIS_AUTO", "DEFAULT_ENABLE_TURN_EMBEDDINGS", diff --git a/tests/unit/aio/services/test_dedup_vector_async.py b/tests/unit/aio/services/test_dedup_vector_async.py index 3667be0..36ee349 100644 --- a/tests/unit/aio/services/test_dedup_vector_async.py +++ b/tests/unit/aio/services/test_dedup_vector_async.py @@ -9,6 +9,14 @@ from azure.cosmos.agent_memory.aio.services.pipeline import AsyncPipelineService +@pytest.fixture(autouse=True) +def _enable_vector_folding(monkeypatch: pytest.MonkeyPatch) -> None: + # DEDUP_VECTOR_ENABLED now defaults to False (add-only); this suite exercises + # the in-place folding path, so enable it. Tests that assert the flag-off + # behavior patch the getter directly and override this. + monkeypatch.setenv("DEDUP_VECTOR_ENABLED", "true") + + def _service() -> AsyncPipelineService: p = AsyncPipelineService.__new__(AsyncPipelineService) p._memories_container = MagicMock() diff --git a/tests/unit/aio/store/test_memory_store.py b/tests/unit/aio/store/test_memory_store.py index 2e7940b..e0626de 100644 --- a/tests/unit/aio/store/test_memory_store.py +++ b/tests/unit/aio/store/test_memory_store.py @@ -61,6 +61,35 @@ async def test_add_upserts_memory_document(): assert body["ttl"] == 2_592_000 +async def test_add_honors_explicit_created_at_string(): + turns = MagicMock() + turns.upsert_item = AsyncMock() + store = AsyncMemoryStore(containers=_containers(turns=turns)) + + await store.add( + user_id="u1", + role="user", + content="hello", + thread_id="t1", + created_at="2024-03-01T12:00:00+00:00", + ) + + body = turns.upsert_item.call_args.kwargs["body"] + assert body["created_at"] == "2024-03-01T12:00:00+00:00" + + +async def test_add_normalizes_datetime_created_at(): + turns = MagicMock() + turns.upsert_item = AsyncMock() + store = AsyncMemoryStore(containers=_containers(turns=turns)) + + dt = datetime(2024, 3, 1, 12, 0, tzinfo=timezone.utc) + await store.add(user_id="u1", role="user", content="hello", thread_id="t1", created_at=dt) + + body = turns.upsert_item.call_args.kwargs["body"] + assert body["created_at"] == dt.isoformat() + + @pytest.mark.parametrize( ("memory_type", "expected_ttl"), [ diff --git a/tests/unit/aio/test_cosmos_memory_client.py b/tests/unit/aio/test_cosmos_memory_client.py index e3dfb3e..4d6adf1 100644 --- a/tests/unit/aio/test_cosmos_memory_client.py +++ b/tests/unit/aio/test_cosmos_memory_client.py @@ -985,3 +985,50 @@ async def test_list_tags_delegates_to_store(): kwargs = container.query_items.call_args.kwargs assert "SELECT VALUE c.tags" in kwargs["query"] assert kwargs["parameters"] == [{"name": "@user_id", "value": "u1"}] + + +class TestAsyncSearchCosmosUnifiedRetrieval: + def _stub_store(self, mem, *, memories, turns=None, turns_raises=False): + store = MagicMock() + store.search = AsyncMock(return_value=list(memories)) + if turns_raises: + store.search_turns = AsyncMock(side_effect=RuntimeError("turns not embedded")) + else: + store.search_turns = AsyncMock(return_value=list(turns or [])) + mem._get_store = MagicMock(return_value=store) + return store + + async def test_default_does_not_search_turns(self): + mem, _ = _connected_client() + store = self._stub_store(mem, memories=[{"content": "fact A", "type": "fact"}]) + + out = await mem.search_cosmos("q", user_id="u1") + + assert out == [{"content": "fact A", "type": "fact"}] + store.search_turns.assert_not_called() + + async def test_include_turns_appends_and_dedups(self): + mem, _ = _connected_client() + self._stub_store( + mem, + memories=[{"content": "fact A", "type": "fact"}], + turns=[ + {"content": "fact A", "type": "turn"}, # dup -> skipped + {"content": "raw dialogue B", "type": "turn"}, + ], + ) + + out = await mem.search_cosmos("q", user_id="u1", include_turns=True) + + assert out == [ + {"content": "fact A", "type": "fact"}, + {"content": "raw dialogue B", "type": "turn"}, + ] + + async def test_include_turns_failure_returns_memories_only(self): + mem, _ = _connected_client() + self._stub_store(mem, memories=[{"content": "fact A", "type": "fact"}], turns_raises=True) + + out = await mem.search_cosmos("q", user_id="u1", include_turns=True) + + assert out == [{"content": "fact A", "type": "fact"}] diff --git a/tests/unit/services/test_dedup_vector.py b/tests/unit/services/test_dedup_vector.py index 9cabf3f..3b8de6c 100644 --- a/tests/unit/services/test_dedup_vector.py +++ b/tests/unit/services/test_dedup_vector.py @@ -4,9 +4,19 @@ from typing import Any from unittest.mock import MagicMock +import pytest + from azure.cosmos.agent_memory.services.pipeline import PipelineService +@pytest.fixture(autouse=True) +def _enable_vector_folding(monkeypatch: pytest.MonkeyPatch) -> None: + # DEDUP_VECTOR_ENABLED now defaults to False (add-only); this suite exercises + # the in-place folding path, so enable it. Tests that assert the flag-off + # behavior patch the getter directly and override this. + monkeypatch.setenv("DEDUP_VECTOR_ENABLED", "true") + + def _make_pipeline() -> PipelineService: p = PipelineService.__new__(PipelineService) p._memories_container = MagicMock() diff --git a/tests/unit/services/test_extract_dry.py b/tests/unit/services/test_extract_dry.py index 97bae44..610b9b5 100644 --- a/tests/unit/services/test_extract_dry.py +++ b/tests/unit/services/test_extract_dry.py @@ -416,3 +416,111 @@ def test_extract_defers_retryable_batch_leaving_turns_unstamped(monkeypatch) -> assert len(out["processed_turn_docs"]) == 2 stats = [u for u in out["updates"] if u.get("op") == "stats" and "deferred_turn_count" in u] assert stats and stats[0]["deferred_turn_count"] == 1 + + +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. + return { + "facts": [ + { + "text": "The agent booked the user on flight TP204 to Lisbon.", + "category": "other", + "source": "agent", + "confidence": 0.95, + "salience": 0.8, + "temporal_context": None, + "tags": ["travel"], + }, + { + "text": "The user prefers window seats.", + "category": "preference", + "source": "user", + "confidence": 0.9, + "salience": 0.7, + "temporal_context": None, + "tags": ["travel"], + }, + { + "text": "The user lives in Seattle.", + "category": "biographical", + "confidence": 0.9, + "salience": 0.7, + "temporal_context": None, + "tags": ["identity"], + }, + ], + "episodic": [], + } + + +def _fact_by_text(facts: list[dict[str, Any]], needle: str) -> dict[str, Any]: + for doc in facts: + if needle in doc["content"]: + return doc + raise AssertionError(f"no fact containing {needle!r}") + + +def test_agent_sourced_fact_is_tagged_and_stamped() -> None: + memories_store = _Store([]) + service = PipelineService( + memories_store, + _SyncChat([_agent_source_response()]), + _SyncEmbeddings(), + containers=_containers_for_store(memories_store, turns_store=_Store([_turn(1)])), + ) + + out = service.extract_memories_dry("u1", "t1") + + agent_fact = _fact_by_text(out["facts"], "booked the user on flight") + assert agent_fact["metadata"]["source"] == "agent" + assert "sys:agent-fact" in agent_fact["tags"] + + user_fact = _fact_by_text(out["facts"], "prefers window seats") + assert user_fact["metadata"]["source"] == "user" + assert "sys:agent-fact" not in user_fact["tags"] + + # Source omitted by the model must default to user (no agent tag). + defaulted = _fact_by_text(out["facts"], "lives in Seattle") + assert defaulted["metadata"]["source"] == "user" + assert "sys:agent-fact" not in defaulted["tags"] + + +@pytest.mark.asyncio +async def test_async_agent_sourced_fact_is_tagged_and_stamped() -> None: + memories_store = _AsyncStore([]) + service = AsyncPipelineService( + memories_store, + _AsyncChat([_agent_source_response()]), + _AsyncEmbeddings(), + containers=_async_containers_for_store(memories_store, turns_store=_AsyncStore([_turn(1)])), + ) + + out = await service.extract_memories_dry("u1", "t1") + + agent_fact = _fact_by_text(out["facts"], "booked the user on flight") + assert agent_fact["metadata"]["source"] == "agent" + assert "sys:agent-fact" in agent_fact["tags"] + + defaulted = _fact_by_text(out["facts"], "lives in Seattle") + assert defaulted["metadata"]["source"] == "user" + assert "sys:agent-fact" not in defaulted["tags"] + + +def test_extraction_transcript_includes_turn_timestamps() -> None: + # The extraction prompt must carry each turn's event time so the LLM can + # resolve relative dates. _turn(i) stamps created_at=2025-01-01T00:0i:00. + chat = _SyncChat([_response()]) + memories_store = _Store([]) + turns_store = _Store([_turn(1)]) + service = PipelineService( + memories_store, + chat, + _SyncEmbeddings(), + containers=_containers_for_store(memories_store, turns_store=turns_store), + ) + + service.extract_memories_dry("u1", "t1") + + prompt_text = json.dumps(chat.messages) + assert "2025-01-01T00:01:00+00:00 | user" in prompt_text diff --git a/tests/unit/services/test_prompty_loader.py b/tests/unit/services/test_prompty_loader.py index daa1ee6..784a601 100644 --- a/tests/unit/services/test_prompty_loader.py +++ b/tests/unit/services/test_prompty_loader.py @@ -56,13 +56,16 @@ def test_loader_prompt_version_is_cached(tmp_path: Path) -> None: def test_all_shipped_prompts_declare_version() -> None: loader = PromptyLoader() - for filename in ( - "extract_memories.prompty", - "dedup.prompty", - "summarize.prompty", - "summarize_update.prompty", - "user_summary.prompty", - "user_summary_update.prompty", - "synthesize_procedural.prompty", - ): - assert loader.prompt_version(filename) == "v1" + # extract_memories bumped to v2 when agent-sourced fact extraction landed; + # the rest remain v1. Every shipped prompt must declare *some* version. + expected = { + "extract_memories.prompty": "v2", + "dedup.prompty": "v1", + "summarize.prompty": "v1", + "summarize_update.prompty": "v1", + "user_summary.prompty": "v1", + "user_summary_update.prompty": "v1", + "synthesize_procedural.prompty": "v1", + } + for filename, version in expected.items(): + assert loader.prompt_version(filename) == version diff --git a/tests/unit/services/test_transcript_metadata.py b/tests/unit/services/test_transcript_metadata.py index e6b9f85..8dec24d 100644 --- a/tests/unit/services/test_transcript_metadata.py +++ b/tests/unit/services/test_transcript_metadata.py @@ -327,3 +327,34 @@ def fake_pipeline(*args: object, **kwargs: object) -> object: mod.AsyncPipelineService = original assert captured.get("transcript_metadata_keys") == ("agent_id",) + + +class TestIncludeTimestamp: + """include_timestamp prefixes lines with the turn's event time so the + extraction LLM can resolve relative time expressions to absolute dates.""" + + def test_timestamp_prefix_flat(self) -> None: + items = [ + {"role": "user", "content": "Hello", "created_at": "2024-06-20T10:00:00+00:00"}, + {"role": "assistant", "content": "Hi", "created_at": "2024-06-20T10:01:00+00:00"}, + ] + out = build_transcript(items, include_timestamp=True) + assert out == ("[2024-06-20T10:00:00+00:00 | user]: Hello\n[2024-06-20T10:01:00+00:00 | assistant]: Hi") + + def test_timestamp_absent_falls_back_to_plain_role(self) -> None: + items = [{"role": "user", "content": "Hello"}] # no created_at + out = build_transcript(items, include_timestamp=True) + assert out == "[user]: Hello" + + def test_timestamp_off_by_default(self) -> None: + items = [{"role": "user", "content": "Hello", "created_at": "2024-06-20T10:00:00+00:00"}] + out = build_transcript(items) + assert out == "[user]: Hello" + + def test_timestamp_prefix_grouped(self) -> None: + items = [ + {"role": "user", "content": "Q", "thread_id": "t1", "created_at": "2024-06-20T10:00:00+00:00"}, + ] + out = build_transcript(items, group_by_thread=True, include_timestamp=True) + assert "[2024-06-20T10:00:00+00:00 | user]: Q" in out + assert "=== Thread t1 ===" in out diff --git a/tests/unit/store/test_memory_store.py b/tests/unit/store/test_memory_store.py index 80a61c9..f96d504 100644 --- a/tests/unit/store/test_memory_store.py +++ b/tests/unit/store/test_memory_store.py @@ -47,6 +47,44 @@ def test_add_upserts_memory_document(): assert body["ttl"] == 2_592_000 +def test_add_defaults_created_at_to_ingestion_time(): + turns = MagicMock() + store = MemoryStore(containers=_containers(turns=turns)) + + store.add(user_id="u1", role="user", content="hello", thread_id="t1") + + body = turns.upsert_item.call_args.kwargs["body"] + # No created_at passed -> defaulted (present, non-empty ISO string). + assert body["created_at"] + + +def test_add_honors_explicit_created_at_string(): + turns = MagicMock() + store = MemoryStore(containers=_containers(turns=turns)) + + store.add( + user_id="u1", + role="user", + content="hello", + thread_id="t1", + created_at="2024-03-01T12:00:00+00:00", + ) + + body = turns.upsert_item.call_args.kwargs["body"] + assert body["created_at"] == "2024-03-01T12:00:00+00:00" + + +def test_add_normalizes_datetime_created_at(): + turns = MagicMock() + store = MemoryStore(containers=_containers(turns=turns)) + + dt = datetime(2024, 3, 1, 12, 0, tzinfo=timezone.utc) + store.add(user_id="u1", role="user", content="hello", thread_id="t1", created_at=dt) + + body = turns.upsert_item.call_args.kwargs["body"] + assert body["created_at"] == dt.isoformat() + + @pytest.mark.parametrize( ("memory_type", "expected_ttl"), [ diff --git a/tests/unit/test_cosmos_memory_client.py b/tests/unit/test_cosmos_memory_client.py index b438103..351b543 100644 --- a/tests/unit/test_cosmos_memory_client.py +++ b/tests/unit/test_cosmos_memory_client.py @@ -602,6 +602,20 @@ def test_add_cosmos(self): assert body["user_id"] == "u1" assert body["role"] == "user" + def test_add_cosmos_threads_explicit_created_at(self): + mem, container = _connected_client() + mem._maybe_auto_trigger = MagicMock() + mem.add_cosmos( + user_id="u1", + role="user", + content="hello", + thread_id="t1", + created_at="2024-03-01T12:00:00+00:00", + ) + + body = mem._turns_container_client.upsert_item.call_args.kwargs["body"] + assert body["created_at"] == "2024-03-01T12:00:00+00:00" + def test_add_cosmos_not_connected(self): mem = _make_client() with pytest.raises(CosmosNotConnectedError): @@ -648,6 +662,62 @@ def test_add_cosmos_swallows_cadence_failure(self): mem._turns_container_client.upsert_item.assert_called_once() +class TestSearchCosmosUnifiedRetrieval: + def _stub_store(self, mem, *, memories, turns=None, turns_raises=False): + store = MagicMock() + store.search.return_value = list(memories) + if turns_raises: + store.search_turns.side_effect = RuntimeError("turns not embedded") + else: + store.search_turns.return_value = list(turns or []) + mem._get_store = MagicMock(return_value=store) + return store + + def test_default_does_not_search_turns(self): + mem, _ = _connected_client() + store = self._stub_store(mem, memories=[{"content": "fact A", "type": "fact"}]) + + out = mem.search_cosmos("q", user_id="u1") + + assert out == [{"content": "fact A", "type": "fact"}] + store.search_turns.assert_not_called() + + def test_include_turns_appends_and_dedups(self): + mem, _ = _connected_client() + self._stub_store( + mem, + memories=[{"content": "fact A", "type": "fact"}], + turns=[ + {"content": "fact A", "type": "turn"}, # dup of memory -> skipped + {"content": "raw dialogue B", "type": "turn"}, + ], + ) + + out = mem.search_cosmos("q", user_id="u1", include_turns=True) + + assert out == [ + {"content": "fact A", "type": "fact"}, + {"content": "raw dialogue B", "type": "turn"}, + ] + + def test_include_turns_failure_returns_memories_only(self): + mem, _ = _connected_client() + self._stub_store(mem, memories=[{"content": "fact A", "type": "fact"}], turns_raises=True) + + out = mem.search_cosmos("q", user_id="u1", include_turns=True) + + assert out == [{"content": "fact A", "type": "fact"}] + + def test_include_turns_without_user_id_skips_turns(self): + mem, _ = _connected_client() + store = self._stub_store(mem, memories=[{"content": "fact A", "type": "fact"}]) + + out = mem.search_cosmos("q", include_turns=True) + + assert out == [{"content": "fact A", "type": "fact"}] + store.search_turns.assert_not_called() + + class TestPushToCosmos: def test_push_to_cosmos(self): mem, container = _connected_client() diff --git a/tests/unit/test_thresholds.py b/tests/unit/test_thresholds.py index d867d71..eb350d0 100644 --- a/tests/unit/test_thresholds.py +++ b/tests/unit/test_thresholds.py @@ -179,9 +179,18 @@ def test_processor_owner_invalid_uses_default(monkeypatch: pytest.MonkeyPatch) - def test_internalized_getters_return_fixed_constants_and_ignore_env(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("EXTRACTION_BATCH_MAX_TOKENS", "999") - monkeypatch.setenv("DEDUP_VECTOR_ENABLED", "false") monkeypatch.setenv("DEDUP_SIM_HIGH", "0.50") assert thresholds.get_extraction_batch_max_tokens() == 7000 - assert thresholds.get_dedup_vector_enabled() is True assert thresholds.get_dedup_sim_high() == 0.97 + + +def test_dedup_vector_enabled_defaults_false_and_reads_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("DEDUP_VECTOR_ENABLED", raising=False) + assert thresholds.get_dedup_vector_enabled() is False + + monkeypatch.setenv("DEDUP_VECTOR_ENABLED", "true") + assert thresholds.get_dedup_vector_enabled() is True + + monkeypatch.setenv("DEDUP_VECTOR_ENABLED", "false") + assert thresholds.get_dedup_vector_enabled() is False From a1d9d23933fccda065687fd3fa5a03148bdb43f3 Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Tue, 14 Jul 2026 10:20:15 -0700 Subject: [PATCH 2/3] Resolving comments --- CHANGELOG.md | 5 +++ Docs/concepts.md | 4 ++ azure/cosmos/agent_memory/_utils.py | 23 ++++++++++ .../agent_memory/aio/cosmos_memory_client.py | 7 +-- .../agent_memory/aio/services/pipeline.py | 32 +++++++++++++- .../agent_memory/aio/store/memory_store.py | 3 +- .../agent_memory/cosmos_memory_client.py | 7 +-- .../prompts/extract_memories.prompty | 2 +- .../cosmos/agent_memory/services/pipeline.py | 32 +++++++++++++- .../cosmos/agent_memory/store/memory_store.py | 3 +- .../aio/services/test_dedup_vector_async.py | 20 +++++++++ tests/unit/services/test_dedup_vector.py | 22 ++++++++++ tests/unit/test_reconcile.py | 9 ++++ tests/unit/test_utils.py | 44 +++++++++++++++++++ 14 files changed, 201 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 35cc968..d90636b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,10 @@ #### Features Added * Write-time in-place deduplication: near-duplicate memories fold into the existing record (same id, newer content) instead of creating a new doc. See [PR:#26](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/26) * Reconciliation now resolves contradictions only, soft-deleting the loser with `superseded_by`; `get_memory_history()` walks that chain. See [PR:#26](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/26) +* Agent-sourced facts: extraction can now capture the assistant's own actions and recommendations (not just user statements), tagged `source=agent` / `sys:agent-fact` so retrieval can include or exclude them. See [PR:#31](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/31) +* Event time on write: `add_cosmos(..., created_at=...)` sets a memory's event time (falls back to ingestion time), enabling time-aware retrieval and temporal reasoning over backfilled conversations. See [PR:#31](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/31) +* Unified retrieval: `search_cosmos(include_turns=True)` blends raw conversation turns into results alongside extracted memories. See [PR:#31](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/31) +* `DEDUP_VECTOR_ENABLED` is now an environment knob (default `false` = add-only) instead of a fixed internal constant. See [PR:#31](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/31) #### Bugs Fixed * Fixed a re-extraction loop that re-extracted the whole conversation every cycle (turns were never stamped `extracted_at` when vector dedup was on). See [PR:#26](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/26) @@ -12,6 +16,7 @@ #### Other Changes * Reworked the extraction prompt (anti-inference, preserve specifics, topic-grouped memories) and simplified the schema to `fact`/`episodic` with fixed fact categories. See [PR:#26](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/26) * Token-bounded extraction batches with per-batch failure isolation; embedding inputs truncated to the model token budget. See [PR:#26](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/26) +* Extraction now sees per-turn timestamps and resolves relative dates ("3 weeks ago") to absolute dates in the fact text (time-range filtering uses the memory's `created_at`). See [PR:#31](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/31) ## [0.2.0b3] (2026-07-08) diff --git a/Docs/concepts.md b/Docs/concepts.md index ee19084..7484ec0 100644 --- a/Docs/concepts.md +++ b/Docs/concepts.md @@ -96,6 +96,10 @@ Facts work especially well for vector search because each fact is stored as a sm 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. + --- ## Processing Pipeline diff --git a/azure/cosmos/agent_memory/_utils.py b/azure/cosmos/agent_memory/_utils.py index d9a97fb..8373e18 100644 --- a/azure/cosmos/agent_memory/_utils.py +++ b/azure/cosmos/agent_memory/_utils.py @@ -108,6 +108,29 @@ def _coerce_datetime_iso(value: Optional[str | datetime]) -> Optional[str]: return value +def normalize_created_at_iso(value: Optional[str | datetime]) -> Optional[str]: + """Normalize a caller-supplied event time to a tz-aware **UTC** ISO-8601 string.""" + if value is None: + return None + if isinstance(value, datetime): + dt = value + else: + text = str(value).strip() + if not text: + raise ValidationError("created_at must be a non-empty ISO-8601 string or datetime") + try: + # ``fromisoformat`` doesn't accept a trailing 'Z' before 3.11 in all + # forms; normalize it to an explicit UTC offset first. + dt = datetime.fromisoformat(text.replace("Z", "+00:00")) + except ValueError as exc: + raise ValidationError(f"created_at is not a valid ISO-8601 timestamp: {value!r}") from exc + if dt.tzinfo is None: + dt = dt.replace(tzinfo=timezone.utc) + else: + dt = dt.astimezone(timezone.utc) + return dt.isoformat() + + def _normalize_for_hash(text: str) -> str: """Lowercase + collapse whitespace for write-time exact-dedup. diff --git a/azure/cosmos/agent_memory/aio/cosmos_memory_client.py b/azure/cosmos/agent_memory/aio/cosmos_memory_client.py index 89676fa..b058426 100644 --- a/azure/cosmos/agent_memory/aio/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/aio/cosmos_memory_client.py @@ -722,8 +722,9 @@ async def search_cosmos( include_turns: bool = False, turn_top_k: Optional[int] = None, ) -> list[dict[str, Any]]: - """Search memories in Cosmos DB using vector similarity.""" - results = await self._get_store().search( + """Search memories using vector similarity; optionally blend raw turns.""" + store = self._get_store() + results = await store.search( search_terms=search_terms, memory_id=memory_id, user_id=user_id, @@ -745,7 +746,7 @@ async def search_cosmos( seen_content = {str(r.get("content") or "").strip() for r in results} try: - turns = await self._get_store().search_turns( + turns = await store.search_turns( search_terms=search_terms, user_id=user_id, thread_id=thread_id, diff --git a/azure/cosmos/agent_memory/aio/services/pipeline.py b/azure/cosmos/agent_memory/aio/services/pipeline.py index fc03034..81369b9 100644 --- a/azure/cosmos/agent_memory/aio/services/pipeline.py +++ b/azure/cosmos/agent_memory/aio/services/pipeline.py @@ -609,7 +609,14 @@ async def extract_memories_dry( seed = _ID_SEED_SEP.join((user_id, thread_id, new_content_hash)) det_id = f"fact_{hashlib.sha256(seed.encode()).hexdigest()[:32]}" topic_tags = build_topic_tags(fact.get("tags", [])) - fact_source = fact.get("source") if fact.get("source") in ("user", "agent") else "user" + raw_source = fact.get("source") + if raw_source is not None and raw_source not in ("user", "agent"): + logger.debug( + "extract_memories: coercing invalid fact source=%r to 'user' user_id=%s", + raw_source, + user_id, + ) + fact_source = raw_source if raw_source in ("user", "agent") else "user" source_tags = ["sys:agent-fact"] if fact_source == "agent" else [] confidence = fact.get("confidence") doc: dict[str, Any] = { @@ -865,10 +872,27 @@ async def _nearest_active_full( return None, 0.0 async def _apply_inplace_update(self, neighbor: dict[str, Any], new_doc: dict[str, Any]) -> bool: - """Async mirror of the sync in-place refresh (recency-wins content+embedding).""" + """Async mirror of the sync in-place refresh (recency-wins content+embedding). + + Folds only within the same ``metadata.source`` (user vs agent); a + cross-source pair returns False so the caller keeps it as a novel ADD, + preventing tag/source desync. + """ from azure.core import MatchConditions from azure.cosmos.exceptions import CosmosAccessConditionFailedError + neighbor_source = (neighbor.get("metadata") or {}).get("source") or "user" + new_source = (new_doc.get("metadata") or {}).get("source") or "user" + if neighbor_source != new_source: + logger.info( + "in-place dedup update skipped (source mismatch neighbor=%s new=%s) " + "target_id=%s; keeping new doc as novel", + neighbor_source, + new_source, + neighbor.get("id"), + ) + return False + try: old_etag = neighbor.get("_etag") updated = dict(neighbor) @@ -1568,11 +1592,15 @@ def _emit_reconcile_outcome( async def _active_memories_for_reconcile(self, user_id: str, memory_type: str, n: int) -> list[dict[str, Any]]: capped_n = top_literal(n, name="reconcile_memories.n") + # Agent-sourced facts (sys:agent-fact) are excluded: they record what the + # agent did/recommended (historical events), not mutable user state, so + # they must never be contradiction-superseded by a later user statement. query = ( f"SELECT TOP {capped_n} * FROM c " "WHERE c.user_id = @user_id " "AND c.type = @memory_type " f"AND {_ACTIVE_DOC_FILTER} " + "AND NOT ARRAY_CONTAINS(c.tags, 'sys:agent-fact') " "ORDER BY c.created_at DESC" ) return await self._query_items( diff --git a/azure/cosmos/agent_memory/aio/store/memory_store.py b/azure/cosmos/agent_memory/aio/store/memory_store.py index 986d486..1b21987 100644 --- a/azure/cosmos/agent_memory/aio/store/memory_store.py +++ b/azure/cosmos/agent_memory/aio/store/memory_store.py @@ -20,6 +20,7 @@ compute_content_hash, extract_keywords, new_id, + normalize_created_at_iso, ) from azure.cosmos.agent_memory.exceptions import ( ConfigurationError, @@ -187,7 +188,7 @@ async def add( if salience is not None: kwargs["salience"] = salience if created_at is not None: - kwargs["created_at"] = created_at.isoformat() if isinstance(created_at, datetime) else created_at + kwargs["created_at"] = normalize_created_at_iso(created_at) if memory_type != "turn": kwargs.setdefault("content_hash", compute_content_hash(content)) kwargs.setdefault("prompt_id", "manual:add") diff --git a/azure/cosmos/agent_memory/cosmos_memory_client.py b/azure/cosmos/agent_memory/cosmos_memory_client.py index 0e26ca7..aaf4859 100644 --- a/azure/cosmos/agent_memory/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/cosmos_memory_client.py @@ -677,8 +677,9 @@ def search_cosmos( include_turns: bool = False, turn_top_k: Optional[int] = None, ) -> list[dict[str, Any]]: - """Search memories in Cosmos DB using vector similarity.""" - results = self._get_store().search( + """Search memories using vector similarity; optionally blend raw turns.""" + store = self._get_store() + results = store.search( search_terms=search_terms, memory_id=memory_id, user_id=user_id, @@ -700,7 +701,7 @@ def search_cosmos( seen_content = {str(r.get("content") or "").strip() for r in results} try: - turns = self._get_store().search_turns( + turns = store.search_turns( search_terms=search_terms, user_id=user_id, thread_id=thread_id, diff --git a/azure/cosmos/agent_memory/prompts/extract_memories.prompty b/azure/cosmos/agent_memory/prompts/extract_memories.prompty index 20caa51..b4f2cc0 100644 --- a/azure/cosmos/agent_memory/prompts/extract_memories.prompty +++ b/azure/cosmos/agent_memory/prompts/extract_memories.prompty @@ -33,7 +33,7 @@ Every memory must be explicitly grounded in the provided content - never inferre ## Speaker Discrimination - Where Memories May Come From -The transcript below is line-tagged: `[user]:` lines are the human's own words, `[agent]:` lines are the agent's response. Both 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"`). +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 speakers 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"`. diff --git a/azure/cosmos/agent_memory/services/pipeline.py b/azure/cosmos/agent_memory/services/pipeline.py index e340c30..3ca1c96 100644 --- a/azure/cosmos/agent_memory/services/pipeline.py +++ b/azure/cosmos/agent_memory/services/pipeline.py @@ -671,7 +671,14 @@ def extract_memories_dry( seed = _ID_SEED_SEP.join((user_id, thread_id, new_content_hash)) det_id = f"fact_{hashlib.sha256(seed.encode()).hexdigest()[:32]}" topic_tags = build_topic_tags(fact.get("tags", [])) - fact_source = fact.get("source") if fact.get("source") in ("user", "agent") else "user" + raw_source = fact.get("source") + if raw_source is not None and raw_source not in ("user", "agent"): + logger.debug( + "extract_memories: coercing invalid fact source=%r to 'user' user_id=%s", + raw_source, + user_id, + ) + fact_source = raw_source if raw_source in ("user", "agent") else "user" source_tags = ["sys:agent-fact"] if fact_source == "agent" else [] confidence = fact.get("confidence") doc: dict[str, Any] = { @@ -979,10 +986,28 @@ def _apply_inplace_update(self, neighbor: dict[str, Any], new_doc: dict[str, Any Recency wins: the neighbor keeps its id / created_at / partition but takes the new doc's content + embedding, unions tags, and takes the max salience/confidence. + + Folds only happen within the same ``metadata.source`` (user vs agent): an + agent action and a user statement are distinct records even when their + embeddings are near-identical, so folding across sources would corrupt + attribution (tag/source desync, content swap). Cross-source pairs return + False and the caller keeps the new doc as a novel ADD. """ from azure.core import MatchConditions from azure.cosmos.exceptions import CosmosAccessConditionFailedError + neighbor_source = (neighbor.get("metadata") or {}).get("source") or "user" + new_source = (new_doc.get("metadata") or {}).get("source") or "user" + if neighbor_source != new_source: + logger.info( + "in-place dedup update skipped (source mismatch neighbor=%s new=%s) " + "target_id=%s; keeping new doc as novel", + neighbor_source, + new_source, + neighbor.get("id"), + ) + return False + try: old_etag = neighbor.get("_etag") updated = dict(neighbor) @@ -1819,11 +1844,16 @@ def _active_memories_for_reconcile(self, user_id: str, memory_type: str, n: int) # physical partitions and matches the dedup prompt's tiebreaker # ("more recent created_at first"). Cosmos's _ts is the last-write # timestamp, which would diverge from created_at after any UPDATE. + # + # Agent-sourced facts (sys:agent-fact) are excluded: they record what the + # agent did/recommended (historical events), not mutable user state, so + # they must never be contradiction-superseded by a later user statement. query = ( f"SELECT TOP {top_literal(n, name='reconcile_memories.n')} * FROM c " "WHERE c.user_id = @user_id " "AND c.type = @memory_type " f"AND {_ACTIVE_DOC_FILTER} " + "AND NOT ARRAY_CONTAINS(c.tags, 'sys:agent-fact') " "ORDER BY c.created_at DESC" ) return list( diff --git a/azure/cosmos/agent_memory/store/memory_store.py b/azure/cosmos/agent_memory/store/memory_store.py index 2ccd8e9..cbcf9b8 100644 --- a/azure/cosmos/agent_memory/store/memory_store.py +++ b/azure/cosmos/agent_memory/store/memory_store.py @@ -18,6 +18,7 @@ compute_content_hash, extract_keywords, new_id, + normalize_created_at_iso, ) from azure.cosmos.agent_memory.exceptions import ( ConfigurationError, @@ -221,7 +222,7 @@ def add( if salience is not None: kwargs["salience"] = salience if created_at is not None: - kwargs["created_at"] = created_at.isoformat() if isinstance(created_at, datetime) else created_at + kwargs["created_at"] = normalize_created_at_iso(created_at) if memory_type != "turn": kwargs.setdefault("content_hash", compute_content_hash(content)) kwargs.setdefault("prompt_id", "manual:add") diff --git a/tests/unit/aio/services/test_dedup_vector_async.py b/tests/unit/aio/services/test_dedup_vector_async.py index 36ee349..122126c 100644 --- a/tests/unit/aio/services/test_dedup_vector_async.py +++ b/tests/unit/aio/services/test_dedup_vector_async.py @@ -261,6 +261,26 @@ async def test_apply_inplace_update_etag_conflict_returns_false(): assert await p._apply_inplace_update(neighbor, new_doc) is False +@pytest.mark.asyncio +async def test_apply_inplace_update_skips_cross_source_fold(): + p = _service() + p._replace_item = AsyncMock() + neighbor = _fact( + "existing-1", "same content", tags=["sys:fact"], metadata={"category": "preference", "source": "user"} + ) + neighbor["_etag"] = "etag-xyz" + new_doc = _fact( + "f-new", + "same content", + embedding=[0.5, 0.5], + tags=["sys:fact", "sys:agent-fact"], + metadata={"category": "other", "source": "agent"}, + ) + + assert await p._apply_inplace_update(neighbor, new_doc) is False + p._replace_item.assert_not_awaited() + + @pytest.mark.asyncio async def test_nearest_active_full_returns_full_doc_and_skips_excluded(): p = _service() diff --git a/tests/unit/services/test_dedup_vector.py b/tests/unit/services/test_dedup_vector.py index 3b8de6c..02beaff 100644 --- a/tests/unit/services/test_dedup_vector.py +++ b/tests/unit/services/test_dedup_vector.py @@ -314,6 +314,28 @@ def test_apply_inplace_update_shorter_restatement_keeps_richer_content() -> None assert p._apply_inplace_update(neighbor, new_doc) is False +def test_apply_inplace_update_skips_cross_source_fold() -> None: + p = _make_pipeline() + neighbor = _doc( + "existing-1", + "same content", + tags=["sys:fact"], + metadata={"category": "preference", "source": "user"}, + ) + neighbor["_etag"] = "etag-xyz" + new_doc = _doc( + "f-new", + "same content", + embedding=[0.5, 0.5], + tags=["sys:fact", "sys:agent-fact"], + metadata={"category": "other", "source": "agent"}, + ) + + assert p._apply_inplace_update(neighbor, new_doc) is False + p._memories_container.replace_item.assert_not_called() + p._memories_container.upsert_item.assert_not_called() + + def test_nearest_active_full_returns_full_doc_and_skips_excluded() -> None: p = _make_pipeline() doc_a = _doc("a", "first") diff --git a/tests/unit/test_reconcile.py b/tests/unit/test_reconcile.py index 81c563c..cd94c3a 100644 --- a/tests/unit/test_reconcile.py +++ b/tests/unit/test_reconcile.py @@ -598,6 +598,15 @@ def test_query_uses_is_null(self): assert "IS_NULL(c.superseded_by)" in sql assert "c.superseded_by = null" not in sql + def test_reconcile_pool_excludes_agent_facts(self): + p = _make_pipeline() + p._container.query_items.return_value = iter([]) + p._run_prompty = MagicMock(return_value=json.dumps({})) + p.reconcile_memories("u1") + call = p._container.query_items.call_args + sql = (call.kwargs.get("query") or call.args[0]) if call else "" + assert "NOT ARRAY_CONTAINS(c.tags, 'sys:agent-fact')" in sql + class TestFactsTextHandlesNullConfidence: """Pool facts with ``confidence=None`` / ``salience=None`` (legacy diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 6c29d95..4295404 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -448,3 +448,47 @@ def test_build_cosmos_user_agent_suffixes_toolkit_behind_custom(): assert result == f"MyApp/1.2.3 {COSMOS_USER_AGENT}" assert result.startswith("MyApp/1.2.3 ") assert result.endswith(COSMOS_USER_AGENT) + + +class TestNormalizeCreatedAtIso: + def test_none_passthrough(self): + from azure.cosmos.agent_memory._utils import normalize_created_at_iso + + assert normalize_created_at_iso(None) is None + + def test_naive_datetime_assumed_utc(self): + from datetime import datetime + + from azure.cosmos.agent_memory._utils import normalize_created_at_iso + + out = normalize_created_at_iso(datetime(2024, 3, 1, 12, 0, 0)) + assert out == "2024-03-01T12:00:00+00:00" + + def test_aware_datetime_converted_to_utc(self): + from datetime import datetime, timedelta, timezone + + from azure.cosmos.agent_memory._utils import normalize_created_at_iso + + # 09:00 at -03:00 == 12:00 UTC + dt = datetime(2024, 3, 1, 9, 0, 0, tzinfo=timezone(timedelta(hours=-3))) + assert normalize_created_at_iso(dt) == "2024-03-01T12:00:00+00:00" + + def test_offsetless_string_assumed_utc(self): + from azure.cosmos.agent_memory._utils import normalize_created_at_iso + + assert normalize_created_at_iso("2024-03-01T12:00:00") == "2024-03-01T12:00:00+00:00" + + def test_z_suffix_and_offset_normalized_to_utc(self): + from azure.cosmos.agent_memory._utils import normalize_created_at_iso + + assert normalize_created_at_iso("2024-03-01T12:00:00Z") == "2024-03-01T12:00:00+00:00" + assert normalize_created_at_iso("2024-03-01T09:00:00-03:00") == "2024-03-01T12:00:00+00:00" + + def test_invalid_string_raises(self): + from azure.cosmos.agent_memory._utils import normalize_created_at_iso + from azure.cosmos.agent_memory.exceptions import ValidationError + + with pytest.raises(ValidationError): + normalize_created_at_iso("not-a-date") + with pytest.raises(ValidationError): + normalize_created_at_iso("") From 2391e0c5c4af395423049d9a1ec41ef3823cb7c8 Mon Sep 17 00:00:00 2001 From: Aayush Kataria Date: Tue, 14 Jul 2026 11:18:45 -0700 Subject: [PATCH 3/3] Resolving comments --- .../prompts/extract_memories.prompty | 2 +- .../services/_pipeline_helpers.py | 31 +++++++++++++++- tests/unit/services/test_extract_dry.py | 22 ++++++++++++ .../unit/services/test_transcript_metadata.py | 36 +++++++++++++++++-- 4 files changed, 86 insertions(+), 5 deletions(-) diff --git a/azure/cosmos/agent_memory/prompts/extract_memories.prompty b/azure/cosmos/agent_memory/prompts/extract_memories.prompty index b4f2cc0..812acd7 100644 --- a/azure/cosmos/agent_memory/prompts/extract_memories.prompty +++ b/azure/cosmos/agent_memory/prompts/extract_memories.prompty @@ -33,7 +33,7 @@ Every memory must be explicitly grounded in the provided content - never inferre ## 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 speakers 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"`). +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"`. diff --git a/azure/cosmos/agent_memory/services/_pipeline_helpers.py b/azure/cosmos/agent_memory/services/_pipeline_helpers.py index 0a88e05..3f9953a 100644 --- a/azure/cosmos/agent_memory/services/_pipeline_helpers.py +++ b/azure/cosmos/agent_memory/services/_pipeline_helpers.py @@ -344,7 +344,7 @@ def build_transcript( keys = _normalize_metadata_keys(metadata_keys) def _line(m: dict[str, Any]) -> str: - role = m.get("role", "unknown") + role = _canonical_speaker(m.get("role", "unknown")) content = m.get("content", "") meta_str = _format_metadata_segment(m.get("metadata", {}), keys) created_at = m.get("created_at") if include_timestamp else None @@ -367,6 +367,35 @@ def _line(m: dict[str, Any]) -> str: return "\n".join(parts) +# Map common role synonyms onto the toolkit's canonical speaker labels. Callers +# may write turns with any role string (e.g. OpenAI's ``assistant``); the +# extraction prompt reasons about ``user`` vs ``agent``, so synonyms are folded +# onto those. ``tool`` and ``system`` are distinct roles and kept as-is; +# unrecognized labels pass through unchanged rather than being misattributed. +_SPEAKER_ALIASES = { + "user": "user", + "human": "user", + "customer": "user", + "end_user": "user", + "person": "user", + "agent": "agent", + "assistant": "agent", + "ai": "agent", + "bot": "agent", + "chatbot": "agent", + "model": "agent", + "copilot": "agent", + "tool": "tool", + "system": "system", +} + + +def _canonical_speaker(role: Any) -> str: + """Normalize a free-form role to the canonical speaker label for prompts.""" + normalized = str(role or "").strip().lower() + 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. diff --git a/tests/unit/services/test_extract_dry.py b/tests/unit/services/test_extract_dry.py index 610b9b5..2ad225a 100644 --- a/tests/unit/services/test_extract_dry.py +++ b/tests/unit/services/test_extract_dry.py @@ -524,3 +524,25 @@ def test_extraction_transcript_includes_turn_timestamps() -> None: prompt_text = json.dumps(chat.messages) assert "2025-01-01T00:01:00+00:00 | user" in prompt_text + + +def test_extraction_transcript_canonicalizes_speaker_role() -> None: + # A turn written with the OpenAI-style role="assistant" must reach the + # extraction prompt as the "agent" speaker, matching the prompt's rules. + chat = _SyncChat([_response()]) + memories_store = _Store([]) + assistant_turn = dict(_turn(1)) + assistant_turn["role"] = "assistant" + turns_store = _Store([assistant_turn]) + service = PipelineService( + memories_store, + chat, + _SyncEmbeddings(), + containers=_containers_for_store(memories_store, turns_store=turns_store), + ) + + service.extract_memories_dry("u1", "t1") + + prompt_text = json.dumps(chat.messages) + assert "| agent]" in prompt_text + assert "| assistant]" not in prompt_text diff --git a/tests/unit/services/test_transcript_metadata.py b/tests/unit/services/test_transcript_metadata.py index 8dec24d..a480457 100644 --- a/tests/unit/services/test_transcript_metadata.py +++ b/tests/unit/services/test_transcript_metadata.py @@ -29,7 +29,7 @@ def test_metadata_omitted_by_default_flat(self) -> None: _turn("assistant", "Hi", model_id="gpt-4o", token_count=42), ] out = build_transcript(items) - assert out == "[user]: Hello\n[assistant]: Hi" + assert out == "[user]: Hello\n[agent]: Hi" assert "metadata" not in out assert "raw_response" not in out assert "tool_calls" not in out @@ -44,7 +44,7 @@ def test_metadata_omitted_by_default_grouped(self) -> None: assert "metadata" not in out assert "=== Thread t1 ===" in out assert "[user]: Q" in out - assert "[assistant]: A" in out + assert "[agent]: A" in out def test_empty_allowlist_is_same_as_default(self) -> None: items = [_turn("user", "x", a=1, b=2)] @@ -339,7 +339,7 @@ def test_timestamp_prefix_flat(self) -> None: {"role": "assistant", "content": "Hi", "created_at": "2024-06-20T10:01:00+00:00"}, ] out = build_transcript(items, include_timestamp=True) - assert out == ("[2024-06-20T10:00:00+00:00 | user]: Hello\n[2024-06-20T10:01:00+00:00 | assistant]: Hi") + assert out == ("[2024-06-20T10:00:00+00:00 | user]: Hello\n[2024-06-20T10:01:00+00:00 | agent]: Hi") def test_timestamp_absent_falls_back_to_plain_role(self) -> None: items = [{"role": "user", "content": "Hello"}] # no created_at @@ -358,3 +358,33 @@ def test_timestamp_prefix_grouped(self) -> None: out = build_transcript(items, group_by_thread=True, include_timestamp=True) assert "[2024-06-20T10:00:00+00:00 | user]: Q" in out assert "=== Thread t1 ===" in out + + +class TestCanonicalSpeaker: + """Roles are normalized to canonical speaker labels unconditionally, so the + extraction prompt sees a stable vocabulary regardless of the caller's role + string. Synonyms fold onto user/agent; ``tool`` and ``system`` are kept.""" + + def test_agent_synonym_normalized(self) -> None: + # OpenAI-style "assistant" must render as the canonical "agent". + assert build_transcript([{"role": "assistant", "content": "Hi"}]) == "[agent]: Hi" + + def test_user_synonym_normalized(self) -> None: + assert build_transcript([{"role": "human", "content": "a"}]) == "[user]: a" + + def test_tool_and_system_preserved(self) -> None: + items = [ + {"role": "tool", "content": "b"}, + {"role": "system", "content": "c"}, + {"role": "agent", "content": "d"}, + ] + out = build_transcript(items) + assert out == "[tool]: b\n[system]: c\n[agent]: d" + + def test_unknown_role_passed_through(self) -> None: + assert build_transcript([{"role": "narrator", "content": "x"}]) == "[narrator]: x" + + def test_normalized_with_timestamp(self) -> None: + items = [{"role": "assistant", "content": "Hi", "created_at": "2024-06-20T10:01:00+00:00"}] + out = build_transcript(items, include_timestamp=True) + assert out == "[2024-06-20T10:01:00+00:00 | agent]: Hi"