mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 18:41:59 +00:00
ThearchyHelios
244 lines
12 KiB
Python
244 lines
12 KiB
Python
"""按需求挑选：表、能力、工作流版本共用同一份挑选提示词（L3），每种对象只给自己的目录。
|
||
|
||
挑选是一次分类：同一个助手模型，关掉思考、温度 0，PICK_TIMEOUT_S 时限。只在主模型写了 look、只给了需求、
|
||
候选又超过 PICK_LIMIT 时才调；失败（拿不到模型、超时、报错、被截断、回复认不出）不挡生成，由查看服务交回目录、让主模型
|
||
点名。没成的原因照实分开说：到了额度上限被截断、回复是空的、写了字却不是约定的格式，不混成一句「格式不对」。
|
||
|
||
演示模型挑不了：pick_model 照实交回原因，不拿它乱挑。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
import math
|
||
import re
|
||
from typing import Any
|
||
|
||
from app.look.base import PICK_LIMIT, PICK_TIMEOUT_S
|
||
from app.providers.budget import BudgetFallback, budget_refused, cut_off, cut_words, truncated_error
|
||
from app.providers.factory import ModelSpec, get_chat_model
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
#: 挑选调用的输出额度。挑出的十几个名字本身用不了几百 token，可经 OpenAI 兼容接入的推理模型（DeepSeek、GLM、Kimi 这一类）
|
||
#: 收不到 thinking=off，思考也算在额度里：原来的 1024 常常想完就没额度写结果，回复是空的或 JSON 只写了一半，被当成
|
||
#: 「不是约定的格式」，每次都把目录交回主模型。和证据裁判、核查（engine/judge.MAX_TOKENS）同一档
|
||
PICK_MAX_TOKENS = 16384
|
||
#: 服务方嫌 PICK_MAX_TOKENS 太大（400 / 422，或 anthropic SDK 在本地就拒了）时退回这个。被嫌大过的按 (接入, 模型, 额度)
|
||
#: 记在进程里（providers/budget.BudgetFallback），和裁判、核查的记法相同、互相认
|
||
PICK_FALLBACK_MAX_TOKENS = 4096
|
||
#: 从正文里找 JSON 时最多试几个起点（每个 { 或 [ 是一个）
|
||
_MAX_STARTS = 200
|
||
#: 有的网关把思考写在正文里，包在 <think> 里：先去掉，不然思考里打的草稿会被当成结果
|
||
_THINK = re.compile(r"<think>.*?</think>", re.S)
|
||
#: 挑选没成的原因（look 事件的 reason，也交回给模型）
|
||
REASON_FORMAT = "模型的回复不是约定的格式，无法识别挑出的对象"
|
||
REASON_EMPTY = "模型没有写出挑选结果（回复是空的）"
|
||
|
||
PICK_SYSTEM_PROMPT = (
|
||
"你在帮工作流助手挑选它要查看的对象。用户这一轮提了一个需求，助手接下来要据此搭工作流；对象太多，没法每个都把"
|
||
"详情给它看，它只会看到你挑出的那些。请从目录里挑出完成这个需求要用到的。\n\n"
|
||
"挑选规则：\n"
|
||
f"- 最多挑 {PICK_LIMIT} 个，按和需求的相关程度从高到低排\n"
|
||
"- 完成需求要一起用的也挑上：要按渠道、门店、日期之类分组或筛选时，把对应的维度表挑上；两张表要靠一张关联表"
|
||
"才能连起来时，关联表也挑上\n"
|
||
"- 标「当前工作流已用到」的，需求还用得着就挑上\n"
|
||
"- 拿不准要不要的，宁可挑上\n"
|
||
"- name 照抄目录里的名字（含 schema 前缀），不要改写、不要编造；source 写它在目录里所属的那一组的名字\n"
|
||
"- 需求和这份目录无关时返回空列表\n\n"
|
||
'只输出 JSON：{"items": [{"source": "所属的组", "name": "名字"}]}'
|
||
)
|
||
PICK_HUMAN_PROMPT = "需求：{need}\n\n{used}{what}目录（{line}；行尾标「推断，未确认」的是推断出来的）：\n\n{index}"
|
||
PICK_USED_PROMPT = "当前工作流已用到：{items}\n\n"
|
||
PICK_SCHEMA = {
|
||
"title": "picked_items",
|
||
"description": "完成需求要用到的对象，按相关程度从高到低排",
|
||
"type": "object",
|
||
"properties": {
|
||
"items": {
|
||
"type": "array",
|
||
"items": {
|
||
"type": "object",
|
||
"properties": {
|
||
"source": {"type": "string", "description": "它在目录里所属的那一组的名字"},
|
||
"name": {"type": "string", "description": "照抄目录里的名字"},
|
||
},
|
||
"required": ["name"],
|
||
},
|
||
},
|
||
},
|
||
"required": ["items"],
|
||
}
|
||
|
||
|
||
class PickFailed(Exception):
|
||
"""挑选没成：原因写给人看（look 事件的 reason），也交回给模型。"""
|
||
|
||
|
||
async def pick_model(spec: ModelSpec) -> tuple[Any | None, str | None]:
|
||
"""挑选用的模型：(模型, None)，拿不到时 (None, 原因)。spec 是助手的模型规格（copilot_model_spec）。
|
||
|
||
在流里拿：请求的会话那时已经不能用了，这里另开一个短会话读接入设置。没配接入时 get_chat_model 静默退回演示模型，
|
||
演示模型挑不了，直接按挑选失败处理。额度按 PICK_MAX_TOKENS，服务方嫌大时退回 PICK_FALLBACK_MAX_TOKENS（要退回时才建）。
|
||
"""
|
||
from app.core.errors import explain
|
||
from app.db.base import SessionLocal
|
||
from app.providers.mock_model import MockChatModel
|
||
|
||
spec = spec.model_copy(update={"thinking": "off", "max_tokens": PICK_MAX_TOKENS, "temperature": 0,
|
||
"timeout": math.ceil(PICK_TIMEOUT_S)})
|
||
try:
|
||
async with SessionLocal() as session:
|
||
model, model_id = await get_chat_model(session, spec)
|
||
except Exception as e: # noqa: BLE001 - 拿不到模型只是挑不了，生成照常
|
||
return None, f"助手使用的模型不可用：{explain(e)[0]}"
|
||
if isinstance(model, MockChatModel):
|
||
return None, "演示模型不能按需求挑选"
|
||
|
||
async def build() -> Any:
|
||
async with SessionLocal() as session:
|
||
return (await get_chat_model(session, spec.model_copy(update={"max_tokens": PICK_FALLBACK_MAX_TOKENS})))[0]
|
||
|
||
return BudgetFallback(model, build=build, limits=(PICK_MAX_TOKENS, PICK_FALLBACK_MAX_TOKENS),
|
||
rejected=budget_refused, key=(spec.provider, model_id, PICK_MAX_TOKENS)), None
|
||
|
||
|
||
def picked(raw: Any) -> list[tuple[str | None, str]] | None:
|
||
"""模型的回复 → [(组名或 None, 名字)]。结构化输出的对象、正文里的 JSON、字符串列表都认；键名认 items / tables、
|
||
name / table（以前挑表时的写法）。对不上格式返回 None。
|
||
|
||
正文里的 JSON 前后可以有说明、围栏，前面可以有 <think> 包着的思考：逐个起点试，第一个对得上形状的算数——
|
||
{"items": [...]} 一类的对象（空列表也算，是「和需求无关」），或者至少有一个名字的列表（正文里顺手写的 [1] 不算）。"""
|
||
if hasattr(raw, "model_dump"):
|
||
raw = raw.model_dump()
|
||
if not isinstance(raw, str):
|
||
return _items(raw)
|
||
text = _THINK.sub("", raw)
|
||
decoder = json.JSONDecoder()
|
||
starts = 0
|
||
for i, ch in enumerate(text):
|
||
if ch not in "{[":
|
||
continue
|
||
starts += 1
|
||
if starts > _MAX_STARTS:
|
||
break
|
||
try:
|
||
value, _ = decoder.raw_decode(text, i)
|
||
except ValueError:
|
||
continue
|
||
got = _items(value)
|
||
if got is not None and (got or isinstance(value, dict)):
|
||
return got
|
||
return None
|
||
|
||
|
||
def _items(raw: Any) -> list[tuple[str | None, str]] | None:
|
||
if isinstance(raw, dict):
|
||
raw = raw.get("items", raw.get("tables"))
|
||
if not isinstance(raw, list):
|
||
return None
|
||
out: list[tuple[str | None, str]] = []
|
||
for item in raw:
|
||
if isinstance(item, str) and item.strip():
|
||
out.append((None, item.strip()))
|
||
elif isinstance(item, dict):
|
||
name = item.get("name", item.get("table"))
|
||
if isinstance(name, str) and name.strip():
|
||
src = item.get("source")
|
||
out.append((src.strip() if isinstance(src, str) and src.strip() else None, name.strip()))
|
||
return out
|
||
|
||
|
||
def _unpack(out: Any) -> tuple[Any, Any]:
|
||
"""with_structured_output(include_raw=True) 交回 {raw, parsed, parsing_error}；测试的假模型直接交回解析好的对象。"""
|
||
if isinstance(out, dict) and "raw" in out and ("parsed" in out or "parsing_error" in out):
|
||
return out.get("raw"), out.get("parsed")
|
||
return None, out
|
||
|
||
|
||
def _in_raw(raw: Any) -> list[Any]:
|
||
"""结构化输出没解析出来时，原始回复里还可能有结果的地方：工具调用的参数、正文。"""
|
||
from app.engine.state import message_text
|
||
|
||
if raw is None:
|
||
return []
|
||
found: list[Any] = [c.get("args") for c in getattr(raw, "tool_calls", None) or [] if isinstance(c, dict)]
|
||
found += [c.get("args") for c in getattr(raw, "invalid_tool_calls", None) or [] if isinstance(c, dict)]
|
||
return [*found, message_text(raw)]
|
||
|
||
|
||
def _cut(model: Any) -> str:
|
||
return f"{cut_words('按需求挑选的模型', getattr(model, 'max_tokens', None))}，没写完挑选结果"
|
||
|
||
|
||
async def _ask(model: Any, messages: list[tuple[str, str]]) -> list[tuple[str | None, str]]:
|
||
"""问一次。先走结构化输出（带原始回复）：解析出来的不对，再从原始回复的工具调用参数、正文里找；不支持结构化输出
|
||
（有的网关不支持工具调用，一调就报错，纯文本却能用，同 catalog._ask_model）或都没找到，再退回纯文本、从正文里找。
|
||
|
||
到了额度上限被截断的照实说是截断（同证据裁判 engine/judge._ask）：结构化调用截断的不再退回纯文本——思考占满额度时
|
||
纯文本一样截断，白花一次额度；写了一半的清单也不采信，分不清最后一个名字写完了没有。
|
||
"""
|
||
from app.engine.state import message_text
|
||
|
||
try:
|
||
raw, parsed = _unpack(await model.with_structured_output(PICK_SCHEMA, include_raw=True).ainvoke(messages))
|
||
except Exception as e: # noqa: BLE001 - 退回纯文本
|
||
if truncated_error(e):
|
||
raise PickFailed(_cut(model)) from e
|
||
else:
|
||
if cut_off(raw):
|
||
raise PickFailed(_cut(model))
|
||
for found in (parsed, *_in_raw(raw)):
|
||
if (got := picked(found)) is not None:
|
||
return got
|
||
try:
|
||
reply = await model.ainvoke(messages)
|
||
except Exception as e:
|
||
if truncated_error(e):
|
||
raise PickFailed(_cut(model)) from e
|
||
raise
|
||
if cut_off(reply):
|
||
raise PickFailed(_cut(model))
|
||
text = message_text(reply)
|
||
if (got := picked(text)) is not None:
|
||
return got
|
||
logger.warning("按需求挑选：回复认不出（正文 %d 字，开头：%r）", len(text), text[:200])
|
||
raise PickFailed(REASON_FORMAT if text.strip() else REASON_EMPTY)
|
||
|
||
|
||
class Picker:
|
||
"""一轮里共用的挑选：模型第一次要用时才拿（多数轮次根本不挑），拿不到的原因记住，不反复拿。"""
|
||
|
||
def __init__(self, spec: ModelSpec | None) -> None:
|
||
self.spec = spec
|
||
self._model: Any | None = None
|
||
self._unavailable: str | None = None if spec is not None else "没有可用的模型"
|
||
self.calls = 0
|
||
|
||
async def _get(self) -> Any:
|
||
if self._model is None and self._unavailable is None:
|
||
self._model, self._unavailable = await pick_model(self.spec)
|
||
if self._model is None:
|
||
raise PickFailed(self._unavailable or "没有可用的模型")
|
||
return self._model
|
||
|
||
async def pick(self, *, need: str, what: str, line: str, index: str,
|
||
used: str = "") -> list[tuple[str | None, str]]:
|
||
"""按需求从目录里挑。what 是对象的类别名（数据表、能力……），line 说目录一行写的是什么。失败抛 PickFailed。"""
|
||
model = await self._get()
|
||
human = PICK_HUMAN_PROMPT.format(need=need.strip(), used=PICK_USED_PROMPT.format(items=used) if used else "",
|
||
what=what, line=line, index=index)
|
||
self.calls += 1
|
||
try:
|
||
return await asyncio.wait_for(_ask(model, [("system", PICK_SYSTEM_PROMPT), ("human", human)]),
|
||
timeout=PICK_TIMEOUT_S)
|
||
except PickFailed:
|
||
raise
|
||
except asyncio.TimeoutError as e:
|
||
raise PickFailed(f"按需求挑选超过 {PICK_TIMEOUT_S:g} 秒未完成") from e
|
||
except Exception as e: # noqa: BLE001 - 挑选失败不挡生成
|
||
from app.core.errors import explain
|
||
|
||
raise PickFailed(f"按需求挑选的模型调用失败：{explain(e)[0]}") from e
|