dekun
123 lines
4.0 KiB
Python
123 lines
4.0 KiB
Python
"""中控 AI：单会话聊天（直到用户点击新开）。"""
|
||
from __future__ import annotations
|
||
|
||
from typing import Any, Optional
|
||
|
||
from hub_ai.attachments import parse_chat_attachments
|
||
from hub_ai.client import generate_text, model_label
|
||
from hub_ai.config import (
|
||
CHAT_CONTEXT_MAX_CHARS,
|
||
CHAT_MAX_HISTORY_TURNS,
|
||
CHAT_MAX_OUTPUT_TOKENS,
|
||
CHAT_SUMMARY_EXCERPT_MAX_CHARS,
|
||
CHAT_TEMPERATURE,
|
||
)
|
||
from hub_ai.context import build_daily_context, format_chat_context_for_chat
|
||
from hub_ai.prompts import CHAT_SYSTEM, build_chat_user_prompt
|
||
from hub_ai.store import (
|
||
append_chat_message,
|
||
create_new_session,
|
||
ensure_active_session,
|
||
get_active_session,
|
||
load_chat_store,
|
||
summary_excerpt_for_chat,
|
||
)
|
||
|
||
|
||
def _history_lines(messages: list[dict], max_turns: int = CHAT_MAX_HISTORY_TURNS) -> str:
|
||
rows = [m for m in (messages or []) if m.get("role") in ("user", "assistant")]
|
||
rows = rows[-max_turns * 2 :]
|
||
lines = []
|
||
for m in rows:
|
||
role = "用户" if m.get("role") == "user" else "搭档"
|
||
content = m.get("content") or ""
|
||
att = m.get("attachments") or []
|
||
if att:
|
||
names = "、".join(str(a.get("name") or "附件") for a in att[:3])
|
||
content = f"{content} [附件: {names}]".strip()
|
||
lines.append(f"{role}：{content}")
|
||
return "\n".join(lines)
|
||
|
||
|
||
def get_chat_state() -> dict[str, Any]:
|
||
store = load_chat_store()
|
||
session = get_active_session()
|
||
return {
|
||
"active_session_id": store.get("active_session_id"),
|
||
"session": session,
|
||
"model": model_label(),
|
||
}
|
||
|
||
|
||
def start_new_chat(*, trading_day: str) -> dict:
|
||
session = create_new_session(trading_day=trading_day)
|
||
return {"ok": True, "session": session, "model": model_label()}
|
||
|
||
|
||
def send_chat_message(
|
||
exchanges: list[dict],
|
||
message: str,
|
||
*,
|
||
trading_day: str | None = None,
|
||
raw_attachments: Optional[list[dict]] = None,
|
||
) -> dict[str, Any]:
|
||
text = (message or "").strip()
|
||
parsed = parse_chat_attachments(raw_attachments or [])
|
||
if parsed.get("errors") and not text and not parsed.get("images_b64"):
|
||
return {"ok": False, "msg": "；".join(parsed["errors"])}
|
||
if not text and not parsed.get("images_b64") and not parsed.get("text_append"):
|
||
return {"ok": False, "msg": "消息不能为空"}
|
||
|
||
user_visible = text
|
||
if parsed.get("text_append"):
|
||
user_visible = (user_visible + "\n\n" + parsed["text_append"]).strip()
|
||
if not user_visible and parsed.get("attachment_note"):
|
||
user_visible = f"（上传了 {parsed['attachment_note']}）"
|
||
|
||
ctx = build_daily_context(exchanges, trading_day=trading_day)
|
||
day = ctx["trading_day"]
|
||
session = ensure_active_session(trading_day=day)
|
||
sid = session["id"]
|
||
history = _history_lines(session.get("messages") or [])
|
||
|
||
append_chat_message(
|
||
sid,
|
||
"user",
|
||
user_visible,
|
||
attachments=parsed.get("attachment_meta") or [],
|
||
)
|
||
|
||
brief_ctx = format_chat_context_for_chat(ctx, max_chars=CHAT_CONTEXT_MAX_CHARS)
|
||
excerpt = summary_excerpt_for_chat(day, max_chars=CHAT_SUMMARY_EXCERPT_MAX_CHARS)
|
||
|
||
user_prompt = build_chat_user_prompt(
|
||
context_text=brief_ctx,
|
||
trading_day=day,
|
||
summary_excerpt=excerpt,
|
||
history_lines=history,
|
||
user_message=text or user_visible,
|
||
attachment_note=str(parsed.get("attachment_note") or ""),
|
||
)
|
||
if parsed.get("text_append"):
|
||
user_prompt += "\n\n【附件正文】\n" + parsed["text_append"]
|
||
|
||
reply = generate_text(
|
||
system=CHAT_SYSTEM,
|
||
user=user_prompt,
|
||
temperature=CHAT_TEMPERATURE,
|
||
images_b64=parsed.get("images_b64") or None,
|
||
max_tokens=CHAT_MAX_OUTPUT_TOKENS,
|
||
)
|
||
if reply.startswith("AI 调用失败"):
|
||
return {"ok": False, "msg": reply, "session_id": sid}
|
||
|
||
session = append_chat_message(sid, "assistant", reply)
|
||
return {
|
||
"ok": True,
|
||
"trading_day": day,
|
||
"session": session,
|
||
"reply": reply,
|
||
"model": model_label(),
|
||
"attachment_warnings": parsed.get("errors") or [],
|
||
}
|