121 lines
4.0 KiB
Python
121 lines
4.0 KiB
Python
"""Baseline memory enrichment: single LLM call → session summary + facts."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from app.features.memory.enrichment import enrich_memory_after_ingest_sync
|
|
from app.features.memory.llm_schemas import EnrichmentPayload, parse_json_payload
|
|
from app.features.memory.models import MemorySource
|
|
from app.features.user.models import User
|
|
|
|
|
|
def test_enrichment_payload_roundtrip() -> None:
|
|
raw = (
|
|
'{"summary":"要点摘要",'
|
|
'"facts":[{"fact_type":"event","subject":"王伟","predicate":"去",'
|
|
'"object_json":{"value":"北京","approximate_era":"1990年代"},'
|
|
'"confidence":0.85,"source_chunk_id":"ch-1"}]}'
|
|
)
|
|
p = parse_json_payload(raw, EnrichmentPayload)
|
|
assert p is not None
|
|
assert p.summary == "要点摘要"
|
|
assert len(p.facts) == 1
|
|
assert p.facts[0].subject == "王伟"
|
|
|
|
|
|
def test_enrich_memory_after_ingest_sync_single_llm_call(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
from app.features.memory import enrichment as mod
|
|
|
|
monkeypatch.setattr("app.core.config.settings.memory_enrichment_enabled", True)
|
|
|
|
invoke_count = {"n": 0}
|
|
|
|
def fake_invoke(llm, prompt, max_tokens, agent):
|
|
invoke_count["n"] += 1
|
|
assert agent == "memory.enrichment_sync"
|
|
return (
|
|
'{"summary":"本轮要点",'
|
|
'"facts":[{"fact_type":"event","subject":"王伟","predicate":"住",'
|
|
'"object_json":{"value":"上海"},"confidence":0.8,"source_chunk_id":"ch1"}]}'
|
|
)
|
|
|
|
monkeypatch.setattr(mod, "invoke_json_object", fake_invoke)
|
|
monkeypatch.setattr(
|
|
mod,
|
|
"list_chunks_for_source_sync",
|
|
lambda s, sid: [SimpleNamespace(id="ch1", content="王伟住在上海。")],
|
|
)
|
|
|
|
summaries: list[dict] = []
|
|
facts: list[dict] = []
|
|
|
|
def capture_summary(session, **kwargs):
|
|
summaries.append(kwargs)
|
|
|
|
def capture_fact(session, **kwargs):
|
|
facts.append(kwargs)
|
|
|
|
monkeypatch.setattr(mod, "create_memory_summary_sync", capture_summary)
|
|
monkeypatch.setattr(mod, "create_memory_fact_sync", capture_fact)
|
|
|
|
class FakeSession:
|
|
def get(self, model, key):
|
|
if model is User and key == "u1":
|
|
return SimpleNamespace(nickname="老王")
|
|
if model is MemorySource and key == "src-1":
|
|
return SimpleNamespace(lineage_json=None)
|
|
return None
|
|
|
|
enrich_memory_after_ingest_sync(FakeSession(), "u1", "src-1", llm=object())
|
|
|
|
assert invoke_count["n"] == 1
|
|
assert len(summaries) == 1
|
|
assert summaries[0]["summary_type"] == "session"
|
|
assert summaries[0]["content"] == "本轮要点"
|
|
assert summaries[0]["source_chunk_ids"] == ["ch1"]
|
|
assert len(facts) == 1
|
|
assert facts[0]["predicate"] == "住"
|
|
assert facts[0]["status"] == "confirmed"
|
|
|
|
|
|
def test_enrich_memory_skips_when_parse_returns_none(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
from app.features.memory import enrichment as mod
|
|
|
|
monkeypatch.setattr("app.core.config.settings.memory_enrichment_enabled", True)
|
|
monkeypatch.setattr(mod, "invoke_json_object", lambda *a, **k: "{not json")
|
|
monkeypatch.setattr(
|
|
mod,
|
|
"list_chunks_for_source_sync",
|
|
lambda s, sid: [SimpleNamespace(id="c1", content="x")],
|
|
)
|
|
called = {"summary": False, "fact": False}
|
|
|
|
monkeypatch.setattr(
|
|
mod,
|
|
"create_memory_summary_sync",
|
|
lambda *a, **k: called.update(summary=True),
|
|
)
|
|
monkeypatch.setattr(
|
|
mod,
|
|
"create_memory_fact_sync",
|
|
lambda *a, **k: called.update(fact=True),
|
|
)
|
|
|
|
class FakeSession:
|
|
def get(self, model, key):
|
|
if model is User and key == "u":
|
|
return None
|
|
if model is MemorySource and key == "s":
|
|
return SimpleNamespace(lineage_json=None)
|
|
return None
|
|
|
|
enrich_memory_after_ingest_sync(FakeSession(), "u", "s", llm=object())
|
|
assert called == {"summary": False, "fact": False}
|