mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(memory): expose memory IDs in auto-tail + memory_list tool + ID-usage guidance
Pre-this-PR the auto-injected memory block rendered rows as `1. <content>`
with no addressable handle. To UPDATE or DELETE a row the model first had
to call memory_search to discover its ID — two round trips, against the
model-as-judge architecture.
This PR adds three tightly-coupled affordances so the model can act on
memory directly:
1. Auto-tail rows now carry the memory ID:
`1. [mem_alpha_001] User prefers Python`
The bracketed token is the canonical ID — same identifier accepted by
memory_update and memory_delete.
2. New `memory_list` tool — chronological browse (vs `memory_search`'s
semantic lookup). Returns recent memories with their IDs. Backend
dispatches to `Backend.list_memories` if available, else falls back
to an empty-query `search_memories`. Caps at 100 entries.
3. ID-usage guidance text appended to the auto-tail block. Tells the
model that bracketed IDs can go straight to memory_update /
memory_delete with no intervening search. The guidance lives in the
user-message tail (never system) — preserves cache-prefix byte
stability (invariant I2).
`memory_update` and `memory_delete` tool descriptions also point at the
[id] block as a valid ID source — keeps tool docs consistent with the
new affordance.
Verification:
- 10/10 tests pass in tests/test_memory_auto_tail.py (incl. 2 new
guidance tests + 2 new ID-format tests)
- 31/31 tests pass in tests/test_memory_handler_native_ops.py (incl. 4
new memory_list dispatch tests + existing assertions updated for the
[id] format change)
- Golden fixtures regenerated for the tool-description copy changes
(tests/fixtures/memory_tool_definitions/{anthropic,openai}.json)
- Live end-to-end test against real Anthropic API
(tests/test_proxy_memory_integration.py::TestMemoryIdAutoTailAndUpdate):
seeded memory → auto-tail → Claude → memory_update with exact ID.
PASSED.
This commit is contained in:
parent
4e1b218544
commit
c62d45eea8
7 changed files with 557 additions and 11 deletions
|
|
@ -182,7 +182,7 @@ The update creates a new version while preserving history, allowing point-in-tim
|
|||
"properties": {
|
||||
"memory_id": {
|
||||
"type": "string",
|
||||
"description": "The unique ID of the memory to update. Obtain this from a memory_search result.",
|
||||
"description": "The unique ID of the memory to update. Take this from the [id] prefix shown in the auto-injected memory block, or from a memory_search / memory_list result.",
|
||||
},
|
||||
"new_content": {
|
||||
"type": "string",
|
||||
|
|
@ -232,7 +232,7 @@ Always provide a reason for deletion to maintain an audit trail.""",
|
|||
"properties": {
|
||||
"memory_id": {
|
||||
"type": "string",
|
||||
"description": "The unique ID of the memory to delete. Obtain this from a memory_search result.",
|
||||
"description": "The unique ID of the memory to delete. Take this from the [id] prefix shown in the auto-injected memory block, or from a memory_search / memory_list result.",
|
||||
},
|
||||
"reason": {
|
||||
"type": "string",
|
||||
|
|
@ -243,6 +243,40 @@ Always provide a reason for deletion to maintain an audit trail.""",
|
|||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "memory_list",
|
||||
"description": """Browse memories without a semantic query — list recent or all memories with their IDs.
|
||||
|
||||
Use this when:
|
||||
- You want to see what's stored without a specific search term
|
||||
- "What do you remember about me / this project?"
|
||||
- "Show me everything you've saved recently"
|
||||
- You need a memory ID for `memory_update` or `memory_delete` but don't have a good search query
|
||||
- You're auditing the memory store (debugging, cleanup, review)
|
||||
|
||||
Differences from `memory_search`:
|
||||
- `memory_search(query)` is SEMANTIC — finds memories similar to a query string
|
||||
- `memory_list()` is CHRONOLOGICAL — returns the most recent memories first
|
||||
- Use `memory_search` when you know what you're looking for; use `memory_list` when you want to browse
|
||||
|
||||
Returns memories in reverse chronological order (newest first). Each entry includes
|
||||
the `memory_id` you'd use to update / delete it.""",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"limit": {
|
||||
"type": "integer",
|
||||
"description": "Maximum number of memories to return (default 10, max 100). Use a smaller number for a quick overview; larger when you need to find a specific memory ID.",
|
||||
"minimum": 1,
|
||||
"maximum": 100,
|
||||
},
|
||||
},
|
||||
"required": [],
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
|
|
@ -388,6 +422,7 @@ MEMORY_TOOLS_OPTIMIZED: list[dict[str, Any]] = [
|
|||
MEMORY_TOOLS[1], # memory_search (unchanged)
|
||||
MEMORY_TOOLS[2], # memory_update (unchanged)
|
||||
MEMORY_TOOLS[3], # memory_delete (unchanged)
|
||||
MEMORY_TOOLS[4], # memory_list (new — chronological browse)
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -70,7 +70,13 @@ class MemoryMode(str, enum.Enum):
|
|||
|
||||
|
||||
# Memory tool names for detection (Headroom's custom tools)
|
||||
MEMORY_TOOL_NAMES = {"memory_save", "memory_search", "memory_update", "memory_delete"}
|
||||
MEMORY_TOOL_NAMES = {
|
||||
"memory_save",
|
||||
"memory_search",
|
||||
"memory_update",
|
||||
"memory_delete",
|
||||
"memory_list",
|
||||
}
|
||||
|
||||
# Anthropic's native memory tool name
|
||||
NATIVE_MEMORY_TOOL_NAME = "memory"
|
||||
|
|
@ -88,6 +94,26 @@ NATIVE_MEMORY_TOOL_TYPE = "memory_20250818"
|
|||
STARTUP_INIT_TIMEOUT_SECONDS = 30.0
|
||||
|
||||
|
||||
def _serialize_created_at(value: Any) -> str | None:
|
||||
"""Best-effort timestamp serialization for tool-result payloads.
|
||||
|
||||
The backend may return ``datetime`` (from a freshly-saved row) or
|
||||
string (from a hydrated SQLite row). Either way the model needs
|
||||
a string to render in chat. Unparseable values → None.
|
||||
"""
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
if hasattr(value, "isoformat"):
|
||||
try:
|
||||
iso = value.isoformat()
|
||||
return iso if isinstance(iso, str) else str(iso)
|
||||
except Exception:
|
||||
return str(value)
|
||||
return str(value)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MemoryConfig:
|
||||
"""Configuration for memory handler.
|
||||
|
|
@ -732,10 +758,17 @@ class MemoryHandler:
|
|||
# post-filter results too).
|
||||
filtered_results = filtered_results[: effective_budget.max_entries]
|
||||
|
||||
# Format as context.
|
||||
# Format as context. Each row prefixes the memory ID so the
|
||||
# model can address it directly (e.g.,
|
||||
# ``memory_update(id, ...)``) without first calling
|
||||
# ``memory_search`` to discover IDs. Pre-this-PR the block
|
||||
# only carried content; the model had to round-trip through
|
||||
# search to do any UPDATE / DELETE on a row visible in the
|
||||
# auto-injected tail.
|
||||
memory_lines = []
|
||||
for i, result in enumerate(filtered_results, 1):
|
||||
memory_lines.append(f"{i}. {result.memory.content}")
|
||||
memory_id = getattr(result.memory, "id", None) or "?"
|
||||
memory_lines.append(f"{i}. [{memory_id}] {result.memory.content}")
|
||||
if hasattr(result, "related_entities") and result.related_entities:
|
||||
entities_str = ", ".join(result.related_entities[:3])
|
||||
memory_lines.append(f" (Related: {entities_str})")
|
||||
|
|
@ -754,7 +787,10 @@ The following information was previously saved in this scope:
|
|||
|
||||
{chr(10).join(memory_lines)}
|
||||
|
||||
Use this context to provide personalized and contextually relevant responses."""
|
||||
Each row begins with an ID in square brackets. To update or delete a row, \
|
||||
pass that ID directly to memory_update or memory_delete — you do not need \
|
||||
to call memory_search first to discover IDs. Use this context to provide \
|
||||
personalized, contextually relevant responses."""
|
||||
|
||||
# Apply the token-budget cap on the formatted block. Pre-this-
|
||||
# PR there was no cap — up to ~4000 tokens could be injected
|
||||
|
|
@ -1010,6 +1046,8 @@ Use this context to provide personalized and contextually relevant responses."""
|
|||
return await self._execute_update(input_data, user_id, provider, request_context)
|
||||
elif tool_name == "memory_delete":
|
||||
return await self._execute_delete(input_data, user_id, request_context)
|
||||
elif tool_name == "memory_list":
|
||||
return await self._execute_list(input_data, user_id, request_context)
|
||||
else:
|
||||
return json.dumps({"error": f"Unknown tool: {tool_name}"})
|
||||
|
||||
|
|
@ -1306,6 +1344,75 @@ Use this context to provide personalized and contextually relevant responses."""
|
|||
}
|
||||
)
|
||||
|
||||
async def _execute_list(
|
||||
self,
|
||||
input_data: dict[str, Any],
|
||||
user_id: str,
|
||||
request_context: RequestContext | None = None,
|
||||
) -> str:
|
||||
"""Execute memory_list tool — chronological browse without semantic query.
|
||||
|
||||
Returns memories in reverse-chronological order (newest first).
|
||||
Different from ``memory_search`` (which needs a semantic query).
|
||||
Use case: the model needs a memory ID for update/delete but
|
||||
doesn't have a good query string to find it.
|
||||
|
||||
Backend dispatch: prefer ``list_memories`` if the backend
|
||||
exposes it; otherwise fall back to an empty-query
|
||||
``search_memories(query="", top_k=limit)`` which most backends
|
||||
treat as "return everything ordered by recency."
|
||||
"""
|
||||
limit = input_data.get("limit", 10)
|
||||
try:
|
||||
limit = max(1, min(100, int(limit)))
|
||||
except (TypeError, ValueError):
|
||||
limit = 10
|
||||
|
||||
await self._ensure_initialized()
|
||||
if not self._backend:
|
||||
return json.dumps({"status": "error", "error": "Memory backend not initialized"})
|
||||
|
||||
backend, _scope, effective_user_id = self._resolve_for_request(user_id, request_context)
|
||||
|
||||
# Prefer a native list_memories if the backend has one (LocalBackend
|
||||
# does); fall back to a recency-keyed search when not available.
|
||||
list_fn = getattr(backend, "list_memories", None)
|
||||
if callable(list_fn):
|
||||
try:
|
||||
results = await list_fn(user_id=effective_user_id, limit=limit)
|
||||
except Exception as e:
|
||||
logger.warning(f"Memory: list_memories failed for user {effective_user_id}: {e}")
|
||||
return json.dumps({"status": "error", "error": str(e)})
|
||||
else:
|
||||
try:
|
||||
results = await backend.search_memories(
|
||||
query="",
|
||||
user_id=effective_user_id,
|
||||
top_k=limit,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Memory: list fallback search failed: {e}")
|
||||
return json.dumps({"status": "error", "error": str(e)})
|
||||
|
||||
entries: list[dict[str, Any]] = []
|
||||
for r in results:
|
||||
mem = getattr(r, "memory", r)
|
||||
entries.append(
|
||||
{
|
||||
"id": getattr(mem, "id", None),
|
||||
"content": getattr(mem, "content", ""),
|
||||
"created_at": _serialize_created_at(getattr(mem, "created_at", None)),
|
||||
}
|
||||
)
|
||||
|
||||
return json.dumps(
|
||||
{
|
||||
"status": "ok",
|
||||
"count": len(entries),
|
||||
"memories": entries,
|
||||
}
|
||||
)
|
||||
|
||||
# =========================================================================
|
||||
# Native Memory Tool (Anthropic's memory_20250818)
|
||||
# =========================================================================
|
||||
|
|
|
|||
|
|
@ -153,7 +153,7 @@
|
|||
"properties": {
|
||||
"memory_id": {
|
||||
"type": "string",
|
||||
"description": "The unique ID of the memory to update. Obtain this from a memory_search result."
|
||||
"description": "The unique ID of the memory to update. Take this from the [id] prefix shown in the auto-injected memory block, or from a memory_search / memory_list result."
|
||||
},
|
||||
"new_content": {
|
||||
"type": "string",
|
||||
|
|
@ -178,7 +178,7 @@
|
|||
"properties": {
|
||||
"memory_id": {
|
||||
"type": "string",
|
||||
"description": "The unique ID of the memory to delete. Obtain this from a memory_search result."
|
||||
"description": "The unique ID of the memory to delete. Take this from the [id] prefix shown in the auto-injected memory block, or from a memory_search / memory_list result."
|
||||
},
|
||||
"reason": {
|
||||
"type": "string",
|
||||
|
|
@ -189,6 +189,22 @@
|
|||
"memory_id"
|
||||
]
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "memory_list",
|
||||
"description": "Browse memories without a semantic query \u2014 list recent or all memories with their IDs.\n\nUse this when:\n- You want to see what's stored without a specific search term\n - \"What do you remember about me / this project?\"\n - \"Show me everything you've saved recently\"\n- You need a memory ID for `memory_update` or `memory_delete` but don't have a good search query\n- You're auditing the memory store (debugging, cleanup, review)\n\nDifferences from `memory_search`:\n- `memory_search(query)` is SEMANTIC \u2014 finds memories similar to a query string\n- `memory_list()` is CHRONOLOGICAL \u2014 returns the most recent memories first\n- Use `memory_search` when you know what you're looking for; use `memory_list` when you want to browse\n\nReturns memories in reverse chronological order (newest first). Each entry includes\nthe `memory_id` you'd use to update / delete it.",
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"limit": {
|
||||
"type": "integer",
|
||||
"description": "Maximum number of memories to return (default 10, max 100). Use a smaller number for a quick overview; larger when you need to find a specific memory ID.",
|
||||
"minimum": 1,
|
||||
"maximum": 100
|
||||
}
|
||||
},
|
||||
"required": []
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
|
|||
|
|
@ -161,7 +161,7 @@
|
|||
"properties": {
|
||||
"memory_id": {
|
||||
"type": "string",
|
||||
"description": "The unique ID of the memory to update. Obtain this from a memory_search result."
|
||||
"description": "The unique ID of the memory to update. Take this from the [id] prefix shown in the auto-injected memory block, or from a memory_search / memory_list result."
|
||||
},
|
||||
"new_content": {
|
||||
"type": "string",
|
||||
|
|
@ -189,7 +189,7 @@
|
|||
"properties": {
|
||||
"memory_id": {
|
||||
"type": "string",
|
||||
"description": "The unique ID of the memory to delete. Obtain this from a memory_search result."
|
||||
"description": "The unique ID of the memory to delete. Take this from the [id] prefix shown in the auto-injected memory block, or from a memory_search / memory_list result."
|
||||
},
|
||||
"reason": {
|
||||
"type": "string",
|
||||
|
|
@ -201,6 +201,25 @@
|
|||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "memory_list",
|
||||
"description": "Browse memories without a semantic query \u2014 list recent or all memories with their IDs.\n\nUse this when:\n- You want to see what's stored without a specific search term\n - \"What do you remember about me / this project?\"\n - \"Show me everything you've saved recently\"\n- You need a memory ID for `memory_update` or `memory_delete` but don't have a good search query\n- You're auditing the memory store (debugging, cleanup, review)\n\nDifferences from `memory_search`:\n- `memory_search(query)` is SEMANTIC \u2014 finds memories similar to a query string\n- `memory_list()` is CHRONOLOGICAL \u2014 returns the most recent memories first\n- Use `memory_search` when you know what you're looking for; use `memory_list` when you want to browse\n\nReturns memories in reverse chronological order (newest first). Each entry includes\nthe `memory_id` you'd use to update / delete it.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"limit": {
|
||||
"type": "integer",
|
||||
"description": "Maximum number of memories to return (default 10, max 100). Use a smaller number for a quick overview; larger when you need to find a specific memory ID.",
|
||||
"minimum": 1,
|
||||
"maximum": 100
|
||||
}
|
||||
},
|
||||
"required": []
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
|
|||
|
|
@ -311,3 +311,132 @@ def test_unknown_provider_raises() -> None:
|
|||
"ctx",
|
||||
provider="bogus", # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Memory IDs in the auto-tail block (new contract for this PR).
|
||||
#
|
||||
# Pre-this-PR the block rendered entries as ``f"{i}. {content}"`` — no ID,
|
||||
# so the model could see "1. fact X" but had no addressable handle on it.
|
||||
# To UPDATE or DELETE that row, the model first had to call
|
||||
# ``memory_search`` to discover its ID. Two round trips for one
|
||||
# operation, against the model-as-judge architecture.
|
||||
#
|
||||
# Post-this-PR the format is ``f"{i}. [{id}] {content}"``. The model
|
||||
# can call ``memory_update('mem_alpha_001', ...)`` directly from a
|
||||
# row it sees in the auto-injected tail.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_auto_tail_block_includes_memory_ids() -> None:
|
||||
"""Each entry in the formatted block carries the memory's ID in
|
||||
square brackets, immediately after the row number. The model uses
|
||||
this to address rows directly (memory_update / memory_delete)
|
||||
without round-tripping through memory_search."""
|
||||
handler = _build_handler()
|
||||
context = asyncio.run(
|
||||
handler.search_and_format_context("alpha", [{"role": "user", "content": "hi"}])
|
||||
)
|
||||
assert context is not None
|
||||
# IDs from the stub backend fixture.
|
||||
assert "[mem_alpha_001]" in context
|
||||
assert "[mem_alpha_002]" in context
|
||||
# Format is row-number then bracketed-id then content.
|
||||
assert "1. [mem_alpha_001] User prefers Python" in context
|
||||
assert "2. [mem_alpha_002] User's timezone" in context
|
||||
|
||||
|
||||
def test_auto_tail_block_id_format_handles_missing_id() -> None:
|
||||
"""Defensive: if the backend returns a memory without an ID (edge
|
||||
case during a migration), the format must not crash. Render with
|
||||
a placeholder so the model sees the row exists but can't address
|
||||
it — calling memory_update("?") will fail cleanly."""
|
||||
|
||||
class _NoIdBackend:
|
||||
async def search_memories(self, **_: Any) -> list[_StubResult]:
|
||||
return [
|
||||
_StubResult(
|
||||
memory=_StubMemory(id=None, content="legacy row", metadata={}), # type: ignore[arg-type]
|
||||
score=0.9,
|
||||
related_entities=[],
|
||||
)
|
||||
]
|
||||
|
||||
config = MemoryConfig(
|
||||
enabled=True,
|
||||
backend="local",
|
||||
inject_context=True,
|
||||
inject_tools=True,
|
||||
top_k=5,
|
||||
min_similarity=0.3,
|
||||
mode=MemoryMode.AUTO_TAIL,
|
||||
)
|
||||
handler = MemoryHandler(config)
|
||||
handler._backend = _NoIdBackend() # type: ignore[assignment]
|
||||
handler._initialized = True
|
||||
|
||||
context = asyncio.run(
|
||||
handler.search_and_format_context("alpha", [{"role": "user", "content": "hi"}])
|
||||
)
|
||||
assert context is not None
|
||||
# Placeholder ID is "?" — no crash; format is preserved.
|
||||
assert "[?]" in context
|
||||
assert "legacy row" in context
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Memory-ID-usage guidance (new contract for this PR).
|
||||
#
|
||||
# Pre-this-PR the auto-tail block closed with a generic line that said
|
||||
# nothing about the [id] prefix. Real Claude could *learn* to use the IDs
|
||||
# when explicitly told in the user prompt (see live integration test in
|
||||
# tests/test_proxy_memory_integration.py), but had no signal in the block
|
||||
# itself that the bracketed token was an addressable handle.
|
||||
#
|
||||
# Post-this-PR the block carries a short guidance line that names the
|
||||
# direct-update / direct-delete affordance. This is the "memory prelude"
|
||||
# referenced in the realignment plan — embedded in the same user-message
|
||||
# tail as the memories themselves, never in system/instructions.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_auto_tail_block_includes_id_usage_guidance() -> None:
|
||||
"""The formatted block tells the model that [id]-prefixed rows can be
|
||||
passed straight to memory_update / memory_delete. Without this the
|
||||
model has to be primed by the user; with it the affordance is
|
||||
self-describing."""
|
||||
handler = _build_handler()
|
||||
context = asyncio.run(
|
||||
handler.search_and_format_context("alpha", [{"role": "user", "content": "hi"}])
|
||||
)
|
||||
assert context is not None
|
||||
# The block names BOTH update and delete so the affordance covers
|
||||
# the two ID-addressable mutations.
|
||||
assert "memory_update" in context
|
||||
assert "memory_delete" in context
|
||||
# And it names the [id] convention so the model maps brackets → IDs.
|
||||
assert "square brackets" in context.lower() or "[id]" in context.lower()
|
||||
|
||||
|
||||
def test_id_usage_guidance_lives_in_user_tail_not_system() -> None:
|
||||
"""Invariant: the guidance text is part of the auto-tail block (which
|
||||
`_append_to_latest_user_tail` writes to the latest user message). It
|
||||
must NEVER be written to the system message — that would invalidate
|
||||
the cache-hot-zone byte-stability invariant (I2)."""
|
||||
handler = _build_handler()
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "tell me something"},
|
||||
]
|
||||
context = asyncio.run(handler.search_and_format_context("alpha", messages))
|
||||
assert context is not None
|
||||
assert "memory_update" in context
|
||||
|
||||
new_messages, _ = MemoryHandler._append_to_latest_user_tail(
|
||||
messages, context, provider="openai"
|
||||
)
|
||||
# System message is byte-stable.
|
||||
assert new_messages[0]["content"] == "You are a helpful assistant."
|
||||
# Guidance only appears in the user tail.
|
||||
assert "memory_update" not in new_messages[0]["content"]
|
||||
assert "memory_update" in new_messages[1]["content"]
|
||||
|
|
|
|||
|
|
@ -962,7 +962,10 @@ async def test_search_and_format_context_and_handle_memory_tool_calls(
|
|||
[{"role": "user", "content": "What food does Alice like?"}],
|
||||
)
|
||||
assert "## Relevant Memories for This User" in context
|
||||
assert "1. Alice likes pizza" in context
|
||||
# Format now includes memory ID in brackets so the model can address
|
||||
# rows directly via memory_update / memory_delete without round-
|
||||
# tripping through memory_search.
|
||||
assert "1. [m1] Alice likes pizza" in context
|
||||
assert "(Related: Alice, pizza)" in context
|
||||
|
||||
backend.raise_on = "search"
|
||||
|
|
@ -1404,3 +1407,103 @@ async def test_extract_tool_calls_and_handle_tool_calls_parse_edges(
|
|||
"openai",
|
||||
)
|
||||
assert results == [{"role": "tool", "tool_call_id": "fc1", "content": "ok:memory_search:{}"}]
|
||||
|
||||
|
||||
# ── memory_list (browse without semantic query) ──────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_list_returns_recent_memories_with_ids(handler: MemoryHandler) -> None:
|
||||
"""memory_list returns memories in reverse-chronological order with
|
||||
IDs, content, and timestamps. Distinct from memory_search — no
|
||||
semantic query required; model can browse to discover IDs."""
|
||||
|
||||
class ListBackend:
|
||||
async def search_memories(self, **kwargs): # noqa: ANN003
|
||||
return []
|
||||
|
||||
async def list_memories(self, *, user_id, limit): # noqa: ANN001, ANN201
|
||||
assert user_id == "u1"
|
||||
assert limit == 5
|
||||
return [
|
||||
make_result("mem_002", "newest fact", created_at="2026-05-19T12:00:00+00:00"),
|
||||
make_result("mem_001", "older fact", created_at="2026-05-18T10:00:00+00:00"),
|
||||
]
|
||||
|
||||
handler._backend = ListBackend() # type: ignore[assignment]
|
||||
handler._initialized = True
|
||||
out = await handler._execute_list({"limit": 5}, "u1")
|
||||
payload = json.loads(out)
|
||||
assert payload["status"] == "ok"
|
||||
assert payload["count"] == 2
|
||||
assert payload["memories"][0]["id"] == "mem_002"
|
||||
assert payload["memories"][0]["content"] == "newest fact"
|
||||
assert payload["memories"][0]["created_at"] == "2026-05-19T12:00:00+00:00"
|
||||
assert payload["memories"][1]["id"] == "mem_001"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_list_falls_back_to_search_when_list_unavailable(
|
||||
handler: MemoryHandler,
|
||||
) -> None:
|
||||
"""Backends without list_memories fall back to an empty-query
|
||||
search — most backends treat that as "return recent." Locks the
|
||||
fallback path so a future backend without list_memories still works."""
|
||||
|
||||
class SearchOnlyBackend:
|
||||
async def search_memories(self, *, query, user_id, top_k, **kwargs): # noqa: ANN001, ANN003
|
||||
assert query == "" # fallback uses empty query
|
||||
assert top_k == 3
|
||||
return [make_result("mem_x", "anything", created_at=None)]
|
||||
|
||||
handler._backend = SearchOnlyBackend() # type: ignore[assignment]
|
||||
handler._initialized = True
|
||||
out = await handler._execute_list({"limit": 3}, "u1")
|
||||
payload = json.loads(out)
|
||||
assert payload["status"] == "ok"
|
||||
assert payload["count"] == 1
|
||||
assert payload["memories"][0]["id"] == "mem_x"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_list_caps_limit_to_1_100(handler: MemoryHandler) -> None:
|
||||
"""Defensive: input limit is clamped to [1, 100]. Protects the
|
||||
backend from a runaway value the model might invent."""
|
||||
|
||||
received_limits: list[int] = []
|
||||
|
||||
class LimitWatcher:
|
||||
async def search_memories(self, **kwargs): # noqa: ANN003
|
||||
return []
|
||||
|
||||
async def list_memories(self, *, user_id, limit): # noqa: ANN001, ANN201
|
||||
received_limits.append(limit)
|
||||
return []
|
||||
|
||||
handler._backend = LimitWatcher() # type: ignore[assignment]
|
||||
handler._initialized = True
|
||||
|
||||
await handler._execute_list({"limit": 99999}, "u1")
|
||||
await handler._execute_list({"limit": 0}, "u1")
|
||||
await handler._execute_list({"limit": "bogus"}, "u1") # type: ignore[dict-item]
|
||||
assert received_limits == [100, 1, 10] # clamped to 100, 1, default 10
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_memory_list_dispatched_via_execute_memory_tool(handler: MemoryHandler) -> None:
|
||||
"""End-to-end: memory_list in MEMORY_TOOL_NAMES means the
|
||||
handler.handle_memory_tool_calls dispatcher routes it correctly."""
|
||||
|
||||
class StubBackend:
|
||||
async def search_memories(self, **kwargs): # noqa: ANN003
|
||||
return []
|
||||
|
||||
async def list_memories(self, *, user_id, limit): # noqa: ANN001, ANN201
|
||||
return [make_result("m1", "fact", created_at=None)]
|
||||
|
||||
handler._backend = StubBackend() # type: ignore[assignment]
|
||||
handler._initialized = True
|
||||
out = await handler._execute_memory_tool("memory_list", {"limit": 5}, "u1", "anthropic")
|
||||
payload = json.loads(out)
|
||||
assert payload["status"] == "ok"
|
||||
assert payload["memories"][0]["id"] == "m1"
|
||||
|
|
|
|||
|
|
@ -376,3 +376,140 @@ class TestMemoryStats:
|
|||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "requests" in data
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def memory_client_global(temp_memory_db):
|
||||
"""Memory-enabled client with GLOBAL storage mode.
|
||||
|
||||
GLOBAL keeps every memory in a single SQLite file regardless of
|
||||
project routing, so tests that pre-seed via direct backend access
|
||||
are guaranteed to share the same DB as the proxy's runtime backend.
|
||||
"""
|
||||
config = ProxyConfig(
|
||||
optimize=False,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
memory_enabled=True,
|
||||
memory_backend="local",
|
||||
memory_db_path=temp_memory_db,
|
||||
memory_inject_tools=True,
|
||||
memory_inject_context=True,
|
||||
memory_top_k=5,
|
||||
memory_storage_mode="global",
|
||||
)
|
||||
app = create_app(config)
|
||||
with TestClient(app) as client:
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.mark.skipif(not os.environ.get("ANTHROPIC_API_KEY"), reason="ANTHROPIC_API_KEY not set")
|
||||
class TestMemoryIdAutoTailAndUpdate:
|
||||
"""End-to-end live: model uses [memory_id] from auto-tail to call memory_update.
|
||||
|
||||
Validates that the IDs we added to the auto-injected memory block
|
||||
(see ``MemoryHandler.search_and_format_context``) are extractable by
|
||||
a real Claude model and can be passed directly to ``memory_update``
|
||||
without an intervening ``memory_search`` round-trip.
|
||||
"""
|
||||
|
||||
def test_model_uses_memory_id_to_call_memory_update(
|
||||
self,
|
||||
memory_client_global,
|
||||
anthropic_api_key,
|
||||
temp_memory_db,
|
||||
):
|
||||
import asyncio
|
||||
|
||||
from headroom.memory.backends.local import LocalBackend, LocalBackendConfig
|
||||
|
||||
user_id = f"test-id-update-{int(time.time())}"
|
||||
|
||||
# Pre-seed a known memory directly via a fresh backend so we
|
||||
# control its content and learn its ID up front. Same db_path
|
||||
# AND same embedder (ONNX) as the proxy backend → both read
|
||||
# the same SQLite file and produce comparable vectors.
|
||||
async def _seed() -> str:
|
||||
backend = LocalBackend(
|
||||
LocalBackendConfig(
|
||||
db_path=temp_memory_db,
|
||||
embedder_backend="onnx",
|
||||
embedder_model="all-MiniLM-L6-v2",
|
||||
vector_dimension=384,
|
||||
)
|
||||
)
|
||||
mem = await backend.save_memory(
|
||||
content="The user's favorite color is blue.",
|
||||
user_id=user_id,
|
||||
)
|
||||
return mem.id
|
||||
|
||||
memory_id = asyncio.run(_seed())
|
||||
assert memory_id, "save_memory should return a usable id"
|
||||
|
||||
# Let the SQLite write + index settle before the proxy reads.
|
||||
time.sleep(0.5)
|
||||
|
||||
# Capture tool calls the proxy executes, so we can assert the
|
||||
# model picked the right tool with the right memory_id.
|
||||
proxy = memory_client_global.app.state.proxy
|
||||
assert proxy.memory_handler is not None
|
||||
recorded: list[dict] = []
|
||||
original_execute = proxy.memory_handler._execute_memory_tool
|
||||
|
||||
async def _capturing_execute(
|
||||
tool_name, input_data, user_id_arg, provider, request_context=None
|
||||
):
|
||||
recorded.append({"tool_name": tool_name, "input": dict(input_data)})
|
||||
return await original_execute(
|
||||
tool_name,
|
||||
input_data,
|
||||
user_id_arg,
|
||||
provider,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
proxy.memory_handler._execute_memory_tool = _capturing_execute # type: ignore[assignment]
|
||||
|
||||
try:
|
||||
response = memory_client_global.post(
|
||||
"/v1/messages",
|
||||
headers={
|
||||
"x-api-key": anthropic_api_key,
|
||||
"anthropic-version": "2023-06-01",
|
||||
"x-headroom-user-id": user_id,
|
||||
},
|
||||
json={
|
||||
"model": "claude-sonnet-4-20250514",
|
||||
"max_tokens": 800,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": (
|
||||
"Quick correction: my favorite color is actually "
|
||||
"green, not blue. Please call the memory_update "
|
||||
"tool to fix the relevant memory in your context. "
|
||||
"The relevant memories block lists each memory's "
|
||||
"ID in square brackets — use that ID for "
|
||||
"memory_id."
|
||||
),
|
||||
}
|
||||
],
|
||||
},
|
||||
)
|
||||
finally:
|
||||
proxy.memory_handler._execute_memory_tool = original_execute # type: ignore[assignment]
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
# The model should have called memory_update at least once.
|
||||
update_calls = [c for c in recorded if c["tool_name"] == "memory_update"]
|
||||
assert update_calls, f"Expected at least one memory_update call. Recorded: {recorded}"
|
||||
|
||||
# And it should reference the exact ID we seeded — i.e. the
|
||||
# model used the [id] from the auto-tail block, not a guess.
|
||||
assert any(c["input"].get("memory_id") == memory_id for c in update_calls), (
|
||||
f"Expected memory_update(memory_id={memory_id!r}); got inputs: "
|
||||
f"{[c['input'] for c in update_calls]}"
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue