crypto_monitor/manual_trading_hub/hub_ai/chat.py

"""中控 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_HISTORY_MAX_CHARS_PER_MSG,
    CHAT_MAX_CONTINUATIONS,
    CHAT_MAX_HISTORY_TURNS,
    CHAT_MAX_OUTPUT_TOKENS,
    CHAT_SUMMARY_EXCERPT_MAX_CHARS,
    CHAT_TEMPERATURE,
    trading_day_reset_hour,
)
from hub_trades_lib import current_trading_day
from hub_ai.context import build_daily_context, format_chat_context_for_chat
from hub_ai.prompts import (
    CHAT_GENERAL_SYSTEM,
    CHAT_SYSTEM,
    build_chat_user_prompt,
    build_general_chat_user_prompt,
)
from hub_ai.store import (
    CHAT_BOT_GENERAL,
    CHAT_BOT_TRADING,
    append_chat_message,
    create_new_session,
    delete_chat_session,
    ensure_active_session,
    get_active_session,
    list_chat_sessions,
    load_chat_store,
    set_active_session,
    summary_excerpt_for_chat,
)


def _is_ai_error_reply(text: str) -> bool:
    t = (text or "").strip()
    return t.startswith("AI 调用失败") or t.startswith("AI 生成失败")


def _history_lines(
    messages: list[dict],
    max_turns: int = CHAT_MAX_HISTORY_TURNS,
    *,
    max_chars_per_msg: int = 1500,
) -> 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 = str(m.get("content") or "").strip()
        if m.get("role") == "assistant" and _is_ai_error_reply(content):
            continue
        att = m.get("attachments") or []
        if att:
            names = "、".join(str(a.get("name") or "附件") for a in att[:3])
            content = f"{content} [附件: {names}]".strip()
        if len(content) > max_chars_per_msg:
            content = content[: max_chars_per_msg - 1].rstrip() + "…"
        lines.append(f"{role}：{content}")
    return "\n".join(lines)


def get_chat_state() -> dict[str, Any]:
    store = load_chat_store()
    session = get_active_session()
    if session:
        session.setdefault("bot_mode", CHAT_BOT_TRADING)
    return {
        "active_session_id": store.get("active_session_id"),
        "session": session,
        "sessions": list_chat_sessions(),
        "model": model_label(),
    }


def start_new_chat(*, trading_day: str, bot_mode: str = CHAT_BOT_TRADING) -> dict:
    session = create_new_session(trading_day=trading_day, bot_mode=bot_mode)
    return {
        "ok": True,
        "session": session,
        "sessions": list_chat_sessions(),
        "model": model_label(),
    }


def switch_chat_session(session_id: str) -> dict[str, Any]:
    session = set_active_session(session_id)
    return {
        "ok": True,
        "session": session,
        "sessions": list_chat_sessions(),
        "model": model_label(),
    }


def remove_chat_session(session_id: str) -> dict[str, Any]:
    deleted, new_active = delete_chat_session(session_id)
    if not deleted:
        return {"ok": False, "msg": "session_not_found"}
    session = get_active_session()
    return {
        "ok": True,
        "active_session_id": new_active,
        "session": session,
        "sessions": list_chat_sessions(),
        "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']}）"

    day = (trading_day or "").strip()[:10] or current_trading_day(
        reset_hour=trading_day_reset_hour()
    )
    session = ensure_active_session(trading_day=day)
    sid = session["id"]
    history = _history_lines(
        session.get("messages") or [],
        max_chars_per_msg=CHAT_HISTORY_MAX_CHARS_PER_MSG,
    )

    append_chat_message(
        sid,
        "user",
        user_visible,
        attachments=parsed.get("attachment_meta") or [],
    )

    bot_mode = (session.get("bot_mode") or CHAT_BOT_TRADING).strip().lower()
    if bot_mode == CHAT_BOT_GENERAL:
        user_prompt = build_general_chat_user_prompt(
            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"]
        system_prompt = CHAT_GENERAL_SYSTEM
    else:
        ctx = build_daily_context(exchanges, trading_day=day)
        day = ctx["trading_day"]
        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"]
        system_prompt = CHAT_SYSTEM

    reply = generate_text(
        system=system_prompt,
        user=user_prompt,
        temperature=CHAT_TEMPERATURE,
        images_b64=parsed.get("images_b64") or None,
        max_tokens=CHAT_MAX_OUTPUT_TOKENS,
        max_continuations=CHAT_MAX_CONTINUATIONS,
    )
    if _is_ai_error_reply(reply):
        return {"ok": False, "msg": reply, "session_id": sid}

    session = append_chat_message(sid, "assistant", reply)
    return {
        "ok": True,
        "trading_day": day,
        "session": session,
        "sessions": list_chat_sessions(),
        "reply": reply,
        "model": model_label(),
        "attachment_warnings": parsed.get("errors") or [],
    }