332 lines
13 KiB
Python
332 lines
13 KiB
Python
"""对话服务：上下文组装 + LLM 调用 + 附件/标题/推荐短语/重新生成。"""
|
||
import asyncio
|
||
import base64
|
||
import json
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
|
||
from sqlalchemy.orm import Session
|
||
|
||
from ..core import llm
|
||
from ..models import Agent, ChatMessage, ChatSession, FileEntry, User
|
||
from ..services.default_config import DEFAULT_UI_CONFIG
|
||
|
||
|
||
def get_session(db: Session, user: User, session_id: int) -> ChatSession:
|
||
session = db.get(ChatSession, session_id)
|
||
if not session or session.user_id != user.id:
|
||
raise LookupError("会话不存在")
|
||
return session
|
||
|
||
|
||
def get_ui_config(db: Session) -> dict:
|
||
"""读取 UI 配置（数据库覆盖默认值）。"""
|
||
from ..models import AppConfig
|
||
row = db.get(AppConfig, "ui")
|
||
if row and row.value:
|
||
try:
|
||
stored = json.loads(row.value)
|
||
return _deep_merge(DEFAULT_UI_CONFIG, stored)
|
||
except json.JSONDecodeError:
|
||
pass
|
||
return DEFAULT_UI_CONFIG
|
||
|
||
|
||
def _deep_merge(base: dict, override: dict) -> dict:
|
||
out = dict(base)
|
||
for k, v in override.items():
|
||
if k in out and isinstance(out[k], dict) and isinstance(v, dict):
|
||
out[k] = _deep_merge(out[k], v)
|
||
else:
|
||
out[k] = v
|
||
return out
|
||
|
||
|
||
def _load_attachments(db: Session, file_ids: list[int]) -> tuple[list[str], list[str]]:
|
||
"""加载附件：返回 (图片 data URL 列表, 文本内容列表)。"""
|
||
images: list[str] = []
|
||
texts: list[str] = []
|
||
for fid in file_ids or []:
|
||
f = db.get(FileEntry, fid)
|
||
if not f:
|
||
continue
|
||
path = Path(f.stored_path)
|
||
if not path.exists():
|
||
continue
|
||
mime = f.mime or ""
|
||
if mime.startswith("image/"):
|
||
b64 = base64.b64encode(path.read_bytes()).decode()
|
||
images.append(f"data:{mime};base64,{b64}")
|
||
elif mime.startswith("text/") or f.filename.lower().endswith((".txt", ".md", ".csv", ".json", ".log", ".py", ".js", ".html", ".xml", ".yml", ".yaml")):
|
||
try:
|
||
content = path.read_text(encoding="utf-8", errors="replace")
|
||
texts.append(f"【附件: {f.filename}】\n{content[:20000]}")
|
||
except Exception:
|
||
texts.append(f"【附件: {f.filename}】（无法读取）")
|
||
else:
|
||
texts.append(f"【附件: {f.filename}】（不支持的类型，仅图片和文本文件可被读取）")
|
||
return images, texts
|
||
|
||
|
||
def build_messages(db: Session, session: ChatSession, user_content: str,
|
||
file_ids: list[int] | None = None, exclude_last: bool = False) -> list[dict]:
|
||
"""组装发给 LLM 的消息：系统提示词(agent) + 最近历史 + 当前问题（支持图片附件）。"""
|
||
messages: list[dict] = []
|
||
|
||
if session.agent_id:
|
||
agent = db.get(Agent, session.agent_id)
|
||
if agent and agent.system_prompt:
|
||
messages.append({"role": "system", "content": agent.system_prompt})
|
||
|
||
history = (
|
||
db.query(ChatMessage)
|
||
.filter(ChatMessage.session_id == session.id)
|
||
.order_by(ChatMessage.id.desc())
|
||
.limit(20)
|
||
.all()
|
||
)
|
||
# 重新生成时排除最后一条 AI 回答
|
||
if exclude_last and history and history[0].role == "assistant":
|
||
history = history[1:]
|
||
for m in reversed(history):
|
||
messages.append({"role": m.role, "content": m.content})
|
||
|
||
images, texts = _load_attachments(db, file_ids or [])
|
||
if images:
|
||
content: list[dict] = []
|
||
if user_content:
|
||
content.append({"type": "text", "text": user_content})
|
||
for url in images:
|
||
content.append({"type": "image_url", "image_url": {"url": url}})
|
||
messages.append({"role": "user", "content": content})
|
||
else:
|
||
text = user_content
|
||
if texts:
|
||
text = (user_content + "\n\n" if user_content else "") + "\n\n".join(texts)
|
||
messages.append({"role": "user", "content": text})
|
||
return messages
|
||
|
||
|
||
def resolve_model(session: ChatSession, agent: Agent | None, requested: str) -> str:
|
||
if requested:
|
||
return requested
|
||
if session.model:
|
||
return session.model
|
||
if agent and agent.model:
|
||
return agent.model
|
||
return llm.settings.DEFAULT_MODEL
|
||
|
||
|
||
def _json_files(fids: list[int]) -> str:
|
||
return json.dumps(fids or [])
|
||
|
||
|
||
async def _auto_title(db: Session, session_id: int, first_content: str, reply_content: str = ""):
|
||
"""后台任务：首个回答后自动生成会话标题（结合首轮问答内容）。"""
|
||
from ..database import SessionLocal
|
||
sdb = SessionLocal()
|
||
try:
|
||
session = sdb.get(ChatSession, session_id)
|
||
if not session or session.title != "新对话":
|
||
return
|
||
ui = get_ui_config(sdb)
|
||
if not ui.get("title_auto_gen", True):
|
||
return
|
||
model = ui.get("title_model") or "deepseek-chat"
|
||
user_part = first_content[:100] or "(无文字，仅图片)"
|
||
reply_part = (reply_content or "")[:120] or "(无回答)"
|
||
prompt = (
|
||
"根据下面这轮对话，生成一个简洁的对话标题（10 字以内，不要标点、不要引号，直接输出标题）：\n"
|
||
f"用户：{user_part}\n助手：{reply_part}"
|
||
)
|
||
try:
|
||
title = await llm.chat_completion(
|
||
[{"role": "user", "content": prompt}], model=model, max_tokens=30, temperature=0.3
|
||
)
|
||
title = title.strip().strip('"「」').strip().replace("\n", " ")[:30]
|
||
if title:
|
||
session.title = title
|
||
sdb.commit()
|
||
from ..api.ws import push_to_session
|
||
push_to_session(session_id, {"type": "title", "title": title})
|
||
except Exception:
|
||
pass
|
||
finally:
|
||
sdb.close()
|
||
|
||
|
||
async def _gen_suggestions(db: Session, session_id: int, reply_id: int):
|
||
"""后台任务：根据最近对话生成 1-3 条用户可能的追问短语。"""
|
||
from ..database import SessionLocal
|
||
sdb = SessionLocal()
|
||
try:
|
||
reply = sdb.get(ChatMessage, reply_id)
|
||
if not reply:
|
||
return
|
||
ui = get_ui_config(sdb)
|
||
sug_cfg = ui.get("suggestions", {})
|
||
if not sug_cfg.get("enabled", True):
|
||
return
|
||
count = max(1, min(3, int(sug_cfg.get("count", 3))))
|
||
model = sug_cfg.get("model") or "deepseek-chat"
|
||
|
||
history = (
|
||
sdb.query(ChatMessage)
|
||
.filter(ChatMessage.session_id == session_id, ChatMessage.id <= reply_id)
|
||
.order_by(ChatMessage.id.desc())
|
||
.limit(6)
|
||
.all()
|
||
)
|
||
transcript = "\n".join(
|
||
f"{'用户' if m.role == 'user' else '助手'}: {(m.content or '')[:200]}" for m in reversed(history)
|
||
)
|
||
prompt = (
|
||
f"以下是最近一段对话：\n{transcript}\n\n"
|
||
f"请以用户的视角，预测用户接下来最可能追问的 {count} 个短语（每个不超过 12 字，"
|
||
f"自然口语化，不要重复）。只输出 JSON 数组，如：[\"追问1\", \"追问2\"]，不要其他内容。"
|
||
)
|
||
try:
|
||
raw = await llm.chat_completion(
|
||
[{"role": "user", "content": prompt}], model=model, max_tokens=200, temperature=0.8
|
||
)
|
||
raw = raw.strip()
|
||
if raw.startswith("```"):
|
||
raw = raw.strip("`")
|
||
if raw.startswith("json"):
|
||
raw = raw[4:]
|
||
items = json.loads(raw)
|
||
if isinstance(items, list):
|
||
items = [str(x)[:20] for x in items[:3]]
|
||
reply.suggestions = json.dumps(items, ensure_ascii=False)
|
||
sdb.commit()
|
||
from ..api.ws import push_to_session
|
||
push_to_session(session_id, {"type": "suggestions", "message_id": reply_id, "items": items})
|
||
except Exception:
|
||
pass
|
||
finally:
|
||
sdb.close()
|
||
|
||
|
||
async def chat_once(db: Session, user: User, session: ChatSession, content: str,
|
||
agent_id: int | None = None, model: str = "", file_ids: list[int] | None = None) -> ChatMessage:
|
||
"""REST 非流式：调用 LLM 并保存双方消息。"""
|
||
if agent_id is not None:
|
||
session.agent_id = agent_id
|
||
agent = db.get(Agent, session.agent_id) if session.agent_id else None
|
||
session.model = resolve_model(session, agent, model)
|
||
|
||
db.add(ChatMessage(session_id=session.id, role="user", content=content,
|
||
file_ids=_json_files(file_ids)))
|
||
db.commit()
|
||
|
||
messages = build_messages(db, session, content, file_ids)
|
||
reply_text, reply_reasoning = await llm.chat_completion_full(
|
||
messages, model=session.model, temperature=agent.temperature if agent else 0.7)
|
||
|
||
reply = ChatMessage(session_id=session.id, role="assistant", content=reply_text,
|
||
reasoning_content=reply_reasoning, model=session.model)
|
||
db.add(reply)
|
||
session.updated_at = datetime.utcnow()
|
||
db.commit()
|
||
db.refresh(reply)
|
||
|
||
_fire_background_jobs(db, session, content, reply.id, reply_text)
|
||
return reply
|
||
|
||
|
||
async def chat_stream(db: Session, user: User, session: ChatSession, content: str,
|
||
agent_id: int | None = None, model: str = "", file_ids: list[int] | None = None):
|
||
"""流式：先存用户消息，逐段产出回复增量，结束后存完整回复。"""
|
||
if agent_id is not None:
|
||
session.agent_id = agent_id
|
||
agent = db.get(Agent, session.agent_id) if session.agent_id else None
|
||
session.model = resolve_model(session, agent, model)
|
||
|
||
db.add(ChatMessage(session_id=session.id, role="user", content=content,
|
||
file_ids=_json_files(file_ids)))
|
||
db.commit()
|
||
|
||
messages = build_messages(db, session, content, file_ids)
|
||
reply = ChatMessage(session_id=session.id, role="assistant", content="", reasoning_content="", model=session.model)
|
||
db.add(reply)
|
||
db.commit()
|
||
db.refresh(reply)
|
||
|
||
parts: list[str] = []
|
||
reasoning_parts: list[str] = []
|
||
try:
|
||
async for kind, delta in llm.chat_completion_stream(
|
||
messages, model=session.model, temperature=agent.temperature if agent else 0.7
|
||
):
|
||
if kind == "reasoning":
|
||
reasoning_parts.append(delta)
|
||
yield ("reasoning", delta), reply.id
|
||
else:
|
||
parts.append(delta)
|
||
yield ("content", delta), reply.id
|
||
except Exception as e:
|
||
reply.content = "".join(parts) or f"（调用失败：{e}）"
|
||
reply.reasoning_content = "".join(reasoning_parts)
|
||
session.updated_at = datetime.utcnow()
|
||
db.commit()
|
||
raise
|
||
else:
|
||
reply.content = "".join(parts)
|
||
reply.reasoning_content = "".join(reasoning_parts)
|
||
session.updated_at = datetime.utcnow()
|
||
db.commit()
|
||
_fire_background_jobs(db, session, content, reply.id, reply.content)
|
||
|
||
|
||
def _fire_background_jobs(db: Session, session: ChatSession, first_content: str, reply_id: int, reply_content: str = ""):
|
||
"""回答完成后：异步生成标题 + 推荐短语（各自独立会话，不阻塞请求）。"""
|
||
asyncio.create_task(_auto_title(db, session.id, first_content, reply_content))
|
||
asyncio.create_task(_gen_suggestions(db, session.id, reply_id))
|
||
|
||
|
||
async def regenerate_stream(db: Session, session: ChatSession):
|
||
"""重新生成最后一条 AI 回答（流式）：先删除旧回答内容，再流式输出。"""
|
||
last = (
|
||
db.query(ChatMessage)
|
||
.filter(ChatMessage.session_id == session.id, ChatMessage.role == "assistant")
|
||
.order_by(ChatMessage.id.desc())
|
||
.first()
|
||
)
|
||
if not last:
|
||
raise LookupError("没有可重新生成的消息")
|
||
|
||
# 先删除旧回答（内容清空，前端同步清空显示）
|
||
last.content = ""
|
||
last.reasoning_content = ""
|
||
last.suggestions = "[]"
|
||
db.commit()
|
||
|
||
agent = db.get(Agent, session.agent_id) if session.agent_id else None
|
||
messages = build_messages(db, session, "", exclude_last=True)
|
||
|
||
parts: list[str] = []
|
||
reasoning_parts: list[str] = []
|
||
try:
|
||
async for kind, delta in llm.chat_completion_stream(
|
||
messages, model=session.model, temperature=agent.temperature if agent else 0.7
|
||
):
|
||
if kind == "reasoning":
|
||
reasoning_parts.append(delta)
|
||
yield ("reasoning", delta), last.id
|
||
else:
|
||
parts.append(delta)
|
||
yield ("content", delta), last.id
|
||
except Exception as e:
|
||
last.content = "".join(parts) or f"（调用失败：{e}）"
|
||
last.reasoning_content = "".join(reasoning_parts)
|
||
last.regenerated += 1
|
||
db.commit()
|
||
raise
|
||
else:
|
||
last.content = "".join(parts)
|
||
last.reasoning_content = "".join(reasoning_parts)
|
||
last.regenerated += 1
|
||
session.updated_at = datetime.utcnow()
|
||
db.commit()
|
||
asyncio.create_task(_gen_suggestions(db, session.id, last.id))
|