my_bookmate/backend/app/llm/hub.py

import json
import os
import re
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import AsyncGenerator

from sqlalchemy.ext.asyncio import AsyncSession

# DeepSeek 部分模型会以「纯文本 content」输出 DSML 工具调用而不是标准 tool_calls：
#   <｜｜DSML｜｜ calls>
#   <｜｜DSML｜｜ invoke name="tool">
#     <｜｜DSML｜｜ parameter name="k" string="true">v</｜｜DSML｜｜ parameter>
#   </｜｜DSML｜｜ invoke>
#   </｜｜DSML｜｜ calls>
# 这里把 DSML 块解析成 OpenAI 兼容的 tool_calls，保证工具循环不中断（模型选型不确定时也能兜底）。
DSML_CALLS_START = "<｜｜DSML｜｜ calls>"
DSML_CALLS_END = "</｜｜DSML｜｜ calls>"
_DSML_CALLS_RE = re.compile(r"<｜｜DSML｜｜ calls>(.*?)</｜｜DSML｜｜ calls>", re.DOTALL)
_DSML_INVOKE_RE = re.compile(r'<｜｜DSML｜｜ invoke name="([^"]+)">(.*?)</｜｜DSML｜｜ invoke>', re.DOTALL)
_DSML_PARAM_RE = re.compile(r'<｜｜DSML｜｜ parameter name="([^"]+)"(?: string="true")?>(.*?)</｜｜DSML｜｜ parameter>', re.DOTALL)


def _parse_dsml_calls(block: str) -> list[dict]:
    """把一段 DSML calls 块解析成 OpenAI 兼容的 tool_calls 列表（参数统一按字符串）。

    id 必须生成真实值：后续第二轮会带 role="tool" 的消息（tool_call_id 指向它），
    空 id 会导致模型 API 拒绝、工具循环第二轮中断。
    """
    calls: list[dict] = []
    m = _DSML_CALLS_RE.search(block)
    if not m:
        return calls
    for n, inv in enumerate(_DSML_INVOKE_RE.finditer(m.group(1))):
        name = inv.group(1).strip()
        args: dict = {}
        for pm in _DSML_PARAM_RE.finditer(inv.group(2)):
            key = pm.group(1).strip()
            args[key] = pm.group(2).strip()
        calls.append({
            "id": f"call_{name}_{n}",
            "type": "function",
            "function": {"name": name, "arguments": json.dumps(args, ensure_ascii=False)},
        })
    return calls


@dataclass
class ChatResponse:
    content: str
    model: str
    prompt_tokens: int = 0
    completion_tokens: int = 0


class LLMProvider(ABC):
    name: str = "base"
    max_context: int = 128_000
    # 是否原生支持 function calling（OpenAI 兼容协议）；Anthropic 暂未接入工具循环
    supports_tools: bool = False

    @abstractmethod
    async def chat(self, messages: list[dict], temperature: float = 0.7) -> ChatResponse:
        ...

    @abstractmethod
    async def chat_stream(self, messages: list[dict], temperature: float = 0.7, tools: list | None = None) -> AsyncGenerator[str, None]:
        ...


class OpenAIProvider(LLMProvider):
    name = "openai"
    max_context = 128_000
    supports_tools = True

    def __init__(self, api_key: str | None = None, base_url: str | None = None, model: str | None = None):
        self.api_key = api_key if api_key is not None else os.getenv("OPENAI_API_KEY", "")
        self.base_url = base_url if base_url is not None else os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
        self.model = model if model is not None else os.getenv("OPENAI_MODEL", "gpt-4o")
        # 流式工具循环的工作态：chat_stream 结束时填充 tool_calls / last_content
        self.tool_calls: list[dict] = []
        self.last_content: str = ""

    async def chat(self, messages: list[dict], temperature: float = 0.7) -> ChatResponse:
        import httpx
        async with httpx.AsyncClient(timeout=120) as client:
            resp = await client.post(
                f"{self.base_url}/chat/completions",
                headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
                json={"model": self.model, "messages": messages, "temperature": temperature},
            )
            data = resp.json()
            if "error" in data:
                raise RuntimeError(f"OpenAI error: {data['error']}")
            choice = data["choices"][0]
            return ChatResponse(
                content=choice["message"]["content"],
                model=data["model"],
                prompt_tokens=data.get("usage", {}).get("prompt_tokens", 0),
                completion_tokens=data.get("usage", {}).get("completion_tokens", 0),
            )

    async def chat_stream(self, messages: list[dict], temperature: float = 0.7, tools: list | None = None) -> AsyncGenerator[str, None]:
        import httpx
        payload = {
            "model": self.model,
            "messages": messages,
            "temperature": temperature,
            "stream": True,
        }
        if tools:
            payload["tools"] = tools
        # 工作态复位：上一轮（若有）的 tool_calls 不残留
        self.tool_calls = []
        self.last_content = ""
        async with httpx.AsyncClient(timeout=300) as client:
            async with client.stream(
                "POST",
                f"{self.base_url}/chat/completions",
                headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
                json=payload,
            ) as resp:
                tool_calls_acc: dict[int, dict] = {}
                dsml_buf = ""
                in_dsml = False
                async for line in resp.aiter_lines():
                    if line.startswith("data: "):
                        data_str = line[6:]
                        if data_str.strip() == "[DONE]":
                            break
                        try:
                            data = json.loads(data_str)
                            # 流末尾可能来一个 usage-only chunk（choices 为空），直接跳过
                            choices = data.get("choices") or []
                            if not choices:
                                continue
                            delta = choices[0].get("delta", {})
                            content = delta.get("content", "")
                            if content:
                                dsml_buf += content
                                if in_dsml:
                                    # 正在 DSML 块内：等结束标记后解析工具调用
                                    end = dsml_buf.find(DSML_CALLS_END)
                                    if end != -1:
                                        block = dsml_buf[: end + len(DSML_CALLS_END)]
                                        self.tool_calls.extend(_parse_dsml_calls(block))
                                        tail = dsml_buf[end + len(DSML_CALLS_END):]
                                        in_dsml = False
                                        dsml_buf = tail
                                        # 块后的残余文本（罕见）按正常内容继续
                                        if dsml_buf:
                                            self.last_content += dsml_buf
                                            yield dsml_buf
                                            dsml_buf = ""
                                else:
                                    idx = dsml_buf.find(DSML_CALLS_START)
                                    if idx != -1:
                                        before = dsml_buf[:idx]
                                        if before:
                                            self.last_content += before
                                            yield before
                                        dsml_buf = dsml_buf[idx:]
                                        in_dsml = True
                                    else:
                                        # 未发现 DSML 标记：yield 除检测窗口外的文本（跨 chunk 检测需保留尾部）
                                        keep = len(DSML_CALLS_START) - 1
                                        if len(dsml_buf) > keep:
                                            head = dsml_buf[:-keep]
                                            self.last_content += head
                                            yield head
                                            dsml_buf = dsml_buf[-keep:]
                            # 流式 tool_calls：每个 index 首个 delta 带 id/name，后续增量拼 arguments
                            for tc in delta.get("tool_calls") or []:
                                idx = tc.get("index", 0)
                                entry = tool_calls_acc.setdefault(idx, {
                                    "id": "", "type": "function", "function": {"name": "", "arguments": ""},
                                })
                                if tc.get("id"):
                                    entry["id"] = tc["id"]
                                fn = tc.get("function") or {}
                                if fn.get("name"):
                                    entry["function"]["name"] = fn["name"]
                                if fn.get("arguments"):
                                    entry["function"]["arguments"] += fn["arguments"]
                        except json.JSONDecodeError:
                            continue
                # 流结束：把检测窗口残留的文本冲出来，避免丢失内容
                if dsml_buf:
                    self.last_content += dsml_buf
                    yield dsml_buf
                if tool_calls_acc:
                    self.tool_calls.extend(tool_calls_acc[i] for i in sorted(tool_calls_acc))


class AnthropicProvider(LLMProvider):
    name = "claude"
    max_context = 200_000

    def __init__(self, api_key: str | None = None, model: str | None = None):
        self.api_key = api_key if api_key is not None else os.getenv("ANTHROPIC_API_KEY", "")
        self.model = model if model is not None else os.getenv("ANTHROPIC_MODEL", "claude-sonnet-4-20250514")

    def _to_anthropic_messages(self, messages: list[dict]) -> tuple[str | None, list[dict]]:
        system_msg = None
        converted = []
        for m in messages:
            if m["role"] == "system":
                system_msg = m["content"]
            else:
                converted.append({"role": m["role"], "content": m["content"]})
        return system_msg, converted

    async def chat(self, messages: list[dict], temperature: float = 0.7) -> ChatResponse:
        import httpx
        system_msg, anthropic_msgs = self._to_anthropic_messages(messages)
        body = {
            "model": self.model,
            "max_tokens": 4096,
            "messages": anthropic_msgs,
            "temperature": temperature,
        }
        if system_msg:
            body["system"] = system_msg
        async with httpx.AsyncClient(timeout=120) as client:
            resp = await client.post(
                "https://api.anthropic.com/v1/messages",
                headers={
                    "x-api-key": self.api_key,
                    "anthropic-version": "2023-06-01",
                    "Content-Type": "application/json",
                },
                json=body,
            )
            data = resp.json()
            if "error" in data:
                raise RuntimeError(f"Anthropic error: {data['error']}")
            content_blocks = data.get("content", [])
            text = "".join(b.get("text", "") for b in content_blocks if b.get("type") == "text")
            return ChatResponse(
                content=text,
                model=data["model"],
                prompt_tokens=data.get("usage", {}).get("input_tokens", 0),
                completion_tokens=data.get("usage", {}).get("output_tokens", 0),
            )

    async def chat_stream(self, messages: list[dict], temperature: float = 0.7, tools: list | None = None) -> AsyncGenerator[str, None]:
        import httpx
        # tools 参数为兼容 ABC 而存在：Anthropic 暂未接入工具循环（supports_tools=False）
        system_msg, anthropic_msgs = self._to_anthropic_messages(messages)
        body = {
            "model": self.model,
            "max_tokens": 4096,
            "messages": anthropic_msgs,
            "temperature": temperature,
            "stream": True,
        }
        if system_msg:
            body["system"] = system_msg
        async with httpx.AsyncClient(timeout=300) as client:
            async with client.stream(
                "POST",
                "https://api.anthropic.com/v1/messages",
                headers={
                    "x-api-key": self.api_key,
                    "anthropic-version": "2023-06-01",
                    "Content-Type": "application/json",
                },
                json=body,
            ) as resp:
                async for line in resp.aiter_lines():
                    if line.startswith("data: "):
                        data_str = line[6:]
                        try:
                            data = json.loads(data_str)
                            if data.get("type") == "content_block_delta":
                                delta = data.get("delta", {})
                                text = delta.get("text", "")
                                if text:
                                    yield text
                        except json.JSONDecodeError:
                            continue


class DeepSeekProvider(OpenAIProvider):
    name = "deepseek"
    max_context = 128_000

    def __init__(self, api_key: str | None = None, model: str | None = None):
        super().__init__(
            api_key=api_key if api_key is not None else os.getenv("DEEPSEEK_API_KEY", ""),
            base_url=os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com/v1"),
            model=model if model is not None else os.getenv("DEEPSEEK_MODEL", "deepseek-chat"),
        )


# ---- 内置模型（.env）与用户自定义模型（ai_models 表）的统一解析 ----

def builtin_models() -> list[dict]:
    """扫描 .env 里已配置 key 的内置模型（无 key 的不返回，前端不会展示废选项）。

    label 跟随 .env 的 model_id 动态生成，不硬编码（如 DEEPSEEK_MODEL=deepseek-flash
    就显示「DeepSeek deepseek-flash」，避免与实际请求的模型对不上）。
    """
    out = []
    if os.getenv("DEEPSEEK_API_KEY"):
        mid = os.getenv("DEEPSEEK_MODEL", "deepseek-chat")
        out.append({
            "name": "deepseek", "label": f"DeepSeek {mid}",
            "kind": "openai_compatible", "base_url": os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com/v1"),
            "api_key_masked": _mask_key(os.getenv("DEEPSEEK_API_KEY")),
            "model_id": mid, "builtin": True,
        })
    if os.getenv("OPENAI_API_KEY"):
        mid = os.getenv("OPENAI_MODEL", "gpt-4o")
        out.append({
            "name": "openai", "label": f"GPT {mid}",
            "kind": "openai_compatible", "base_url": os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1"),
            "api_key_masked": _mask_key(os.getenv("OPENAI_API_KEY")),
            "model_id": mid, "builtin": True,
        })
    if os.getenv("ANTHROPIC_API_KEY"):
        mid = os.getenv("ANTHROPIC_MODEL", "claude-sonnet-4-20250514")
        out.append({
            "name": "claude", "label": f"Claude {mid}",
            "kind": "anthropic", "base_url": "https://api.anthropic.com/v1",
            "api_key_masked": _mask_key(os.getenv("ANTHROPIC_API_KEY")),
            "model_id": mid, "builtin": True,
        })
    if os.getenv("OPENROUTER_API_KEY"):
        mid = os.getenv("OPENROUTER_MODEL", "openrouter/auto")
        out.append({
            "name": "openrouter", "label": f"OpenRouter {mid}",
            "kind": "openai_compatible", "base_url": "https://openrouter.ai/api/v1",
            "api_key_masked": _mask_key(os.getenv("OPENROUTER_API_KEY")),
            "model_id": mid, "builtin": True,
        })
    return out


def _mask_key(key: str | None) -> str:
    if not key:
        return ""
    if len(key) <= 8:
        return "****"
    return f"{key[:4]}****{key[-4:]}"


def _builtin_provider(entry: dict) -> LLMProvider:
    """内置模型：key 从对应 .env 变量读取（builtin_models 只暴露掩码，不回传明文）。"""
    name = entry["name"]
    if name == "deepseek":
        return DeepSeekProvider()
    if name == "claude":
        return AnthropicProvider()
    if name == "openrouter":
        return OpenAIProvider(
            api_key=os.getenv("OPENROUTER_API_KEY", ""),
            base_url="https://openrouter.ai/api/v1",
            model=os.getenv("OPENROUTER_MODEL", "openrouter/auto"),
        )
    return OpenAIProvider()  # openai


def get_provider(name: str) -> LLMProvider | None:
    """兼容旧的同步入口（仅内置模型）。"""
    for m in builtin_models():
        if m["name"] == name:
            return _builtin_provider(m)
    return None


async def resolve_provider(db: AsyncSession, name: str) -> LLMProvider | None:
    """按模型名解析 provider：内置名走 .env 默认构造，自定义走 ai_models 表；找不到返回 None。"""
    from sqlalchemy import select
    from app.models.ai_model import AiModel

    for m in builtin_models():
        if m["name"] == name:
            return _builtin_provider(m)
    result = await db.execute(select(AiModel).where(AiModel.name == name))
    row = result.scalars().first()
    if row is None:
        return None
    if row.kind == "anthropic":
        return AnthropicProvider(api_key=row.api_key, model=row.model_id)
    return OpenAIProvider(api_key=row.api_key, base_url=row.base_url, model=row.model_id)