import json import unittest from types import SimpleNamespace from unittest.mock import AsyncMock, patch import httpx from fastapi import HTTPException from db.models import ModelResource, TestRun from models import ( AssistantConfig, RuntimeModelResource, RuntimeTool, RuntimeToolBinding, ) from routes.test_runs import create_prompt_improvement from services.test_runs.prompt_improvement import ( PROMPT_IMPROVEMENT_SNAPSHOT_KEY, PromptImprovementAgent, PromptImprovementError, _GeneratedImprovement, build_prompt_improvement_context, build_prompt_improvement_snapshot, public_test_run_config, ) from test_schemas import PromptImprovementOut def evaluator_resource() -> RuntimeModelResource: return RuntimeModelResource( id="model_eval", name="测试评估模型", capability="LLM", interface_type="openai-llm", values={ "apiUrl": "https://model.example/v1", "modelId": "judge-model", }, secrets={"apiKey": "model-secret"}, ) def failed_result() -> dict: return { "startup": { "completed": True, "greeting": "您好", "assistantOutputs": ["请问需要什么帮助？"], "runtimeRecords": [{"privateNoise": "must-not-be-forwarded"}], }, "turns": [ { "id": "turn_1", "index": 0, "userInput": "查询订单", "assistantReply": "请提供订单号", "toolCalls": [ { "functionName": "query_order", "argumentsJson": '{"orderId":"A1"}', "resultJson": '{"status":"ok"}', "outcome": "success", } ], "runtimeRecords": [{"privateNoise": "must-not-be-forwarded"}], "evaluations": [ { "label": "回复 · 关键词", "status": "fail", "expected": "应包含：查询成功", "actual": "请提供订单号", "reason": "没有告知查询结果", } ], } ], "overallCriteria": [], } def missing_tool_result() -> dict: return { "startup": {"completed": True}, "turns": [ { "id": "turn_photo", "index": 0, "userInput": "1", "assistantReply": "现在开始拍照，请拍车辆前方。", "toolCalls": [], "evaluations": [ { "id": "behavior_photo", "kind": "tool_call", "label": "工具 · set_photo_button_visible", "status": "fail", "expected": "应调用 set_photo_button_visible 1 次以上", "actual": "实际调用 0 次 set_photo_button_visible", "reason": "工具调用次数不符合要求。", } ], } ], "overallCriteria": [], } def missing_tool_case() -> dict: return { "name": "进入拍照阶段", "turns": [ { "id": "turn_photo", "userInput": "1", "behaviors": [ { "id": "behavior_photo", "type": "tool_call", "toolId": "tool_photo", "functionName": "set_photo_button_visible", "expectation": "called", "minCalls": 1, "maxCalls": None, "paramAssertions": [ { "name": "visible", "matchMode": "exact", "value": "true", } ], } ], } ], } def missing_tool_snapshot() -> dict: return { "assistantType": "prompt", "runtimeMode": "pipeline", "prompt": "信息确认后立即调用 `set_photo_button_visible(true)`。", "greeting": "您好", "dynamicVariableNames": [], "tools": [ { "toolId": "tool_photo", "functionName": "set_photo_button_visible", "type": "client", "parameters": [{"name": "visible", "type": "boolean"}], } ], } class PromptImprovementSnapshotTests(unittest.TestCase): def test_snapshot_keeps_model_visible_tool_settings_without_secrets(self): config = AssistantConfig( type="prompt", runtimeMode="pipeline", prompt="服务 {{customer_name}}", greeting="您好", dynamic_variable_definitions={"customer_name": {"type": "string"}}, startup={ "openingMode": "confirmation", "entryMode": "wait_user", "openingMessage": { "title": "提示", "message": "请确认", "confirmLabel": "确认", }, "actions": [ { "phase": "opening", "required": True, "toolBinding": { "toolId": "tool_order", "argumentBindings": { "token": { "source": "literal", "value": "literal-secret", } }, }, } ], }, tools=[ RuntimeTool( id="tool_order", name="查询订单", function_name="query_order", type="http", description="按订单号查询", definition={ "type": "http", "config": { "url": "https://private.example/orders", "headers": {"Authorization": "header-secret"}, "parameters": [ { "name": "order_id", "type": "string", "location": "body", "description": "订单号", "required": True, } ], }, }, secrets={"token": "tool-secret"}, ) ], tool_bindings=[ RuntimeToolBinding( tool_id="tool_order", argument_bindings={ "order_id": {"source": "model"}, "token": {"source": "literal", "value": "binding-secret"}, }, result_assignments={"status": "order_status"}, ) ], ) snapshot = build_prompt_improvement_snapshot(config) serialized = json.dumps(snapshot, ensure_ascii=False) self.assertEqual(snapshot["prompt"], "服务 {{customer_name}}") self.assertEqual(snapshot["tools"][0]["functionName"], "query_order") self.assertEqual(snapshot["tools"][0]["parameters"][0]["name"], "order_id") self.assertEqual( snapshot["tools"][0]["binding"]["argumentSources"]["token"], {"source": "literal"}, ) for secret in ( "private.example", "header-secret", "tool-secret", "literal-secret", "binding-secret", ): self.assertNotIn(secret, serialized) def test_internal_snapshot_is_hidden_from_public_run_config(self): public = public_test_run_config( { "timeoutSecs": 60, PROMPT_IMPROVEMENT_SNAPSHOT_KEY: {"prompt": "secret context"}, } ) self.assertEqual(public, {"timeoutSecs": 60}) def test_context_contains_failure_trajectory_but_not_runtime_records(self): context = build_prompt_improvement_context( { "assistantType": "prompt", "runtimeMode": "pipeline", "prompt": "你是订单助手", "greeting": "您好", "startup": {}, "tools": [], "dynamicVariableNames": [], }, { "name": "订单查询", "description": "查询已存在订单", "contextTurns": [{"role": "user", "content": "之前的消息"}], "turns": [{"id": "turn_1", "userInput": "查询订单"}], }, failed_result(), ) self.assertIn("你是订单助手", context) self.assertIn("没有告知查询结果", context) self.assertIn("query_order", context) self.assertNotIn("must-not-be-forwarded", context) def test_context_keeps_structured_missing_tool_assertion(self): context = json.loads( build_prompt_improvement_context( missing_tool_snapshot(), missing_tool_case(), missing_tool_result(), ) ) failure = context["failedEvaluations"][0] self.assertEqual(failure["kind"], "tool_call") self.assertEqual(failure["turnIndex"], 1) self.assertEqual(failure["actualCallCount"], 0) self.assertEqual( failure["expectedBehavior"]["functionName"], "set_photo_button_visible", ) self.assertEqual( failure["expectedBehavior"]["paramAssertions"][0], {"name": "visible", "matchMode": "exact", "value": "true"}, ) class PromptImprovementAgentTests(unittest.IsolatedAsyncioTestCase): async def test_valid_change_preserves_existing_variable(self): agent = PromptImprovementAgent(evaluator_resource()) generated = _GeneratedImprovement.model_validate( { "diagnosis": "缺少查询后的回复要求", "promptChange": { "after": "服务 {{customer_name}}。查询完成后说明结果。", "reason": "补充回复规则", "evidence": ["实际回复没有查询结果"], }, "manualRecommendations": [], } ) with patch.object(agent, "_request", AsyncMock(return_value=generated)): result = await agent.improve( snapshot={ "prompt": "服务 {{customer_name}}。", "greeting": "您好", "dynamicVariableNames": ["customer_name"], "tools": [], }, case_snapshot={"name": "订单", "turns": []}, result=failed_result(), ) self.assertEqual(result.prompt_change.before, "服务 {{customer_name}}。") self.assertIn("说明结果", result.prompt_change.after) async def test_missing_tool_noop_is_retried_with_focused_instruction(self): agent = PromptImprovementAgent(evaluator_resource()) before = missing_tool_snapshot()["prompt"] noop = _GeneratedImprovement.model_validate( { "diagnosis": "Prompt 已经要求调用工具", "promptChange": { "after": before, "reason": "无需修改", "evidence": ["已有调用要求"], }, } ) repaired = _GeneratedImprovement.model_validate( { "diagnosis": "助手用文字进入下一阶段，漏掉实际工具调用", "promptChange": { "after": ( f"{before}\n调用成功前不得输出拍照提示，" "不得用自然语言代替工具调用。" ), "reason": "将工具调用设为进入下一阶段的前置门槛", "evidence": ["实际调用 0 次 set_photo_button_visible"], }, } ) request = AsyncMock(side_effect=[noop, repaired]) with patch.object(agent, "_request", request): result = await agent.improve( snapshot=missing_tool_snapshot(), case_snapshot=missing_tool_case(), result=missing_tool_result(), ) self.assertEqual(request.await_count, 2) retry_instruction = request.await_args_list[1].kwargs["retry_instruction"] self.assertIn("实际调用 0 次", retry_instruction) self.assertIn("调用成功前", result.prompt_change.after) async def test_missing_tool_uses_safe_fallback_after_two_noop_results(self): agent = PromptImprovementAgent(evaluator_resource()) before = missing_tool_snapshot()["prompt"] noop = _GeneratedImprovement.model_validate( { "diagnosis": "Prompt 已包含工具调用要求", "promptChange": { "after": before, "reason": "保持原文", "evidence": ["已有调用要求"], }, } ) with patch.object( agent, "_request", AsyncMock(side_effect=[noop, noop]), ): result = await agent.improve( snapshot=missing_tool_snapshot(), case_snapshot=missing_tool_case(), result=missing_tool_result(), ) self.assertIn("# 工具调用执行约束", result.prompt_change.after) self.assertIn("不得用自然语言", result.prompt_change.after) self.assertIn("没有实际调用", result.diagnosis) self.assertEqual( result.prompt_change.evidence, [ "应调用 set_photo_button_visible 1 次以上；" "实际调用 0 次 set_photo_button_visible；" "工具调用次数不符合要求。" ], ) async def test_unknown_variable_is_rejected(self): agent = PromptImprovementAgent(evaluator_resource()) generated = _GeneratedImprovement.model_validate( { "diagnosis": "需要修改", "promptChange": { "after": "服务 {{unknown_name}}。", "reason": "错误建议", "evidence": [], }, } ) with patch.object(agent, "_request", AsyncMock(return_value=generated)): with self.assertRaisesRegex(PromptImprovementError, "未知动态变量"): await agent.improve( snapshot={ "prompt": "服务用户。", "greeting": "您好", "dynamicVariableNames": [], "tools": [], }, case_snapshot={"name": "订单", "turns": []}, result=failed_result(), ) async def test_non_prompt_root_cause_can_return_manual_recommendation(self): agent = PromptImprovementAgent(evaluator_resource()) generated = _GeneratedImprovement.model_validate( { "diagnosis": "工具描述与测试预期不一致", "promptChange": None, "manualRecommendations": [ { "category": "tool", "title": "检查工具描述", "detail": "明确工具的适用条件。", } ], } ) with patch.object(agent, "_request", AsyncMock(return_value=generated)): result = await agent.improve( snapshot={ "prompt": "服务用户。", "greeting": "您好", "dynamicVariableNames": [], "tools": [], }, case_snapshot={"name": "订单", "turns": []}, result=failed_result(), ) self.assertIsNone(result.prompt_change) self.assertEqual(result.manual_recommendations[0].category, "tool") async def test_model_timeout_is_reported_as_gateway_timeout(self): agent = PromptImprovementAgent(evaluator_resource()) fake_client = AsyncMock() fake_client.__aenter__.return_value = fake_client fake_client.post.side_effect = httpx.ReadTimeout("timeout") with patch( "services.test_runs.prompt_improvement.httpx.AsyncClient", return_value=fake_client, ): with self.assertRaises(PromptImprovementError) as raised: await agent._request("{}") self.assertEqual(raised.exception.status_code, 504) class FakeScalarResult: def __init__(self, values): self._values = values def scalars(self): return self def all(self): return self._values class FakePromptImprovementSession: def __init__(self, run, run_case, evaluator): self.run = run self.run_case = run_case self.evaluator = evaluator async def get(self, model, identity): if model is TestRun and identity == self.run.id: return self.run if model is ModelResource and identity == self.evaluator.id: return self.evaluator return None async def execute(self, _query): return FakeScalarResult([self.run_case]) class PromptImprovementRouteTests(unittest.IsolatedAsyncioTestCase): def make_session(self, *, status="fail", run_status="completed"): snapshot = { "assistantType": "prompt", "runtimeMode": "pipeline", "prompt": "你是订单助手", "dynamicVariableNames": [], } run = SimpleNamespace( id="run_1", status=run_status, config={ "evaluatorModelResourceId": "model_eval", PROMPT_IMPROVEMENT_SNAPSHOT_KEY: snapshot, }, ) run_case = SimpleNamespace( status=status, position=0, case_snapshot={"name": "订单查询", "turns": []}, result=failed_result(), ) evaluator = SimpleNamespace( id="model_eval", name="测试评估模型", capability="LLM", interface_type="openai-llm", values={"apiUrl": "https://model.example/v1", "modelId": "judge"}, secrets={"apiKey": "secret"}, support_image_input=False, enabled=True, ) return FakePromptImprovementSession(run, run_case, evaluator) async def test_route_returns_generated_result_without_writing_assistant(self): expected = PromptImprovementOut.model_validate( { "diagnosis": "缺少结果说明", "promptChange": { "before": "你是订单助手", "after": "你是订单助手。查询后说明结果。", "reason": "补充规则", "evidence": [], }, "evaluatorName": "测试评估模型", } ) session = self.make_session() with patch.object( PromptImprovementAgent, "improve", AsyncMock(return_value=expected), ): result = await create_prompt_improvement("run_1", session=session) self.assertEqual(result, expected) async def test_route_rejects_execution_error(self): session = self.make_session(status="error") with self.assertRaises(HTTPException) as raised: await create_prompt_improvement("run_1", session=session) self.assertEqual(raised.exception.status_code, 409) if __name__ == "__main__": unittest.main()