262 lines
10 KiB
Python
262 lines
10 KiB
Python
"""Python role workers for task-card execution."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from pathlib import Path
|
|
from typing import Callable, Protocol, Any
|
|
|
|
from scripts.runtime.roles import RoleDefinition, RuntimeProfile
|
|
from scripts.runtime.skills import SkillRegistry
|
|
from scripts.runtime.tasks import TaskCard, validate_packet
|
|
from scripts.runtime.sources import append_packet_sources
|
|
|
|
|
|
class ChatClient(Protocol):
|
|
def chat_complete(self, **kwargs) -> str:
|
|
...
|
|
|
|
|
|
class SearchProvider(Protocol):
|
|
def search(self, *, query: str, route: str, num_results: int) -> list[dict[str, Any]]:
|
|
...
|
|
|
|
|
|
class ProjectSearchProvider:
|
|
"""Thin adapter over the project-owned search client."""
|
|
|
|
def __init__(self, *, strict_specialized: bool = True) -> None:
|
|
from scripts.lib.search_client import SearchClient
|
|
|
|
self.client = SearchClient(strict_specialized=strict_specialized)
|
|
|
|
def search(self, *, query: str, route: str, num_results: int) -> list[dict[str, Any]]:
|
|
if route == "scholar":
|
|
hits = self.client.scholar(query, num_results=num_results, year_low=2020)
|
|
elif route == "patents":
|
|
hits = self.client.patents(query, num_results=num_results)
|
|
elif route == "news":
|
|
hits = self.client.news(query, num_results=num_results, time_range="y")
|
|
else:
|
|
hits = self.client.search(query, num_results=num_results)
|
|
return [
|
|
{
|
|
"title": hit.title,
|
|
"url": hit.url,
|
|
"snippet": hit.snippet,
|
|
"route": route,
|
|
}
|
|
for hit in hits
|
|
]
|
|
|
|
def close(self) -> None:
|
|
self.client.close()
|
|
|
|
|
|
def _extract_json_object(text: str) -> dict:
|
|
stripped = text.strip()
|
|
if stripped.startswith("```"):
|
|
stripped = stripped.strip("`")
|
|
if stripped.startswith("json"):
|
|
stripped = stripped[4:].strip()
|
|
start = stripped.find("{")
|
|
end = stripped.rfind("}")
|
|
if start == -1 or end == -1 or end < start:
|
|
raise ValueError("worker response does not contain a JSON object")
|
|
return json.loads(stripped[start : end + 1])
|
|
|
|
|
|
def _safe_source_stem(task_id: str) -> str:
|
|
return re.sub(r"[^a-zA-Z0-9]+", "_", task_id).strip("_").lower()
|
|
|
|
|
|
def build_search_context(
|
|
card: TaskCard,
|
|
search_provider: SearchProvider,
|
|
*,
|
|
num_results_per_route: int = 5,
|
|
) -> dict[str, Any]:
|
|
candidate_sources: list[dict[str, Any]] = []
|
|
routes_used: list[str] = []
|
|
source_stem = _safe_source_stem(card.task_id)
|
|
idx = 1
|
|
query = " ".join(card.questions)
|
|
for route in card.search_routes:
|
|
routes_used.append(route)
|
|
hits = search_provider.search(query=query, route=route, num_results=num_results_per_route)
|
|
for hit in hits:
|
|
candidate_sources.append(
|
|
{
|
|
"id": f"src_{source_stem}_{idx:03d}",
|
|
"title": hit.get("title", ""),
|
|
"url": hit.get("url", ""),
|
|
"snippet": hit.get("snippet", ""),
|
|
"route": hit.get("route", route),
|
|
"tier": "Tier 2",
|
|
"score": 6,
|
|
}
|
|
)
|
|
idx += 1
|
|
return {"routes_used": routes_used, "candidate_sources": candidate_sources}
|
|
|
|
|
|
def build_packet_user_prompt(card: TaskCard, search_context: dict[str, Any] | None = None) -> str:
|
|
context = search_context or {"routes_used": [], "candidate_sources": []}
|
|
return (
|
|
"请根据以下 task card 产出一个证据包 JSON。\n"
|
|
"正式结论、summary、open_questions 用中文；英文原文摘录、source title、DOI/URL 可以保留英文。\n"
|
|
"必须主动包含 counter_evidence，且所有引用的 source_id 必须出现在 source_ids 中。\n\n"
|
|
"只能使用 candidate_sources 中的来源，不得编造 URL、DOI、trial ID 或 source_id。\n"
|
|
"输出 JSON 必须包含 sources 字段，且 sources 只能来自 candidate_sources。\n\n"
|
|
f"{json.dumps(card.to_dict(), ensure_ascii=False, indent=2)}\n\n"
|
|
f"Search context:\n{json.dumps(context, ensure_ascii=False, indent=2)}\n\n"
|
|
"只输出 JSON，不要输出 Markdown 解释。"
|
|
)
|
|
|
|
|
|
def build_packet_repair_prompt(
|
|
*,
|
|
card: TaskCard,
|
|
raw_response: str,
|
|
error: Exception,
|
|
search_context: dict[str, Any] | None = None,
|
|
) -> str:
|
|
context = search_context or {"routes_used": [], "candidate_sources": []}
|
|
return (
|
|
"请修复上一次 evidence packet 输出，使其成为合法且通过 schema 校验的 JSON。\n"
|
|
"只输出 JSON 对象，不要输出 Markdown、解释或代码块。\n"
|
|
"保留中文主写作；英文只允许出现在来源标题、URL、DOI、原文摘录或检索笔记中。\n"
|
|
"不得编造 candidate_sources 以外的来源、URL、DOI、trial ID 或 source_id。\n\n"
|
|
f"Schema error:\n{error}\n\n"
|
|
f"Task card:\n{json.dumps(card.to_dict(), ensure_ascii=False, indent=2)}\n\n"
|
|
f"Search context:\n{json.dumps(context, ensure_ascii=False, indent=2)}\n\n"
|
|
f"Previous raw response:\n{raw_response[:12000]}"
|
|
)
|
|
|
|
|
|
class PacketWorker:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
role: RoleDefinition,
|
|
client: ChatClient,
|
|
search_provider: SearchProvider | None = None,
|
|
skill_registry: SkillRegistry | None = None,
|
|
num_results_per_route: int = 5,
|
|
) -> None:
|
|
self.role = role
|
|
self.client = client
|
|
self.search_provider = search_provider
|
|
self.skill_registry = skill_registry or SkillRegistry()
|
|
self.num_results_per_route = num_results_per_route
|
|
|
|
def _system_prompt(self) -> str:
|
|
skill_texts = []
|
|
for name in self.role.skills:
|
|
try:
|
|
skill_texts.append(f"# Skill: {name}\n\n{self.skill_registry.read(name)}")
|
|
except FileNotFoundError:
|
|
skill_texts.append(f"# Skill: {name}\n\n[missing skill: {name}]")
|
|
return (
|
|
"你是 Deep Research v0.20 Python runtime 的证据包 worker。\n"
|
|
"你的唯一任务是把一个 task card 转换为结构化 evidence packet。\n"
|
|
"遵循中文主写作原则；不要写章节正文；不要编造 URL、DOI、trial ID 或 source_id。\n\n"
|
|
"搜索只能走项目 Python search gateway 或调用方提供的 search_context；不要直接使用 Tavily MCP、browser MCP、平台 web search 或任何需要用户权限确认的外部搜索工具。\n\n"
|
|
+ "\n\n".join(skill_texts)
|
|
)
|
|
|
|
def run(self, card: TaskCard) -> dict:
|
|
search_context = None
|
|
if self.search_provider:
|
|
search_context = build_search_context(
|
|
card,
|
|
self.search_provider,
|
|
num_results_per_route=self.num_results_per_route,
|
|
)
|
|
raw = self.client.chat_complete(
|
|
model=self.role.model,
|
|
system=self._system_prompt(),
|
|
user=build_packet_user_prompt(card, search_context),
|
|
temperature=self.role.temperature,
|
|
max_tokens=self.role.max_tokens,
|
|
tag=f"packet:{card.task_id}",
|
|
)
|
|
try:
|
|
packet = _extract_json_object(raw)
|
|
validate_packet(packet)
|
|
return packet
|
|
except Exception as error:
|
|
repaired = self.client.chat_complete(
|
|
model=self.role.model,
|
|
system=self._system_prompt(),
|
|
user=build_packet_repair_prompt(
|
|
card=card,
|
|
raw_response=raw,
|
|
error=error,
|
|
search_context=search_context,
|
|
),
|
|
temperature=0,
|
|
max_tokens=self.role.max_tokens,
|
|
tag=f"packet-repair:{card.task_id}",
|
|
)
|
|
packet = _extract_json_object(repaired)
|
|
validate_packet(packet)
|
|
return packet
|
|
|
|
|
|
def _write_packet_error(project_root: Path, card: TaskCard, error: Exception) -> None:
|
|
path = project_root / "phase2" / "packet_errors" / f"{card.task_id}.json"
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
payload = {
|
|
"task_id": card.task_id,
|
|
"status": "failed",
|
|
"error": str(error),
|
|
"output_packet": card.output_packet,
|
|
}
|
|
path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
|
|
|
|
|
def run_packet_workers(
|
|
*,
|
|
project_root: Path,
|
|
cards: list[TaskCard],
|
|
runtime: RuntimeProfile,
|
|
client_factory: Callable[[RoleDefinition], ChatClient],
|
|
search_provider_factory: Callable[[], SearchProvider] | None = None,
|
|
workers: int,
|
|
) -> int:
|
|
role = runtime.role_for_task("evidence_packet")
|
|
max_workers = max(1, min(workers, role.max_concurrency))
|
|
|
|
def run_one(card: TaskCard) -> tuple[TaskCard, dict | None, Exception | None]:
|
|
search_provider = search_provider_factory() if search_provider_factory else None
|
|
try:
|
|
worker = PacketWorker(role=role, client=client_factory(role), search_provider=search_provider)
|
|
return card, worker.run(card), None
|
|
except Exception as error:
|
|
return card, None, error
|
|
finally:
|
|
close = getattr(search_provider, "close", None)
|
|
if close:
|
|
close()
|
|
|
|
written = 0
|
|
with ThreadPoolExecutor(max_workers=max_workers) as pool:
|
|
futures = [pool.submit(run_one, card) for card in cards]
|
|
for future in as_completed(futures):
|
|
card, packet, error = future.result()
|
|
if error is not None:
|
|
_write_packet_error(project_root, card, error)
|
|
continue
|
|
if packet is None:
|
|
_write_packet_error(project_root, card, RuntimeError("packet worker returned no packet"))
|
|
continue
|
|
path = project_root / card.output_packet
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(json.dumps(packet, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
|
append_packet_sources(project_root / "phase2" / "sources.jsonl", packet)
|
|
written += 1
|
|
return written
|