app/adapters/openai_chat.py

"""OpenAI /v1/chat/completions 兼容接口（流式 + 非流式 + tool calls + usage + v1 key 校验）。"""
from __future__ import annotations

import json
import time
import uuid
from collections.abc import AsyncIterator
from typing import Any

from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

from app.adapters import extract_user_prompt, normalize_model, upstream_id_for
from app.deps import get_client, verify_api_key
from app.orchestrator import stream_with_retry
from app.tokens import estimate_tokens, first_usage
from app.tools import (
    ToolCallStreamParser,
    ToolDef,
    missing_required,
    parse_tool_calls,
    strip_tool_calls,
)

router = APIRouter()


class ChatMessage(BaseModel):
    role: str
    content: Any = None
    name: str | None = None
    tool_call_id: str | None = None
    reasoning_content: str | None = None  # DeepSeek / OpenAI o-series 兼容
    model_config = {"extra": "allow"}


class ChatCompletionRequest(BaseModel):
    model: str | None = None
    messages: list[ChatMessage]
    stream: bool = False
    tools: list[dict[str, Any]] | None = None
    model_config = {"extra": "ignore"}


def _completion_id() -> str:
    return f"chatcmpl-{uuid.uuid4().hex[:24]}"


def _now() -> int:
    return int(time.time())


def _sse(obj: dict) -> bytes:
    return f"data: {json.dumps(obj, ensure_ascii=False)}\n\n".encode()


def _usage_obj(u: Any, prompt: str, completion: str) -> dict[str, Any]:
    """OpenAI usage：上游真实 usage 优先，否则 token 估算（CJK 感知 + tiktoken 兜底）。

    若有 thinking_tokens，附带 ``completion_tokens_details.reasoning_tokens``
    （o-series / DeepSeek 兼容；reasoning 是 completion 的子集明细，不重复加总）。
    """
    if u.input_tokens or u.output_tokens or u.thinking_tokens:
        prompt_tokens = int(u.input_tokens or 0)
        completion_tokens = int(u.output_tokens or 0)
        usage: dict[str, Any] = {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": prompt_tokens + completion_tokens,
        }
        if u.thinking_tokens:
            usage["completion_tokens_details"] = {"reasoning_tokens": int(u.thinking_tokens)}
        return usage
    p, c = estimate_tokens(prompt), estimate_tokens(completion)
    return {"prompt_tokens": p, "completion_tokens": c, "total_tokens": p + c}


def _build_prompt(req: ChatCompletionRequest, model: str) -> tuple[str, list[ToolDef]]:
    tools = [ToolDef.from_openai(t.get("function", t)) for t in (req.tools or [])]
    base_prompt = extract_user_prompt(
        [m.model_dump() for m in req.messages], model_id=model)
    return base_prompt, tools


async def _collect(client: Any, prompt: str, tools: list[ToolDef],
                   model_id: str | None = None) -> tuple[str, str, list]:
    parts: list[str] = []
    thinking_parts: list[str] = []
    usages: list = []
    try:
        agen = stream_with_retry(client, prompt, tools, model_id=model_id)
        async for ir in agen:
            if ir.kind == "error":
                raise HTTPException(status_code=502, detail=ir.error)
            if ir.kind == "text" and ir.text:
                parts.append(ir.text)
            if ir.kind == "thinking" and ir.thinking:
                thinking_parts.append(ir.thinking)
            if ir.usage_delta:
                usages.append(ir.usage_delta)
            if ir.kind == "finish":
                break
    except HTTPException:
        raise
    except Exception as e:  # noqa: BLE001
        # 上游账号级错误（503 换号失败 / token 过期等）统一转 502 给客户端明确信息
        raise HTTPException(status_code=502, detail=getattr(e, "detail", None) or str(e)) from e
    return "".join(parts), "".join(thinking_parts), usages


async def _gen_stream(client: Any, prompt: str, tools: list[ToolDef],
                      model: str, model_id: str | None = None) -> AsyncIterator[bytes]:
    cid, created = _completion_id(), _now()

    def chunk(delta: dict, finish: str | None = None) -> dict:
        return {"id": cid, "object": "chat.completion.chunk", "created": created, "model": model,
                "choices": [{"index": 0, "delta": delta, "finish_reason": finish}]}

    yield _sse(chunk({"role": "assistant"}))
    parts: list[str] = []
    thinking_parts: list[str] = []
    usages: list = []
    saw_tool_calls = False
    # 增量解析：围栏可能跨 IREvent 分片到达（模型把 <tool_call> 当正文输出时），
    # 逐事件 parse_tool_calls 对半截围栏束手无策会泄漏正文；状态机逐 token 累积解决。
    # known_names 传入时非客户端工具（上游注入的 AnumaSearchMCP-* 等）的围栏被直接丢弃。
    tc = ToolCallStreamParser(known_names={t.name for t in tools} if tools else None)
    tool_index = 0  # 全局递增：跨事件/跨批次编号，客户端按 index 累积 tool_calls

    def emit_tool_call(c: Any) -> bytes:
        nonlocal tool_index
        out = _sse(chunk({"tool_calls": [{
            "index": tool_index, "id": c.id, "type": "function",
            "function": {"name": c.name,
                         "arguments": json.dumps(c.arguments, ensure_ascii=False)},
        }]}))
        tool_index += 1
        return out

    agen = stream_with_retry(client, prompt, tools, model_id=model_id)
    while True:
        try:
            ir = await agen.__anext__()
        except StopAsyncIteration:
            break
        except Exception as e:  # noqa: BLE001
            # 上游抛错（账号失效 503 / token 过期等）时流已开始：发 error 帧干净收尾，
            # 避免 Starlette "response already started" 掐断连接导致客户端挂死。
            detail = getattr(e, "detail", None) or str(e)
            yield _sse({**chunk({}), "error": {"message": detail}})
            yield _sse(chunk({}, finish="stop"))
            yield b"data: [DONE]\n\n"
            return
        if ir.kind == "error":
            yield _sse({**chunk({}), "error": {"message": ir.error or "unknown"}})
            yield _sse(chunk({}, finish="stop"))
            yield b"data: [DONE]\n\n"
            return
        if ir.kind == "text" and ir.text:
            parts.append(ir.text)
            for kind, value in tc.feed(ir.text):
                if kind == "tool":
                    # 客户端未声明 tools 或调用缺必填字段（模型发的 arguments:{} 空壳）→ 静默丢弃
                    if tools and not missing_required(value, tools):
                        saw_tool_calls = True
                        yield emit_tool_call(value)
                elif value:
                    yield _sse(chunk({"content": value}))
        if ir.kind == "thinking" and ir.thinking:
            thinking_parts.append(ir.thinking)
            yield _sse(chunk({"reasoning_content": ir.thinking}))
        if ir.usage_delta:
            usages.append(ir.usage_delta)
        if ir.kind == "finish":
            break
    for kind, value in tc.finish():  # 流结束收尾：未闭合围栏 / 残留文本
        if kind == "tool":
            if tools and not missing_required(value, tools):
                saw_tool_calls = True
                yield emit_tool_call(value)
        elif value:
            yield _sse(chunk({"content": value}))

    full_text = "".join(parts)
    finish_reason = "tool_calls" if saw_tool_calls else "stop"
    yield _sse(chunk({}, finish=finish_reason) | {
        "usage": _usage_obj(first_usage(usages), prompt, full_text),
    })
    yield b"data: [DONE]\n\n"


@router.post("/v1/chat/completions")
async def chat_completions(
    req: ChatCompletionRequest,
    client: Any = Depends(get_client),
    _: None = Depends(verify_api_key),
) -> Any:
    model = normalize_model(req.model)
    prompt, tools = _build_prompt(req, model)
    model_id = upstream_id_for(model)

    if req.stream:
        return StreamingResponse(
            _gen_stream(client, prompt, tools, model, model_id),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "X-Accel-Buffering": "no",
                "Connection": "keep-alive",
            },
        )

    full_text, thinking_text, usages = await _collect(client, prompt, tools, model_id)
    # 围栏（含上游注入工具调用）一律剥离，不允许泄漏进正文
    message: dict[str, Any] = {"role": "assistant", "content": strip_tool_calls(full_text)}
    if thinking_text:
        message["reasoning_content"] = thinking_text
    finish_reason = "stop"
    if tools:
        known = {t.name for t in tools}
        calls = [c for c in parse_tool_calls(full_text, known_names=known)
                 if c.name in known and not missing_required(c, tools)]
        if calls:
            finish_reason = "tool_calls"
            message["content"] = None
            message["tool_calls"] = [{
                "id": c.id, "type": "function",
                "function": {"name": c.name, "arguments": json.dumps(c.arguments, ensure_ascii=False)},
            } for c in calls]
    return {
        "id": _completion_id(), "object": "chat.completion", "created": _now(), "model": model,
        "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}],
        "usage": _usage_obj(first_usage(usages), prompt, full_text),
    }