156 lines
5.7 KiB
Python
156 lines
5.7 KiB
Python
"""快速新对话的首条消息自动命名。
|
||
|
||
这是平台 UI 元数据调用，不进入 agent loop、不改 working_dir。调用前后以
|
||
tasks.auto_title_pending + auto_title_version 为闸：人工 PATCH name 会清闸，清空
|
||
会递增版本，在途标题结果因此既不能覆盖用户命名，也不能跨会话轮次写回。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
from typing import Any, Optional
|
||
from uuid import UUID
|
||
|
||
from sqlalchemy import select, update
|
||
|
||
from .agent_builder import ROOT, load_config
|
||
from .capabilities import ModelCapabilities
|
||
from .llm import LLM
|
||
from .storage import session_scope
|
||
from .storage.models import Task
|
||
from .storage.usage import record_chat_usage
|
||
|
||
|
||
_TITLE_PROMPT = """请给下面这条新对话消息生成一个简短、具体的中文对话标题。
|
||
|
||
规则：
|
||
1. 只输出标题本身，不要解释、引号、前缀或 Markdown。
|
||
2. 6～18 个汉字为宜，必要的英文缩写、材料牌号和数字可以保留。
|
||
3. 概括用户真正要做的事，不要使用“新对话”“咨询问题”“帮我处理”等空泛标题。
|
||
4. 不得包含斜杠或反斜杠。
|
||
|
||
用户消息：
|
||
{content}
|
||
"""
|
||
|
||
_PREFIX_RE = re.compile(r"^(?:标题|对话标题)\s*[:：]\s*", re.I)
|
||
_ATTACHMENT_LINE_RE = re.compile(r"^\[用户上传的(?:参考图|文件)\]\s+\S.*$")
|
||
|
||
|
||
def is_attachment_only_message(content: str) -> bool:
|
||
"""消息是否只包含前端注入的附件路径标记。"""
|
||
lines = [line.strip() for line in (content or "").splitlines() if line.strip()]
|
||
return bool(lines) and all(_ATTACHMENT_LINE_RE.fullmatch(line) for line in lines)
|
||
|
||
|
||
def clean_generated_title(raw: str, user_message: str) -> str:
|
||
"""把模型输出收敛为 validate_task_name 可接受的短标题，失败时按用户消息降级。"""
|
||
first = next((line.strip() for line in (raw or "").splitlines() if line.strip()), "")
|
||
title = _PREFIX_RE.sub("", first).strip(" \t\r\n\"'“”‘’`#*")
|
||
title = re.sub(r"[/\\\x00]+", "·", title)
|
||
title = re.sub(r"\s+", " ", title).strip(" .。,:：;；-—_")
|
||
if not title or title == "新对话":
|
||
fallback = next(
|
||
(line.strip() for line in (user_message or "").splitlines() if line.strip()),
|
||
"新对话",
|
||
)
|
||
title = re.sub(r"[/\\\x00]+", "·", fallback)
|
||
title = re.sub(r"\s+", " ", title).strip(" .。,:：;；-—_")
|
||
return (title or "新对话")[:24]
|
||
|
||
|
||
def generate_task_title(
|
||
*,
|
||
task_id: UUID,
|
||
user_id: UUID,
|
||
user_message: str,
|
||
model_profile: str,
|
||
) -> Optional[str]:
|
||
"""若 task 仍待自动命名，调用一次模型并原子写标题；任何失败均不影响主 run。"""
|
||
with session_scope() as s:
|
||
state = s.execute(
|
||
select(Task.auto_title_pending, Task.auto_title_version).where(
|
||
Task.task_id == task_id,
|
||
Task.user_id == user_id,
|
||
)
|
||
).first()
|
||
if state is None or not state.auto_title_pending:
|
||
return None
|
||
title_version = state.auto_title_version
|
||
|
||
response: Any = None
|
||
title: Optional[str] = None
|
||
caps: Optional[ModelCapabilities] = None
|
||
try:
|
||
cfg = load_config()
|
||
profile = model_profile or cfg["default_model"]
|
||
caps = ModelCapabilities.load(profile, ROOT / cfg["models_dir"])
|
||
response = LLM(caps).chat(
|
||
messages=[{
|
||
"role": "user",
|
||
"content": _TITLE_PROMPT.format(content=user_message[:3000]),
|
||
}],
|
||
tools=None,
|
||
max_retries=2,
|
||
)
|
||
choices = getattr(response, "choices", None) or []
|
||
raw = ((choices[0].message.content if choices else "") or "").strip()
|
||
title = clean_generated_title(raw, user_message)
|
||
except Exception as e:
|
||
print(
|
||
f"[task_title] generate failed task={task_id}: "
|
||
f"{type(e).__name__}: {e}",
|
||
flush=True,
|
||
)
|
||
|
||
# 一次性消费 pending。WHERE pending=true 是与人工 PATCH name 的竞态闸：
|
||
# 用户先改名时 PATCH 已清 false，此处 rowcount=0，不覆盖。
|
||
with session_scope() as s:
|
||
values: dict[str, Any] = {
|
||
"auto_title_pending": False,
|
||
"title_source": "auto",
|
||
}
|
||
if title:
|
||
values["name"] = title
|
||
result = s.execute(
|
||
update(Task)
|
||
.where(
|
||
Task.task_id == task_id,
|
||
Task.user_id == user_id,
|
||
Task.auto_title_pending.is_(True),
|
||
Task.auto_title_version == title_version,
|
||
)
|
||
.values(**values)
|
||
)
|
||
applied = bool(result.rowcount)
|
||
if response is not None and caps is not None:
|
||
usage = getattr(response, "usage", None)
|
||
try:
|
||
record_chat_usage(
|
||
task_id=task_id,
|
||
user_id=user_id,
|
||
message_id=None,
|
||
model_profile=f"{caps.family}.{caps.variant}",
|
||
prompt_tokens=getattr(usage, "prompt_tokens", 0) or 0,
|
||
completion_tokens=getattr(usage, "completion_tokens", 0) or 0,
|
||
input_cny_per_mtoken=caps.input_cny_per_mtoken,
|
||
output_cny_per_mtoken=caps.output_cny_per_mtoken,
|
||
response=response,
|
||
kind="task_title",
|
||
)
|
||
except Exception:
|
||
pass
|
||
return title if applied else None
|
||
|
||
|
||
def generate_task_title_safe(**kwargs: Any) -> Optional[str]:
|
||
"""后台辅助任务入口：DB/配置层异常也必须被吞掉，主对话永不受标题影响。"""
|
||
try:
|
||
return generate_task_title(**kwargs)
|
||
except Exception as e:
|
||
print(
|
||
f"[task_title] auxiliary failed task={kwargs.get('task_id')}: "
|
||
f"{type(e).__name__}: {e}",
|
||
flush=True,
|
||
)
|
||
return None
|