server.py

"""OpenAI 兼容入口（带 tool-call 透传）
链路：调用方(OpenAI SDK, 带 tools) → 本服务 → ima 对话(协议见 ima_upstream.py)
     tool_calls 由 tool_bridge 从模型输出解析还原；工具由调用方执行后带 role=tool 回传。
跑法: uvicorn server:app --host 0.0.0.0 --port 8300
"""
from __future__ import annotations

import json
import re
import os
import time
import uuid

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, StreamingResponse

from config import UPSTREAM_MODEL
from ima_upstream import ImaUpstream
from tool_bridge import messages_to_question, extract_tool_calls, openai_response

app = FastAPI(title="ima OpenAI-compatible downstream")
up = ImaUpstream()

MOCK = os.getenv("MOCK_UPSTREAM", "") == "1"


async def _upstream_answer(question: str, model: str, has_tools: bool = False) -> str:
    """上游应答（copilot-only）：model 指定什么传什么，无自动路由。
    MOCK=1 走模拟；票据过期（600001）→ TokenExpiredError 由调用方转 401；
    模型不在列表 → ValueError 由调用方转 400。
    has_tools=True 时若模型未产出 tool_calls（部分模型会把注入协议当攻击拒绝或幻觉），
    自动换新会话重试一次——自动优选按时段路由不同模型，遵循度有波动。"""
    if MOCK:
        from mock_upstream import mock_answer
        return mock_answer(question, model)
    r = await up.ask_once(question, model=model)
    if has_tools and "tool_calls" not in r["text"]:
        r = await up.ask_once(question, model=model)
    return r["text"]


def _err(msg: str, code: int = 502):
    return JSONResponse({"error": {"message": msg, "type": "upstream_error"}}, status_code=code)


@app.post("/v1/chat/completions")
async def chat_completions(req: Request):
    body = await req.json()
    model = body.get("model") or UPSTREAM_MODEL or "ima-copilot"
    stream = body.get("stream", False)
    messages = body.get("messages", [])
    tools = body.get("tools") or None

    # OpenAI 语义 → ima question（system 注入 + 历史序列化 + tool 结果回注）
    question = messages_to_question(messages, tools)

    if not stream:
        try:
            answer = await _upstream_answer(question, model, has_tools=bool(tools))
        except Exception as e:
            from ima_upstream import TokenExpiredError
            if isinstance(e, ValueError):
                return JSONResponse({"error": {"message": str(e), "type": "invalid_model"}}, status_code=400)
            if isinstance(e, TokenExpiredError):
                return JSONResponse(
                    {"error": {"message": f"票据已过期，请重新从网页版提取 x-ima-cookie：{e}",
                               "type": "token_expired"}}, status_code=401)
            return _err(str(e))
        clean, tool_calls = extract_tool_calls(answer)
        return JSONResponse(openai_response(model, clean or None, tool_calls))

    # ---- 流式：聚合 ima 流，按 OpenAI chunk 语义输出 ----
    async def sse():
        chat_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"
        created = int(time.time())

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

        buf = []
        try:
            answer = await _upstream_answer(question, model)
            for seg in (answer[i:i+64] for i in range(0, len(answer), 64)):
                buf.append(seg)
                yield f"data: {json.dumps(chunk({'content': seg}), ensure_ascii=False)}\n\n"
        except Exception as e:
            yield f"data: {json.dumps({'error': {'message': str(e)}})}\n\n"
            yield "data: [DONE]\n\n"
            return

        # 流结束：检查是否 tool call（是则撤回已发的正文，补 tool_calls chunk）
        answer = "".join(buf)
        clean, tool_calls = extract_tool_calls(answer)
        if tool_calls:
            yield f"data: {json.dumps(chunk({'role': 'assistant', 'content': None}))}\n\n"
            for i, tc in enumerate(tool_calls):
                yield f"data: {json.dumps(chunk({'tool_calls': [{
                    'index': i, 'id': tc['id'], 'type': 'function',
                    'function': {'name': tc['function']['name'],
                                 'arguments': tc['function']['arguments']}}]}), ensure_ascii=False)}\n\n"
            yield f"data: {json.dumps(chunk({}, finish='tool_calls'))}\n\n"
        else:
            yield f"data: {json.dumps(chunk({}, finish='stop'))}\n\n"
        yield "data: [DONE]\n\n"

    return StreamingResponse(sse(), media_type="text/event-stream")


@app.get("/v1/models")
async def models():
    """动态返回可指定模型列表（普通通道 scene=0，qa 传法实测生效的 4 个）"""
    try:
        models = await up.get_models(scene=0)

        def multiplier(m):
            """积分倍率：从 model_token_copywriting（如 "×0.2"）提取数值"""
            t = m.get("model_token_copywriting") or ""
            mm = re.search(r'([\d.]+)', t)
            return float(mm.group(1)) if mm else None

        return {"object": "list", "data": [
            {"id": m.get("model_id"), "object": "model",
             "owned_by": "ima",
             "model_name": m.get("model_name"),
             "description": m.get("describe", ""),
             "token_multiplier": multiplier(m),       # 积分倍率（None=默认档 ×1.0）
             "token_copywriting": m.get("model_token_copywriting", ""),
             "is_default": bool(m.get("is_default"))}
            for m in models]}
    except Exception as e:
        return JSONResponse({"error": {"message": str(e)}}, status_code=502)


@app.get("/healthz")
async def healthz():
    return {"ok": True, "upstream": up.base}