96 lines
4.2 KiB
Python
96 lines
4.2 KiB
Python
|
import unittest
|
||
|
|
|
||
|
|
from sqlmodel import SQLModel, Session, create_engine
|
||
|
|
|
||
|
|
from app.agents.prompting import (
|
||
|
|
DEFAULT_OUTPUT_SCHEMA,
|
||
|
|
render_prompt,
|
||
|
|
validate_schema,
|
||
|
|
)
|
||
|
|
from app.main import (create_agent, create_prompt_template, delete_prompt_template,
|
||
|
|
list_prompt_templates, update_prompt_template)
|
||
|
|
from app.models import Agent, AgentPromptTemplate
|
||
|
|
|
||
|
|
|
||
|
|
class PromptingTests(unittest.TestCase):
|
||
|
|
def setUp(self):
|
||
|
|
self.engine = create_engine("sqlite://")
|
||
|
|
SQLModel.metadata.create_all(self.engine)
|
||
|
|
self.db = Session(self.engine)
|
||
|
|
|
||
|
|
def tearDown(self):
|
||
|
|
self.db.close()
|
||
|
|
|
||
|
|
def test_render_prompt_keeps_sections_and_replaces_known_slots(self):
|
||
|
|
prompt = render_prompt(
|
||
|
|
{
|
||
|
|
"name": "测试模板",
|
||
|
|
"role_prompt": "你是{agent_name}，专门负责{agent_role}。",
|
||
|
|
"task_template": "标题：{title}\n标签：{content_tag}",
|
||
|
|
"constraints": ["只输出正文"],
|
||
|
|
"examples": [{"input": "旧输入", "output": "旧输出"}],
|
||
|
|
"output_schema": DEFAULT_OUTPUT_SCHEMA,
|
||
|
|
"style": "大白话",
|
||
|
|
"version": 3,
|
||
|
|
},
|
||
|
|
{"agent_name": "心法 Agent", "agent_role": "production", "title": "关系沟通",
|
||
|
|
"content_tag": "心法", "time_range": "00:00-05:00"},
|
||
|
|
rules="只引用已提供材料",
|
||
|
|
knowledge="<source>真实素材</source>",
|
||
|
|
runtime_rules="Runtime 硬规则",
|
||
|
|
)
|
||
|
|
self.assertIn("心法 Agent", prompt)
|
||
|
|
self.assertIn("关系沟通", prompt)
|
||
|
|
self.assertIn("旧输入", prompt)
|
||
|
|
self.assertIn("只引用已提供材料", prompt)
|
||
|
|
self.assertIn("Runtime 硬规则", prompt)
|
||
|
|
self.assertNotIn("{agent_name}", prompt)
|
||
|
|
|
||
|
|
def test_validate_schema_requires_fields_and_types(self):
|
||
|
|
ok, error = validate_schema(
|
||
|
|
{"core_concepts": "沟通", "techniques": "倾听", "products_ref": "无",
|
||
|
|
"golden_lines": "先听再说", "emotion_tone": "温和", "ending_hook": "继续练习"},
|
||
|
|
DEFAULT_OUTPUT_SCHEMA,
|
||
|
|
)
|
||
|
|
self.assertTrue(ok)
|
||
|
|
self.assertEqual(error, "")
|
||
|
|
invalid, reason = validate_schema({"core_concepts": ""}, DEFAULT_OUTPUT_SCHEMA)
|
||
|
|
self.assertFalse(invalid)
|
||
|
|
self.assertIn("缺少必填字段", reason)
|
||
|
|
invalid_type, reason = validate_schema(
|
||
|
|
{"core_concepts": [], "techniques": "倾听", "products_ref": "无", "golden_lines": "先听再说",
|
||
|
|
"emotion_tone": "温和", "ending_hook": "继续练习"}, DEFAULT_OUTPUT_SCHEMA)
|
||
|
|
self.assertFalse(invalid_type)
|
||
|
|
self.assertIn("不是字符串", reason)
|
||
|
|
|
||
|
|
def test_prompt_template_api_versions_and_delete_protection(self):
|
||
|
|
created = create_prompt_template({
|
||
|
|
"name": "可测试模板", "role_prompt": "你是{agent_name}",
|
||
|
|
"task_template": "标题：{title}", "constraints": ["只输出正文"],
|
||
|
|
"examples": [{"input": "A", "output": "B"}],
|
||
|
|
"output_schema": {"type": "object", "required": ["result"]},
|
||
|
|
}, self.db)
|
||
|
|
template = self.db.get(AgentPromptTemplate, created["id"])
|
||
|
|
self.assertIsNotNone(template)
|
||
|
|
self.assertEqual(template.version, 1)
|
||
|
|
updated = update_prompt_template(template.id, {"style": "大白话"}, self.db)
|
||
|
|
self.assertEqual(updated["version"], 2)
|
||
|
|
listed = list_prompt_templates(self.db)
|
||
|
|
self.assertEqual(listed[0]["agent_count"], 0)
|
||
|
|
agent = Agent(name="绑定模板", prompt_template_id=template.id)
|
||
|
|
self.db.add(agent); self.db.commit()
|
||
|
|
with self.assertRaises(Exception):
|
||
|
|
delete_prompt_template(template.id, self.db)
|
||
|
|
|
||
|
|
def test_new_agent_is_bound_to_enabled_default_template(self):
|
||
|
|
default = AgentPromptTemplate(name="默认绑定模板", builtin=True, enabled=True)
|
||
|
|
self.db.add(default); self.db.commit(); self.db.refresh(default)
|
||
|
|
created = create_agent({"name": "自动绑定 Agent", "role": "qc"}, self.db)
|
||
|
|
agent = self.db.get(Agent, created["id"])
|
||
|
|
self.assertIsNotNone(agent)
|
||
|
|
self.assertEqual(agent.prompt_template_id, default.id)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|