351 lines
12 KiB
Python
351 lines
12 KiB
Python
"""Focused persistence tests for project-scoped cross-session memory."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
import pytest
|
|
from sqlalchemy import delete, event, func, select
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from deerflow.persistence.analysis_project_memory import (
|
|
AnalysisProjectMemoryItemRow,
|
|
AnalysisProjectMemoryJobRow,
|
|
AnalysisProjectMemoryStateRow,
|
|
AnalysisProjectMemoryStore,
|
|
)
|
|
from deerflow.persistence.analysis_projects import AnalysisProjectRow
|
|
|
|
|
|
async def _sql_store(tmp_path, *, project_id: str = "project-1", owner: str = "user-1"):
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'project-memory.db'}")
|
|
|
|
@event.listens_for(engine.sync_engine, "connect")
|
|
def _foreign_keys(dbapi_connection, _connection_record): # noqa: ANN001
|
|
cursor = dbapi_connection.cursor()
|
|
cursor.execute("PRAGMA foreign_keys=ON")
|
|
cursor.close()
|
|
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(AnalysisProjectRow.__table__.create)
|
|
await connection.run_sync(AnalysisProjectMemoryStateRow.__table__.create)
|
|
await connection.run_sync(AnalysisProjectMemoryItemRow.__table__.create)
|
|
await connection.run_sync(AnalysisProjectMemoryJobRow.__table__.create)
|
|
sf = async_sessionmaker(engine, expire_on_commit=False)
|
|
now = datetime.now(UTC)
|
|
async with sf() as session:
|
|
session.add(
|
|
AnalysisProjectRow(
|
|
id=project_id,
|
|
user_id=owner,
|
|
name="Project",
|
|
mode="qa",
|
|
background="",
|
|
knowledge_base_ids=[],
|
|
assistant_knowledge_base_ids=[],
|
|
created_at=now,
|
|
updated_at=now,
|
|
)
|
|
)
|
|
await session.commit()
|
|
return AnalysisProjectMemoryStore(sf), sf, engine
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_memory_fallback_supersedes_items_and_builds_snapshot() -> None:
|
|
store = AnalysisProjectMemoryStore()
|
|
state = await store.get_state("p1", "u1")
|
|
assert state is not None
|
|
assert state["enabled"] is True
|
|
assert await store.get_state("p1", "another-user") is None
|
|
|
|
first = await store.upsert_item(
|
|
"p1",
|
|
"u1",
|
|
logical_key="deadline",
|
|
kind="constraint",
|
|
content="Deadline is Friday",
|
|
source_thread_id="thread-1",
|
|
source_message_ids=["message-1"],
|
|
)
|
|
second = await store.upsert_item(
|
|
"p1",
|
|
"u1",
|
|
logical_key="deadline",
|
|
kind="constraint",
|
|
content="Deadline is Monday",
|
|
source_thread_id="thread-2",
|
|
)
|
|
|
|
assert first["id"] != second["id"]
|
|
all_items = await store.list_items("p1", "u1", status=None)
|
|
assert {item["status"] for item in all_items} == {"active", "superseded"}
|
|
old = next(item for item in all_items if item["id"] == first["id"])
|
|
assert old["superseded_by"] == second["id"]
|
|
|
|
state = await store.get_state("p1", "u1")
|
|
assert state is not None
|
|
assert state["item_count"] == 1
|
|
assert state["snapshot"]["summary"] == "已沉淀 1 条项目记忆。"
|
|
assert state["snapshot"]["statistics"]["by_kind"]["constraint"] == 1
|
|
assert state["snapshot"]["constraints"] == ["Deadline is Monday"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_recall_uses_lexical_fallback_and_vector_identity_gate() -> None:
|
|
store = AnalysisProjectMemoryStore()
|
|
await store.upsert_item(
|
|
"p1",
|
|
"u1",
|
|
logical_key="budget",
|
|
kind="fact",
|
|
content="Budget is 500 yuan",
|
|
importance=0.2,
|
|
embedding=[1.0, 0.0],
|
|
embedding_model="encoder-a",
|
|
embedding_fingerprint="a-v1",
|
|
)
|
|
await store.upsert_item(
|
|
"p1",
|
|
"u1",
|
|
logical_key="venue",
|
|
kind="fact",
|
|
content="Venue is Shanghai",
|
|
importance=0.9,
|
|
embedding=[0.0, 1.0],
|
|
embedding_model="encoder-b",
|
|
embedding_fingerprint="b-v1",
|
|
)
|
|
|
|
lexical = await store.recall("p1", "u1", query="What is the budget?")
|
|
assert lexical["items"][0]["logical_key"] == "budget"
|
|
assert all(item["recall_signals"]["vector"] == 0 for item in lexical["items"])
|
|
|
|
wrong_model = await store.recall(
|
|
"p1",
|
|
"u1",
|
|
query="unrelated",
|
|
query_embedding=[1.0, 0.0],
|
|
query_embedding_model="encoder-b",
|
|
query_embedding_fingerprint="b-v1",
|
|
)
|
|
budget = next(item for item in wrong_model["items"] if item["logical_key"] == "budget")
|
|
assert budget["recall_signals"]["vector"] == 0
|
|
assert budget["recall_signals"]["vector_compatible"] is False
|
|
|
|
matching = await store.recall(
|
|
"p1",
|
|
"u1",
|
|
query="unrelated",
|
|
query_embedding=[1.0, 0.0],
|
|
query_embedding_model="encoder-a",
|
|
query_embedding_fingerprint="a-v1",
|
|
)
|
|
budget = next(item for item in matching["items"] if item["logical_key"] == "budget")
|
|
assert budget["recall_signals"]["vector"] == 1.0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_apply_extraction_snapshot_is_order_independent() -> None:
|
|
first = AnalysisProjectMemoryStore()
|
|
second = AnalysisProjectMemoryStore()
|
|
deltas = [
|
|
{
|
|
"logical_key": "decision-b",
|
|
"kind": "decision",
|
|
"content": "Use option B",
|
|
"importance": 0.7,
|
|
},
|
|
{
|
|
"logical_key": "fact-a",
|
|
"kind": "fact",
|
|
"content": "The audience is internal",
|
|
"importance": 0.8,
|
|
},
|
|
]
|
|
for item in deltas:
|
|
result = await first.apply_extraction(
|
|
"p1", "u1", items=[item], snapshot={"summary": "untrusted"}
|
|
)
|
|
assert result["applied"] is True
|
|
for item in reversed(deltas):
|
|
await second.apply_extraction(
|
|
"p1", "u1", items=[item], snapshot={"summary": "different"}
|
|
)
|
|
|
|
first_state = await first.get_state("p1", "u1")
|
|
second_state = await second.get_state("p1", "u1")
|
|
assert first_state is not None and second_state is not None
|
|
assert first_state["snapshot"] == second_state["snapshot"]
|
|
assert first_state["snapshot"]["summary"] == "已沉淀 2 条项目记忆。"
|
|
assert "untrusted" not in str(first_state["snapshot"])
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sql_job_idempotency_lease_retry_and_new_success_recovers_ready(tmp_path) -> None:
|
|
store, _sf, engine = await _sql_store(tmp_path)
|
|
try:
|
|
await store.mark_rebuild("project-1", "user-1")
|
|
first, created = await store.enqueue_job(
|
|
"project-1",
|
|
"user-1",
|
|
source_hash="hash-1",
|
|
extractor_version="v1",
|
|
payload={"source_message_ids": ["m1"]},
|
|
source_thread_id="thread-1",
|
|
max_attempts=1,
|
|
)
|
|
duplicate, duplicate_created = await store.enqueue_job(
|
|
"project-1",
|
|
"user-1",
|
|
source_hash="hash-1",
|
|
extractor_version="v1",
|
|
payload={"source_message_ids": ["m1"]},
|
|
source_thread_id="thread-1",
|
|
)
|
|
assert created is True
|
|
assert duplicate_created is False
|
|
assert duplicate["id"] == first["id"]
|
|
assert (await store.get_state("project-1", "user-1"))["status"] == "rebuilding"
|
|
|
|
lease_until = datetime.now(UTC) + timedelta(minutes=1)
|
|
[claimed] = await store.claim_jobs(
|
|
lease_owner="worker-1", lease_until=lease_until, limit=1
|
|
)
|
|
assert claimed["attempt"] == 1
|
|
failed = await store.fail_job(
|
|
claimed["id"], lease_owner="worker-1", error="temporary failure"
|
|
)
|
|
assert failed is not None and failed["status"] == "failed"
|
|
assert (await store.get_state("project-1", "user-1"))["status"] == "error"
|
|
|
|
await store.enqueue_job(
|
|
"project-1",
|
|
"user-1",
|
|
source_hash="hash-2",
|
|
extractor_version="v1",
|
|
payload={"source_message_ids": ["m2"]},
|
|
source_thread_id="thread-2",
|
|
)
|
|
[claimed] = await store.claim_jobs(
|
|
lease_owner="worker-2", lease_until=lease_until, limit=1
|
|
)
|
|
applied = await store.apply_extraction(
|
|
"project-1",
|
|
"user-1",
|
|
items=[{"logical_key": "new", "kind": "fact", "content": "New fact"}],
|
|
source_thread_id="thread-2",
|
|
job_id=claimed["id"],
|
|
lease_owner="worker-2",
|
|
)
|
|
assert applied["applied"] is True
|
|
completed = await store.complete_job(
|
|
claimed["id"], lease_owner="worker-2", result={"count": 1}
|
|
)
|
|
assert completed is not None and completed["status"] == "completed"
|
|
assert (await store.get_state("project-1", "user-1"))["status"] == "ready"
|
|
finally:
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_invalidate_and_clear_fence_running_jobs(tmp_path) -> None:
|
|
store, _sf, engine = await _sql_store(tmp_path)
|
|
try:
|
|
await store.upsert_item(
|
|
"project-1",
|
|
"user-1",
|
|
logical_key="old",
|
|
kind="fact",
|
|
content="Old fact",
|
|
source_thread_id="thread-1",
|
|
source_message_ids=["m1"],
|
|
)
|
|
job, _ = await store.enqueue_job(
|
|
"project-1",
|
|
"user-1",
|
|
source_hash="turn-1",
|
|
extractor_version="v1",
|
|
payload={"source_message_ids": ["m1"]},
|
|
source_thread_id="thread-1",
|
|
)
|
|
[claimed] = await store.claim_jobs(
|
|
lease_owner="worker", lease_until=datetime.now(UTC) + timedelta(minutes=1)
|
|
)
|
|
assert claimed["id"] == job["id"]
|
|
|
|
invalidated = await store.invalidate_source(
|
|
"project-1",
|
|
"user-1",
|
|
source_thread_id="thread-1",
|
|
source_message_ids=["m1"],
|
|
)
|
|
assert invalidated["invalidated_items"] == 1
|
|
assert invalidated["invalidated_jobs"] == 1
|
|
assert invalidated["state"]["status"] == "rebuilding"
|
|
fenced = await store.apply_extraction(
|
|
"project-1",
|
|
"user-1",
|
|
items=[{"logical_key": "stale", "kind": "fact", "content": "Stale fact"}],
|
|
source_thread_id="thread-1",
|
|
job_id=job["id"],
|
|
lease_owner="worker",
|
|
)
|
|
assert fenced["applied"] is False
|
|
assert await store.list_items("project-1", "user-1") == []
|
|
|
|
second, _ = await store.enqueue_job(
|
|
"project-1",
|
|
"user-1",
|
|
source_hash="turn-2",
|
|
extractor_version="v1",
|
|
payload={"source_message_ids": ["m2"]},
|
|
source_thread_id="thread-2",
|
|
)
|
|
[second_claimed] = await store.claim_jobs(
|
|
lease_owner="worker", lease_until=datetime.now(UTC) + timedelta(minutes=1)
|
|
)
|
|
assert second_claimed["id"] == second["id"]
|
|
await store.clear("project-1", "user-1")
|
|
fenced = await store.apply_extraction(
|
|
"project-1",
|
|
"user-1",
|
|
items=[{"logical_key": "stale-2", "kind": "fact", "content": "Stale again"}],
|
|
job_id=second["id"],
|
|
lease_owner="worker",
|
|
)
|
|
assert fenced["applied"] is False
|
|
assert await store.list_items("project-1", "user-1") == []
|
|
finally:
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sql_project_delete_cascades_all_memory_rows(tmp_path) -> None:
|
|
store, sf, engine = await _sql_store(tmp_path)
|
|
try:
|
|
await store.upsert_item(
|
|
"project-1", "user-1", logical_key="fact", kind="fact", content="Remember me"
|
|
)
|
|
await store.enqueue_job(
|
|
"project-1",
|
|
"user-1",
|
|
source_hash="hash",
|
|
extractor_version="v1",
|
|
payload={},
|
|
)
|
|
async with sf() as session:
|
|
await session.execute(delete(AnalysisProjectRow).where(AnalysisProjectRow.id == "project-1"))
|
|
await session.commit()
|
|
async with sf() as session:
|
|
counts = []
|
|
for model in (
|
|
AnalysisProjectMemoryStateRow,
|
|
AnalysisProjectMemoryItemRow,
|
|
AnalysisProjectMemoryJobRow,
|
|
):
|
|
counts.append((await session.execute(select(func.count()).select_from(model))).scalar_one())
|
|
assert counts == [0, 0, 0]
|
|
finally:
|
|
await engine.dispose()
|