routeopt_agent/llm_client.py

from __future__ import annotations

import json
import os
from dataclasses import dataclass
from typing import Any

import requests


POLLINATIONS_OPENAI_URL = "https://text.pollinations.ai/openai"


@dataclass(frozen=True)
class LLMConfig:
    provider: str
    url: str
    model: str
    api_key: str | None = None
    ignore_proxy: bool = True


DEFAULT_MODELS = {
    "pollinations": "openai-fast",
    "groq": "openai/gpt-oss-20b",
    "siliconflow": "THUDM/GLM-Z1-9B-0414",
    "gemini": "gemini-2.5-flash",
    "openrouter": "qwen/qwen3-32b:free",
}


OPENAI_COMPATIBLE_URLS = {
    "pollinations": POLLINATIONS_OPENAI_URL,
    "groq": "https://api.groq.com/openai/v1/chat/completions",
    "siliconflow": "https://api.siliconflow.cn/v1/chat/completions",
    "gemini": "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
    "openrouter": "https://openrouter.ai/api/v1/chat/completions",
}


API_KEY_ENVS = {
    "groq": "GROQ_API_KEY",
    "siliconflow": "SILICONFLOW_API_KEY",
    "gemini": "GEMINI_API_KEY",
    "openrouter": "OPENROUTER_API_KEY",
}


class LLMError(RuntimeError):
    pass


def chat_completion(
    messages: list[dict[str, Any]],
    tools: list[dict[str, Any]] | None = None,
    temperature: float = 0.2,
    timeout: int = 45,
) -> dict[str, Any]:
    config = get_llm_config()
    payload: dict[str, Any] = {
        "model": config.model,
        "messages": messages,
        "temperature": temperature,
        "stream": False,
    }
    if tools:
        payload["tools"] = tools
        payload["tool_choice"] = "auto"

    headers = {"Content-Type": "application/json"}
    if config.api_key:
        headers["Authorization"] = f"Bearer {config.api_key}"
    if config.provider == "openrouter":
        headers["HTTP-Referer"] = os.getenv("OPENROUTER_SITE_URL", "http://localhost:7860")
        headers["X-Title"] = os.getenv("OPENROUTER_APP_NAME", "RouteOpt Agent")

    session = requests.Session()
    session.trust_env = not config.ignore_proxy

    try:
        response = session.post(config.url, json=payload, headers=headers, timeout=timeout)
        response.raise_for_status()
        data = response.json()
    except Exception as exc:
        raise LLMError(f"LLM 请求失败（provider={config.provider}, model={config.model}）：{exc}") from exc

    try:
        return data["choices"][0]["message"]
    except Exception as exc:
        raise LLMError(f"LLM 返回格式异常：{json.dumps(data, ensure_ascii=False)[:300]}") from exc


def extract_tool_arguments(message: dict[str, Any], tool_name: str) -> dict[str, Any] | None:
    for tool_call in message.get("tool_calls") or []:
        function = tool_call.get("function") or {}
        if function.get("name") != tool_name:
            continue
        raw_arguments = function.get("arguments") or "{}"
        if isinstance(raw_arguments, dict):
            return raw_arguments
        try:
            return json.loads(raw_arguments)
        except json.JSONDecodeError:
            return None
    return None


def get_llm_config() -> LLMConfig:
    provider = normalize_provider(os.getenv("LLM_PROVIDER", "auto"))
    if provider == "auto":
        provider = first_configured_provider() or "pollinations"

    if provider == "custom":
        url = required_env("OPENAI_COMPATIBLE_BASE_URL")
        api_key = os.getenv("LLM_API_KEY") or os.getenv("OPENAI_API_KEY")
        model = os.getenv("LLM_MODEL") or os.getenv("OPENAI_COMPATIBLE_MODEL") or required_env("LLM_MODEL")
    else:
        url = os.getenv(f"{provider.upper()}_BASE_URL") or OPENAI_COMPATIBLE_URLS[provider]
        api_key_env = API_KEY_ENVS.get(provider)
        api_key = os.getenv(api_key_env) if api_key_env else None
        provider_model_env = os.getenv(f"{provider.upper()}_MODEL")
        legacy_pollinations_model = os.getenv("POLLINATIONS_MODEL") if provider == "pollinations" else None
        model = (
            os.getenv("LLM_MODEL")
            or provider_model_env
            or legacy_pollinations_model
            or DEFAULT_MODELS[provider]
        )

    ignore_proxy = os.getenv("LLM_IGNORE_PROXY", "1").strip().lower() not in {"0", "false", "no"}
    return LLMConfig(provider=provider, url=url, model=model, api_key=api_key, ignore_proxy=ignore_proxy)


def normalize_provider(provider: str) -> str:
    provider = provider.strip().lower()
    aliases = {
        "pollinations.ai": "pollinations",
        "pollination": "pollinations",
        "silicon": "siliconflow",
        "sf": "siliconflow",
        "google": "gemini",
        "google-gemini": "gemini",
        "openai-compatible": "custom",
    }
    provider = aliases.get(provider, provider)
    allowed = {"auto", "pollinations", "groq", "siliconflow", "gemini", "openrouter", "custom"}
    if provider not in allowed:
        raise LLMError(f"未知 LLM_PROVIDER={provider}，可选：{', '.join(sorted(allowed))}")
    return provider


def first_configured_provider() -> str | None:
    for provider, env_name in API_KEY_ENVS.items():
        if os.getenv(env_name):
            return provider
    if os.getenv("OPENAI_COMPATIBLE_BASE_URL"):
        return "custom"
    return None


def required_env(name: str) -> str:
    value = os.getenv(name)
    if not value:
        raise LLMError(f"缺少环境变量 {name}")
    return value


def extract_json_from_text(text: str) -> dict[str, Any] | None:
    text = text.strip()
    if not text:
        return None
    if text.startswith("```"):
        text = text.strip("`")
        text = text.removeprefix("json").strip()
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    start = text.find("{")
    end = text.rfind("}")
    if start >= 0 and end > start:
        try:
            return json.loads(text[start : end + 1])
        except json.JSONDecodeError:
            return None
    return None