265 lines
12 KiB
Python
265 lines
12 KiB
Python
"""记忆抽取集成测试（真实 MySQL）。
|
||
|
||
P2 之后记忆链路的语义变化：
|
||
|
||
- 记忆键来自**受控词表**（如 `preference:risk_level`），不再是 `conversation.{session_id}`；
|
||
- 记忆内容是模型抽取出的**结构化值**，而不是用户原文整句。
|
||
|
||
抽取依赖模型端点，而本地 `model_endpoint_config` 为 0 行，端点缺失时按设计失败关闭。
|
||
因此本文件注入确定性替身抽取器，只验证与模型供应商无关的契约：正文回查、幂等边界、
|
||
客户隔离、证据落库与清理。
|
||
"""
|
||
|
||
from datetime import UTC, datetime
|
||
from uuid import uuid4
|
||
|
||
import pytest
|
||
from sqlalchemy import delete, select
|
||
|
||
from app.infrastructure.db import SessionFactory
|
||
from app.model.conversation import ConversationMessage
|
||
from app.model.memory import MemoryConflict, MemoryEvidence, MemoryUnit
|
||
from app.model.platform import AgentRun, DomainEventOutbox, OutboxDelivery, RequestIdempotency
|
||
from app.service.memory_extraction_service import ExtractedMemory
|
||
from app.service.memory_service import MemoryService
|
||
from app.worker.memory_extraction_worker import MemoryExtractionWorker
|
||
from app.worker.outbox_worker import OutboxWorker
|
||
|
||
KEY_PREFIX = "it-memory-"
|
||
MEMORY_KEY = "preference:risk_level"
|
||
EXTRACTED_VALUE = "稳健型"
|
||
USER_FACT = f"{KEY_PREFIX}我偏好低风险稳健型基金"
|
||
OTHER_FACT = f"{KEY_PREFIX}本人可承受中等风险"
|
||
THIRD_FACT = f"{KEY_PREFIX}本人只做货币基金"
|
||
PLAN_FACT = f"{KEY_PREFIX}我计划两年内买房"
|
||
ASSISTANT_REPLY = "已为您记录该偏好。"
|
||
|
||
|
||
class StubExtractor:
|
||
"""确定性替身抽取器：记录收到的正文，返回固定抽取结果，不发网络请求。"""
|
||
|
||
def __init__(self, value: str = EXTRACTED_VALUE) -> None:
|
||
self.value = value
|
||
self.messages: list[str] = []
|
||
|
||
async def extract(
|
||
self, *, message: str, agent_type: str = "customer_service"
|
||
) -> ExtractedMemory:
|
||
del agent_type
|
||
self.messages.append(message)
|
||
return ExtractedMemory(
|
||
memory_key=MEMORY_KEY,
|
||
value=self.value,
|
||
memory_type="preference",
|
||
confidence=0.9,
|
||
)
|
||
|
||
|
||
class Case:
|
||
def __init__(self, customer_id: int, run_id: str, session_id: str) -> None:
|
||
self.customer_id = customer_id
|
||
self.run_id = run_id
|
||
self.session_id = session_id
|
||
|
||
@property
|
||
def memory_key(self) -> str:
|
||
"""抽取结果对应的受控键；同一客户的记忆键与客户一一对应，便于隔离断言。"""
|
||
return MEMORY_KEY
|
||
|
||
|
||
def new_case(offset: int) -> Case:
|
||
return Case(
|
||
customer_id=uuid4().int % 10**14 + offset * 10**14,
|
||
run_id=str(uuid4()),
|
||
session_id=f"{KEY_PREFIX}{uuid4().hex[:24]}",
|
||
)
|
||
|
||
|
||
async def seed(case: Case, *, user_fact: str) -> int:
|
||
"""建立 run 到用户消息与助手消息的真实链路，事件只携带 message_id。"""
|
||
now = datetime.now(UTC).replace(tzinfo=None)
|
||
async with SessionFactory() as session, session.begin():
|
||
request_message = ConversationMessage(
|
||
session_id=case.session_id, customer_id=case.customer_id, portal="api", role="user",
|
||
content=user_fact, trace_id=str(uuid4()), created_at=now,
|
||
)
|
||
session.add(request_message)
|
||
await session.flush()
|
||
result_message = ConversationMessage(
|
||
session_id=case.session_id, customer_id=case.customer_id, portal="agent",
|
||
role="assistant", content=ASSISTANT_REPLY, trace_id=str(uuid4()), created_at=now,
|
||
)
|
||
session.add(result_message)
|
||
await session.flush()
|
||
idempotency = RequestIdempotency(
|
||
user_id=case.customer_id, session_id=case.session_id, agent_type="customer_service",
|
||
idempotency_key=str(uuid4()), request_hash=uuid4().hex * 2, trace_id=str(uuid4()),
|
||
status="completed", expire_at=now, created_at=now, updated_at=now,
|
||
)
|
||
session.add(idempotency)
|
||
await session.flush()
|
||
session.add(AgentRun(
|
||
run_id=case.run_id, idempotency_id=idempotency.id, session_id=case.session_id,
|
||
user_id=case.customer_id, agent_type="customer_service", trace_id=str(uuid4()),
|
||
request_message_id=request_message.id, result_message_id=result_message.id,
|
||
status="succeeded", attempt_count=1, created_at=now, updated_at=now,
|
||
))
|
||
return result_message.id
|
||
|
||
|
||
async def enqueue(case: Case, message_id: int, *, event_id: str) -> None:
|
||
now = datetime.now(UTC).replace(tzinfo=None)
|
||
async with SessionFactory() as session, session.begin():
|
||
session.add(DomainEventOutbox(
|
||
event_id=event_id, event_type="memory.extraction_requested",
|
||
aggregate_type="agent_run", aggregate_id=case.run_id, trace_id=str(uuid4()),
|
||
payload={"run_id": case.run_id, "message_id": message_id,
|
||
"customer_id": case.customer_id},
|
||
status="pending", retry_count=0, occurred_at=now, created_at=now, updated_at=now,
|
||
))
|
||
|
||
|
||
async def consume_once(case: Case, extractor: StubExtractor) -> bool:
|
||
async with SessionFactory() as session:
|
||
worker = OutboxWorker(session, {
|
||
"memory.extraction_requested": MemoryExtractionWorker(
|
||
session, extractor=extractor
|
||
).handle,
|
||
})
|
||
return await worker.publish_one(aggregate_id=case.run_id)
|
||
|
||
|
||
async def realtime_seen_count(extractor: StubExtractor, fact: str) -> int:
|
||
return sum(1 for message in extractor.messages if message == fact)
|
||
|
||
|
||
async def cleanup(cases: list[Case]) -> None:
|
||
run_ids = [case.run_id for case in cases]
|
||
session_ids = [case.session_id for case in cases]
|
||
customer_ids = [case.customer_id for case in cases]
|
||
async with SessionFactory() as session, session.begin():
|
||
message_ids = list(await session.scalars(select(ConversationMessage.id).where(
|
||
ConversationMessage.session_id.in_(session_ids))))
|
||
memory_ids = list(await session.scalars(select(MemoryUnit.id).where(
|
||
MemoryUnit.customer_id.in_(customer_ids))))
|
||
# 测试证据只可能挂在本次测试的会话消息或本次测试客户的记忆上。
|
||
await session.execute(delete(MemoryEvidence).where(
|
||
(MemoryEvidence.source_record_id.in_([str(item) for item in message_ids]))
|
||
| (MemoryEvidence.memory_id.in_(memory_ids))))
|
||
await session.execute(delete(MemoryConflict).where(
|
||
(MemoryConflict.left_memory_id.in_(memory_ids))
|
||
| (MemoryConflict.right_memory_id.in_(memory_ids))))
|
||
await session.execute(delete(MemoryUnit).where(
|
||
MemoryUnit.customer_id.in_(customer_ids)))
|
||
event_ids = select(DomainEventOutbox.event_id).where(
|
||
DomainEventOutbox.aggregate_id.in_(run_ids))
|
||
await session.execute(delete(OutboxDelivery).where(OutboxDelivery.event_id.in_(event_ids)))
|
||
await session.execute(delete(DomainEventOutbox).where(
|
||
DomainEventOutbox.aggregate_id.in_(run_ids)))
|
||
await session.execute(delete(AgentRun).where(AgentRun.run_id.in_(run_ids)))
|
||
await session.execute(delete(RequestIdempotency).where(
|
||
RequestIdempotency.session_id.in_(session_ids)))
|
||
await session.execute(delete(ConversationMessage).where(
|
||
ConversationMessage.session_id.in_(session_ids)))
|
||
|
||
|
||
@pytest.mark.integration
|
||
@pytest.mark.asyncio
|
||
async def test_duplicate_consumption_produces_single_memory() -> None:
|
||
case = new_case(9)
|
||
extractor = StubExtractor()
|
||
event_id = str(uuid4())
|
||
try:
|
||
message_id = await seed(case, user_fact=USER_FACT)
|
||
await enqueue(case, message_id, event_id=event_id)
|
||
|
||
assert await consume_once(case, extractor)
|
||
# 重复投递同一事件：事件回到 pending 后再次消费。
|
||
async with SessionFactory() as session, session.begin():
|
||
event = await session.scalar(select(DomainEventOutbox).where(
|
||
DomainEventOutbox.event_id == event_id))
|
||
assert event is not None
|
||
event.status = "pending"
|
||
event.published_at = None
|
||
assert await consume_once(case, extractor)
|
||
|
||
async with SessionFactory() as session:
|
||
memories = list(await session.scalars(select(MemoryUnit).where(
|
||
MemoryUnit.customer_id == case.customer_id,
|
||
MemoryUnit.memory_key == case.memory_key)))
|
||
assert len(memories) == 1
|
||
# 落库的是抽取结果，不是用户原文整句。
|
||
assert memories[0].content == EXTRACTED_VALUE
|
||
assert memories[0].content != USER_FACT
|
||
assert memories[0].memory_type == "preference"
|
||
assert memories[0].evidence_count <= 1
|
||
evidence = list(await session.scalars(select(MemoryEvidence).where(
|
||
MemoryEvidence.idempotency_key == f"memory.extraction_requested:{event_id}")))
|
||
assert len(evidence) == 1
|
||
assert evidence[0].memory_id == memories[0].id
|
||
assert evidence[0].source_record_id is not None
|
||
finally:
|
||
await cleanup([case])
|
||
|
||
|
||
@pytest.mark.integration
|
||
@pytest.mark.asyncio
|
||
async def test_memory_is_not_recallable_across_customers() -> None:
|
||
owner = new_case(8)
|
||
other = new_case(7)
|
||
# 两个客户用不同的抽取值，才能用内容区分"谁记住了什么"——受控键相同，
|
||
# 单靠 memory_key 无法判定归属。
|
||
owner_extractor = StubExtractor("稳健型")
|
||
other_extractor = StubExtractor("激进型")
|
||
try:
|
||
owner_message = await seed(owner, user_fact=OTHER_FACT)
|
||
other_message = await seed(other, user_fact=THIRD_FACT)
|
||
await enqueue(owner, owner_message, event_id=str(uuid4()))
|
||
await enqueue(other, other_message, event_id=str(uuid4()))
|
||
assert await consume_once(owner, owner_extractor)
|
||
assert await consume_once(other, other_extractor)
|
||
|
||
async with SessionFactory() as session:
|
||
service = MemoryService(session)
|
||
owner_memories = await service.recall(owner.customer_id)
|
||
other_memories = await service.recall(other.customer_id)
|
||
assert [item.memory_key for item in owner_memories] == [owner.memory_key]
|
||
assert [item.memory_key for item in other_memories] == [other.memory_key]
|
||
assert all(item.customer_id == owner.customer_id for item in owner_memories)
|
||
assert all(item.customer_id == other.customer_id for item in other_memories)
|
||
assert owner_memories[0].content == "稳健型"
|
||
assert other_memories[0].content == "激进型"
|
||
# 客户范围过滤：即使知道他人的内容，也查不到挂在他人名下的行。
|
||
assert await session.scalar(select(MemoryUnit.id).where(
|
||
MemoryUnit.customer_id == owner.customer_id,
|
||
MemoryUnit.content == "激进型")) is None
|
||
assert await session.scalar(select(MemoryUnit.id).where(
|
||
MemoryUnit.customer_id == other.customer_id,
|
||
MemoryUnit.content == "稳健型")) is None
|
||
finally:
|
||
await cleanup([owner, other])
|
||
|
||
|
||
@pytest.mark.integration
|
||
@pytest.mark.asyncio
|
||
async def test_worker_reads_user_content_not_assistant_reply() -> None:
|
||
case = new_case(6)
|
||
extractor = StubExtractor()
|
||
try:
|
||
message_id = await seed(case, user_fact=PLAN_FACT)
|
||
await enqueue(case, message_id, event_id=str(uuid4()))
|
||
assert await consume_once(case, extractor)
|
||
|
||
# 抽取器收到的必须是库中的用户消息正文，而不是助手回复。
|
||
assert extractor.messages == [PLAN_FACT]
|
||
assert ASSISTANT_REPLY not in extractor.messages[0]
|
||
|
||
async with SessionFactory() as session:
|
||
memory = await session.scalar(select(MemoryUnit).where(
|
||
MemoryUnit.customer_id == case.customer_id,
|
||
MemoryUnit.memory_key == case.memory_key))
|
||
assert memory is not None
|
||
assert memory.content == EXTRACTED_VALUE
|
||
finally:
|
||
await cleanup([case])
|