qihuaiflow/backend/python/app/core/llm_client.py

"""Shared, standards-based LLM client for web chat + IM bot.

平台红线(07-24)：模型层以成熟/开源标准为设计基准。本模块使用官方 OpenAI SDK
(`openai.AsyncOpenAI`)——它是 OpenAI 兼容端点的事实标准客户端。平台本身已在讲
OpenAI 协议(`/v1/chat/completions`)，因此零协议改造成本，并天然获得：
  * 原生 SSE 流式(`stream=True`)；
  * 结构化的 tool_calls / usage 解析；
  * 统一的代理(trust_env=False + 显式 proxy)与超时处理。

本模块取代 chat_service.py 与 im_llm.py 中两套手写 httpx 循环，使 web chat 与
IM bot 共用**同一模型层**、同一套 failover/代理/超时语义——这是「各端能力统一对齐」
红线在模型层的落地。

API:
  achat_complete(api_url, api_key, request_body, *, proxy=None, timeout=120, stream=False)
    - stream=False -> dict，形状对齐上游 `data` 响应：
        {"choices":[{"message":{"content","tool_calls","reasoning_content"}}], "usage":{...}}
    - stream=True  -> async generator，逐块产出 ("content"|"reasoning", text) 增量
"""
from __future__ import annotations

import httpx
from openai import AsyncOpenAI, OpenAI


def _base_url(api_url: str) -> str:
    """从完整 /v1/chat/completions 或 /v1/embeddings 地址推导 OpenAI SDK 所需的 base_url。"""
    u = (api_url or "").rstrip()
    for suf in ("/v1/chat/completions", "/chat/completions", "/v1/embeddings", "/embeddings"):
        if u.endswith(suf):
            return u[: -len(suf)]
    return u


def build_openai_client(api_url, api_key, *, proxy=None, timeout=120) -> AsyncOpenAI:
    base = _base_url(api_url)
    # 始终 trust_env=False：避免容器 HTTP_PROXY 对国内端点(siliconflow 等)造成 502；
    # 代理经 proxy 显式传入（AGNES 等），未传则直连。
    http_client = httpx.AsyncClient(proxy=proxy, trust_env=False)
    return AsyncOpenAI(
        base_url=base,
        api_key=api_key or "EMPTY",
        timeout=timeout,
        http_client=http_client,
    )


def _create_kwargs(request_body: dict) -> dict:
    kw = {
        "model": request_body.get("model"),
        "messages": request_body.get("messages", []),
        "stream": False,
    }
    if request_body.get("temperature") is not None:
        kw["temperature"] = request_body["temperature"]
    if request_body.get("top_p") is not None:
        kw["top_p"] = request_body["top_p"]
    if request_body.get("max_tokens") is not None:
        kw["max_tokens"] = request_body["max_tokens"]
    if request_body.get("stop") is not None:
        kw["stop"] = request_body["stop"]
    tools = request_body.get("tools")
    if tools:
        kw["tools"] = tools
        kw["tool_choice"] = request_body.get("tool_choice", "auto")
    return kw


def _normalize_tool_calls(tcs):
    if not tcs:
        return None
    out = []
    for tc in tcs:
        out.append(
            {
                "id": getattr(tc, "id", None),
                "type": "function",
                "function": {
                    "name": getattr(tc.function, "name", None),
                    "arguments": getattr(tc.function, "arguments", "") or "",
                },
            }
        )
    return out


async def achat_complete(api_url, api_key, request_body, *, proxy=None, timeout=120, stream=False):
    """统一的 OpenAI 兼容聊天补全。详见模块 docstring。"""
    if stream:

        async def _gen():
            client = build_openai_client(api_url, api_key, proxy=proxy, timeout=timeout)
            try:
                gen = await client.chat.completions.create(
                    **_create_kwargs(request_body), stream=True
                )
                async for chunk in gen:
                    if not chunk.choices:
                        continue
                    delta = chunk.choices[0].delta
                    rc = getattr(delta, "reasoning_content", None)
                    if rc:
                        yield ("reasoning", rc)
                    ct = getattr(delta, "content", None)
                    if ct:
                        yield ("content", ct)
            finally:
                await client.close()

        return _gen()

    client = build_openai_client(api_url, api_key, proxy=proxy, timeout=timeout)
    async with client:
        resp = await client.chat.completions.create(**_create_kwargs(request_body))
        msg = resp.choices[0].message
        return {
            "choices": [
                {
                    "message": {
                        "content": msg.content or "",
                        "tool_calls": _normalize_tool_calls(getattr(msg, "tool_calls", None)),
                        "reasoning_content": getattr(msg, "reasoning_content", None),
                    }
                }
            ],
            "usage": (resp.usage.model_dump() if resp.usage else {}),
        }


def build_openai_client_sync(api_url, api_key, *, proxy=None, timeout=120) -> OpenAI:
    """同步客户端工厂（供 IM bot 等同步上下文复用同一标准模型层）。"""
    base = _base_url(api_url)
    http_client = httpx.Client(proxy=proxy, trust_env=False)
    return OpenAI(base_url=base, api_key=api_key or "EMPTY", timeout=timeout, http_client=http_client)


def chat_complete(api_url, api_key, request_body, *, proxy=None, timeout=120):
    """同步版（与 achat_complete 行为/返回形状完全一致），供非 async 上下文使用。"""
    client = build_openai_client_sync(api_url, api_key, proxy=proxy, timeout=timeout)
    with client:
        resp = client.chat.completions.create(**_create_kwargs(request_body))
        msg = resp.choices[0].message
        return {
            "choices": [
                {
                    "message": {
                        "content": msg.content or "",
                        "tool_calls": _normalize_tool_calls(getattr(msg, "tool_calls", None)),
                        "reasoning_content": getattr(msg, "reasoning_content", None),
                    }
                }
            ],
            "usage": (resp.usage.model_dump() if resp.usage else {}),
        }


# ---------------------------------------------------------------------------
# 嵌入（Embeddings）：OpenAI 兼容 /v1/embeddings，收敛散落的裸 httpx 嵌入调用
# ---------------------------------------------------------------------------
def _embed_kwargs(input, model):
    return {"model": model, "input": input if isinstance(input, list) else [input]}


async def aembed(api_url, api_key, input, model, *, proxy=None, timeout=60) -> list:
    """统一的 OpenAI 兼容文本嵌入（异步）。返回 list[list[float]]。

    供 pgvector_search.get_embedding 等 embeddings 路径复用，把散落的裸 httpx
    嵌入调用收敛到同一标准模型层（与聊天共用代理/超时/failover 语义）。
    """
    client = build_openai_client(api_url, api_key, proxy=proxy, timeout=timeout)
    async with client:
        resp = await client.embeddings.create(**_embed_kwargs(input, model))
        return [d.embedding for d in resp.data]


def embed(api_url, api_key, input, model, *, proxy=None, timeout=60) -> list:
    """同步版嵌入（与非 async 上下文复用同一标准模型层）。"""
    client = build_openai_client_sync(api_url, api_key, proxy=proxy, timeout=timeout)
    with client:
        resp = client.embeddings.create(**_embed_kwargs(input, model))
        return [d.embedding for d in resp.data]