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"])