sync: 鍚屾灞€鍩熺綉鏈€鏂颁唬鐮?(source master @ b69568e4)
This commit is contained in:
gozeros committed
1 parent
4f7f07504a
commit
8bc8b0e082
1651 files changed
+278599
-11416
No files matched your search
@@ -0,0 +1,350 @@
|
||||
"""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()
|
||||
Reference in new issue
Block a user