api_client.py

"""api_client：模型调用层（OpenAI 兼容端点 + Ollama 原生 /api/chat）。

统一入口：
- _chat(group_name, messages, ...) —— 按 config.xml 的 select 模式选模型并调用，返回文本或 None；
- _create(model, group, messages, ...) —— 直接调指定模型，返回 _Completion。

响应统一为 _Completion（content + 规范化 tool_calls），
调用失败/端点拒绝时内部自动降级重试（去额外字段 / 去工具 / 去图片），返回 None。
"""

import json
from typing import List, Optional

from openai import OpenAI

import config as engine_config

# 模型未配置 <temperature> 时的默认采样温度
DEFAULT_TEMPERATURE = 0.7


# ---------- API Key ----------

def _api_key_for(model: engine_config.ModelConfig) -> str:
    """取请求使用的 API Key（openai 库不允许 api_key 为空）。

    优先级：本地端点（Ollama/localhost）→ 占位 key（服务端不校验）；
    模型级 <apikey> → 全局 <apiKey> → 环境变量 QQPILOT_API_KEY → "none"。
    """
    if model.is_local:
        return "ollama"
    if model.api_key:
        return model.api_key
    cfg = engine_config.get_config()
    if cfg.api_key:
        return cfg.api_key
    return  "none"


# ---------- 图片降级辅助 ----------

def _has_images(native_messages: List[dict]) -> bool:
    """native 消息里是否含图片（顶层 images 数组）。"""
    return any(m.get("images") for m in native_messages)


def _strip_images(native_messages: List[dict]) -> List[dict]:
    """去掉 native 消息中的 images，保留文本（模型不支持多模态时回退）。"""
    return [{k: v for k, v in m.items() if k != "images"} for m in native_messages]


def _strip_openai_images(messages: List[dict]) -> List[dict]:
    """去掉 OpenAI 兼容消息 content 数组中的 image_url 分段（模型不支持多模态时回退）。"""
    out: List[dict] = []
    for m in messages:
        c = m.get("content")
        if not isinstance(c, list):
            out.append(dict(m))
            continue
        parts = [p for p in c if not (isinstance(p, dict) and p.get("type") == "image_url")]
        if not parts:
            continue  # 全是图片 → 整条丢弃
        nm = dict(m)
        if len(parts) == 1:
            nm["content"] = parts[0].get("text", "") if isinstance(parts[0], dict) else str(parts[0])
        else:
            nm["content"] = parts
        out.append(nm)
    return out


# ---------- 统一响应 ----------

class _Completion:
    """统一两种 API 的完成响应：content（文本） + tool_calls（规范化结构）。

    tool_calls 规范化为 list[{"function": {"name": str, "arguments": str}}]，
    arguments 保证是 JSON 字符串（原生 API 可能返回对象）。
    """
    __slots__ = ("content", "tool_calls")

    def __init__(self, content: Optional[str] = None, tool_calls: Optional[list] = None):
        self.content = content
        self.tool_calls = tool_calls


def _normalize_tool_calls(tool_calls) -> Optional[list]:
    """把 openai 库对象 / Ollama 原生 dict 的 tool_calls 规范化为统一 dict 结构。"""
    if not tool_calls:
        return None
    out = []
    for tc in tool_calls:
        if isinstance(tc, dict):
            fn = tc.get("function") or {}
            name = fn.get("name", "")
            args = fn.get("arguments", "{}")
        else:  # openai 库对象（ChatCompletionMessageToolCall）
            fn = tc.function
            name = fn.name
            args = fn.arguments
        if not isinstance(args, str):
            args = json.dumps(args, ensure_ascii=False)
        out.append({"function": {"name": name, "arguments": args}})
    return out or None


# ---------- 消息格式转换 ----------

def _to_native_messages(messages: List[dict]) -> List[dict]:
    """OpenAI 兼容 messages → Ollama 原生 /api/chat 格式。

    多模态 content 数组拆为 text + 顶层 images（纯 base64，不带 data: 前缀）。
    """
    out: List[dict] = []
    for m in messages:
        content = m.get("content")
        if isinstance(content, str):
            out.append(dict(m))
            continue
        text_parts: List[str] = []
        images: List[str] = []
        for part in content or []:
            if not isinstance(part, dict):
                continue
            ptype = part.get("type")
            if ptype == "text":
                text_parts.append(part.get("text", "") or "")
            elif ptype == "image_url":
                url = part.get("image_url", {}).get("url", "") if isinstance(part.get("image_url"), dict) else ""
                if url.startswith("data:"):
                    images.append(url.split(",", 1)[1] if "," in url else url)
                else:
                    images.append(url)
        nm = {"role": m.get("role", "user"), "content": "".join(text_parts)}
        if images:
            nm["images"] = images
        out.append(nm)
    return out


def _native_url(model: engine_config.ModelConfig) -> str:
    """原生 API 地址：base_url 去掉尾部 /v1 后拼 /api/chat。"""
    base = model.base_url.rstrip("/")
    if base.endswith("/v1"):
        base = base[:-3]
    return base + "/api/chat"


def DebugPrintJson(messages) -> None:
    """调试用：打印 native 格式的 messages（默认关闭）。"""
    return
    print(json.dumps(_to_native_messages(messages), ensure_ascii=False, indent=4))


# ---------- 调用入口 ----------

def _create(model: engine_config.ModelConfig, group: engine_config.ModelGroup,
            messages: List[dict],
            max_tokens: Optional[int] = None,
            temperature: Optional[float] = None,
            tools: Optional[list] = None,
            tool_choice: Optional[dict] = None) -> Optional[_Completion]:
    """调用单个模型（按 use_native 分发到 Ollama 原生 API 或 OpenAI 兼容端点）。

    temperature 为 None 时使用模型自身配置（<model><temperature>），
    模型也未配置则用 DEFAULT_TEMPERATURE。
    inject（组级+模型级）合并进请求体：native 模式下 think:false 等思考控制参数生效，
    openai 兼容模式下无效（仅当端点支持时透传）。
    """
    if temperature is None:
        temperature = model.temperature if model.temperature is not None else DEFAULT_TEMPERATURE
    if model.use_native:
        return _create_native(model, group, messages, max_tokens, temperature, tools, tool_choice)
    return _create_openai(model, group, messages, max_tokens, temperature, tools, tool_choice)


def _chat(group_name: str, messages: List[dict],
          max_tokens: Optional[int] = None,
          temperature: Optional[float] = None) -> Optional[str]:
    """调用指定模型组（按 select 模式选模型），返回文本或 None。

    temperature: None（默认）→ 使用模型自身配置（<model><temperature>），
                 模型也未配置 → DEFAULT_TEMPERATURE。显式传值可覆盖。
    """
    cfg = engine_config.get_config()
    group = cfg.group(group_name)
    if group is None or not group.models:
        print(f"[api_client] 模型组 {group_name} 未配置，跳过")
        return None
    model = cfg.select_model(group)
    resp = _create(model, group, messages, max_tokens, temperature)
    if resp is None:
        return None
    content = resp.content
    return content.strip() if content else None


# ---------- OpenAI 兼容端点 ----------

def _create_openai(model: engine_config.ModelConfig, group: engine_config.ModelGroup,
                   messages: List[dict],
                   max_tokens: Optional[int] = None,
                   temperature: Optional[float] = None,
                   tools: Optional[list] = None,
                   tool_choice: Optional[dict] = None) -> Optional[_Completion]:
    """OpenAI 兼容端点（/v1/chat/completions）。请求体被拒（400）时去掉额外字段与工具重试。"""
    cfg = engine_config.get_config()
    extra: dict = {}
    extra.update(group.inject)
    extra.update(model.inject)
    if max_tokens is not None:
        extra["max_tokens"] = max_tokens
    if temperature is not None:
        extra["temperature"] = temperature

    try:
        client = OpenAI(
            api_key=_api_key_for(model),
            base_url=model.base_url,
            timeout=cfg.timeout,
            max_retries=1,
        )
        kwargs = {"model": model.name, "messages": messages, "extra_body": extra}
        if tools:
            kwargs["tools"] = tools
        if tool_choice:
            kwargs["tool_choice"] = tool_choice
        resp = client.chat.completions.create(**kwargs)
        DebugPrintJson(messages)

        message = resp.choices[0].message
        return _Completion(content=message.content, tool_calls=_normalize_tool_calls(message.tool_calls))
    except Exception as e:
        if getattr(e, "status_code", None) == 400:
            err_str = str(e).lower()
            retry_messages = messages
            if "multimodal" in err_str or "does not support" in err_str:
                print(f"[api_client] openai 模型不支持多模态，去掉图片重试")
                retry_messages = _strip_openai_images(messages)
            else:
                print(f"[api_client] openai 请求体被拒（{e}），去掉额外参数与工具重试")
            try:
                client = OpenAI(
                    api_key=_api_key_for(model),
                    base_url=model.base_url,
                    timeout=cfg.timeout,
                    max_retries=1,
                )
                kwargs = {"model": model.name, "messages": retry_messages}
                DebugPrintJson(messages)

                if max_tokens is not None:
                    kwargs["max_tokens"] = max_tokens
                if temperature is not None:
                    kwargs["temperature"] = temperature
                resp = client.chat.completions.create(**kwargs)
                message = resp.choices[0].message
                return _Completion(content=message.content,
                                   tool_calls=_normalize_tool_calls(message.tool_calls))
            except Exception as e2:
                print(f"[api_client] openai 重试仍失败：{e2}")
                return None
        print(f"[api_client] openai 调用失败：{e}")
        return None


# ---------- Ollama 原生 /api/chat ----------

def _create_native(model: engine_config.ModelConfig, group: engine_config.ModelGroup,
                   messages: List[dict],
                   max_tokens: Optional[int] = None,
                   temperature: Optional[float] = None,
                   tools: Optional[list] = None,
                   tool_choice: Optional[dict] = None) -> Optional[_Completion]:
    """Ollama 原生 /api/chat。think:false 等思考控制参数在这里生效。"""
    cfg = engine_config.get_config()
    try:
        import httpx
    except ImportError:
        print("[api_client] native 模式需要 httpx")
        return None

    extra: dict = {}
    extra.update(group.inject)
    extra.update(model.inject)

    body: dict = {"model": model.name, "messages": _to_native_messages(messages), "stream": False}
    DebugPrintJson(messages)

    body.update(extra)  # think:false 等顶层参数
    options: dict = {}
    if max_tokens is not None:
        options["num_predict"] = max_tokens
    if temperature is not None:
        options["temperature"] = temperature
    if options:
        body["options"] = options
    if tools:
        body["tools"] = tools
        body["tool_choice"] = tool_choice if tool_choice else "auto"

    def _post(payload: dict):
        return httpx.post(_native_url(model), json=payload, timeout=cfg.timeout)

    try:
        r = _post(body)
        if r.status_code != 200:
            # 400：先尝试去掉 extra 与 tools 重试；若错误提示不支持多模态，则再去掉图片
            if r.status_code == 400:
                native_msgs = _to_native_messages(messages)
                err_text = r.text.lower()
                no_mm = "multimodal" in err_text or "does not support" in err_text
                if no_mm and _has_images(native_msgs):
                    print(f"[api_client] native 模型不支持多模态，去掉图片重试")
                    retry = {"model": model.name,
                             "messages": _strip_images(native_msgs), "stream": False}
                    if options:
                        retry["options"] = options
                    r = _post(retry)
                    if r.status_code != 200:
                        print(f"[api_client] native 去图重试仍失败：{r.status_code} {r.text[:200]}")
                        return None
                elif extra or tools:
                    print(f"[api_client] native 请求体被拒（{r.status_code}），去掉额外参数与工具重试")
                    retry = {"model": model.name,
                             "messages": native_msgs, "stream": False}
                    if options:
                        retry["options"] = options
                    r = _post(retry)
                    if r.status_code != 200:
                        print(f"[api_client] native 重试仍失败：{r.status_code} {r.text[:200]}")
                        return None
                else:
                    print(f"[api_client] native 调用失败：{r.status_code} {r.text[:300]}")
                    return None
            else:
                print(f"[api_client] native 调用失败：{r.status_code} {r.text[:300]}")
                return None
        d = r.json()
        msg = d.get("message") or {}
        return _Completion(content=msg.get("content"),
                           tool_calls=_normalize_tool_calls(msg.get("tool_calls")))
    except Exception as e:
        print(f"[api_client] native 调用失败：{e}")
        return None