fix: return metadata in recall responses (#680)
* fix: return metadata in recall responses (#674) Metadata stored during retain was never retrieved during recall. Add metadata to all SQL SELECT queries, the RetrievalResult dataclass, ScoredResult.to_dict(), and MemoryFact construction in the recall pipeline. * test: add metadata round-trip test for retain→recall Replace placeholder metadata test with one that actually passes metadata via retain_batch_async and asserts it is returned on recall. * fix: parse metadata JSON string from database in MemoryFact asyncpg may return JSONB columns as strings. Add a field_validator to MemoryFact.metadata to handle JSON string deserialization.
This commit is contained in:
parent
f0f0d554f2
commit
0bcbf8491b
7 changed files with 46 additions and 22 deletions
|
|
@ -3234,7 +3234,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||||
source_rows = await sf_conn.fetch(
|
source_rows = await sf_conn.fetch(
|
||||||
f"""
|
f"""
|
||||||
SELECT id, text, fact_type, context, occurred_start, occurred_end,
|
SELECT id, text, fact_type, context, occurred_start, occurred_end,
|
||||||
mentioned_at, document_id, chunk_id, tags
|
mentioned_at, document_id, chunk_id, tags, metadata
|
||||||
FROM {fq_table("memory_units")}
|
FROM {fq_table("memory_units")}
|
||||||
WHERE id = ANY($1::uuid[])
|
WHERE id = ANY($1::uuid[])
|
||||||
""",
|
""",
|
||||||
|
|
@ -3255,6 +3255,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||||
occurred_end=r["occurred_end"].isoformat() if r["occurred_end"] else None,
|
occurred_end=r["occurred_end"].isoformat() if r["occurred_end"] else None,
|
||||||
mentioned_at=r["mentioned_at"].isoformat() if r["mentioned_at"] else None,
|
mentioned_at=r["mentioned_at"].isoformat() if r["mentioned_at"] else None,
|
||||||
document_id=r["document_id"],
|
document_id=r["document_id"],
|
||||||
|
metadata=r["metadata"],
|
||||||
chunk_id=str(r["chunk_id"]) if r["chunk_id"] else None,
|
chunk_id=str(r["chunk_id"]) if r["chunk_id"] else None,
|
||||||
tags=r["tags"] or None,
|
tags=r["tags"] or None,
|
||||||
)
|
)
|
||||||
|
|
@ -3333,6 +3334,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||||
occurred_end=result_dict.get("occurred_end"),
|
occurred_end=result_dict.get("occurred_end"),
|
||||||
mentioned_at=result_dict.get("mentioned_at"),
|
mentioned_at=result_dict.get("mentioned_at"),
|
||||||
document_id=result_dict.get("document_id"),
|
document_id=result_dict.get("document_id"),
|
||||||
|
metadata=result_dict.get("metadata"),
|
||||||
chunk_id=result_dict.get("chunk_id"),
|
chunk_id=result_dict.get("chunk_id"),
|
||||||
tags=result_dict.get("tags"),
|
tags=result_dict.get("tags"),
|
||||||
source_fact_ids=source_fact_ids_by_obs.get(result_id) if include_source_facts else None,
|
source_fact_ids=source_fact_ids_by_obs.get(result_id) if include_source_facts else None,
|
||||||
|
|
|
||||||
|
|
@ -8,7 +8,7 @@ API stability even if internal models change.
|
||||||
|
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from pydantic import BaseModel, ConfigDict, Field
|
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||||||
|
|
||||||
# Valid fact types for recall operations (excludes 'opinion' which is deprecated)
|
# Valid fact types for recall operations (excludes 'opinion' which is deprecated)
|
||||||
VALID_RECALL_FACT_TYPES = frozenset(["world", "experience", "observation"])
|
VALID_RECALL_FACT_TYPES = frozenset(["world", "experience", "observation"])
|
||||||
|
|
@ -159,6 +159,19 @@ class MemoryFact(BaseModel):
|
||||||
mentioned_at: str | None = Field(None, description="ISO format date when the fact was mentioned/learned")
|
mentioned_at: str | None = Field(None, description="ISO format date when the fact was mentioned/learned")
|
||||||
document_id: str | None = Field(None, description="ID of the document this memory belongs to")
|
document_id: str | None = Field(None, description="ID of the document this memory belongs to")
|
||||||
metadata: dict[str, str] | None = Field(None, description="User-defined metadata")
|
metadata: dict[str, str] | None = Field(None, description="User-defined metadata")
|
||||||
|
|
||||||
|
@field_validator("metadata", mode="before")
|
||||||
|
@classmethod
|
||||||
|
def parse_metadata(cls, v: Any) -> dict[str, str] | None:
|
||||||
|
"""Parse metadata from JSON string if needed (asyncpg may return JSONB as str)."""
|
||||||
|
if v is None:
|
||||||
|
return None
|
||||||
|
if isinstance(v, str):
|
||||||
|
import json
|
||||||
|
|
||||||
|
return json.loads(v)
|
||||||
|
return v
|
||||||
|
|
||||||
chunk_id: str | None = Field(
|
chunk_id: str | None = Field(
|
||||||
None, description="ID of the chunk this fact was extracted from (format: bank_id_document_id_chunk_index)"
|
None, description="ID of the chunk this fact was extracted from (format: bank_id_document_id_chunk_index)"
|
||||||
)
|
)
|
||||||
|
|
|
||||||
|
|
@ -231,7 +231,7 @@ class BFSGraphRetriever(GraphRetriever):
|
||||||
f"""
|
f"""
|
||||||
SELECT mu.id, mu.text, mu.context, mu.occurred_start, mu.occurred_end,
|
SELECT mu.id, mu.text, mu.context, mu.occurred_start, mu.occurred_end,
|
||||||
mu.mentioned_at, mu.fact_type,
|
mu.mentioned_at, mu.fact_type,
|
||||||
mu.document_id, mu.chunk_id, mu.tags,
|
mu.document_id, mu.chunk_id, mu.tags, mu.metadata,
|
||||||
ml.weight, ml.link_type, ml.from_unit_id
|
ml.weight, ml.link_type, ml.from_unit_id
|
||||||
FROM {fq_table("memory_links")} ml
|
FROM {fq_table("memory_links")} ml
|
||||||
JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id
|
JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id
|
||||||
|
|
|
||||||
|
|
@ -449,7 +449,7 @@ async def fetch_memory_units_by_ids(
|
||||||
rows = await conn.fetch(
|
rows = await conn.fetch(
|
||||||
f"""
|
f"""
|
||||||
SELECT id, text, context, event_date, occurred_start, occurred_end,
|
SELECT id, text, context, event_date, occurred_start, occurred_end,
|
||||||
mentioned_at, fact_type, document_id, chunk_id, tags
|
mentioned_at, fact_type, document_id, chunk_id, tags, metadata
|
||||||
FROM {fq_table("memory_units")}
|
FROM {fq_table("memory_units")}
|
||||||
WHERE id = ANY($1::uuid[])
|
WHERE id = ANY($1::uuid[])
|
||||||
AND fact_type = $2
|
AND fact_type = $2
|
||||||
|
|
|
||||||
|
|
@ -148,7 +148,7 @@ async def retrieve_semantic_bm25_combined(
|
||||||
|
|
||||||
cols = (
|
cols = (
|
||||||
"id, text, context, event_date, occurred_start, occurred_end, mentioned_at, "
|
"id, text, context, event_date, occurred_start, occurred_end, mentioned_at, "
|
||||||
"fact_type, document_id, chunk_id, tags"
|
"fact_type, document_id, chunk_id, tags, metadata"
|
||||||
)
|
)
|
||||||
table = fq_table("memory_units")
|
table = fq_table("memory_units")
|
||||||
|
|
||||||
|
|
@ -343,7 +343,7 @@ async def retrieve_temporal_combined(
|
||||||
{groups_clause}
|
{groups_clause}
|
||||||
),
|
),
|
||||||
sim_ranked AS (
|
sim_ranked AS (
|
||||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.metadata,
|
||||||
1 - (mu.embedding <=> $1::vector) AS similarity,
|
1 - (mu.embedding <=> $1::vector) AS similarity,
|
||||||
ROW_NUMBER() OVER (PARTITION BY mu.fact_type ORDER BY mu.embedding <=> $1::vector) AS sim_rn
|
ROW_NUMBER() OVER (PARTITION BY mu.fact_type ORDER BY mu.embedding <=> $1::vector) AS sim_rn
|
||||||
FROM date_ranked dr
|
FROM date_ranked dr
|
||||||
|
|
@ -351,7 +351,7 @@ async def retrieve_temporal_combined(
|
||||||
WHERE dr.rn <= 50
|
WHERE dr.rn <= 50
|
||||||
AND (1 - (mu.embedding <=> $1::vector)) >= $6
|
AND (1 - (mu.embedding <=> $1::vector)) >= $6
|
||||||
)
|
)
|
||||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, document_id, chunk_id, tags, similarity
|
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, document_id, chunk_id, tags, metadata, similarity
|
||||||
FROM sim_ranked
|
FROM sim_ranked
|
||||||
WHERE sim_rn <= 10
|
WHERE sim_rn <= 10
|
||||||
""",
|
""",
|
||||||
|
|
@ -449,7 +449,7 @@ async def retrieve_temporal_combined(
|
||||||
# bank_id on memory_units lets the planner use idx_memory_units_bank_fact_type.
|
# bank_id on memory_units lets the planner use idx_memory_units_bank_fact_type.
|
||||||
neighbors = await conn.fetch(
|
neighbors = await conn.fetch(
|
||||||
f"""
|
f"""
|
||||||
SELECT src.from_unit_id, mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
SELECT src.from_unit_id, mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.metadata,
|
||||||
l.weight, l.link_type,
|
l.weight, l.link_type,
|
||||||
1 - (mu.embedding <=> $1::vector) AS similarity
|
1 - (mu.embedding <=> $1::vector) AS similarity
|
||||||
FROM unnest($2::uuid[]) AS src(from_unit_id)
|
FROM unnest($2::uuid[]) AS src(from_unit_id)
|
||||||
|
|
|
||||||
|
|
@ -47,6 +47,7 @@ class RetrievalResult:
|
||||||
document_id: str | None = None
|
document_id: str | None = None
|
||||||
chunk_id: str | None = None
|
chunk_id: str | None = None
|
||||||
tags: list[str] | None = None # Visibility scope tags
|
tags: list[str] | None = None # Visibility scope tags
|
||||||
|
metadata: dict[str, str] | None = None # User-provided metadata
|
||||||
|
|
||||||
# Retrieval-specific scores (only one will be set depending on retrieval method)
|
# Retrieval-specific scores (only one will be set depending on retrieval method)
|
||||||
similarity: float | None = None # Semantic retrieval
|
similarity: float | None = None # Semantic retrieval
|
||||||
|
|
@ -70,6 +71,7 @@ class RetrievalResult:
|
||||||
document_id=row.get("document_id"),
|
document_id=row.get("document_id"),
|
||||||
chunk_id=row.get("chunk_id"),
|
chunk_id=row.get("chunk_id"),
|
||||||
tags=row.get("tags"),
|
tags=row.get("tags"),
|
||||||
|
metadata=row.get("metadata"),
|
||||||
similarity=row.get("similarity"),
|
similarity=row.get("similarity"),
|
||||||
bm25_score=row.get("bm25_score"),
|
bm25_score=row.get("bm25_score"),
|
||||||
activation=row.get("activation"),
|
activation=row.get("activation"),
|
||||||
|
|
@ -153,6 +155,7 @@ class ScoredResult:
|
||||||
"document_id": self.retrieval.document_id,
|
"document_id": self.retrieval.document_id,
|
||||||
"chunk_id": self.retrieval.chunk_id,
|
"chunk_id": self.retrieval.chunk_id,
|
||||||
"tags": self.retrieval.tags,
|
"tags": self.retrieval.tags,
|
||||||
|
"metadata": self.retrieval.metadata,
|
||||||
"semantic_similarity": self.retrieval.similarity,
|
"semantic_similarity": self.retrieval.similarity,
|
||||||
"bm25_score": self.retrieval.bm25_score,
|
"bm25_score": self.retrieval.bm25_score,
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -814,34 +814,36 @@ async def test_context_with_batch(memory, request_context):
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_metadata_storage_and_retrieval(memory, request_context):
|
async def test_metadata_storage_and_retrieval(memory, request_context):
|
||||||
"""
|
"""
|
||||||
Test that user-defined metadata is preserved.
|
Test that user-defined metadata passed during retain is returned on recall.
|
||||||
Metadata allows arbitrary key-value data to be stored with facts.
|
Metadata allows arbitrary key-value data to be stored with facts.
|
||||||
"""
|
"""
|
||||||
bank_id = f"test_metadata_{datetime.now(timezone.utc).timestamp()}"
|
bank_id = f"test_metadata_{datetime.now(timezone.utc).timestamp()}"
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# Store content with custom metadata
|
|
||||||
custom_metadata = {
|
custom_metadata = {
|
||||||
"source": "slack",
|
"source": "slack",
|
||||||
"channel": "engineering",
|
"channel": "engineering",
|
||||||
"importance": "high",
|
"importance": "high",
|
||||||
"tags": "product,launch"
|
|
||||||
}
|
}
|
||||||
|
|
||||||
# Note: retain_async doesn't directly support metadata parameter
|
# Use retain_batch_async which supports the metadata parameter
|
||||||
# Metadata would need to be supported in the API layer
|
unit_ids_list = await memory.retain_batch_async(
|
||||||
# For now, we test that the system handles content without errors
|
|
||||||
unit_ids = await memory.retain_async(
|
|
||||||
bank_id=bank_id,
|
bank_id=bank_id,
|
||||||
content="The product launch is scheduled for March 1st.",
|
contents=[
|
||||||
context="planning meeting",
|
{
|
||||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
"content": "The product launch is scheduled for March 1st.",
|
||||||
|
"context": "planning meeting",
|
||||||
|
"event_date": datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||||
|
"metadata": custom_metadata,
|
||||||
|
}
|
||||||
|
],
|
||||||
request_context=request_context,
|
request_context=request_context,
|
||||||
)
|
)
|
||||||
|
|
||||||
assert len(unit_ids) > 0, "Should create memory units"
|
assert len(unit_ids_list) > 0, "Should create memory units"
|
||||||
|
assert len(unit_ids_list[0]) > 0, "Should have at least one unit ID"
|
||||||
|
|
||||||
# Recall to verify storage worked
|
# Recall and verify metadata is returned
|
||||||
result = await memory.recall_async(
|
result = await memory.recall_async(
|
||||||
bank_id=bank_id,
|
bank_id=bank_id,
|
||||||
query="When is the product launch?",
|
query="When is the product launch?",
|
||||||
|
|
@ -853,8 +855,12 @@ async def test_metadata_storage_and_retrieval(memory, request_context):
|
||||||
|
|
||||||
assert len(result.results) > 0, "Should recall stored facts"
|
assert len(result.results) > 0, "Should recall stored facts"
|
||||||
|
|
||||||
print("✓ Successfully stored and retrieved facts")
|
# Verify metadata is present on recalled facts
|
||||||
print(" (Note: Metadata support depends on API implementation)")
|
fact = result.results[0]
|
||||||
|
assert fact.metadata is not None, "Metadata should not be null on recall"
|
||||||
|
assert fact.metadata.get("source") == "slack"
|
||||||
|
assert fact.metadata.get("channel") == "engineering"
|
||||||
|
assert fact.metadata.get("importance") == "high"
|
||||||
|
|
||||||
finally:
|
finally:
|
||||||
await memory.delete_bank(bank_id, request_context=request_context)
|
await memory.delete_bank(bank_id, request_context=request_context)
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue