diff --git a/CHANGELOG.md b/CHANGELOG.md index d90636b..915b979 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,20 +3,22 @@ ## [0.3.0b1] (Unreleased) #### 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) +* Searchable summaries: the summaries container now carries a vector index, and `search_summaries()` / `search_cosmos(include_summaries=True)` vector-search across a user's thread + user summaries. See [PR:#34](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/34) * 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) +* 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) #### 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) #### Other Changes +* Extraction prompt v3: named attribution (use real names, not a generic "the user", so multi-speaker turns don't blur), stronger temporal grounding (resolve relative dates against the turn's time, never the current clock), and meaning/completeness rules. See [PR:#34](https://github.com/AzureCosmosDB/AgentMemoryToolkit/pull/34) +* 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) * 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/azure/cosmos/agent_memory/_utils.py b/azure/cosmos/agent_memory/_utils.py index 8373e18..e3cafae 100644 --- a/azure/cosmos/agent_memory/_utils.py +++ b/azure/cosmos/agent_memory/_utils.py @@ -118,10 +118,10 @@ def normalize_created_at_iso(value: Optional[str | datetime]) -> Optional[str]: text = str(value).strip() if not text: raise ValidationError("created_at must be a non-empty ISO-8601 string or datetime") + if text.endswith("Z"): + text = text[:-1] + "+00:00" 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")) + dt = datetime.fromisoformat(text) 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: diff --git a/azure/cosmos/agent_memory/aio/cosmos_memory_client.py b/azure/cosmos/agent_memory/aio/cosmos_memory_client.py index b058426..089c6a9 100644 --- a/azure/cosmos/agent_memory/aio/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/aio/cosmos_memory_client.py @@ -40,25 +40,6 @@ logger = get_logger(__name__) -_SUMMARIES_INDEXING_POLICY = { - "indexingMode": "consistent", - "automatic": True, - "includedPaths": [{"path": "/*"}], - "excludedPaths": [ - {"path": "/embedding/?"}, - {"path": "/source_memory_ids/*"}, - {"path": "/supersedes_ids/*"}, - {"path": '/"_etag"/?'}, - ], - "compositeIndexes": [ - [ - {"path": "/user_id", "order": "ascending"}, - {"path": "/thread_id", "order": "ascending"}, - {"path": "/version", "order": "descending"}, - ] - ], -} - def _log_auto_trigger_task_failure(task: "asyncio.Task[Any]") -> None: if not task.cancelled() and (exc := task.exception()) is not None: @@ -345,15 +326,6 @@ async def create_memory_store( vector_index_type=_resolve_vector_index_type(vector_index_type), ) vec_policy, idx_policy, ft_policy = _container_policies(**_policy_kwargs) - # Turns always carry the vector index (primed for search) but skip the - # salience composite index, which only procedural synthesis needs. The - # turns vector index is pinned to quantizedFlat so search_turns() works - # on accounts without the DiskANN capability, independent of the - # memories container's configured vector_index_type. - turns_vec_policy, turns_idx_policy, turns_ft_policy = _container_policies( - **{**_policy_kwargs, "vector_index_type": "quantizedFlat"}, - include_salience_composite=False, - ) self._memories_container_client = await db.create_container_if_not_exists( **_build_container_kwargs( container_id=self._cosmos_container, @@ -365,6 +337,11 @@ async def create_memory_store( full_text_policy=ft_policy, ) ) + + turns_vec_policy, turns_idx_policy, turns_ft_policy = _container_policies( + **{**_policy_kwargs, "vector_index_type": "quantizedFlat"}, + include_salience_composite=False, + ) self._turns_container_client = await db.create_container_if_not_exists( **_build_container_kwargs( container_id=self._cosmos_turns_container, @@ -377,13 +354,26 @@ async def create_memory_store( ) ) logger.info("Created turns container: %s/%s", self._cosmos_database, self._cosmos_turns_container) + summaries_vec_policy, summaries_idx_policy, summaries_ft_policy = _container_policies( + **{**_policy_kwargs, "vector_index_type": "quantizedFlat"}, + include_salience_composite=False, + ) + summaries_idx_policy["compositeIndexes"] = [ + [ + {"path": "/user_id", "order": "ascending"}, + {"path": "/thread_id", "order": "ascending"}, + {"path": "/version", "order": "descending"}, + ] + ] self._summaries_container_client = await db.create_container_if_not_exists( **_build_container_kwargs( container_id=self._cosmos_summaries_container, partition_key=partition_key, offer_throughput=offer, default_ttl=-1, - indexing_policy=_SUMMARIES_INDEXING_POLICY, + indexing_policy=summaries_idx_policy, + vector_embedding_policy=summaries_vec_policy, + full_text_policy=summaries_ft_policy, ) ) logger.info( @@ -721,8 +711,11 @@ async def search_cosmos( created_before: Optional[str | datetime] = None, include_turns: bool = False, turn_top_k: Optional[int] = None, + include_summaries: bool = False, + summary_top_k: Optional[int] = None, ) -> list[dict[str, Any]]: - """Search memories using vector similarity; optionally blend raw turns.""" + """Search memories using vector similarity, with optional summary / raw-turn + blending. See the sync client for details.""" store = self._get_store() results = await store.search( search_terms=search_terms, @@ -741,31 +734,73 @@ async def search_cosmos( created_after=created_after, created_before=created_before, ) - if not include_turns or not user_id: + if not user_id: return results seen_content = {str(r.get("content") or "").strip() for r in results} - try: - turns = await 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) + if include_summaries: + try: + summaries = await store.search_summaries( + search_terms=search_terms, + user_id=user_id, + thread_id=thread_id, + top_k=summary_top_k if summary_top_k is not None else top_k, + exclude_tags=exclude_tags, + ) + except Exception as exc: # noqa: BLE001 + logger.warning("search_cosmos: include_summaries search failed (%s); skipping summaries", exc) + summaries = [] + # Order is facts -> summaries -> raw turns: append summaries after the + # relevance-ranked memory hits (and before turns). + for s in summaries: + content = str(s.get("content") or "").strip() + if content and content not in seen_content: + seen_content.add(content) + results.append(s) + + if include_turns: + try: + turns = await 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) + turns = [] + for turn in turns: + content = str(turn.get("content") or "").strip() + if content and content not in seen_content: + seen_content.add(content) + results.append(turn) return results + async def search_summaries( + self, + search_terms: str, + user_id: str, + thread_id: Optional[str] = None, + top_k: int = 5, + tags_all: Optional[list[str]] = None, + tags_any: Optional[list[str]] = None, + exclude_tags: Optional[list[str]] = None, + ) -> list[dict[str, Any]]: + """Vector-search the user's thread + user summaries (see ``AsyncMemoryStore.search_summaries``).""" + return await self._get_store().search_summaries( + search_terms=search_terms, + user_id=user_id, + thread_id=thread_id, + top_k=top_k, + tags_all=tags_all, + tags_any=tags_any, + exclude_tags=exclude_tags, + ) + async def search_turns( self, search_terms: str, diff --git a/azure/cosmos/agent_memory/aio/store/memory_store.py b/azure/cosmos/agent_memory/aio/store/memory_store.py index 1b21987..fce795b 100644 --- a/azure/cosmos/agent_memory/aio/store/memory_store.py +++ b/azure/cosmos/agent_memory/aio/store/memory_store.py @@ -989,6 +989,51 @@ async def search_turns( partition_key=partition_key, ) + async def search_summaries( + self, + search_terms: Optional[str] = None, + user_id: Optional[str] = None, + thread_id: Optional[str] = None, + top_k: int = 5, + tags_all: Optional[list[str]] = None, + tags_any: Optional[list[str]] = None, + exclude_tags: Optional[list[str]] = None, + *, + query: Optional[str] = None, + ) -> list[dict[str, Any]]: + """Vector-search the summaries container (thread_summary + user_summary). + + Searches across *all* of the user's summary docs at once - the single + user_summary and every per-thread thread_summary - ranked by relevance. + ``user_id`` is required. Pass ``thread_id`` to narrow to one thread. + """ + if not user_id: + raise ValidationError("user_id is required for search_summaries") + terms = require_search_terms(search_terms, query) + top = top_literal(top_k, name="top_k") + query_vector = await self._embed(terms) + keywords = extract_keywords(terms) + + qb = _QueryBuilder() + qb.add_filter("c.user_id", "@user_id", user_id) + qb.add_filter("c.thread_id", "@thread_id", thread_id) + add_tag_filters(qb, tags_all=tags_all, tags_any=tags_any, exclude_tags=exclude_tags) + + sql = build_search_sql(qb=qb, top=top, keyword_count=len(keywords), include_superseded=False) + parameters = qb.get_parameters() + parameters.append({"name": "@embedding", "value": query_vector}) + for i, kw in enumerate(keywords): + parameters.append({"name": f"@kw{i}", "value": kw}) + + partition_key, _ = query_scope(user_id, thread_id) + logger.debug("AsyncMemoryStore.search_summaries query: %s", sql) + return await self.query( + sql, + parameters, + container_key=ContainerKey.SUMMARIES, + partition_key=partition_key, + ) + async def search_episodic( self, user_id: str, diff --git a/azure/cosmos/agent_memory/cosmos_memory_client.py b/azure/cosmos/agent_memory/cosmos_memory_client.py index aaf4859..c1ea414 100644 --- a/azure/cosmos/agent_memory/cosmos_memory_client.py +++ b/azure/cosmos/agent_memory/cosmos_memory_client.py @@ -37,25 +37,6 @@ logger = get_logger(__name__) -_SUMMARIES_INDEXING_POLICY = { - "indexingMode": "consistent", - "automatic": True, - "includedPaths": [{"path": "/*"}], - "excludedPaths": [ - {"path": "/embedding/?"}, - {"path": "/source_memory_ids/*"}, - {"path": "/supersedes_ids/*"}, - {"path": '/"_etag"/?'}, - ], - "compositeIndexes": [ - [ - {"path": "/user_id", "order": "ascending"}, - {"path": "/thread_id", "order": "ascending"}, - {"path": "/version", "order": "descending"}, - ] - ], -} - class CosmosMemoryClient(_BaseMemoryClient): """Manages local and Cosmos-backed memories via store and service layers.""" @@ -342,13 +323,29 @@ def create_memory_store( ) ) logger.info("Created turns container: %s/%s", self._cosmos_database, self._cosmos_turns_container) + # Summaries carry the vector + full-text index too, so summaries are + # semantically searchable (search_summaries). Keep the point-lookup + # composite (user_id, thread_id, version) get_*_summary relies on. + summaries_vec_policy, summaries_idx_policy, summaries_ft_policy = _container_policies( + **{**_policy_kwargs, "vector_index_type": "quantizedFlat"}, + include_salience_composite=False, + ) + summaries_idx_policy["compositeIndexes"] = [ + [ + {"path": "/user_id", "order": "ascending"}, + {"path": "/thread_id", "order": "ascending"}, + {"path": "/version", "order": "descending"}, + ] + ] self._summaries_container_client = db.create_container_if_not_exists( **_build_container_kwargs( container_id=self._cosmos_summaries_container, partition_key=partition_key, offer_throughput=offer, default_ttl=-1, - indexing_policy=_SUMMARIES_INDEXING_POLICY, + indexing_policy=summaries_idx_policy, + vector_embedding_policy=summaries_vec_policy, + full_text_policy=summaries_ft_policy, ) ) logger.info( @@ -676,8 +673,14 @@ def search_cosmos( created_before: Optional[str | datetime] = None, include_turns: bool = False, turn_top_k: Optional[int] = None, + include_summaries: bool = False, + summary_top_k: Optional[int] = None, ) -> list[dict[str, Any]]: - """Search memories using vector similarity; optionally blend raw turns.""" + """Search memories using vector similarity, with optional summary / raw-turn blending. + + ``include_summaries`` / ``include_turns`` prepend matching summaries / + append matching raw turns (best-effort); see Docs/concepts.md. + """ store = self._get_store() results = store.search( search_terms=search_terms, @@ -696,31 +699,75 @@ def search_cosmos( created_after=created_after, created_before=created_before, ) - if not include_turns or not user_id: + + if not user_id: return results seen_content = {str(r.get("content") or "").strip() for r in results} - try: - turns = 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) + if include_summaries: + try: + summaries = store.search_summaries( + search_terms=search_terms, + user_id=user_id, + thread_id=thread_id, + top_k=summary_top_k if summary_top_k is not None else top_k, + exclude_tags=exclude_tags, + ) + except Exception as exc: # noqa: BLE001 + logger.warning("search_cosmos: include_summaries search failed (%s); skipping summaries", exc) + summaries = [] + # Order is facts -> summaries -> raw turns: append summaries after the + # relevance-ranked memory hits (and before turns) so they neither jump + # the ranking nor evict facts under context-window truncation. + for s in summaries: + content = str(s.get("content") or "").strip() + if content and content not in seen_content: + seen_content.add(content) + results.append(s) + + if include_turns: + try: + turns = 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) + turns = [] + for turn in turns: + content = str(turn.get("content") or "").strip() + if content and content not in seen_content: + seen_content.add(content) + results.append(turn) return results + def search_summaries( + self, + search_terms: str, + user_id: str, + thread_id: Optional[str] = None, + top_k: int = 5, + tags_all: Optional[list[str]] = None, + tags_any: Optional[list[str]] = None, + exclude_tags: Optional[list[str]] = None, + ) -> list[dict[str, Any]]: + """Vector-search the user's thread + user summaries (see ``MemoryStore.search_summaries``).""" + return self._get_store().search_summaries( + search_terms=search_terms, + user_id=user_id, + thread_id=thread_id, + top_k=top_k, + tags_all=tags_all, + tags_any=tags_any, + exclude_tags=exclude_tags, + ) + def search_turns( self, search_terms: str, diff --git a/azure/cosmos/agent_memory/prompts/extract_memories.prompty b/azure/cosmos/agent_memory/prompts/extract_memories.prompty index 812acd7..d5b7cb4 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: v2 +version: v3 description: Extract facts and episodic memories from a conversation or provided content. model: apiType: chat @@ -85,9 +85,26 @@ stamped `2024-06-20` saying "I flew to Tokyo 3 weeks ago" → `text: "The user f `temporal_context: "3 weeks ago"`. If no line timestamp is present, keep the relative expression as-is in both fields. Never invent a date when there is no anchor. +**Two clocks - never confuse them.** A line's `[ | ...]` prefix is the fact's **event time**: +when it was actually said. The moment you are running now (processing time) is unrelated and may be +months or years later. Resolve every relative expression ("last year", "3 weeks ago", "recently", +"next month") **only** against the event time of the line it appears on - **never** against today / +the current processing date. A fact whose "last year" is computed from today instead of from when it +was said is silently wrong and stays wrong forever. + +**Grounding rules for time:** +- **Relative → absolute, never absolute → vague.** Convert "yesterday", "last week", "3 weeks ago" + into concrete dates using the event time. But never do the reverse: an exact "18 days", "June 3rd", + or "every other Tuesday" must survive verbatim - do not soften it to "some time" or "a while". +- **Preserve exact durations, frequencies, and counts** as stated: "18 days" stays "18 days" (not + "about 3 weeks"), "twice a week" stays "twice a week" (not "regularly"). +- **Keep the relationship, not just the endpoint.** "The deadline moved from March 1 to March 15" + records both the change and both dates - not merely "the deadline is March 15". + ### Fact Formatting Rules - Each fact must be self-contained and intelligible without context - no pronouns like "it" or "they" without antecedents -- Write in third person ("The user...", "The project...") +- **Name people when the conversation names them.** Write in third person (never "I" / "you"), but use a person's actual name once it is given - the user's own name if stated, and always for third parties (family, colleagues, friends: "Elena", "Sara", "the user's manager Raj"). Fall back to "the user" only when no name is available. In multi-speaker transcripts this is essential: attribute each fact to the specific person who stated or is described by it (by name), never a generic "the user" - otherwise two speakers blur into one and later retrieval cannot tell who did what. +- **Preserve the exact meaning - direction and polarity matter.** Read each statement carefully before recording its meaning. "The user hasn't shipped the release yet" ≠ "shipped the release". "The user moved *off* Datadog" ≠ "uses Datadog". "The user used to run marathons" ≠ "runs marathons" (no longer). Getting the direction of a statement backwards is worse than not extracting it. - **Preserve every specific detail stated in the source: exact dates, times, durations, numbers, quantities, amounts, proper nouns, names, brands, product/model names, and named locations or events.** NEVER generalize a specific into a category - keep "June 3rd" (not "a date"), "7 days" (not "about a week"), "Fitbit Versa 3" (not "a smartwatch"), "$4,500" (not "some money"). A fact that drops the specific detail is useless for later recall - the specific IS the memory. - **Group by topic - one contextually rich memory per distinct topic or event, not one per claim.** Capture the whole coherent picture (the core fact plus its immediate context, qualifiers, and every specific) in a single self-contained memory instead of shattering it into fragments. Split into separate memories only when the content spans genuinely distinct topics (e.g. career vs. family vs. a trip) or when one topic is so detail-dense that a single memory would run past ~3 sentences - then split along natural sub-groupings (e.g. one memory per enemy type in a stat block), never by shaving off individual details. Do NOT merge two unrelated topics to save space, and do NOT split one topic just to make memories smaller. - **Capture transitions as one memory - the new state AND what it replaced.** When the user changes, switches, replaces, upgrades, or stops something in favor of something else, the link between old and new is essential context: record both together. "The user switched from almond milk to oat milk after developing an almond sensitivity" - not two disconnected facts. If the change is explicitly temporary or a trial ("for a month", "trying out"), capture that too. Keeping old→new in one memory prevents the two states from later surfacing as contradictory-looking facts. @@ -164,6 +181,8 @@ Do NOT extract: - Raw agent reasoning or intermediate steps that did not produce a confirmed outcome - Memories that are only meaningful within the context of this conversation and have no future reference value +**Extract the content, not the act of sharing.** When the user pastes or references a document, spec, dataset, case, stat block, or any reference material, extract the concrete facts *inside* it - the parties, dates, figures, findings, named items, quantities - each as its own memory. Do NOT produce a memory that merely describes the gesture ("The user shared a contract summary", "The user asked to remember a spec"). A memory of "the user pasted a case" is useless; the case's actual facts are the memory. The same applies to structured data: preserve every count and value ("4 mummies, AC 11, 45 HP"), splitting a detail-dense block into a few focused memories rather than compressing the numbers away. + ### Scoped intents are NOT facts A statement that only applies inside a particular trip, project, event, session, or other bounded context is an **episodic** memory, not a fact. Examples that look like facts but are NOT: - "For this Paris trip, I want luxury accommodations." → episodic (`scope_type=trip`, `scope_value=Paris`) @@ -187,6 +206,13 @@ A fact is a standing claim that holds outside any specific context. The test: if 2. Is the outcome clearly stated, not assumed? 3. Is the lesson genuinely transferable, or too specific to this one situation? +**Before finalizing - scan the whole transcript.** Re-read the entire conversation, including the +**middle and end**, not only the opening. A common failure is capturing the first prominent topic +thoroughly and skimming the rest ("first-topic dominance"). If a single message raised several +distinct topics (e.g. a job change, a health note, and a trip), make sure each distinct topic is +represented. This is about not *missing* a stated fact - it is NOT a directive to inflate the count +or extract filler; the precision and grounding rules above still apply. + --- ## Few-Shot Examples @@ -461,6 +487,83 @@ The only user-stated fact is the request itself (a scoped intent → episodic). } ``` +### Example 9: Reference material - extract the content, not the act of sharing + +**Conversation:** +> User: "Remember this case for me. Bajimaya v Reward Homes [2021] NSWCATAP 297 - construction began in 2014, the contract was signed in 2015 with completion due by October 2015, and the owner received the keys in December 2016 with defects. The tribunal found the builder breached the contract." + +Extract the actual facts *inside* the shared material - dates, parties, findings - each as its own memory. Do NOT record "The user shared a case summary." + +**Output:** +```json +{ + "facts": [ + { + "text": "In Bajimaya v Reward Homes [2021] NSWCATAP 297, construction began in 2014, the contract was signed in 2015, and completion was due by October 2015.", + "category": "other", + "source": "user", + "confidence": 0.95, + "salience": 0.6, + "temporal_context": null, + "tags": ["topic:legal", "topic:construction"] + }, + { + "text": "In Bajimaya v Reward Homes, the owner received the keys in December 2016 and the tribunal found the builder had breached the contract.", + "category": "other", + "source": "user", + "confidence": 0.95, + "salience": 0.6, + "temporal_context": "December 2016", + "tags": ["topic:legal", "topic:construction"] + } + ], + "episodic": [] +} +``` + +### Example 10: Multi-topic message + named people - one memory per distinct topic + +**Conversation:** +> User: "Quick update - my sister Priya just moved to Portland, I got promoted to team lead last week, and my daughter Sara started aerial yoga on Tuesdays." + +Three unrelated topics in one message. Extract each separately, and use the actual names (Priya, Sara) rather than collapsing everyone into "the user". + +**Output:** +```json +{ + "facts": [ + { + "text": "The user's sister Priya recently moved to Portland.", + "category": "biographical", + "source": "user", + "confidence": 0.95, + "salience": 0.6, + "temporal_context": null, + "tags": ["topic:family"] + }, + { + "text": "The user was promoted to team lead last week.", + "category": "biographical", + "source": "user", + "confidence": 0.95, + "salience": 0.8, + "temporal_context": "last week", + "tags": ["topic:career"] + }, + { + "text": "The user's daughter Sara started aerial yoga on Tuesdays.", + "category": "biographical", + "source": "user", + "confidence": 0.95, + "salience": 0.6, + "temporal_context": null, + "tags": ["topic:family"] + } + ], + "episodic": [] +} +``` + --- ## Output Format diff --git a/azure/cosmos/agent_memory/store/memory_store.py b/azure/cosmos/agent_memory/store/memory_store.py index cbcf9b8..7bf2072 100644 --- a/azure/cosmos/agent_memory/store/memory_store.py +++ b/azure/cosmos/agent_memory/store/memory_store.py @@ -1028,6 +1028,54 @@ def search_turns( cross_partition=cross_partition, ) + def search_summaries( + self, + search_terms: Optional[str] = None, + user_id: Optional[str] = None, + thread_id: Optional[str] = None, + top_k: int = 5, + tags_all: Optional[list[str]] = None, + tags_any: Optional[list[str]] = None, + exclude_tags: Optional[list[str]] = None, + *, + query: Optional[str] = None, + ) -> list[dict[str, Any]]: + """Vector-search the summaries container (thread_summary + user_summary). + + Searches across *all* of the user's summary docs at once - the single + user_summary and every per-thread thread_summary - ranked by relevance, + so a holistic ("summarize everything"), a session-scoped, or a + topic-scoped summary question all surface the right summary. ``user_id`` + is required. Pass ``thread_id`` to narrow to one thread's summary. + """ + if not user_id: + raise ValidationError("user_id is required for search_summaries") + terms = require_search_terms(search_terms, query) + top = top_literal(top_k, name="top_k") + query_vector = self._embed(terms) + keywords = extract_keywords(terms) + + qb = _QueryBuilder() + qb.add_filter("c.user_id", "@user_id", user_id) + qb.add_filter("c.thread_id", "@thread_id", thread_id) + add_tag_filters(qb, tags_all=tags_all, tags_any=tags_any, exclude_tags=exclude_tags) + + sql = build_search_sql(qb=qb, top=top, keyword_count=len(keywords), include_superseded=False) + parameters = qb.get_parameters() + parameters.append({"name": "@embedding", "value": query_vector}) + for i, kw in enumerate(keywords): + parameters.append({"name": f"@kw{i}", "value": kw}) + + partition_key, cross_partition = query_scope(user_id, thread_id) + logger.debug("MemoryStore.search_summaries query: %s", sql) + return self.query( + sql, + parameters, + container_key=ContainerKey.SUMMARIES, + partition_key=partition_key, + cross_partition=cross_partition, + ) + def search_episodic( self, user_id: str, diff --git a/tests/unit/aio/test_cosmos_memory_client.py b/tests/unit/aio/test_cosmos_memory_client.py index 4d6adf1..098fcae 100644 --- a/tests/unit/aio/test_cosmos_memory_client.py +++ b/tests/unit/aio/test_cosmos_memory_client.py @@ -375,8 +375,9 @@ async def test_create_memory_store_with_counter_container(self): assert lease_call.kwargs["id"] == "leases" assert lease_call.kwargs["offer_throughput"].auto_scale_max_throughput == 1000 assert summaries_call.kwargs["id"] == "memories_summaries" - assert "vector_embedding_policy" not in summaries_call.kwargs - assert "full_text_policy" not in summaries_call.kwargs + assert "vector_embedding_policy" in summaries_call.kwargs + assert "full_text_policy" in summaries_call.kwargs + assert summaries_call.kwargs["indexing_policy"]["vectorIndexes"][0]["path"] == "/embedding" assert summaries_call.kwargs["indexing_policy"]["compositeIndexes"][0][-1] == { "path": "/version", "order": "descending", @@ -1032,3 +1033,52 @@ async def test_include_turns_failure_returns_memories_only(self): out = await mem.search_cosmos("q", user_id="u1", include_turns=True) assert out == [{"content": "fact A", "type": "fact"}] + + async def test_include_summaries_appends_after_facts_and_dedups(self): + mem, _ = _connected_client() + store = self._stub_store(mem, memories=[{"content": "fact A", "type": "fact"}]) + store.search_summaries = AsyncMock( + return_value=[ + {"content": "session overview", "type": "thread_summary"}, + {"content": "fact A", "type": "user_summary"}, # dup -> skipped + ] + ) + + out = await mem.search_cosmos("q", user_id="u1", include_summaries=True) + + assert out == [ + {"content": "fact A", "type": "fact"}, + {"content": "session overview", "type": "thread_summary"}, + ] + + async def test_facts_summaries_turns_ordering(self): + mem, _ = _connected_client() + store = self._stub_store( + mem, + memories=[{"content": "fact A", "type": "fact"}], + turns=[{"content": "raw turn C", "type": "turn"}], + ) + store.search_summaries = AsyncMock(return_value=[{"content": "summary B", "type": "thread_summary"}]) + + out = await mem.search_cosmos("q", user_id="u1", include_summaries=True, include_turns=True) + + assert [r["content"] for r in out] == ["fact A", "summary B", "raw turn C"] + + async def test_include_summaries_failure_returns_memories_only(self): + mem, _ = _connected_client() + store = self._stub_store(mem, memories=[{"content": "fact A", "type": "fact"}]) + store.search_summaries = AsyncMock(side_effect=RuntimeError("no summaries index")) + + out = await mem.search_cosmos("q", user_id="u1", include_summaries=True) + + assert out == [{"content": "fact A", "type": "fact"}] + + async def test_search_summaries_delegates_to_store(self): + mem, _ = _connected_client() + store = self._stub_store(mem, memories=[]) + store.search_summaries = AsyncMock(return_value=[{"content": "s", "type": "user_summary"}]) + + out = await mem.search_summaries("q", user_id="u1") + + assert out == [{"content": "s", "type": "user_summary"}] + store.search_summaries.assert_awaited_once() diff --git a/tests/unit/services/test_prompty_loader.py b/tests/unit/services/test_prompty_loader.py index 784a601..76c46da 100644 --- a/tests/unit/services/test_prompty_loader.py +++ b/tests/unit/services/test_prompty_loader.py @@ -59,7 +59,7 @@ def test_all_shipped_prompts_declare_version() -> None: # 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", + "extract_memories.prompty": "v3", "dedup.prompty": "v1", "summarize.prompty": "v1", "summarize_update.prompty": "v1", diff --git a/tests/unit/store/test_memory_store.py b/tests/unit/store/test_memory_store.py index f96d504..aaf9ff5 100644 --- a/tests/unit/store/test_memory_store.py +++ b/tests/unit/store/test_memory_store.py @@ -634,6 +634,34 @@ def test_search_turns_queries_turns_container(): assert "VectorDistance(c.embedding, @embedding)" in sql +def test_search_summaries_queries_summaries_container_scoped_to_user(): + summaries = MagicMock() + summaries.query_items.return_value = [] + memories = MagicMock() + memories.query_items.return_value = [] + embeddings = MagicMock() + embeddings.generate.return_value = [0.1, 0.2] + store = MemoryStore( + containers=_containers(summaries=summaries, memories=memories), + embeddings_client=embeddings, + ) + + # No thread_id: fans across the user's summary partitions (user + all threads). + store.search_summaries(search_terms="summarize", user_id="u1") + + summaries.query_items.assert_called_once() + memories.query_items.assert_not_called() + kwargs = summaries.query_items.call_args.kwargs + assert "VectorDistance(c.embedding, @embedding)" in kwargs["query"] + assert {"name": "@user_id", "value": "u1"} in kwargs["parameters"] + + +def test_search_summaries_requires_user_id(): + store = MemoryStore(containers=_containers(), embeddings_client=MagicMock()) + with pytest.raises(ValidationError, match="user_id is required"): + store.search_summaries(search_terms="x") + + def test_search_turns_scopes_to_single_partition_with_thread_id(): turns = MagicMock() turns.query_items.return_value = [] diff --git a/tests/unit/test_cosmos_memory_client.py b/tests/unit/test_cosmos_memory_client.py index 351b543..eade659 100644 --- a/tests/unit/test_cosmos_memory_client.py +++ b/tests/unit/test_cosmos_memory_client.py @@ -473,8 +473,11 @@ def test_create_memory_store_with_custom_dimensions(self): assert lease_call.kwargs["id"] == "leases" assert lease_call.kwargs["offer_throughput"].auto_scale_max_throughput == 1000 assert summaries_call.kwargs["id"] == "memories_summaries" - assert "vector_embedding_policy" not in summaries_call.kwargs - assert "full_text_policy" not in summaries_call.kwargs + # Summaries now carry the vector + full-text index (semantic search over + # summaries), while keeping the (user_id, thread_id, version) composite. + assert "vector_embedding_policy" in summaries_call.kwargs + assert "full_text_policy" in summaries_call.kwargs + assert summaries_call.kwargs["indexing_policy"]["vectorIndexes"][0]["path"] == "/embedding" assert summaries_call.kwargs["indexing_policy"]["compositeIndexes"][0][-1] == { "path": "/version", "order": "descending", @@ -717,6 +720,54 @@ def test_include_turns_without_user_id_skips_turns(self): assert out == [{"content": "fact A", "type": "fact"}] store.search_turns.assert_not_called() + def test_include_summaries_appends_after_facts_and_dedups(self): + mem, _ = _connected_client() + store = self._stub_store(mem, memories=[{"content": "fact A", "type": "fact"}]) + store.search_summaries.return_value = [ + {"content": "session overview", "type": "thread_summary"}, + {"content": "fact A", "type": "user_summary"}, # dup of memory -> skipped + ] + + out = mem.search_cosmos("q", user_id="u1", include_summaries=True) + + # Order is facts -> summaries. + assert out == [ + {"content": "fact A", "type": "fact"}, + {"content": "session overview", "type": "thread_summary"}, + ] + + def test_facts_summaries_turns_ordering(self): + mem, _ = _connected_client() + store = self._stub_store( + mem, + memories=[{"content": "fact A", "type": "fact"}], + turns=[{"content": "raw turn C", "type": "turn"}], + ) + store.search_summaries.return_value = [{"content": "summary B", "type": "thread_summary"}] + + out = mem.search_cosmos("q", user_id="u1", include_summaries=True, include_turns=True) + + assert [r["content"] for r in out] == ["fact A", "summary B", "raw turn C"] + + def test_include_summaries_failure_returns_memories_only(self): + mem, _ = _connected_client() + store = self._stub_store(mem, memories=[{"content": "fact A", "type": "fact"}]) + store.search_summaries.side_effect = RuntimeError("no summaries index") + + out = mem.search_cosmos("q", user_id="u1", include_summaries=True) + + assert out == [{"content": "fact A", "type": "fact"}] + + def test_search_summaries_delegates_to_store(self): + mem, _ = _connected_client() + store = self._stub_store(mem, memories=[]) + store.search_summaries.return_value = [{"content": "s", "type": "user_summary"}] + + out = mem.search_summaries("q", user_id="u1") + + assert out == [{"content": "s", "type": "user_summary"}] + store.search_summaries.assert_called_once() + class TestPushToCosmos: def test_push_to_cosmos(self): diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 4295404..8b462af 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -484,6 +484,12 @@ def test_z_suffix_and_offset_normalized_to_utc(self): 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_date_only_with_and_without_z_coerce_to_utc_midnight(self): + from azure.cosmos.agent_memory._utils import normalize_created_at_iso + + assert normalize_created_at_iso("2024-03-01") == "2024-03-01T00:00:00+00:00" + assert normalize_created_at_iso("2024-03-01Z") == "2024-03-01T00: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