zhifu-platform/tests/knowledge_memory_model_acceptance.py

import asyncio
from uuid import UUID

import pytest
from sqlalchemy import select
from test_after_sales import scenario  # noqa: F401
from test_knowledge_memory import MEMORY, cleanup_vectors, document

from zhifu.core.errors import BusinessError
from zhifu.modules.agent_runtime.config import AgentSettings
from zhifu.modules.agent_runtime.models import AgentRun, ModelAttempt
from zhifu.modules.agent_runtime.worker import execute_one as execute_agent
from zhifu.modules.memory.models import MemoryJob
from zhifu.modules.memory.worker import execute_one as execute_memory


async def test_actual_model_retrieves_published_document(scenario):  # noqa: F811
    path = await document(scenario)
    app, call = scenario["app"], scenario["call"]
    try:
        response = await call("POST", path + "/publish", "agent", {"expected_version": 1})
        assert response.status_code == 200, response.text
        response = await call(
            "POST",
            "/openapi/v1/agent/runs",
            payload={
                "store_id": str(scenario["store"]),
                "message": "请查询店铺的键盘清洁说明，告诉我能否直接水洗，并标明资料来源。",
            },
        )
        assert response.status_code == 202, response.text
        run_id = UUID(response.json()["run_id"])
        assert await execute_agent(
            app.state.sessions, app.state.state, AgentSettings(), app.state.knowledge.search
        )
        response = await call("GET", f"/openapi/v1/agent/runs/{run_id}")
        assert response.json()["status"] == "COMPLETED", response.text
        answer = response.json()["answer"]
        assert "水洗" in answer and "键盘清洁说明" in answer
        async with app.state.sessions() as session:
            run = await session.get(AgentRun, run_id.bytes)
            calls = [
                call["name"]
                for message in run.messages
                for call in message["data"].get("tool_calls", [])
            ]
            assert "search_knowledge" in calls
    finally:
        await cleanup_vectors(scenario["tenant"])


async def test_actual_langmem_extraction_and_cross_conversation(scenario):  # noqa: F811
    app, call = scenario["app"], scenario["call"]
    response = await call(
        "PUT", MEMORY + "/consent", payload={"expected_revision": 0, "enabled": True}
    )
    assert response.status_code == 200, response.text
    response = await call(
        "POST", "/openapi/v1/conversations", payload={"store_id": str(scenario["store"])}
    )
    conversation = response.json()["conversation_id"]
    response = await call(
        "POST",
        f"/openapi/v1/conversations/{conversation}/messages",
        payload={
            "expected_control_version": 1,
            "text": "以后请一直用英文与我交流，回复尽量简短，通过文字联系我。",
        },
    )
    assert response.status_code == 201, response.text
    assert await execute_agent(app.state.sessions, app.state.state, AgentSettings())
    assert await execute_memory(app.state.sessions, AgentSettings())
    response = await call("GET", MEMORY)
    assert response.status_code == 200, response.text
    content = response.json()
    assert content["content"]["communication"] == {
        "language": "en-US",
        "response_style": "concise",
        "contact_channel": "text",
    }
    response = await call(
        "POST",
        "/openapi/v1/agent/runs",
        payload={
            "store_id": str(scenario["store"]),
            "message": "你好，请和我打个招呼。",
        },
    )
    assert response.status_code == 202, response.text
    run_id = UUID(response.json()["run_id"])
    assert await execute_agent(app.state.sessions, app.state.state, AgentSettings())
    response = await call("GET", f"/openapi/v1/agent/runs/{run_id}")
    answer = response.json()["answer"]
    assert len(answer) > 10 and not any("\u4e00" <= char <= "\u9fff" for char in answer), (
        response.text
    )
    await call("DELETE", MEMORY + f"?expected_revision={content['revision']}")
    async with app.state.sessions() as session:
        runs = (
            await session.scalars(
                select(AgentRun).where(
                    AgentRun.tenant_id == scenario["tenant"].bytes,
                )
            )
        ).all()
        assert all(run.status == "COMPLETED" for run in runs)


async def test_clear_during_actual_extraction(scenario):  # noqa: F811
    app, call = scenario["app"], 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},
    )
    run_id = response.json()["payload"]["run_id"]
    await call("POST", f"/openapi/v1/agent/runs/{run_id}/cancel")
    task = asyncio.create_task(execute_memory(app.state.sessions, AgentSettings()))
    try:
        async with asyncio.timeout(10):
            while True:
                async with app.state.sessions() as session:
                    job = await session.scalar(
                        select(MemoryJob).where(
                            MemoryJob.tenant_id == scenario["tenant"].bytes,
                        )
                    )
                if job.model_name:
                    break
                await asyncio.sleep(0.02)
        response = await call("DELETE", MEMORY + "?expected_revision=1")
        assert response.status_code == 200, response.text
        assert await task
    finally:
        if not task.done():
            task.cancel()
            await asyncio.gather(task, return_exceptions=True)
    assert (await call("GET", MEMORY)).json()["content"] == {}
    async with app.state.sessions() as session:
        job = await session.get(MemoryJob, job.id)
        assert job.status == "CANCELLED" and job.usage


async def test_revoke_during_actual_answer(scenario):  # noqa: F811
    path = await document(scenario)
    app, call = scenario["app"], scenario["call"]
    await call("POST", path + "/publish", "agent", {"expected_version": 1})
    response = await call(
        "POST",
        "/openapi/v1/agent/runs",
        payload={
            "store_id": str(scenario["store"]),
            "message": "请查询店铺资料，详细解释键盘清洁步骤，必须注明资料标题和版本。",
        },
    )
    run_id = UUID(response.json()["run_id"])
    task = asyncio.create_task(execute_agent(app.state.sessions, app.state.state, AgentSettings()))
    try:
        async with asyncio.timeout(20):
            while True:
                async with app.state.sessions() as session:
                    attempts = (
                        await session.scalars(
                            select(ModelAttempt).where(
                                ModelAttempt.run_id == run_id.bytes,
                            )
                        )
                    ).all()
                if len(attempts) >= 2:
                    break
                await asyncio.sleep(0.02)
        response = await call("POST", path + "/revoke", "agent", {"expected_version": 1})
        assert response.status_code == 200, response.text
        with pytest.raises(BusinessError, match="引用资料已变化"):
            await task
        response = await call("GET", f"/openapi/v1/agent/runs/{run_id}")
        assert response.json()["status"] == "FAILED" and response.json()["answer"] is None
    finally:
        if not task.done():
            task.cancel()
            await asyncio.gather(task, return_exceptions=True)
        await cleanup_vectors(scenario["tenant"])