638 lines
28 KiB
Python
638 lines
28 KiB
Python
"""Messages 路由:历史/目录/发消息起 run/取消/清空/润色/bg proc/SSE。
|
||
|
||
drain 背压与 Web 特有模型门控在这里；run 抢占/调度收口在
|
||
web/run_lifecycle.py，BG worker 本体在 web/runs.py。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import time
|
||
from typing import Any, Optional
|
||
from uuid import UUID
|
||
|
||
from fastapi import Depends, HTTPException
|
||
from fastapi.responses import StreamingResponse
|
||
from sqlalchemy import select, update
|
||
|
||
from core.storage import session_scope
|
||
from core.storage.models import Message, Task, User
|
||
|
||
from ..broker import broker
|
||
from ..common import (
|
||
assert_owns_task,
|
||
iso,
|
||
outline_snippet,
|
||
sse_event,
|
||
)
|
||
from ..model_gate import (
|
||
FALLBACK_MODEL_PROFILE,
|
||
list_image_variants,
|
||
list_video_variants,
|
||
model_allowed_for_user,
|
||
resolve_image_model,
|
||
resolve_model_profile,
|
||
resolve_video_model,
|
||
skill_pinned_profiles,
|
||
)
|
||
from ..run_lifecycle import (
|
||
RunScheduleError,
|
||
RunTaskBusy,
|
||
RunTaskNotFound,
|
||
claim_run_with_message,
|
||
schedule_claimed_run,
|
||
)
|
||
from ..schemas import MessageRequest, OptimizePromptRequest
|
||
from ..userfiles import load_user_root
|
||
|
||
|
||
def _proc_view(m: dict) -> dict:
|
||
"""core.procs.list_* 的条目 → API 形态。elapsed:running 现算,finished 用
|
||
exit_code mtime 定格(前端秒数展示与工具卡同体验)。"""
|
||
from core import procs as _procs
|
||
d = m.get("_dir")
|
||
st, ec = m.get("_status"), m.get("_exit_code")
|
||
created = float(m.get("created_ts") or 0) or None
|
||
now = time.time()
|
||
if st == "finished" and d is not None:
|
||
fin = _procs.finished_at(d) or created or now
|
||
elapsed = max(0.0, fin - (created or fin))
|
||
else:
|
||
elapsed = max(0.0, now - created) if created else 0.0
|
||
return {
|
||
"task_id": m.get("task_id"),
|
||
"proc_id": m.get("proc_id"),
|
||
"kind": m.get("kind"),
|
||
"command": m.get("command") or "",
|
||
"status": st,
|
||
"exit_code": ec,
|
||
"killed": bool(m.get("killed")),
|
||
"created_at": m.get("created_at"),
|
||
"elapsed_s": int(elapsed),
|
||
"timeout_s": m.get("timeout_s"),
|
||
}
|
||
|
||
|
||
def register_message_routes(app, *, require_user) -> None:
|
||
@app.get("/v1/tasks/{task_id}/messages", tags=["messages"])
|
||
def list_messages(
|
||
task_id: str,
|
||
limit: Optional[int] = None,
|
||
before_idx: Optional[int] = None,
|
||
after_idx: Optional[int] = None,
|
||
user_id: UUID = Depends(require_user),
|
||
):
|
||
"""task 历史消息(idx 升序);LiteLLM 原 payload 透传给前端,自行渲染。
|
||
|
||
分页(双向窗口):
|
||
- 不传 limit → 升序全量返回(向后兼容旧前端),两个 has_more 都 false。
|
||
- 传 limit(默认)→ 取**尾部**最近 limit 条(idx desc + limit 再 reverse 回升序)。
|
||
- 传 before_idx → 只取 idx < before_idx 的更早部分(向上翻页)。
|
||
- 传 after_idx → 只取 idx > after_idx 的更新部分(向下翻页;从目录跳到旧消息后用)。
|
||
响应恒含 has_more(窗口之前是否还有更早)+ has_more_after(窗口之后是否还有更新)。
|
||
"""
|
||
try:
|
||
tid = UUID(task_id)
|
||
except ValueError:
|
||
raise HTTPException(404, f"invalid task id: {task_id!r}")
|
||
with session_scope() as s:
|
||
assert_owns_task(s, tid, user_id)
|
||
cols = (
|
||
Message.idx, Message.payload, Message.tokens_in,
|
||
Message.tokens_out, Message.model_profile, Message.created_at,
|
||
Message.artifact_refs,
|
||
)
|
||
if limit is None:
|
||
# 旧行为:升序全量
|
||
rows = s.execute(
|
||
select(*cols).where(Message.task_id == tid).order_by(Message.idx)
|
||
).all()
|
||
elif after_idx is not None:
|
||
# 向下窗口:正序取 idx > after_idx 的最早 limit 条
|
||
rows = list(s.execute(
|
||
select(*cols)
|
||
.where(Message.task_id == tid, Message.idx > after_idx)
|
||
.order_by(Message.idx).limit(limit)
|
||
).all())
|
||
else:
|
||
# 尾部 / 向上窗口:倒序取 limit 条,再翻回升序
|
||
q = select(*cols).where(Message.task_id == tid)
|
||
if before_idx is not None:
|
||
q = q.where(Message.idx < before_idx)
|
||
rows = list(s.execute(q.order_by(Message.idx.desc()).limit(limit)).all())
|
||
rows.reverse()
|
||
# 窗口两端外是否还有(供前端顶/底 sentinel 决定要不要继续补)
|
||
has_more = False
|
||
has_more_after = False
|
||
if rows:
|
||
first_idx = rows[0].idx
|
||
last_idx = rows[-1].idx
|
||
has_more = s.execute(
|
||
select(Message.idx)
|
||
.where(Message.task_id == tid, Message.idx < first_idx)
|
||
.limit(1)
|
||
).first() is not None
|
||
has_more_after = s.execute(
|
||
select(Message.idx)
|
||
.where(Message.task_id == tid, Message.idx > last_idx)
|
||
.limit(1)
|
||
).first() is not None
|
||
return {
|
||
"has_more": has_more,
|
||
"has_more_after": has_more_after,
|
||
"messages": [
|
||
{
|
||
"idx": r.idx,
|
||
"payload": dict(r.payload),
|
||
"tokens_in": r.tokens_in,
|
||
"tokens_out": r.tokens_out,
|
||
"model_profile": r.model_profile, # 0006:assistant 行非空,标产生该 msg 的模型
|
||
"created_at": iso(r.created_at),
|
||
"artifact_refs": r.artifact_refs,
|
||
}
|
||
for r in rows
|
||
]
|
||
}
|
||
|
||
@app.get("/v1/tasks/{task_id}/outline", tags=["messages"])
|
||
def task_outline(task_id: str, user_id: UUID = Depends(require_user)):
|
||
"""消息目录:全部 user 轮次的 {idx, snippet}(idx 升序),供右侧圆点轨道导航。
|
||
|
||
只取 role=user 的 idx + content 首行片段,不回传整 payload(轻量,长任务也快);
|
||
走 (task_id, idx) 索引按 task 收窄,role 过滤为残余条件。前端点圆点 → 已加载则
|
||
scrollIntoView,未加载则用 before_idx 拉居中窗口再定位。
|
||
"""
|
||
try:
|
||
tid = UUID(task_id)
|
||
except ValueError:
|
||
raise HTTPException(404, f"invalid task id: {task_id!r}")
|
||
with session_scope() as s:
|
||
assert_owns_task(s, tid, user_id)
|
||
rows = s.execute(
|
||
select(Message.idx, Message.payload["content"].astext)
|
||
.where(
|
||
Message.task_id == tid,
|
||
Message.payload["role"].astext == "user",
|
||
)
|
||
.order_by(Message.idx)
|
||
).all()
|
||
return {
|
||
"items": [
|
||
{"idx": r[0], "snippet": outline_snippet(r[1])} for r in rows
|
||
]
|
||
}
|
||
|
||
@app.post("/v1/tasks/{task_id}/messages", status_code=202, tags=["messages"])
|
||
async def post_message(
|
||
task_id: str,
|
||
body: MessageRequest,
|
||
user_id: UUID = Depends(require_user),
|
||
):
|
||
"""发消息 + 起 BG run。返 `{events_url}`,客户端立刻订阅 SSE 拿流式。
|
||
|
||
单活 run:`SELECT … FOR UPDATE` 锁 task 行 + 活跃状态检查 + 标 running,
|
||
全收进一个事务挡住"用户连点 send 两条消息"导致两个 BG 线程争 `messages.idx`。
|
||
tasks.run_status in ('running','cancelling') → 409;'error' 走起新 run 时清掉
|
||
(跟 ok / cancelled 一样视为可重启)。
|
||
"""
|
||
try:
|
||
tid = UUID(task_id)
|
||
except ValueError:
|
||
raise HTTPException(404, f"invalid task id: {task_id!r}")
|
||
# 关停 drain 期:拒新 run,带 Retry-After 让客户端退避重试(部署窗口背压)。
|
||
if getattr(app.state, "draining", None) is not None and app.state.draining.is_set():
|
||
raise HTTPException(
|
||
503, "server is restarting; retry shortly",
|
||
headers={"Retry-After": "3"},
|
||
)
|
||
content = (body.content or "").strip()
|
||
if not content:
|
||
raise HTTPException(400, "empty content")
|
||
# 快捷指令展开(与渠道入口共用 core/shortcuts.py):整条精确命中触发词 → 换成完整
|
||
# 指令。在起 run 之前、落库之前展开,模型看到的就是完整指令(不进上下文、不问模型)。
|
||
from core.agent_builder import resolve_workspace as _resolve_ws
|
||
from core import shortcuts as _shortcuts
|
||
_ws = await asyncio.to_thread(_resolve_ws, None)
|
||
content, _sc_hit = await asyncio.to_thread(_shortcuts.expand, _ws, user_id, content)
|
||
if _sc_hit:
|
||
print(f"[shortcut] {str(user_id)[:8]} '{_sc_hit}' expanded")
|
||
# 所有可能返回 4xx 的 variant 校验必须发生在事务写入之前。否则 task 已提交
|
||
# running 后才发现参数非法，会留下一个实际上没有 worker 的假活跃任务。
|
||
image_variant = resolve_image_model(body.image_model, user_id=user_id)
|
||
video_variant = resolve_video_model(body.video_model, user_id=user_id)
|
||
def _prepare_claim(s, task):
|
||
values: dict = {}
|
||
# 档位门控:存量 task 的模型已不在用户档位内(如管理员下调了档位)→ 本次起
|
||
# 持久落回 flash(基线必含),UI 下拉随之显示 flash。不报错、不打断会话历史,
|
||
# 符合"老 task 下次发消息直接切 flash"。当前 task 模型仍在档内则原样不动。
|
||
# plan/role 用同一 session 读(避免在 FOR UPDATE 事务里再开嵌套 session)。
|
||
cur_profile = task.model_profile or ""
|
||
if cur_profile:
|
||
from core.model_access import is_allowed
|
||
urow = s.execute(
|
||
select(User.plan, User.role).where(User.user_id == user_id)
|
||
).first()
|
||
plan = (urow.plan if urow else "") or ""
|
||
role = (urow.role if urow else "user") or "user"
|
||
# skill 定向模型豁免:档外但属内置 skill 定向(建 task 默认 / load_skill
|
||
# 热切写入)→ 不降,产品决策放行。开关关闭时集合为空 → 恢复降级,存量
|
||
# 定向 task 自然落回 flash。
|
||
if not is_allowed(cur_profile, plan, role) and cur_profile not in skill_pinned_profiles():
|
||
fb_profile, fb_model_id = resolve_model_profile(FALLBACK_MODEL_PROFILE)
|
||
values["model_profile"] = fb_profile
|
||
values["model"] = fb_model_id
|
||
return values, {
|
||
"title_profile": values.get("model_profile", cur_profile),
|
||
"should_auto_title": bool(task.auto_title_pending),
|
||
}
|
||
|
||
try:
|
||
claim = claim_run_with_message(
|
||
tid, user_id, content, prepare=_prepare_claim,
|
||
)
|
||
except RunTaskNotFound:
|
||
raise HTTPException(404, f"task not found: {tid}")
|
||
except RunTaskBusy as e:
|
||
raise HTTPException(
|
||
409,
|
||
f"task already has an active run (status={e.status}); "
|
||
f"wait for it to finish or cancel",
|
||
)
|
||
|
||
try:
|
||
schedule_claimed_run(
|
||
app,
|
||
tid,
|
||
user_id,
|
||
content,
|
||
image_variant=image_variant,
|
||
video_variant=video_variant,
|
||
)
|
||
except RunScheduleError:
|
||
raise HTTPException(
|
||
500,
|
||
"message persisted, but background scheduling failed; retry this task",
|
||
)
|
||
title_profile = claim.metadata["title_profile"]
|
||
should_auto_title = claim.metadata["should_auto_title"]
|
||
if should_auto_title:
|
||
from core.task_title import is_attachment_only_message
|
||
should_auto_title = not is_attachment_only_message(content)
|
||
# 快速入口只在首条消息时 pending=true。辅助标题与主 run 并行，不阻塞
|
||
# TTFT/SSE；纯附件消息保留 pending，等下一条自然语言再命名，避免标题暴露
|
||
# `[用户上传的文件]` 等内部标记。generate_task_title 内再次查闸并用条件 UPDATE
|
||
# 防人工改名竞态。
|
||
if should_auto_title:
|
||
from core.task_title import generate_task_title_safe
|
||
title_task = asyncio.create_task(asyncio.to_thread(
|
||
generate_task_title_safe,
|
||
task_id=tid,
|
||
user_id=user_id,
|
||
user_message=content,
|
||
model_profile=title_profile,
|
||
))
|
||
app.state.aux_tasks.add(title_task)
|
||
title_task.add_done_callback(app.state.aux_tasks.discard)
|
||
return {"events_url": f"/v1/tasks/{tid}/events"}
|
||
|
||
@app.post("/v1/tasks/{task_id}/cancel", status_code=202, tags=["tasks"])
|
||
def cancel_task(
|
||
task_id: str,
|
||
user_id: UUID = Depends(require_user),
|
||
):
|
||
"""向当前 task 的活跃 run 发协作式 cancel 信号。
|
||
- 单活 run 形态下"取消当前活动"语义无歧义;客户端只需 task_id
|
||
- 校验 task 归属 user;否则 404
|
||
- tasks.run_status 不是 `running` → 409(idle / cancelling / error 都不能 cancel)
|
||
- 标 `cancelling`(过渡态),BG 线程 loop 在 stream chunk 间 + 工具调用之间 poll 看见即退;
|
||
退出后 finally 写终态(正常→idle,异常→error)
|
||
- LLM 走 streaming,cancel 延迟 ~ 单 chunk 间隔(100ms 级)
|
||
"""
|
||
try:
|
||
tid = UUID(task_id)
|
||
except ValueError:
|
||
raise HTTPException(404, f"invalid task id: {task_id!r}")
|
||
with session_scope() as s:
|
||
row = s.execute(
|
||
select(Task.run_status, Task.title_source)
|
||
.where(Task.task_id == tid, Task.user_id == user_id)
|
||
.with_for_update()
|
||
).first()
|
||
if row is None:
|
||
raise HTTPException(404, f"task not found: {tid}")
|
||
if row.run_status != "running":
|
||
raise HTTPException(
|
||
409,
|
||
f"task not running (run_status={row.run_status}); cannot cancel",
|
||
)
|
||
s.execute(
|
||
update(Task).where(Task.task_id == tid).values(run_status="cancelling")
|
||
)
|
||
broker.request_cancel(tid)
|
||
return {"ok": True, "task_id": str(tid), "run_status": "cancelling"}
|
||
|
||
# ───────────── Background procs(bg proc,DESIGN §8.12)─────────────
|
||
|
||
@app.get("/v1/procs", tags=["tasks"])
|
||
async def list_user_procs(user_id: UUID = Depends(require_user)):
|
||
"""当前用户全部 task 的后台进程(shell/run_python background=true 启动)。
|
||
|
||
用户级而非 task 级:前端全局轮询,切到别的 task 也能收到"后台任务完成"提示。
|
||
纯文件系统读取(user_root/.zcbot_procs),无 DB;docker backend 的 running
|
||
探测走 docker inspect,放 to_thread 防塞 event loop。
|
||
"""
|
||
anchor = load_user_root(user_id)
|
||
from core import procs as _procs
|
||
items = await asyncio.to_thread(_procs.list_all_procs, anchor)
|
||
return {"procs": [_proc_view(m) for m in items]}
|
||
|
||
@app.post("/v1/tasks/{task_id}/procs/{proc_id}/kill", tags=["tasks"])
|
||
async def kill_task_proc(
|
||
task_id: str,
|
||
proc_id: str,
|
||
user_id: UUID = Depends(require_user),
|
||
):
|
||
"""强制终止一个后台进程(host 杀进程树 / docker rm -f 容器)。幂等:已结束返 ok。"""
|
||
try:
|
||
tid = UUID(task_id)
|
||
except ValueError:
|
||
raise HTTPException(404, f"invalid task id: {task_id!r}")
|
||
with session_scope() as s:
|
||
assert_owns_task(s, tid, user_id)
|
||
anchor = load_user_root(user_id)
|
||
from core import procs as _procs
|
||
d = _procs.proc_dir(anchor, str(tid), proc_id)
|
||
if d is None or not d.is_dir():
|
||
raise HTTPException(404, f"background proc not found: {proc_id!r}")
|
||
meta = _procs.read_meta(d)
|
||
if meta is None:
|
||
raise HTTPException(404, f"background proc metadata missing: {proc_id!r}")
|
||
msg = await asyncio.to_thread(_procs.kill_proc, meta, d)
|
||
return {"ok": True, "proc_id": proc_id, "message": msg}
|
||
|
||
@app.post("/v1/tasks/{task_id}/clear", tags=["messages"])
|
||
def clear_messages(task_id: str, user_id: UUID = Depends(require_user)):
|
||
"""清空当前 task 全部 messages,token 累计 / cost / run_error 归零。
|
||
|
||
同 working_dir 下的 FS 文件不动(沿用 task delete 的"FS 视图可重生"心智 —
|
||
中间产物保留,模型重起对话时可继续基于已有素材推进)。
|
||
usage_events 不动:那是用户级账户级用量记账,不该被对话清理影响。
|
||
|
||
- 活跃 run(running / cancelling)期间拒绝:409(先 cancel)
|
||
- error 状态可清:顺手 run_status='idle' + run_error=None
|
||
- 跨 user → 404
|
||
"""
|
||
try:
|
||
tid = UUID(task_id)
|
||
except ValueError:
|
||
raise HTTPException(404, f"invalid task id: {task_id!r}")
|
||
from sqlalchemy import delete as _delete
|
||
from ..common import task_dict
|
||
with session_scope() as s:
|
||
row = s.execute(
|
||
select(Task.run_status, Task.title_source)
|
||
.where(Task.task_id == tid, Task.user_id == user_id)
|
||
.with_for_update()
|
||
).first()
|
||
if row is None:
|
||
raise HTTPException(404, f"task not found: {tid}")
|
||
if row.run_status in ("running", "cancelling"):
|
||
raise HTTPException(
|
||
409,
|
||
f"task has an active run (status={row.run_status}); "
|
||
f"cancel it first",
|
||
)
|
||
s.execute(_delete(Message).where(Message.task_id == tid))
|
||
reset_values: dict[str, Any] = {
|
||
"tokens_prompt": 0,
|
||
"tokens_completion": 0,
|
||
"cost_cny": 0,
|
||
"run_status": "idle",
|
||
"run_error": None,
|
||
"next_message_idx": 0,
|
||
# 全删后 idx 从 0 重起,base 必须归零否则 load 窗口起点悬空(0019);
|
||
# 摘要一并清,否则清空后的对话还会被注入旧前情摘要(0021)
|
||
"context_base_idx": 0,
|
||
"context_summary": None,
|
||
}
|
||
# 自动标题描述的是当前对话主题：清空后回到草稿名，下一条消息重新命名。
|
||
# manual/fixed 分别代表用户显式命名与渠道/调度固定名，均不得覆盖。
|
||
if row.title_source == "auto":
|
||
reset_values.update(
|
||
name="新对话",
|
||
auto_title_pending=True,
|
||
auto_title_version=Task.auto_title_version + 1,
|
||
)
|
||
s.execute(
|
||
update(Task).where(Task.task_id == tid).values(**reset_values)
|
||
)
|
||
task_row = s.execute(select(Task).where(Task.task_id == tid)).scalar_one()
|
||
d = task_dict(task_row, n_messages=0)
|
||
return d
|
||
|
||
@app.post("/v1/tasks/{task_id}/optimize_prompt", tags=["messages"])
|
||
def optimize_prompt(
|
||
task_id: str,
|
||
body: OptimizePromptRequest,
|
||
user_id: UUID = Depends(require_user),
|
||
):
|
||
"""用 task 当前 model 润色用户草稿 prompt;返回优化后的文本。
|
||
|
||
- 同步调用(短文本,3-5s),非 stream
|
||
- 不入 messages 表;**不**累计到 tasks.tokens_prompt/completion(顶栏数字保持
|
||
只反映主对话)。usage_events 单独写一行 kind="prompt_optimize",方便对账
|
||
+ 按 kind GROUP BY 评估"这个按钮值不值"
|
||
- 不与主对话 run 互斥(它不写 messages,无 idx 竞争)— 用户在 LLM 流式
|
||
回复期间也可润色下一条草稿
|
||
- image_model 影响 meta-prompt 里给 LLM 的下游 tool 提示;不动 DB
|
||
"""
|
||
from decimal import Decimal
|
||
from core.agent_builder import load_config
|
||
from core.capabilities import ModelCapabilities
|
||
from core.llm import LLM
|
||
from core.paths import ROOT
|
||
from core.storage.models import UsageEvent
|
||
from core.storage.usage import USD_TO_CNY
|
||
|
||
try:
|
||
tid = UUID(task_id)
|
||
except ValueError:
|
||
raise HTTPException(404, f"invalid task id: {task_id!r}")
|
||
text = (body.text or "").strip()
|
||
if not text:
|
||
raise HTTPException(400, "empty text")
|
||
if len(text) > 4000:
|
||
raise HTTPException(400, "text too long (>4000 chars)")
|
||
|
||
with session_scope() as s:
|
||
row = s.execute(
|
||
select(Task.model_profile)
|
||
.where(Task.task_id == tid, Task.user_id == user_id)
|
||
).first()
|
||
if row is None:
|
||
raise HTTPException(404, f"task not found: {tid}")
|
||
task_model_profile = row.model_profile or ""
|
||
|
||
cfg = load_config()
|
||
chosen_profile = task_model_profile or cfg["default_model"]
|
||
# 档位门控:task 存量模型已不在用户档位内 → 润色也落回 flash(与 send 路径一致,
|
||
# 不持久改 task,仅本次润色调用降级)。
|
||
if chosen_profile and not model_allowed_for_user(chosen_profile, user_id):
|
||
chosen_profile = FALLBACK_MODEL_PROFILE
|
||
try:
|
||
caps = ModelCapabilities.load(chosen_profile, ROOT / cfg["models_dir"])
|
||
except (FileNotFoundError, ValueError) as e:
|
||
raise HTTPException(500, f"invalid task model_profile {chosen_profile!r}: {e}")
|
||
|
||
# 收集下游 tool 上下文:对话模型 display_name + 当前选中 image/video variant 元数据
|
||
chat_model_display = caps.display_name or chosen_profile
|
||
image_variant_hint = ""
|
||
img_variant = (body.image_model or "").strip()
|
||
if img_variant:
|
||
for k, v in list_image_variants():
|
||
if k == img_variant:
|
||
name = v.get("display_name") or k
|
||
sz = v.get("default_size")
|
||
size_seg = f"默认尺寸 {sz}" if sz else "输出尺寸固定不可调"
|
||
image_variant_hint = (
|
||
f"\n下游生图工具:{name}({size_seg},支持中英文 prompt,"
|
||
f"擅长写实/插画/构图描述)。若用户意图涉及画面/封面/插图,"
|
||
f"润色后的文本要给出适合该模型的画面细节(主体/风格/光线/构图)。"
|
||
)
|
||
video_variant_hint = ""
|
||
vid_variant = (body.video_model or "").strip()
|
||
if vid_variant:
|
||
for k, v in list_video_variants():
|
||
if k == vid_variant:
|
||
name = v.get("display_name") or k
|
||
res = v.get("default_resolution") or "720p"
|
||
dur = v.get("default_duration") or 5
|
||
video_variant_hint = (
|
||
f"\n下游生视频工具:{name}(默认 {res} / {dur}s / 16:9)。"
|
||
f"若用户意图涉及视频/动画/动起来,润色后的文本要补全 "
|
||
f"主体在做什么(运动)+ 镜头怎么动 + 场景 + 风格,而非静态画面描述。"
|
||
)
|
||
|
||
meta_prompt = (
|
||
f"你的任务是润色用户输入的草稿,使之成为一个清晰、完整、可执行的 prompt。\n"
|
||
f"当前对话模型:{chat_model_display}。{image_variant_hint}{video_variant_hint}\n\n"
|
||
f"规则:\n"
|
||
f"1. 只输出润色后的文本本身,不要任何解释、前后缀、引号、markdown 代码块包裹\n"
|
||
f"2. 保留用户原始语言(中文/英文)\n"
|
||
f"3. 补全模糊点(主体、目标、风格、约束),但不要无中生有改变用户意图\n"
|
||
f"4. 长度合理 — 简短诉求润色后也应当简洁,不要堆砌\n\n"
|
||
f"用户草稿:\n{text}"
|
||
)
|
||
|
||
llm = LLM(caps)
|
||
try:
|
||
response = llm.chat(
|
||
messages=[{"role": "user", "content": meta_prompt}],
|
||
tools=None,
|
||
)
|
||
except Exception as e:
|
||
raise HTTPException(502, f"llm call failed: {type(e).__name__}: {e}")
|
||
|
||
try:
|
||
optimized = (response.choices[0].message.content or "").strip()
|
||
except Exception:
|
||
raise HTTPException(502, "llm response missing content")
|
||
if not optimized:
|
||
raise HTTPException(502, "llm returned empty optimization")
|
||
|
||
usage = getattr(response, "usage", None)
|
||
prompt_tokens = int(getattr(usage, "prompt_tokens", 0) or 0)
|
||
completion_tokens = int(getattr(usage, "completion_tokens", 0) or 0)
|
||
try:
|
||
from litellm import completion_cost
|
||
cost_usd_raw = completion_cost(completion_response=response)
|
||
cost_usd = Decimal(str(cost_usd_raw)) if cost_usd_raw else Decimal("0")
|
||
except Exception:
|
||
cost_usd = Decimal("0")
|
||
cost_cny = (cost_usd * USD_TO_CNY).quantize(Decimal("0.000001"))
|
||
|
||
try:
|
||
with session_scope() as s:
|
||
s.add(UsageEvent(
|
||
user_id=user_id,
|
||
task_id=tid,
|
||
message_id=None,
|
||
kind="prompt_optimize",
|
||
model_profile=chosen_profile,
|
||
units={
|
||
"tokens_in": prompt_tokens,
|
||
"tokens_out": completion_tokens,
|
||
"usd_to_cny": float(USD_TO_CNY),
|
||
"image_model_hint": img_variant or "",
|
||
"video_model_hint": vid_variant or "",
|
||
},
|
||
cost_cny=cost_cny,
|
||
))
|
||
except Exception as e:
|
||
# 记账失败不阻塞返结果 — 用户拿到润色文本要紧,事后人工补
|
||
print(f"[optimize_prompt] usage record failed: {type(e).__name__}: {e}", flush=True)
|
||
|
||
return {
|
||
"optimized": optimized,
|
||
"model_profile": chosen_profile,
|
||
"tokens_in": prompt_tokens,
|
||
"tokens_out": completion_tokens,
|
||
"cost_cny": float(cost_cny),
|
||
}
|
||
|
||
@app.get("/v1/tasks/{task_id}/events", tags=["tasks"])
|
||
async def stream_events(
|
||
task_id: str,
|
||
user_id: UUID = Depends(require_user),
|
||
):
|
||
"""SSE 流。订阅当前 task 的活动 event(单活 run 形态下无歧义)。
|
||
事件类型:run_start / llm_start / text / reasoning / tool_call /
|
||
tool_result / llm_end / cancelled / error / done。data 是 JSON dict
|
||
(已剔除 `type` 字段,移到 event 名)。
|
||
"""
|
||
try:
|
||
tid = UUID(task_id)
|
||
except ValueError:
|
||
raise HTTPException(404, f"invalid task id: {task_id!r}")
|
||
with session_scope() as s:
|
||
assert_owns_task(s, tid, user_id)
|
||
run_status = s.execute(
|
||
select(Task.run_status).where(Task.task_id == tid)
|
||
).scalar_one()
|
||
|
||
# 重连保护:若 task 不在活跃态(进程重启 / reaper 已收尾 / 自然结束),
|
||
# 直接吐 done 关流。否则 broker 进程内队列空,客户端会无限挂在 ping 上。
|
||
is_active = run_status in ("running", "cancelling")
|
||
|
||
async def gen():
|
||
yield b": connected\nretry: 3000\n\n"
|
||
if not is_active:
|
||
yield sse_event("done", {})
|
||
return
|
||
q = broker.subscribe(tid)
|
||
try:
|
||
while True:
|
||
try:
|
||
ev = await asyncio.wait_for(q.get(), timeout=30.0)
|
||
except asyncio.TimeoutError:
|
||
yield b": ping\n\n"
|
||
continue
|
||
ev_type = ev.get("type", "msg")
|
||
payload = {k: v for k, v in ev.items() if k != "type"}
|
||
yield sse_event(ev_type, payload)
|
||
if ev_type in ("done", "error"):
|
||
break
|
||
except asyncio.CancelledError:
|
||
pass # 客户端断开,静默退
|
||
finally:
|
||
broker.unsubscribe(tid, q)
|
||
|
||
return StreamingResponse(
|
||
gen(),
|
||
media_type="text/event-stream",
|
||
headers={
|
||
"Cache-Control": "no-cache",
|
||
"Connection": "keep-alive",
|
||
"X-Accel-Buffering": "no",
|
||
},
|
||
)
|