Agent/agent/llm/openrouter.py

"""
OpenRouter Provider

使用 OpenRouter API 调用各种模型（包括 Claude Sonnet 4.5）
支持 OpenAI 兼容的 API 格式

OpenRouter 转发多种模型，需要根据实际模型处理不同的 usage 格式：
- OpenAI 模型: prompt_tokens, completion_tokens, completion_tokens_details.reasoning_tokens
- Claude 模型: input_tokens, output_tokens, cache_creation_input_tokens, cache_read_input_tokens
- DeepSeek 模型: prompt_tokens, completion_tokens, reasoning_tokens
"""

import os
import json
import httpx
from typing import List, Dict, Any, Optional

from .usage import TokenUsage, create_usage_from_response
from .pricing import calculate_cost


def _detect_provider_from_model(model: str) -> str:
    """根据模型名称检测提供商"""
    model_lower = model.lower()
    if model_lower.startswith("anthropic/") or "claude" in model_lower:
        return "anthropic"
    elif model_lower.startswith("openai/") or model_lower.startswith("gpt") or model_lower.startswith("o1") or model_lower.startswith("o3"):
        return "openai"
    elif model_lower.startswith("deepseek/") or "deepseek" in model_lower:
        return "deepseek"
    elif model_lower.startswith("google/") or "gemini" in model_lower:
        return "gemini"
    else:
        return "openai"  # 默认使用 OpenAI 格式


def _parse_openrouter_usage(usage: Dict[str, Any], model: str) -> TokenUsage:
    """
    解析 OpenRouter 返回的 usage

    OpenRouter 会根据底层模型返回不同格式的 usage
    """
    provider = _detect_provider_from_model(model)

    # OpenRouter 通常返回 OpenAI 格式，但可能包含额外字段
    if provider == "anthropic":
        # Claude 模型可能有缓存字段
        return TokenUsage(
            input_tokens=usage.get("prompt_tokens") or usage.get("input_tokens", 0),
            output_tokens=usage.get("completion_tokens") or usage.get("output_tokens", 0),
            cache_creation_tokens=usage.get("cache_creation_input_tokens", 0),
            cache_read_tokens=usage.get("cache_read_input_tokens", 0),
        )
    elif provider == "deepseek":
        # DeepSeek 可能有 reasoning_tokens
        return TokenUsage(
            input_tokens=usage.get("prompt_tokens", 0),
            output_tokens=usage.get("completion_tokens", 0),
            reasoning_tokens=usage.get("reasoning_tokens", 0),
        )
    else:
        # OpenAI 格式（包括 o1/o3 的 reasoning_tokens）
        reasoning = 0
        if details := usage.get("completion_tokens_details"):
            reasoning = details.get("reasoning_tokens", 0)

        return TokenUsage(
            input_tokens=usage.get("prompt_tokens", 0),
            output_tokens=usage.get("completion_tokens", 0),
            reasoning_tokens=reasoning,
        )


async def openrouter_llm_call(
    messages: List[Dict[str, Any]],
    model: str = "anthropic/claude-sonnet-4.5",
    tools: Optional[List[Dict]] = None,
    **kwargs
) -> Dict[str, Any]:
    """
    OpenRouter LLM 调用函数

    Args:
        messages: OpenAI 格式消息列表
        model: 模型名称（如 "anthropic/claude-sonnet-4.5"）
        tools: OpenAI 格式工具定义
        **kwargs: 其他参数（temperature, max_tokens 等）

    Returns:
        {
            "content": str,
            "tool_calls": List[Dict] | None,
            "prompt_tokens": int,
            "completion_tokens": int,
            "finish_reason": str,
            "cost": float
        }
    """
    api_key = os.getenv("OPEN_ROUTER_API_KEY")
    if not api_key:
        raise ValueError("OPEN_ROUTER_API_KEY environment variable not set")

    base_url = "https://openrouter.ai/api/v1"
    endpoint = f"{base_url}/chat/completions"

    # 构建请求
    payload = {
        "model": model,
        "messages": messages,
    }

    # 添加可选参数
    if tools:
        payload["tools"] = tools

    if "temperature" in kwargs:
        payload["temperature"] = kwargs["temperature"]
    if "max_tokens" in kwargs:
        payload["max_tokens"] = kwargs["max_tokens"]

    # OpenRouter 特定参数
    headers = {
        "Authorization": f"Bearer {api_key}",
        "HTTP-Referer": "https://github.com/your-repo",  # 可选，用于统计
        "X-Title": "Agent Framework",  # 可选，显示在 OpenRouter dashboard
    }

    # 调用 API
    async with httpx.AsyncClient(timeout=120.0) as client:
        try:
            response = await client.post(endpoint, json=payload, headers=headers)
            response.raise_for_status()
            result = response.json()

        except httpx.HTTPStatusError as e:
            error_body = e.response.text
            print(f"[OpenRouter] Error {e.response.status_code}: {error_body}")
            raise
        except Exception as e:
            print(f"[OpenRouter] Request failed: {e}")
            raise

    # 解析响应（OpenAI 格式）
    choice = result["choices"][0] if result.get("choices") else {}
    message = choice.get("message", {})

    content = message.get("content", "")
    tool_calls = message.get("tool_calls")
    finish_reason = choice.get("finish_reason")  # stop, length, tool_calls, content_filter 等

    # 提取 usage（完整版，根据模型类型解析）
    raw_usage = result.get("usage", {})
    usage = _parse_openrouter_usage(raw_usage, model)

    # 计算费用
    cost = calculate_cost(model, usage)

    return {
        "content": content,
        "tool_calls": tool_calls,
        "prompt_tokens": usage.input_tokens,
        "completion_tokens": usage.output_tokens,
        "reasoning_tokens": usage.reasoning_tokens,
        "cache_creation_tokens": usage.cache_creation_tokens,
        "cache_read_tokens": usage.cache_read_tokens,
        "finish_reason": finish_reason,
        "cost": cost,
        "usage": usage,  # 完整的 TokenUsage 对象
    }


def create_openrouter_llm_call(
    model: str = "anthropic/claude-sonnet-4.5"
):
    """
    创建 OpenRouter LLM 调用函数

    Args:
        model: 模型名称
            - "anthropic/claude-sonnet-4.5"
            - "anthropic/claude-opus-4.5"
            - "openai/gpt-4o"
            等等

    Returns:
        异步 LLM 调用函数
    """
    async def llm_call(
        messages: List[Dict[str, Any]],
        model: str = model,
        tools: Optional[List[Dict]] = None,
        **kwargs
    ) -> Dict[str, Any]:
        return await openrouter_llm_call(messages, model, tools, **kwargs)

    return llm_call