841 lines
28 KiB
Python
841 lines
28 KiB
Python
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import re
|
||
import sys
|
||
from copy import deepcopy
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from common.io_utils import append_jsonl, read_json, read_yaml, write_json
|
||
from common.logging_utils import now_iso
|
||
from common.path_utils import resolve_cli_path, resolve_from_file
|
||
from common.schemas import make_ai_message, validate_ai_message, validate_messages,validate_execution_plan
|
||
|
||
|
||
PARSE_ERROR_CONTENT = "模型输出解析失败，无法生成有效工具调用或最终回答。"
|
||
_MODEL_CACHE: dict[tuple[str, ...], tuple[Any, Any]] = {}
|
||
|
||
|
||
def _load_model_config(model_config: str | Path) -> tuple[Path, dict]:
|
||
path = Path(model_config).resolve()
|
||
config = read_yaml(path)
|
||
if not isinstance(config, dict):
|
||
raise ValueError("model.yaml must contain an object")
|
||
return path, config
|
||
|
||
|
||
def _artifact_paths(artifact_dir: str | Path, stem: str | None) -> tuple[Path, Path, Path]:
|
||
directory = Path(artifact_dir)
|
||
prefix = f"{stem}_" if stem else ""
|
||
return (
|
||
directory / f"{prefix}raw_model_output.json",
|
||
directory / f"{prefix}ai_message.json",
|
||
directory / "llm_run_log.jsonl",
|
||
)
|
||
|
||
|
||
def _extract_tool_result(message: dict) -> dict:
|
||
try:
|
||
result = json.loads(message["content"])
|
||
except (KeyError, json.JSONDecodeError, TypeError) as exc:
|
||
raise ValueError("ToolMessage content is not a SkillResult JSON string") from exc
|
||
if not isinstance(result, dict):
|
||
raise ValueError("ToolMessage content must decode to an object")
|
||
return result
|
||
|
||
|
||
def _three_points(text: str) -> list[str]:
|
||
parts = [part.strip(" \t\r\n。") for part in re.split(r"\n+|(?<=[。！？!?])", text) if part.strip()]
|
||
points = []
|
||
for part in parts:
|
||
if part not in points:
|
||
points.append(part)
|
||
if len(points) == 3:
|
||
break
|
||
while len(points) < 3:
|
||
points.append("工具结果未提供更多可提取内容")
|
||
return points
|
||
|
||
|
||
def _mock_generate(messages: list[dict]) -> dict:
|
||
tool_messages = [message for message in messages if message.get("role") == "tool"]
|
||
if not tool_messages:
|
||
return make_ai_message(
|
||
"",
|
||
[
|
||
{
|
||
"id": "call_001",
|
||
"name": "file_reader",
|
||
"args": {"path": "docs/agent_intro.txt", "max_chars": 2000},
|
||
}
|
||
],
|
||
)
|
||
latest = tool_messages[-1]
|
||
result = _extract_tool_result(latest)
|
||
if latest.get("status") != "success" or result.get("status") != "success":
|
||
error = result.get("error") or {}
|
||
detail = error.get("message", "未知工具错误") if isinstance(error, dict) else str(error)
|
||
return make_ai_message(f"工具调用失败，无法完成请求：{detail}", [])
|
||
output = result.get("output") or {}
|
||
content = output.get("content") if isinstance(output, dict) else None
|
||
if not isinstance(content, str) or not content.strip():
|
||
content = json.dumps(output, ensure_ascii=False)
|
||
points = _three_points(content)
|
||
answer = "三条中文要点如下：\n" + "\n".join(f"{index}. {point}" for index, point in enumerate(points, 1))
|
||
return make_ai_message(answer, [])
|
||
|
||
|
||
def _parse_tool_calls_fragment(raw_text: str, original_error: json.JSONDecodeError) -> dict:
|
||
markers = ['"tool_calls":[', '\\"tool_calls\\":[']
|
||
marker_index = -1
|
||
marker = ""
|
||
for item in markers:
|
||
marker_index = raw_text.find(item)
|
||
if marker_index != -1:
|
||
marker = item
|
||
break
|
||
if marker_index == -1:
|
||
raise original_error
|
||
array_start = marker_index + marker.index("[")
|
||
array_end = raw_text.rfind("]")
|
||
if array_end < array_start:
|
||
raise ValueError("model output contains tool_calls marker but no closing array")
|
||
array_text = raw_text[array_start : array_end + 1]
|
||
try:
|
||
tool_calls = json.loads(array_text)
|
||
except json.JSONDecodeError:
|
||
tool_calls = json.loads(array_text.replace('\\"', '"'))
|
||
if not isinstance(tool_calls, list) or not tool_calls:
|
||
raise original_error
|
||
return {"content": "", "tool_calls": tool_calls}
|
||
|
||
|
||
def _parse_json_with_backtick_tail(raw_text: str, original_error: json.JSONDecodeError) -> dict:
|
||
text = raw_text.strip()
|
||
try:
|
||
candidate, end_index = json.JSONDecoder().raw_decode(text)
|
||
except json.JSONDecodeError:
|
||
raise original_error
|
||
trailing = text[end_index:].strip()
|
||
if trailing and set(trailing) <= {"`"}:
|
||
return candidate
|
||
raise original_error
|
||
|
||
|
||
def _parse_execution_plan(raw_text: str) -> dict:
|
||
"""
|
||
将模型生成的 JSON 文本解析为轻量执行计划。
|
||
|
||
计划不是 AIMessage，不包含 content 或 tool_calls，
|
||
因此不能复用 _parse_model_output。
|
||
"""
|
||
try:
|
||
plan = json.loads(raw_text.strip())
|
||
except json.JSONDecodeError as exc:
|
||
plan = _parse_json_with_backtick_tail(raw_text, exc)
|
||
|
||
validate_execution_plan(plan)
|
||
return plan
|
||
|
||
|
||
def _candidate_to_message(candidate: dict) -> tuple[dict, dict]:
|
||
if not isinstance(candidate, dict):
|
||
raise ValueError("model output JSON must be an object")
|
||
expected_keys = {"content", "tool_calls"}
|
||
unknown_keys = set(candidate) - expected_keys
|
||
if unknown_keys:
|
||
raise ValueError(f"model output JSON contains unknown keys: {', '.join(sorted(unknown_keys))}")
|
||
message = {
|
||
"role": "assistant",
|
||
"content": candidate.get("content", ""),
|
||
"tool_calls": candidate.get("tool_calls", []),
|
||
}
|
||
validate_ai_message(message)
|
||
has_content = bool(message["content"].strip())
|
||
has_tool_calls = bool(message["tool_calls"])
|
||
if has_content == has_tool_calls:
|
||
raise ValueError("model output must contain either final content or tool calls, but not both")
|
||
parsed_candidate = {"content": message["content"], "tool_calls": message["tool_calls"]}
|
||
return parsed_candidate, message
|
||
|
||
|
||
def _parse_model_output(raw_text: str) -> tuple[dict, dict]:
|
||
try:
|
||
candidate = json.loads(raw_text.strip())
|
||
except json.JSONDecodeError as exc:
|
||
try:
|
||
candidate = _parse_json_with_backtick_tail(raw_text, exc)
|
||
except json.JSONDecodeError:
|
||
candidate = _parse_tool_calls_fragment(raw_text, exc)
|
||
return _candidate_to_message(candidate)
|
||
|
||
|
||
def _dtype_value(torch_module: Any, configured: str) -> Any:
|
||
if configured == "auto":
|
||
return "auto"
|
||
mapping = {
|
||
"bfloat16": torch_module.bfloat16,
|
||
"float16": torch_module.float16,
|
||
"float32": torch_module.float32,
|
||
}
|
||
if configured not in mapping:
|
||
raise ValueError(f"unsupported torch_dtype: {configured}")
|
||
return mapping[configured]
|
||
|
||
|
||
def _model_cache_key(
|
||
model_path: Path,
|
||
tokenizer_path: Path,
|
||
local_only: bool,
|
||
trust_remote_code: bool,
|
||
dtype: Any,
|
||
device_map: Any,
|
||
max_memory: Any,
|
||
) -> tuple[str, ...]:
|
||
try:
|
||
device_map_key = json.dumps(device_map, sort_keys=True, separators=(",", ":"))
|
||
except TypeError:
|
||
device_map_key = repr(device_map)
|
||
try:
|
||
max_memory_key = json.dumps(max_memory, sort_keys=True, separators=(",", ":"))
|
||
except TypeError:
|
||
max_memory_key = repr(max_memory)
|
||
return (
|
||
str(model_path),
|
||
str(tokenizer_path),
|
||
str(local_only),
|
||
str(trust_remote_code),
|
||
str(dtype),
|
||
device_map_key,
|
||
max_memory_key,
|
||
)
|
||
|
||
|
||
def _load_model_bundle(
|
||
auto_model: Any,
|
||
auto_tokenizer: Any,
|
||
model_path: Path,
|
||
tokenizer_path: Path,
|
||
local_only: bool,
|
||
trust_remote_code: bool,
|
||
dtype: Any,
|
||
device_map: Any,
|
||
max_memory: Any,
|
||
) -> tuple[Any, Any]:
|
||
cache_key = _model_cache_key(
|
||
model_path,
|
||
tokenizer_path,
|
||
local_only,
|
||
trust_remote_code,
|
||
dtype,
|
||
device_map,
|
||
max_memory,
|
||
)
|
||
cached = _MODEL_CACHE.get(cache_key)
|
||
if cached is not None:
|
||
print("model_cache=hit", file=sys.stderr, flush=True)
|
||
return cached
|
||
|
||
print("model_cache=miss", file=sys.stderr, flush=True)
|
||
tokenizer = auto_tokenizer.from_pretrained(
|
||
str(tokenizer_path),
|
||
local_files_only=local_only,
|
||
trust_remote_code=trust_remote_code,
|
||
)
|
||
model = auto_model.from_pretrained(
|
||
str(model_path),
|
||
local_files_only=local_only,
|
||
trust_remote_code=trust_remote_code,
|
||
dtype=dtype,
|
||
device_map=device_map,
|
||
max_memory=max_memory,
|
||
)
|
||
_MODEL_CACHE[cache_key] = (tokenizer, model)
|
||
return tokenizer, model
|
||
|
||
|
||
def _build_prompt_messages(messages: list[dict], tools_schema: list[dict],include_tools_schema: bool = True) -> list[dict]:
|
||
prompt_messages = deepcopy(messages)
|
||
format_instruction = (
|
||
"IMPORTANT OUTPUT FORMAT:\n"
|
||
"You must return exactly one valid JSON object.\n"
|
||
"Do not output markdown.\n"
|
||
"Do not output explanations.\n"
|
||
"Do not output code fences or backticks.\n"
|
||
'The first output character must be "{" and the last output character must be "}".\n\n'
|
||
"Valid schema A:\n"
|
||
'{"content":"final answer text","tool_calls":[]}\n\n'
|
||
"Valid schema B:\n"
|
||
'{"content":"","tool_calls":[{"id":"call_001","name":"file_reader",'
|
||
'"args":{"path":"docs/agent_intro.txt","max_chars":2000}}]}\n\n'
|
||
"The top-level keys must be exactly:\n"
|
||
"- content: string\n"
|
||
"- tool_calls: array\n\n"
|
||
"Never put tool_calls inside content.\n"
|
||
'Never output {"content":"tool_calls": ...}.'
|
||
)
|
||
envelope_reminder = (
|
||
"IMPORTANT OUTPUT FORMAT: Output the JSON object now. "
|
||
'Your first output character must be "{" and your last output character must be "}". '
|
||
"Never output a backtick, Markdown, a code block, an explanation, or text outside the JSON. "
|
||
'Use exactly the top-level keys "content" (string) and "tool_calls" (array). '
|
||
"Choose exactly one schema: final content with an empty tool_calls array, or empty content with tool calls. "
|
||
'Never put tool_calls inside content. Never output {"content":"tool_calls": ...}.'
|
||
)
|
||
|
||
system_instruction = "\n\n" + format_instruction
|
||
|
||
if include_tools_schema:
|
||
system_instruction = (
|
||
"\n\nAvailable tools JSON schema:\n"
|
||
+ json.dumps(tools_schema, ensure_ascii=False)
|
||
+ "\n"
|
||
+ format_instruction
|
||
)
|
||
|
||
if prompt_messages and prompt_messages[0].get("role") == "system":
|
||
prompt_messages[0]["content"] += system_instruction
|
||
else:
|
||
prompt_messages.insert(0, {"role": "system", "content": system_instruction.strip()})
|
||
|
||
for message in reversed(prompt_messages):
|
||
if message.get("role") == "user":
|
||
message["content"] += "\n\n" + envelope_reminder
|
||
break
|
||
if prompt_messages[-1].get("role") == "tool":
|
||
prompt_messages.append(
|
||
{
|
||
"role": "user",
|
||
"content": (
|
||
envelope_reminder
|
||
+ " The latest ToolMessage already contains a tool result. If it provides the requested "
|
||
'information, answer with schema A now and set "tool_calls" to exactly []. Do not repeat the '
|
||
"completed tool call."
|
||
),
|
||
}
|
||
)
|
||
return prompt_messages
|
||
|
||
|
||
def _prompt_json_generate(config_path: Path, config: dict, messages: list[dict], tools_schema: list[dict]) -> str:
|
||
try:
|
||
import torch
|
||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
except ImportError as exc:
|
||
raise RuntimeError("prompt_json mode requires requirements-llm.txt") from exc
|
||
model_config = config.get("model", {})
|
||
generation_config = config.get("generation", {})
|
||
model_setting = model_config.get("model_name_or_path")
|
||
tokenizer_setting = model_config.get("tokenizer_name_or_path", model_setting)
|
||
if not isinstance(model_setting, str) or not isinstance(tokenizer_setting, str):
|
||
raise ValueError("model_name_or_path and tokenizer_name_or_path are required")
|
||
model_path = resolve_from_file(model_setting, config_path)
|
||
tokenizer_path = resolve_from_file(tokenizer_setting, config_path)
|
||
if not model_path.exists() or not tokenizer_path.exists():
|
||
raise FileNotFoundError(f"local model path does not exist: {model_path}")
|
||
local_only = bool(model_config.get("local_files_only", True))
|
||
trust_remote_code = bool(model_config.get("trust_remote_code", False))
|
||
dtype = _dtype_value(torch, str(model_config.get("torch_dtype", "auto")))
|
||
tokenizer, model = _load_model_bundle(
|
||
AutoModelForCausalLM,
|
||
AutoTokenizer,
|
||
model_path,
|
||
tokenizer_path,
|
||
local_only,
|
||
trust_remote_code,
|
||
dtype,
|
||
model_config.get("device_map", "auto"),
|
||
model_config.get("max_memory"),
|
||
)
|
||
tool_calling_config = config.get("tool_calling", {})
|
||
schema_injection = tool_calling_config.get("schema_injection", "prompt")
|
||
|
||
if schema_injection not in {"prompt", "native"}:
|
||
raise ValueError(
|
||
"tool_calling.schema_injection must be prompt or native"
|
||
)
|
||
|
||
template_kwargs = {}
|
||
|
||
if schema_injection == "prompt":
|
||
# 原有方案：schema 转成文字后拼接到 system prompt。
|
||
prompt_messages = _build_prompt_messages(
|
||
messages,
|
||
tools_schema,
|
||
include_tools_schema=True,
|
||
)
|
||
|
||
else:
|
||
# 新方案：schema 不写入 prompt，通过 tokenizer 的 tools 参数注入。
|
||
prompt_messages = _build_prompt_messages(
|
||
messages,
|
||
tools_schema,
|
||
include_tools_schema=False,
|
||
)
|
||
|
||
# 内部临时移除B3 的 x-returns 字段，
|
||
native_tools = deepcopy(tools_schema)
|
||
|
||
for tool in native_tools:
|
||
function = tool.get("function")
|
||
if isinstance(function, dict):
|
||
function.pop("x-returns", None)
|
||
|
||
template_kwargs["tools"] = native_tools
|
||
|
||
inputs = tokenizer.apply_chat_template(
|
||
prompt_messages,
|
||
tokenize=True,
|
||
add_generation_prompt=True,
|
||
return_tensors="pt",
|
||
return_dict=True,
|
||
enable_thinking=False,
|
||
**template_kwargs,
|
||
)
|
||
device = next(model.parameters()).device
|
||
inputs = inputs.to(device)
|
||
input_length = inputs["input_ids"].shape[-1]
|
||
options = {
|
||
"max_new_tokens": int(generation_config.get("max_new_tokens", 1024)),
|
||
"do_sample": bool(generation_config.get("do_sample", False)),
|
||
}
|
||
with torch.no_grad():
|
||
generated = model.generate(**inputs, **options)
|
||
new_tokens = generated[0][input_length:]
|
||
return tokenizer.decode(new_tokens, skip_special_tokens=True)
|
||
|
||
|
||
def _build_plan_prompt_messages(
|
||
messages: list[dict],
|
||
tools_schema: list[dict],
|
||
) -> list[dict]:
|
||
"""
|
||
为复杂任务生成轻量执行计划。
|
||
这里只决定是否需要计划及计划步骤；
|
||
不生成 tool_calls，不执行工具。
|
||
"""
|
||
prompt_messages = deepcopy(messages)
|
||
|
||
plan_instruction = (
|
||
"\n\nPLAN MODE\n"
|
||
"Return exactly one JSON object and nothing else.\n"
|
||
"Choose exactly one mode:\n"
|
||
'DIRECT: {"use_plan":false,"steps":[]}\n'
|
||
'PLAN: {"use_plan":true,"steps":[...]}.\n\n'
|
||
"Use DIRECT when the task needs at most one tool call.\n"
|
||
"Use PLAN only when the task needs two or more dependent tool calls.\n\n"
|
||
"Available tools JSON schema:\n"
|
||
+ json.dumps(tools_schema, ensure_ascii=False)
|
||
+ "\n\n"
|
||
"For PLAN mode, each step must contain "
|
||
'"step_id", "goal", and "suggested_tool".\n'
|
||
"step_id starts at 1 and increases by 1. "
|
||
"suggested_tool must be an available tool name or null. "
|
||
"The final answer step uses null. Do not include tool arguments.\n\n"
|
||
"Example direct output:\n"
|
||
'{"use_plan":false,"steps":[]}\n\n'
|
||
"Example plan output:\n"
|
||
'{"use_plan":true,"steps":['
|
||
'{"step_id":1,"goal":"search related documents","suggested_tool":"local_file_search"},'
|
||
'{"step_id":2,"goal":"read the most relevant document","suggested_tool":"file_reader"},'
|
||
'{"step_id":3,"goal":"answer based on tool results","suggested_tool":null}'
|
||
"]}\n\n"
|
||
"Final check: if use_plan is false, steps must be exactly []."
|
||
)
|
||
|
||
if prompt_messages and prompt_messages[0].get("role") == "system":
|
||
prompt_messages[0]["content"] += plan_instruction
|
||
else:
|
||
prompt_messages.insert(
|
||
0,
|
||
{
|
||
"role": "system",
|
||
"content": plan_instruction.strip(),
|
||
},
|
||
)
|
||
|
||
for message in reversed(prompt_messages):
|
||
if message.get("role") == "user":
|
||
message["content"] += (
|
||
"\n\nReturn the execution-plan JSON now. "
|
||
"Do not answer the task itself yet."
|
||
)
|
||
break
|
||
|
||
return prompt_messages
|
||
|
||
|
||
def _plan_prompt_json_generate(
|
||
config_path: Path,
|
||
config: dict,
|
||
messages: list[dict],
|
||
tools_schema: list[dict],
|
||
) -> str:
|
||
"""
|
||
使用与 prompt_json 相同的模型、配置和完整 tools_schema，
|
||
但生成内容是 execution plan，而不是普通 AIMessage。
|
||
"""
|
||
try:
|
||
import torch
|
||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
except ImportError as exc:
|
||
raise RuntimeError("prompt_json mode requires requirements-llm.txt") from exc
|
||
|
||
model_config = config.get("model", {})
|
||
generation_config = config.get("generation", {})
|
||
|
||
model_setting = model_config.get("model_name_or_path")
|
||
tokenizer_setting = model_config.get("tokenizer_name_or_path", model_setting)
|
||
|
||
if not isinstance(model_setting, str) or not isinstance(tokenizer_setting, str):
|
||
raise ValueError("model_name_or_path and tokenizer_name_or_path are required")
|
||
|
||
model_path = resolve_from_file(model_setting, config_path)
|
||
tokenizer_path = resolve_from_file(tokenizer_setting, config_path)
|
||
|
||
if not model_path.exists() or not tokenizer_path.exists():
|
||
raise FileNotFoundError(f"local model path does not exist: {model_path}")
|
||
|
||
local_only = bool(model_config.get("local_files_only", True))
|
||
trust_remote_code = bool(model_config.get("trust_remote_code", False))
|
||
dtype = _dtype_value(torch, str(model_config.get("torch_dtype", "auto")))
|
||
|
||
tokenizer, model = _load_model_bundle(
|
||
AutoModelForCausalLM,
|
||
AutoTokenizer,
|
||
model_path,
|
||
tokenizer_path,
|
||
local_only,
|
||
trust_remote_code,
|
||
dtype,
|
||
model_config.get("device_map", "auto"),
|
||
model_config.get("max_memory"),
|
||
)
|
||
|
||
# 这里仍然传入完整 tools_schema，只是提示词换成 Plan 提示词。
|
||
prompt_messages = _build_plan_prompt_messages(messages, tools_schema)
|
||
|
||
inputs = tokenizer.apply_chat_template(
|
||
prompt_messages,
|
||
tokenize=True,
|
||
add_generation_prompt=True,
|
||
return_tensors="pt",
|
||
return_dict=True,
|
||
enable_thinking=False,
|
||
)
|
||
|
||
device = next(model.parameters()).device
|
||
inputs = inputs.to(device)
|
||
|
||
input_length = inputs["input_ids"].shape[-1]
|
||
|
||
options = {
|
||
"max_new_tokens": int(generation_config.get("max_new_tokens", 1024)),
|
||
"do_sample": bool(generation_config.get("do_sample", False)),
|
||
}
|
||
|
||
with torch.no_grad():
|
||
generated = model.generate(**inputs, **options)
|
||
|
||
new_tokens = generated[0][input_length:]
|
||
|
||
return tokenizer.decode(new_tokens, skip_special_tokens=True)
|
||
|
||
|
||
def generate_execution_plan(
|
||
model_config: str,
|
||
messages: list[dict],
|
||
tools_schema: list[dict],
|
||
mode: str = "prompt_json",
|
||
artifact_dir: str | None = None,
|
||
artifact_stem: str | None = None,
|
||
) -> dict:
|
||
"""
|
||
生成轻量执行计划。
|
||
|
||
返回的 execution_plan 只描述是否需要多步执行及建议工具，
|
||
不直接产生 AIMessage，也不执行工具。
|
||
"""
|
||
config_path, config = _load_model_config(model_config)
|
||
messages = validate_messages(deepcopy(messages))
|
||
|
||
if not isinstance(tools_schema, list):
|
||
raise ValueError("tools_schema must be an array")
|
||
|
||
generated_at = now_iso()
|
||
backend = (
|
||
"mock"
|
||
if mode == "mock"
|
||
else config.get("model", {}).get("backend", "transformers")
|
||
)
|
||
|
||
if mode == "mock":
|
||
# mock 仅用于验证接口、解析和产物保存。
|
||
# 不虚构复杂计划，真实计划判断交给 prompt_json 模式。
|
||
execution_plan = {
|
||
"use_plan": False,
|
||
"steps": [],
|
||
}
|
||
raw_text = json.dumps(execution_plan, ensure_ascii=False)
|
||
status = "success"
|
||
error = None
|
||
|
||
elif mode == "prompt_json":
|
||
raw_text = _plan_prompt_json_generate(
|
||
config_path,
|
||
config,
|
||
messages,
|
||
tools_schema,
|
||
)
|
||
|
||
try:
|
||
execution_plan = _parse_execution_plan(raw_text)
|
||
status = "success"
|
||
error = None
|
||
except Exception as exc:
|
||
# 解析失败时不伪造 use_plan=false，
|
||
# 让调用方明确知道这是计划生成失败。
|
||
execution_plan = None
|
||
status = "error"
|
||
error = {
|
||
"type": type(exc).__name__,
|
||
"message": str(exc),
|
||
}
|
||
|
||
else:
|
||
raise ValueError("mode must be mock or prompt_json")
|
||
|
||
raw_record = {
|
||
"kind": "execution_plan",
|
||
"mode": mode,
|
||
"backend": backend,
|
||
"raw_text": raw_text,
|
||
"execution_plan": execution_plan,
|
||
"status": status,
|
||
"error": error,
|
||
"generated_at": generated_at,
|
||
}
|
||
|
||
if artifact_dir:
|
||
directory = Path(artifact_dir)
|
||
prefix = f"{artifact_stem}_" if artifact_stem else ""
|
||
|
||
raw_path = directory / f"{prefix}plan_raw_model_output.json"
|
||
plan_path = directory / f"{prefix}execution_plan.json"
|
||
log_path = directory / "llm_run_log.jsonl"
|
||
|
||
write_json(raw_record, raw_path)
|
||
|
||
if execution_plan is not None:
|
||
write_json(execution_plan, plan_path)
|
||
|
||
append_jsonl(
|
||
{
|
||
"timestamp": generated_at,
|
||
"kind": "execution_plan",
|
||
"mode": mode,
|
||
"status": status,
|
||
"raw_output_path": str(raw_path),
|
||
"execution_plan_path": (
|
||
str(plan_path) if execution_plan is not None else None
|
||
),
|
||
"error": error,
|
||
},
|
||
log_path,
|
||
)
|
||
|
||
return {
|
||
"execution_plan": execution_plan,
|
||
"status": status,
|
||
"error": error,
|
||
}
|
||
|
||
|
||
def generate_agent_decision(
|
||
model_config: str,
|
||
messages: list[dict],
|
||
tools_schema: list[dict],
|
||
mode: str = "prompt_json",
|
||
artifact_dir: str | None = None,
|
||
artifact_stem: str | None = None,
|
||
) -> dict:
|
||
"""
|
||
由模型自主判断任务是否需要多步计划。
|
||
|
||
简单任务：
|
||
execution_mode = "direct"
|
||
直接返回普通 AIMessage。
|
||
|
||
复杂任务：
|
||
execution_mode = "plan"
|
||
返回 execution_plan，后续由 B1 逐步执行。
|
||
"""
|
||
plan_result = generate_execution_plan(
|
||
model_config=model_config,
|
||
messages=messages,
|
||
tools_schema=tools_schema,
|
||
mode=mode,
|
||
artifact_dir=artifact_dir,
|
||
artifact_stem=artifact_stem,
|
||
)
|
||
|
||
if plan_result["status"] != "success":
|
||
return {
|
||
"execution_mode": None,
|
||
"execution_plan": None,
|
||
"ai_message": None,
|
||
"status": "error",
|
||
"error": plan_result["error"],
|
||
}
|
||
|
||
execution_plan = plan_result["execution_plan"]
|
||
|
||
if execution_plan is None:
|
||
return {
|
||
"execution_mode": None,
|
||
"execution_plan": None,
|
||
"ai_message": None,
|
||
"status": "error",
|
||
"error": {
|
||
"type": "RuntimeError",
|
||
"message": "execution plan is missing after successful plan generation",
|
||
},
|
||
}
|
||
|
||
# 模型认为是复杂任务：暂不生成普通 AIMessage。
|
||
# 后续 B1 根据计划步骤推进工具调用。
|
||
if execution_plan["use_plan"]:
|
||
return {
|
||
"execution_mode": "plan",
|
||
"execution_plan": execution_plan,
|
||
"ai_message": None,
|
||
"status": "success",
|
||
"error": None,
|
||
}
|
||
|
||
# 模型认为是简单任务：沿用原来的直接工具调用逻辑。
|
||
ai_result = generate_ai_message(
|
||
model_config=model_config,
|
||
messages=messages,
|
||
tools_schema=tools_schema,
|
||
mode=mode,
|
||
artifact_dir=artifact_dir,
|
||
artifact_stem=artifact_stem,
|
||
)
|
||
|
||
return {
|
||
"execution_mode": "direct",
|
||
"execution_plan": execution_plan,
|
||
"ai_message": ai_result["ai_message"],
|
||
"status": ai_result["status"],
|
||
"error": ai_result["error"],
|
||
}
|
||
|
||
|
||
def generate_ai_message(
|
||
model_config: str,
|
||
messages: list[dict],
|
||
tools_schema: list[dict],
|
||
mode: str = "prompt_json",
|
||
artifact_dir: str | None = None,
|
||
artifact_stem: str | None = None,
|
||
) -> dict:
|
||
config_path, config = _load_model_config(model_config)
|
||
messages = validate_messages(deepcopy(messages))
|
||
if not isinstance(tools_schema, list):
|
||
raise ValueError("tools_schema must be an array")
|
||
generated_at = now_iso()
|
||
backend = "mock" if mode == "mock" else config.get("model", {}).get("backend", "transformers")
|
||
if mode == "mock":
|
||
ai_message = _mock_generate(messages)
|
||
raw_text = json.dumps({"content": ai_message["content"], "tool_calls": ai_message["tool_calls"]}, ensure_ascii=False)
|
||
parsed_candidate = {"content": ai_message["content"], "tool_calls": ai_message["tool_calls"]}
|
||
status = "success"
|
||
error = None
|
||
elif mode == "prompt_json":
|
||
raw_text = _prompt_json_generate(config_path, config, messages, tools_schema)
|
||
try:
|
||
parsed_candidate, ai_message = _parse_model_output(raw_text)
|
||
status = "success"
|
||
error = None
|
||
except Exception as exc:
|
||
parsed_candidate = None
|
||
ai_message = make_ai_message(PARSE_ERROR_CONTENT, [])
|
||
status = "error"
|
||
error = {"type": type(exc).__name__, "message": str(exc)}
|
||
else:
|
||
raise ValueError("mode must be mock or prompt_json")
|
||
raw_record = {
|
||
"mode": mode,
|
||
"backend": backend,
|
||
"raw_text": raw_text,
|
||
"parsed_candidate": parsed_candidate,
|
||
"status": status,
|
||
"error": error,
|
||
"generated_at": generated_at,
|
||
}
|
||
if artifact_dir:
|
||
raw_path, message_path, log_path = _artifact_paths(artifact_dir, artifact_stem)
|
||
write_json(raw_record, raw_path)
|
||
write_json(ai_message, message_path)
|
||
append_jsonl(
|
||
{
|
||
"timestamp": generated_at,
|
||
"mode": mode,
|
||
"status": status,
|
||
"raw_output_path": str(raw_path),
|
||
"ai_message_path": str(message_path),
|
||
"error": error,
|
||
},
|
||
log_path,
|
||
)
|
||
return {
|
||
"ai_message": ai_message,
|
||
"status": status,
|
||
"error": error,
|
||
}
|
||
|
||
|
||
def build_parser() -> argparse.ArgumentParser:
|
||
parser = argparse.ArgumentParser(description="Generate one AIMessage with a local or mock LLM.")
|
||
parser.add_argument("--model_config", required=True)
|
||
parser.add_argument("--messages", required=True)
|
||
parser.add_argument("--tools_schema", required=True)
|
||
parser.add_argument("--mode", choices=["mock", "prompt_json"], required=True)
|
||
parser.add_argument("--outdir", required=True)
|
||
return parser
|
||
|
||
|
||
def main(argv: list[str] | None = None) -> int:
|
||
args = build_parser().parse_args(argv)
|
||
try:
|
||
outdir = resolve_cli_path(args.outdir)
|
||
decision = generate_agent_decision(
|
||
model_config=str(resolve_cli_path(args.model_config)),
|
||
messages=read_json(resolve_cli_path(args.messages)),
|
||
tools_schema=read_json(resolve_cli_path(args.tools_schema)),
|
||
mode=args.mode,
|
||
artifact_dir=str(outdir),
|
||
)
|
||
|
||
decision_path = outdir / "agent_decision.json"
|
||
write_json(decision, decision_path)
|
||
|
||
if decision["status"] != "success":
|
||
error = decision["error"] or {}
|
||
raise RuntimeError(
|
||
error.get("message", "agent decision generation failed")
|
||
)
|
||
|
||
print(decision_path)
|
||
return 0
|
||
except Exception as exc:
|
||
print(f"fatal: {type(exc).__name__}: {exc}", file=sys.stderr)
|
||
return 1
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|