zhifu-platform/tests/test_knowledge_memory.py

import asyncio
from uuid import uuid4

import pytest
from sqlalchemy import select
from test_after_sales import scenario  # noqa: F401

from zhifu.modules.identity.access_models import StoreGrant
from zhifu.modules.knowledge.maintenance import clean
from zhifu.modules.knowledge.vector import connection
from zhifu.modules.memory.models import MemoryJob, MemoryProfile
from zhifu.modules.memory.service import read, snapshot
from zhifu.modules.memory.worker import claim

pytestmark = pytest.mark.integration
MEMORY = "/openapi/v1/me/memory"


async def document(context):
    async with context["app"].state.sessions.begin() as session:
        session.add(
            StoreGrant(
                tenant_id=context["tenant"].bytes,
                principal_id=context["accounts"]["agent"].bytes,
                store_id=context["store"].bytes,
                permission="knowledge:write",
            )
        )
    path = f"/admin/v1/stores/{context['store']}/knowledge/{uuid4()}"
    response = await context["call"](
        "PUT",
        path,
        "agent",
        {
            "expected_version": 0,
            "title": "键盘清洁说明",
            "text": "清洁机械键盘前，请断开电源，使用干布擦拭表面。禁止直接水洗。",
        },
    )
    assert response.status_code == 200, response.text
    return path


async def cleanup_vectors(tenant):
    def remove():
        with connection() as (client, collection):
            client.delete(collection, filter=f'tenant == "{tenant.hex}"')

    await asyncio.to_thread(remove)


async def test_publication_and_revocation(scenario):  # noqa: F811
    path = await document(scenario)
    query_path = f"/openapi/v1/stores/{scenario['store']}/knowledge/search"
    body = {"query": "键盘可以用水清洗吗"}
    try:
        assert (await scenario["call"]("POST", query_path, payload=body)).json()["items"] == []
        response = await scenario["call"](
            "POST", path + "/publish", "agent", {"expected_version": 1}
        )
        assert response.status_code == 200, response.text
        result = await scenario["call"]("POST", query_path, payload=body)
        assert result.status_code == 200, result.text
        assert "禁止直接水洗" in result.json()["items"][0]["text"]
        assert (await scenario["call"]("POST", query_path, "foreign", body)).status_code == 404
        assert (
            await scenario["call"]("POST", path + "/revoke", payload={"expected_version": 1})
        ).status_code == 403
        response = await scenario["call"](
            "POST", path + "/revoke", "agent", {"expected_version": 1}
        )
        assert response.status_code == 200, response.text
        assert (await scenario["call"]("POST", query_path, payload=body)).json()["items"] == []
        assert (
            await scenario["call"]("POST", path + "/publish", "agent", {"expected_version": 1})
        ).status_code == 409
        await clean(scenario["app"].state.sessions)
        with connection() as (client, collection):
            assert client.query(collection, filter=f'tenant == "{scenario["tenant"].hex}"') == []
    finally:
        await cleanup_vectors(scenario["tenant"])


async def test_memory_consent_isolation_and_clear(scenario):  # noqa: F811
    call = scenario["call"]
    assert (await call("GET", MEMORY)).json()["consent"] is False
    body = {"expected_revision": 0, "communication": {"response_style": "concise"}}
    assert (await call("PUT", MEMORY, payload=body)).status_code == 403
    response = await call(
        "PUT", MEMORY + "/consent", payload={"expected_revision": 0, "enabled": True}
    )
    assert response.status_code == 200, response.text
    body["expected_revision"] = 1
    response = await call("PUT", MEMORY, payload=body)
    assert response.status_code == 200, response.text
    assert (await call("GET", MEMORY)).json()["content"]["communication"] == {
        "response_style": "concise"
    }
    assert (await call("GET", MEMORY, "other")).json()["content"] == {}
    assert (await call("GET", MEMORY, "agent")).status_code == 403
    assert (await call("PUT", MEMORY, payload=body)).status_code == 409
    response = await call("DELETE", MEMORY + "?expected_revision=2")
    assert response.status_code == 200, response.text
    assert (await call("GET", MEMORY)).json()["content"] == {}


async def test_memory_clear_invalidates_queued_extraction(scenario):  # noqa: F811
    call = scenario["call"]
    await call("PUT", MEMORY + "/consent", payload={"expected_revision": 0, "enabled": True})
    response = await call(
        "POST", "/openapi/v1/conversations", payload={"store_id": str(scenario["store"])}
    )
    path = "/openapi/v1/conversations/" + response.json()["conversation_id"]
    response = await call(
        "POST",
        path + "/messages",
        payload={
            "text": "以后请简短回复",
            "expected_control_version": 1,
        },
    )
    assert response.status_code == 201, response.text
    await call("DELETE", MEMORY + "?expected_revision=1")
    assert await claim(scenario["app"].state.sessions) is None
    async with scenario["app"].state.sessions() as session:
        job = await session.scalar(
            select(MemoryJob).where(MemoryJob.tenant_id == scenario["tenant"].bytes)
        )
        assert job.status == "CANCELLED"
        row = await session.get(
            MemoryProfile, (scenario["tenant"].bytes, scenario["accounts"]["customer"].bytes)
        )
        assert await read(row.tenant_id, row.subject_id, snapshot(row)) == {}