Rewrite consolidation prompt rules to produce clean, single-facet observations: - One observation per distinct facet (count, named entity, relationship) - Match updates by entity/facet, not topic similarity - No computation — never infer/calculate values not explicitly stated - Cascade state changes to all affected observations - Preserve event history (sold, died, moved) — conservative deletes - Include dates on state changes when available - Keep observations concise — no cross-facet narrative bloat Add test_horse_observations.py exercising a realistic sequence of retain operations (farm with horses being named, sold, dying) and verifying that observations track history correctly and mental models can synthesize them.
352 lines
13 KiB
Python
352 lines
13 KiB
Python
"""Test observation tracking for a sequence of horse-related memories.
|
|
|
|
This test retains a series of facts about horses on a farm and inspects
|
|
how observations track the evolving state over time, with full prompt debugging.
|
|
"""
|
|
|
|
import json
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from hindsight_api.config import _get_raw_config
|
|
from hindsight_api.engine.consolidation import consolidator as consolidator_mod
|
|
from hindsight_api.engine.memory_engine import MemoryEngine
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def enable_observations():
|
|
"""Enable observations for all tests in this module."""
|
|
config = _get_raw_config()
|
|
original_value = config.enable_observations
|
|
config.enable_observations = True
|
|
yield
|
|
config.enable_observations = original_value
|
|
|
|
|
|
@dataclass
|
|
class _ActionLog:
|
|
text: str
|
|
source_fact_ids: list[str] = field(default_factory=list)
|
|
observation_id: str = ""
|
|
|
|
|
|
@dataclass
|
|
class _ConsolidationResponse:
|
|
creates: list[_ActionLog] = field(default_factory=list)
|
|
updates: list[_ActionLog] = field(default_factory=list)
|
|
deletes: list[_ActionLog] = field(default_factory=list)
|
|
|
|
|
|
@dataclass
|
|
class _ConsolidationDebugEntry:
|
|
facts: str
|
|
observations_text: str
|
|
response: _ConsolidationResponse
|
|
|
|
|
|
# Store prompts/responses for debugging
|
|
_debug_log: list[_ConsolidationDebugEntry] = []
|
|
|
|
|
|
def _fact_line(m: dict[str, Any]) -> str:
|
|
text = f"[{m['id']}] {m['text']}"
|
|
temporal_parts = []
|
|
if m.get("occurred_start"):
|
|
temporal_parts.append(f"occurred_start={m['occurred_start']}")
|
|
if m.get("occurred_end"):
|
|
temporal_parts.append(f"occurred_end={m['occurred_end']}")
|
|
if m.get("mentioned_at"):
|
|
temporal_parts.append(f"mentioned_at={m['mentioned_at']}")
|
|
if temporal_parts:
|
|
text += f" ({', '.join(temporal_parts)})"
|
|
return text
|
|
|
|
|
|
async def _instrumented_consolidate(
|
|
original_fn: Any,
|
|
*,
|
|
llm_config: Any,
|
|
memories: list[dict[str, Any]],
|
|
union_observations: Any,
|
|
union_source_facts: Any,
|
|
config: Any = None,
|
|
remaining_observation_slots: int | None = None,
|
|
max_observations_per_scope: int = -1,
|
|
) -> Any:
|
|
"""Wrapper that captures the prompt and response for debugging."""
|
|
if union_observations:
|
|
obs_list = consolidator_mod._build_observations_for_llm(union_observations, union_source_facts)
|
|
observations_text = json.dumps(obs_list, indent=2)
|
|
else:
|
|
observations_text = "[]"
|
|
|
|
facts_lines = "\n".join(_fact_line(m) for m in memories)
|
|
|
|
result = await original_fn(
|
|
llm_config=llm_config,
|
|
memories=memories,
|
|
union_observations=union_observations,
|
|
union_source_facts=union_source_facts,
|
|
config=config,
|
|
remaining_observation_slots=remaining_observation_slots,
|
|
max_observations_per_scope=max_observations_per_scope,
|
|
)
|
|
|
|
_debug_log.append(_ConsolidationDebugEntry(
|
|
facts=facts_lines,
|
|
observations_text=observations_text,
|
|
response=_ConsolidationResponse(
|
|
creates=[_ActionLog(text=c.text, source_fact_ids=c.source_fact_ids) for c in result.creates],
|
|
updates=[
|
|
_ActionLog(text=u.text, observation_id=u.observation_id, source_fact_ids=u.source_fact_ids)
|
|
for u in result.updates
|
|
],
|
|
deletes=[_ActionLog(text="", observation_id=d.observation_id) for d in result.deletes],
|
|
),
|
|
))
|
|
|
|
return result
|
|
|
|
|
|
def _print_consolidation_debug(entry: _ConsolidationDebugEntry, index: int) -> None:
|
|
"""Print a single consolidation LLM call for debugging."""
|
|
print(f"\n --- LLM Call #{index} ---")
|
|
print(" FACTS sent to LLM:")
|
|
for line in entry.facts.split("\n"):
|
|
print(f" {line}")
|
|
print("\n EXISTING OBSERVATIONS sent to LLM:")
|
|
obs_data = json.loads(entry.observations_text)
|
|
if obs_data:
|
|
for obs in obs_data:
|
|
src_summary = ""
|
|
if obs.get("source_memories"):
|
|
src_texts = [sm["text"] for sm in obs["source_memories"]]
|
|
src_summary = f" (sources: {src_texts})"
|
|
print(f" [{obs['id'][:8]}..] proof={obs.get('proof_count', '?')}: {obs['text']}{src_summary}")
|
|
else:
|
|
print(" (none)")
|
|
|
|
resp = entry.response
|
|
print("\n LLM RESPONSE:")
|
|
if resp.creates:
|
|
for c in resp.creates:
|
|
print(f" CREATE: \"{c.text}\" (from facts: {[fid[:8] + '..' for fid in c.source_fact_ids]})")
|
|
if resp.updates:
|
|
for u in resp.updates:
|
|
print(
|
|
f" UPDATE [{u.observation_id[:8]}..]: \"{u.text}\""
|
|
f" (from facts: {[fid[:8] + '..' for fid in u.source_fact_ids]})"
|
|
)
|
|
if resp.deletes:
|
|
for d in resp.deletes:
|
|
print(f" DELETE [{d.observation_id[:8]}..]")
|
|
if not resp.creates and not resp.updates and not resp.deletes:
|
|
print(" (no actions)")
|
|
|
|
|
|
def _parse_history(hist: Any) -> list[str]:
|
|
"""Parse observation history from DB (may be list of dicts or JSON strings)."""
|
|
if not hist:
|
|
return []
|
|
parsed = hist if isinstance(hist, list) else json.loads(hist)
|
|
prev_texts = []
|
|
for h in parsed:
|
|
if isinstance(h, str):
|
|
h = json.loads(h)
|
|
prev_texts.append(h.get("previous_text", "?"))
|
|
return prev_texts
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_horse_farm_observation_history(memory: MemoryEngine, request_context: Any) -> None:
|
|
"""Retain a sequence of horse facts and inspect how observations evolve."""
|
|
bank_id = f"test-horses-{uuid.uuid4().hex[:8]}"
|
|
|
|
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
|
|
|
messages = [
|
|
"I have a farm.",
|
|
"I have 2 horses.",
|
|
"I have a horse named Daisy.",
|
|
"I have a horse named Buttercup.",
|
|
"I sold Buttercup.",
|
|
"I now have 1 horse.",
|
|
"I have 5 horses on my farm.",
|
|
"I have a horse named Midnight.",
|
|
"I have horses named Midnight and Shadow.",
|
|
"I have horses named Shadow and Twister.",
|
|
"I am sad to report that Shadow has died.",
|
|
]
|
|
|
|
# Monkey-patch to intercept consolidation LLM calls
|
|
_original_consolidate = consolidator_mod._consolidate_batch_with_llm
|
|
|
|
async def _patched(**kwargs: Any) -> Any:
|
|
return await _instrumented_consolidate(_original_consolidate, **kwargs)
|
|
|
|
consolidator_mod._consolidate_batch_with_llm = _patched
|
|
_debug_log.clear()
|
|
|
|
try:
|
|
for i, content in enumerate(messages):
|
|
print(f"\n{'='*80}")
|
|
print(f"RETAIN #{i+1}: {content}")
|
|
print(f"{'='*80}")
|
|
|
|
log_start = len(_debug_log)
|
|
|
|
await memory.retain_async(
|
|
bank_id=bank_id,
|
|
content=content,
|
|
request_context=request_context,
|
|
)
|
|
await memory.wait_for_background_tasks()
|
|
|
|
for j, entry in enumerate(_debug_log[log_start:]):
|
|
_print_consolidation_debug(entry, j + 1)
|
|
|
|
# Dump current observations
|
|
pool = await memory._get_pool()
|
|
async with pool.acquire() as conn:
|
|
observations = await conn.fetch(
|
|
"""
|
|
SELECT id, text, proof_count, source_memory_ids, history
|
|
FROM memory_units
|
|
WHERE bank_id = $1 AND fact_type = 'observation'
|
|
ORDER BY created_at
|
|
""",
|
|
bank_id,
|
|
)
|
|
|
|
print(f"\n CURRENT OBSERVATIONS ({len(observations)}):")
|
|
for obs in observations:
|
|
prev_texts = _parse_history(obs["history"])
|
|
hist_str = f" (was: {' -> '.join(prev_texts)})" if prev_texts else ""
|
|
print(f" [{str(obs['id'])[:8]}..] proof={obs['proof_count']}: {obs['text']}{hist_str}")
|
|
finally:
|
|
consolidator_mod._consolidate_batch_with_llm = _original_consolidate
|
|
|
|
# Final summary
|
|
print(f"\n{'='*80}")
|
|
print("FINAL STATE")
|
|
print(f"{'='*80}")
|
|
pool = await memory._get_pool()
|
|
async with pool.acquire() as conn:
|
|
observations = await conn.fetch(
|
|
"""
|
|
SELECT id, text, proof_count, source_memory_ids, history
|
|
FROM memory_units
|
|
WHERE bank_id = $1 AND fact_type = 'observation'
|
|
ORDER BY created_at
|
|
""",
|
|
bank_id,
|
|
)
|
|
print(f"\nFinal observations ({len(observations)}):")
|
|
for obs in observations:
|
|
prev_texts = _parse_history(obs["history"])
|
|
if prev_texts:
|
|
chain = prev_texts + [obs["text"]]
|
|
print(f" - [proof={obs['proof_count']}] {obs['text']}")
|
|
print(f" evolution: {' -> '.join(chain)}")
|
|
else:
|
|
print(f" - [proof={obs['proof_count']}] {obs['text']}")
|
|
|
|
# Create a mental model to synthesize the observations
|
|
print(f"\n{'='*80}")
|
|
print("MENTAL MODEL")
|
|
print(f"{'='*80}")
|
|
|
|
# Patch reflect _execute_tool to log tool inputs/outputs
|
|
from hindsight_api.engine.reflect import agent as reflect_agent_mod
|
|
|
|
_original_execute = reflect_agent_mod._execute_tool
|
|
|
|
async def _logging_execute(tool_name: str, args: dict[str, Any], *a: Any, **kw: Any) -> dict[str, Any]:
|
|
result = await _original_execute(tool_name, args, *a, **kw)
|
|
normalized = reflect_agent_mod._normalize_tool_name(tool_name)
|
|
print(f"\n [REFLECT TOOL] {normalized}(args={args})")
|
|
if isinstance(result, dict):
|
|
if "observations" in result:
|
|
print(f" Observations returned ({result.get('count', '?')}, freshness={result.get('freshness', '?')}):")
|
|
for obs in result.get("observations", []):
|
|
print(f" - [proof={obs.get('proof_count', '?')}] {obs.get('text', '?')}")
|
|
if "memories" in result:
|
|
print(f" Memories returned ({result.get('count', '?')}):")
|
|
for mem in result.get("memories", []):
|
|
chunk = mem.get("chunk_text", "")
|
|
chunk_preview = f" | chunk: {chunk[:80]}..." if chunk else ""
|
|
print(f" - [{mem.get('fact_type', '?')}] {mem.get('text', '?')}{chunk_preview}")
|
|
if "mental_models" in result:
|
|
print(f" Mental models returned ({result.get('count', '?')}):")
|
|
for mm_item in result.get("mental_models", []):
|
|
print(f" - {mm_item.get('name', '?')}: {str(mm_item.get('content', '?'))[:120]}")
|
|
if "error" in result:
|
|
print(f" ERROR: {result['error']}")
|
|
return result
|
|
|
|
reflect_agent_mod._execute_tool = _logging_execute
|
|
|
|
source_query = (
|
|
"Produce a structured summary of all animals on the farm. Include:\n"
|
|
"1. A chronological timeline of events (acquisitions, sales, deaths) with dates\n"
|
|
"2. The list of all known horse names and their current status (alive, sold, died)\n"
|
|
"3. The current number of horses on the farm, accounting for all events\n"
|
|
"Reason step by step from the facts. If a horse died or was sold, subtract from the count."
|
|
)
|
|
|
|
try:
|
|
mm = await memory.create_mental_model(
|
|
bank_id=bank_id,
|
|
name="Farm Animals",
|
|
source_query=source_query,
|
|
content="(initial — awaiting refresh)",
|
|
request_context=request_context,
|
|
)
|
|
|
|
refreshed = await memory.refresh_mental_model(
|
|
bank_id=bank_id,
|
|
mental_model_id=mm["id"],
|
|
request_context=request_context,
|
|
)
|
|
content = refreshed["content"]
|
|
finally:
|
|
reflect_agent_mod._execute_tool = _original_execute
|
|
print(f"\nMental model content:\n{content}")
|
|
|
|
reflect_resp = refreshed.get("reflect_response")
|
|
if reflect_resp and isinstance(reflect_resp, str) and reflect_resp.strip():
|
|
try:
|
|
reflect_resp = json.loads(reflect_resp)
|
|
except json.JSONDecodeError:
|
|
reflect_resp = None
|
|
if isinstance(reflect_resp, dict):
|
|
based_on = reflect_resp.get("based_on", [])
|
|
if based_on:
|
|
print("\nBased on:")
|
|
for item in based_on:
|
|
if isinstance(item, str):
|
|
try:
|
|
item = json.loads(item)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
print(f" - [{item.get('fact_type', '?')}] {item.get('text', '?')}")
|
|
|
|
# Verify the mental model captures key facts
|
|
content_lower = content.lower()
|
|
|
|
for name in ["daisy", "buttercup", "midnight", "shadow", "twister"]:
|
|
assert name in content_lower, f"Mental model should mention {name}. Got:\n{content}"
|
|
|
|
assert "sold" in content_lower or "sale" in content_lower, (
|
|
f"Mental model should mention Buttercup was sold. Got:\n{content}"
|
|
)
|
|
|
|
assert "died" in content_lower or "passed" in content_lower or "death" in content_lower, (
|
|
f"Mental model should mention Shadow's death. Got:\n{content}"
|
|
)
|
|
|
|
# Cleanup
|
|
await memory.delete_bank(bank_id, request_context=request_context)
|