Files
yawan-workbench/server/tests/test_planner_versions.py
T

190 lines
8.6 KiB
Python
Raw Normal View History

import shutil
import tempfile
import unittest
from pathlib import Path
from sqlmodel import Session, SQLModel, create_engine, select
from app import planner
from app.agents import runner
from app.rag import dedup
from app.rag.pipeline import iter_lightweight_chunks
from app.models import (
Agent,
Campaign,
CampaignPlanVersion,
KnowledgeBase,
KnowledgeFile,
KnowledgeUsage,
Segment,
)
def parsed_segments():
return {
"status": "ok",
"segments": [
{"title": "开场建立信任", "content_tag": "开场", "time_range": "00:00-00:05",
"mindmap_node": "开场:欢迎与今天主题", "product_ref": ""},
{"title": "核心心法讲解", "content_tag": "心法", "time_range": "00:05-00:10",
"mindmap_node": "主体:核心心法", "product_ref": ""},
],
"pending_confirms": [],
"notes": [],
"provider": "test",
"model": "mock-model",
}
class PlannerVersionTests(unittest.TestCase):
def setUp(self):
self.engine = create_engine("sqlite://")
SQLModel.metadata.create_all(self.engine)
self.db = Session(self.engine)
self.tempdir = Path(tempfile.mkdtemp(prefix="yawan-planner-kb-"))
source_dir = self.tempdir / "kb"
source_dir.mkdir(parents=True, exist_ok=True)
materials = {
"开场素材.txt": (
"开场信任欢迎主题素材。开场要先建立信任,用温和欢迎语说明今天主题。"
"开场阶段必须让听众感到安全、被看见、愿意继续听下去。"
"欢迎语不能夸大承诺,也不能制造焦虑,只能围绕信任和陪伴展开。"
) * 3,
"心法素材.txt": (
"核心心法讲解素材。主体部分解释核心心法,用日常语言说明练习步骤和身体感受。"
"心法讲解要回到呼吸、放松、觉察这些可验证体验,不写素材之外的功效。"
"每个结论都必须能在核心心法素材中找到来源。"
) * 3,
}
for file_name, body in materials.items():
path = source_dir / file_name
path.write_text(body, encoding="utf-8")
# 禁用本地向量:chunk_count 为 0,Planner 必须从登记源文件临时切分并预演。
self.kb = KnowledgeBase(
name="测试知识库", index_status="disabled", file_count=len(materials),
chunk_count=0, source_dir=str(source_dir), organization_id=1,
owner_user_id=0, visibility="global")
self.agent = Agent(name="测试生产 Agent", role="production", status="active",
content_tags=["开场", "心法", "收尾"], kb_ids=[1])
self.db.add_all([self.kb, self.agent])
self.db.commit()
self.db.refresh(self.kb)
self.db.refresh(self.agent)
self.agent.kb_ids = [self.kb.id]
self.db.add(self.agent)
for file_name in materials:
path = source_dir / file_name
self.db.add(KnowledgeFile(
kb_id=self.kb.id, file_name=file_name,
content_hash=dedup.file_hash(path), revision=1))
self.campaign = Campaign(
name="计划版本测试", mindmap_text="开场 -> 心法", objective="完成测试直播",
success_criteria="只使用知识库", duration_minutes=10, status="draft",
owner_user_id=0,
)
self.db.add(self.campaign)
self.db.commit()
self.db.refresh(self.campaign)
self._original_parse = planner.parse_mindmap
planner.parse_mindmap = lambda *args, **kwargs: parsed_segments()
def tearDown(self):
planner.parse_mindmap = self._original_parse
self.db.close()
self.engine.dispose()
shutil.rmtree(self.tempdir, ignore_errors=True)
def test_dag_cycle_is_rejected(self):
with self.assertRaisesRegex(ValueError, "环"):
planner.topo_waves(["n1", "n2"], [{"from": "n1", "to": "n2"}, {"from": "n2", "to": "n1"}])
def test_planner_creates_validated_version_and_approval_materializes_segments(self):
result = planner.run_plan_job(self.db, self.campaign.id)
self.assertEqual(result["status"], "waiting_review")
self.db.refresh(self.campaign)
self.assertEqual(self.campaign.status, "plan_review")
version = self.db.get(CampaignPlanVersion, result["plan_version_id"])
self.assertEqual(version.version, 1)
self.assertEqual(len(version.plan["nodes"]), 2)
self.assertEqual(version.plan["edges"], [{"from": "n1", "to": "n2"}])
self.assertEqual(version.validation["waves"], [["n1"], ["n2"]])
self.assertTrue(version.validation["ok"])
approved = planner.approve_plan(self.db, version.id, reviewer_user_id=1)
self.db.refresh(self.campaign)
self.assertEqual(approved.status, "approved")
self.assertEqual(self.campaign.status, "parsed")
self.assertEqual(self.campaign.approved_plan_version_id, approved.id)
self.assertEqual(self.campaign.canvas_snapshot["plan_version_id"], approved.id)
segments = self.db.exec(select(Segment).where(Segment.campaign_id == self.campaign.id)).all()
self.assertEqual([s.node_id for s in segments], ["n1", "n2"])
self.assertEqual([s.agent_id for s in segments], [self.agent.id, self.agent.id])
self.assertEqual([s.content_tag for s in segments], ["开场", "心法"])
def test_worker_rejects_legacy_campaign_without_approved_plan_in_dsh_mode(self):
self.campaign.status = "parsed"
self.db.add(self.campaign)
self.db.commit()
outcome = runner.run_campaign(self.db, self.campaign.id)
self.assertEqual(outcome["status"], "error")
self.assertIn("人工批准", outcome["error"])
def test_worker_rejects_canvas_that_drifted_from_approved_version(self):
result = planner.run_plan_job(self.db, self.campaign.id)
planner.approve_plan(self.db, result["plan_version_id"], reviewer_user_id=1)
self.db.refresh(self.campaign)
self.campaign.status = "running"
snapshot = dict(self.campaign.canvas_snapshot)
snapshot["plan_version_id"] = 999999
self.campaign.canvas_snapshot = snapshot
self.db.add(self.campaign)
self.db.commit()
outcome = runner.run_campaign(self.db, self.campaign.id)
self.assertEqual(outcome["status"], "error")
self.assertIn("画布快照", outcome["error"])
def test_empty_knowledge_base_fails_validation_and_never_materializes_segments(self):
self.kb.file_count = 0
self.kb.chunk_count = 0
self.db.add(self.kb)
self.db.commit()
result = planner.run_plan_job(self.db, self.campaign.id)
self.assertEqual(result["status"], "failed")
self.db.refresh(self.campaign)
self.assertEqual(self.campaign.status, "plan_failed")
version = self.db.get(CampaignPlanVersion, result["plan_version_id"])
self.assertTrue(version.validation["errors"])
self.assertEqual(self.db.exec(select(Segment)).all(), [])
with self.assertRaisesRegex(ValueError, "不能批准"):
planner.approve_plan(self.db, version.id)
def test_consumed_lightweight_chunks_block_new_plan_before_worker_runs(self):
chunks = iter_lightweight_chunks(self.db, [self.kb.id])
self.assertEqual(len(chunks), 2)
for item in chunks:
self.db.add(KnowledgeUsage(
agent_id=self.agent.id, kb_id=self.kb.id, chunk_id=item["id"],
file_name=item["meta"].get("file", ""), excerpt=item["text"][:300],
campaign_id=999, segment_id=999, status="consumed"))
self.db.commit()
result = planner.run_plan_job(self.db, self.campaign.id)
self.assertEqual(result["status"], "failed")
version = self.db.get(CampaignPlanVersion, result["plan_version_id"])
codes = {item["code"] for item in version.validation["errors"]}
self.assertTrue({"knowledge_exhausted", "knowledge_mismatch"} & codes)
self.assertEqual(self.db.exec(select(Segment)).all(), [])
def test_unassigned_agent_is_blocking_error(self):
self.agent.status = "disabled"
self.db.add(self.agent)
self.db.commit()
result = planner.run_plan_job(self.db, self.campaign.id)
version = self.db.get(CampaignPlanVersion, result["plan_version_id"])
codes = {item["code"] for item in version.validation["errors"]}
self.assertIn("unassigned_agent", codes)
if __name__ == "__main__":
unittest.main()