server.py

#!/usr/bin/env python3
"""
SmartRouter API 服务
基于 aiohttp 的异步透明代理
完整支持 OpenAI /v1/chat/completions 协议
"""

import json
import os
import sys
from typing import Union
from aiohttp import web
import aiohttp

from router import route, get_route_info, CONFIG, get_all_known_models

client_session = None


async def get_client_session():
    global client_session
    if client_session is None or client_session.closed:
        client_session = aiohttp.ClientSession()
    return client_session


async def handle_root(request: web.Request) -> web.Response:
    # 从 CONFIG 获取所有已知模型
    all_models = get_all_known_models()

    return web.json_response(
        {
            "service": "SmartRouter",
            "version": "1.0.0",
            "mode": "transparent-proxy",
            "backend": CONFIG["upstream"]["url"],
            "models": all_models,
            "endpoints": ["/v1/chat/completions", "/v1/models", "/route"],
        },
        dumps=lambda x: json.dumps(x, ensure_ascii=False),
    )


async def handle_models(request: web.Request) -> web.Response:
    # 基础特殊模型
    special_models = [
        {
            "id": "auto",
            "object": "model",
            "created": 1677610602,
            "owned_by": "smart-router",
        },
        {
            "id": "smart-auto",
            "object": "model",
            "created": 1677610602,
            "owned_by": "smart-router",
        },
        {
            "id": "smart-premium",
            "object": "model",
            "created": 1677610602,
            "owned_by": "smart-router",
        },
    ]

    # 从 CONFIG["tiers"] 收集所有模型（去重）
    seen_models = set()
    tier_models = []
    for tier_name, tier_config in CONFIG["tiers"].items():
        for model in tier_config["models"]:
            if model not in seen_models:
                seen_models.add(model)
                tier_models.append(
                    {
                        "id": model,
                        "object": "model",
                        "created": 1677610602,
                        "owned_by": "smart-router",
                    }
                )

    models_data = {
        "object": "list",
        "data": special_models + tier_models,
    }

    return web.json_response(
        models_data, dumps=lambda x: json.dumps(x, ensure_ascii=False)
    )


async def handle_route(request: web.Request) -> web.Response:
    try:
        data = await request.json()
    except:
        data = {}

    messages = data.get("messages", [])
    model = data.get("model", "auto")

    if not messages:
        prompt = data.get("prompt", "")
        system_prompt = data.get("system_prompt", "")
        if prompt:
            messages = []
            if system_prompt:
                messages.append({"role": "system", "content": system_prompt})
            messages.append({"role": "user", "content": prompt})

    session = await get_client_session()
    result = await get_route_info(messages, model, data, session)
    return web.json_response(result, dumps=lambda x: json.dumps(x, ensure_ascii=False))


async def handle_chat_completions(
    request: web.Request,
) -> Union[web.StreamResponse, web.Response]:
    try:
        data = await request.json()
    except Exception as e:
        return web.json_response(
            {
                "error": {
                    "message": f"Invalid JSON body: {str(e)}",
                    "type": "invalid_request_error",
                }
            },
            status=400,
            dumps=lambda x: json.dumps(x, ensure_ascii=False),
        )

    messages = data.get("messages", [])
    if not messages:
        return web.json_response(
            {
                "error": {
                    "message": "messages field is required",
                    "type": "invalid_request_error",
                }
            },
            status=400,
            dumps=lambda x: json.dumps(x, ensure_ascii=False),
        )

    requested_model = data.get("model", "auto")

    # 获取 session 用于异步路由调用
    session = await get_client_session()

    # route() 现在是 async，传入完整 data 和 session
    actual_model = await route(messages, requested_model, data, session)

    forward_data = dict(data)
    forward_data["model"] = actual_model

    # 从 CONFIG 获取上游配置
    upstream_url = CONFIG["upstream"]["url"]
    upstream_key = CONFIG["upstream"]["key"]

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {upstream_key}",
    }

    is_stream = data.get("stream", False)

    try:
        async with session.post(
            upstream_url, json=forward_data, headers=headers
        ) as upstream_resp:
            if upstream_resp.status != 200:
                error_body = await upstream_resp.text()
                try:
                    error_json = json.loads(error_body)
                    return web.json_response(error_json, status=upstream_resp.status)
                except:
                    return web.json_response(
                        {
                            "error": {
                                "message": error_body,
                                "type": "upstream_error",
                                "code": upstream_resp.status,
                            }
                        },
                        status=upstream_resp.status,
                        dumps=lambda x: json.dumps(x, ensure_ascii=False),
                    )

            if is_stream:
                response = web.StreamResponse(
                    status=200,
                    reason="OK",
                    headers={
                        "Content-Type": "text/event-stream",
                        "Cache-Control": "no-cache",
                        "Connection": "keep-alive",
                        "X-SmartRouter-Model": actual_model,
                    },
                )
                await response.prepare(request)

                async for line in upstream_resp.content:
                    await response.write(line)

                await response.write_eof()
                return response
            else:
                response_body = await upstream_resp.read()
                resp = web.Response(
                    status=200,
                    body=response_body,
                    headers={
                        "Content-Type": "application/json",
                        "X-SmartRouter-Model": actual_model,
                    },
                )
                return resp

    except aiohttp.ClientError as e:
        return web.json_response(
            {
                "error": {
                    "message": f"Upstream connection error: {str(e)}",
                    "type": "connection_error",
                }
            },
            status=502,
            dumps=lambda x: json.dumps(x, ensure_ascii=False),
        )
    except Exception as e:
        return web.json_response(
            {
                "error": {
                    "message": f"Internal error: {str(e)}",
                    "type": "internal_error",
                }
            },
            status=500,
            dumps=lambda x: json.dumps(x, ensure_ascii=False),
        )


def create_app() -> web.Application:
    app = web.Application()
    app.router.add_get("/", handle_root)
    app.router.add_get("/v1/models", handle_models)
    app.router.add_get("/models", handle_models)
    app.router.add_post("/v1/chat/completions", handle_chat_completions)
    app.router.add_post("/chat/completions", handle_chat_completions)
    app.router.add_post("/route", handle_route)
    return app


def print_startup_info(port: int):
    print("=" * 60)
    print("SmartRouter 透明代理服务启动")
    print("=" * 60)
    print(f"地址: http://localhost:{port}")
    print(f"后端: {CONFIG['upstream']['url']}")
    print()
    print("端点:")
    print("  GET  /                     - 服务信息")
    print("  GET  /v1/models            - 模型列表")
    print("  POST /v1/chat/completions  - OpenAI 兼容接口 (透明代理)")
    print("  POST /route                - 路由决策信息")
    print()
    print("特殊模型名称:")
    print("  auto / smart-auto          - 自动路由 (根据内容选择模型)")
    print("  smart-premium              - 使用 premium tier 模型")
    print()
    print("Tier 配置:")
    for tier_name, tier_config in CONFIG["tiers"].items():
        default_model = tier_config["models"][0] if tier_config["models"] else "无"
        description = tier_config.get("description", "")
        print(f"  {tier_name}: {default_model} ({description})")
    print()
    print(f"分类器模型: {CONFIG['classifier']['model']}")
    print(f"快速规则: {len(CONFIG['fast_rules'])} 条")
    print()


async def on_startup(app: web.Application):
    pass


async def on_cleanup(app: web.Application):
    global client_session
    if client_session and not client_session.closed:
        await client_session.close()


def main():
    port = int(sys.argv[1]) if len(sys.argv) > 1 else CONFIG["server"]["port"]
    print_startup_info(port)

    app = create_app()
    app.on_startup.append(on_startup)
    app.on_cleanup.append(on_cleanup)

    web.run_app(app, host="0.0.0.0", port=port, print=None)


if __name__ == "__main__":
    main()