Skip to content

Adding support for Temporal Grounding and searchable summaries #34

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
aayush3011 merged 3 commits into from
Jul 23, 2026
Merged

Adding support for Temporal Grounding and searchable summaries #34

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
bdb1ff1
Improving extraction prompt
aayush3011 Jul 23, 2026
7187220
Updating changelog
aayush3011 Jul 23, 2026
048c7b3
Resolving comments
aayush3011 Jul 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 5 additions & 3 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
6 changes: 3 additions & 3 deletions azure/cosmos/agent_memory/_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
135 changes: 85 additions & 50 deletions azure/cosmos/agent_memory/aio/cosmos_memory_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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(
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down
45 changes: 45 additions & 0 deletions azure/cosmos/agent_memory/aio/store/memory_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Loading
Loading