graphrag/layers/universal_llm.py

"""
Universal LLM Layer — LiteLLM-powered multi-provider support
=============================================================
Supports 12 providers through a single interface:
OpenAI, Anthropic, Gemini, Mistral, Cohere, Ollama (local),
OpenRouter, Groq, xAI, Together AI, HuggingFace, DeepSeek

Uses LiteLLM for unified API, falls back to direct OpenAI SDK
if LiteLLM is not installed.
"""
import json
import logging
import os
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)

# ── Provider Registry ─────────────────────────────────

PROVIDERS = {
    "openai": {
        "name": "OpenAI",
        "litellm_prefix": "openai",
        "default_model": "gpt-4o-mini",
        "api_key_env": "OPENAI_API_KEY",
        "cost_input": 0.00015, "cost_output": 0.0006,
    },
    "anthropic": {
        "name": "Anthropic Claude",
        "litellm_prefix": "anthropic",
        "default_model": "claude-sonnet-4-20250514",
        "api_key_env": "ANTHROPIC_API_KEY",
        "cost_input": 0.003, "cost_output": 0.015,
    },
    "gemini": {
        "name": "Google Gemini",
        "litellm_prefix": "gemini",
        "default_model": "gemini-2.0-flash",
        "api_key_env": "GEMINI_API_KEY",
        "cost_input": 0.0001, "cost_output": 0.0004,
    },
    "mistral": {
        "name": "Mistral AI",
        "litellm_prefix": "mistral",
        "default_model": "mistral-large-latest",
        "api_key_env": "MISTRAL_API_KEY",
        "cost_input": 0.002, "cost_output": 0.006,
    },
    "cohere": {
        "name": "Cohere",
        "litellm_prefix": "cohere_chat",
        "default_model": "command-r-plus",
        "api_key_env": "COHERE_API_KEY",
        "cost_input": 0.0025, "cost_output": 0.01,
    },
    "ollama": {
        "name": "Ollama (Local)",
        "litellm_prefix": "ollama_chat",
        "default_model": "llama3.2",
        "api_key_env": "",
        "api_base": "http://localhost:11434",
        "cost_input": 0, "cost_output": 0,
        "is_local": True,
    },
    "openrouter": {
        "name": "OpenRouter",
        "litellm_prefix": "openrouter",
        "default_model": "meta-llama/llama-3.3-70b-instruct",
        "api_key_env": "OPENROUTER_API_KEY",
        "cost_input": 0.0004, "cost_output": 0.0004,
    },
    "groq": {
        "name": "Groq",
        "litellm_prefix": "groq",
        "default_model": "llama-3.3-70b-versatile",
        "api_key_env": "GROQ_API_KEY",
        "cost_input": 0.00059, "cost_output": 0.00079,
    },
    "xai": {
        "name": "xAI Grok",
        "litellm_prefix": "xai",
        "default_model": "grok-3-mini",
        "api_key_env": "XAI_API_KEY",
        "cost_input": 0.0003, "cost_output": 0.0005,
    },
    "together": {
        "name": "Together AI",
        "litellm_prefix": "together_ai",
        "default_model": "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
        "api_key_env": "TOGETHER_API_KEY",
        "cost_input": 0.00088, "cost_output": 0.00088,
    },
    "huggingface": {
        "name": "HuggingFace Inference",
        "litellm_prefix": "huggingface",
        "default_model": "meta-llama/Llama-3.3-70B-Instruct",
        "api_key_env": "HF_TOKEN",
        "cost_input": 0, "cost_output": 0,
    },
    "deepseek": {
        "name": "DeepSeek",
        "litellm_prefix": "deepseek",
        "default_model": "deepseek-chat",
        "api_key_env": "DEEPSEEK_API_KEY",
        "cost_input": 0.00014, "cost_output": 0.00028,
    },
}


@dataclass
class LLMResponse:
    """Universal LLM response."""
    content: str = ""
    input_tokens: int = 0
    output_tokens: int = 0
    total_tokens: int = 0
    latency_ms: float = 0.0
    cost_usd: float = 0.0
    model: str = ""
    provider: str = ""


class UniversalLLM:
    """
    Universal LLM client supporting 12 providers.
    Uses LiteLLM when available, falls back to OpenAI SDK.
    """

    def __init__(self, provider: str = "openai", model: str = None,
                 api_key: str = None, api_base: str = None):
        self.provider_id = provider
        self.provider_config = PROVIDERS.get(provider, PROVIDERS["openai"])
        self.model = model or self.provider_config["default_model"]
        self._api_key = api_key
        self._api_base = api_base
        self._litellm = None
        self._openai_client = None
        self._anthropic_client = None

    def initialize(self):
        """Initialize the appropriate SDK."""
        # Try LiteLLM first (universal)
        try:
            import litellm
            self._litellm = litellm
            litellm.drop_params = True  # auto-drop unsupported params
            logger.info(f"LiteLLM initialized for {self.provider_id}/{self.model}")
            return
        except ImportError:
            pass

        # Fall back to direct SDK
        if self.provider_id == "anthropic":
            try:
                from anthropic import Anthropic
                key = self._api_key or os.getenv(self.provider_config["api_key_env"], "")
                self._anthropic_client = Anthropic(api_key=key)
                logger.info(f"Anthropic SDK initialized: {self.model}")
                return
            except ImportError:
                pass

        # OpenAI SDK (works for OpenAI, Ollama, Groq, Together, etc.)
        try:
            from openai import OpenAI
            api_key_env = self.provider_config.get("api_key_env", "")
            key = self._api_key or os.getenv(api_key_env, "") or "ollama"
            base = self._api_base or self.provider_config.get("api_base", "")

            base_urls = {
                "openai": "https://api.openai.com/v1",
                "gemini": "https://generativelanguage.googleapis.com/v1beta/openai/",
                "mistral": "https://api.mistral.ai/v1",
                "cohere": "https://api.cohere.ai/compatibility/v1",
                "ollama": "http://localhost:11434/v1",
                "openrouter": "https://openrouter.ai/api/v1",
                "groq": "https://api.groq.com/openai/v1",
                "xai": "https://api.x.ai/v1",
                "together": "https://api.together.xyz/v1",
                "huggingface": "https://api-inference.huggingface.co/v1",
                "deepseek": "https://api.deepseek.com/v1",
            }
            base_url = base or base_urls.get(self.provider_id, "https://api.openai.com/v1")

            self._openai_client = OpenAI(base_url=base_url, api_key=key)
            logger.info(f"OpenAI-compat SDK initialized for {self.provider_id}: {base_url}")
        except ImportError:
            logger.warning("No SDK available. Install: pip install openai litellm anthropic")

    def generate(self, messages: List[Dict[str, str]],
                 temperature: float = 0, max_tokens: int = 1024,
                 json_mode: bool = False) -> LLMResponse:
        """Generate a response using the configured provider."""
        start = time.perf_counter()
        cost_in = self.provider_config.get("cost_input", 0)
        cost_out = self.provider_config.get("cost_output", 0)

        # ── LiteLLM path ──────────────────────────────
        if self._litellm:
            return self._call_litellm(messages, temperature, max_tokens, json_mode, start, cost_in, cost_out)

        # ── Anthropic direct path ─────────────────────
        if self._anthropic_client:
            return self._call_anthropic(messages, temperature, max_tokens, start, cost_in, cost_out)

        # ── OpenAI-compat path ────────────────────────
        if self._openai_client:
            return self._call_openai(messages, temperature, max_tokens, json_mode, start, cost_in, cost_out)

        # ── Mock fallback ─────────────────────────────
        return LLMResponse(
            content="[No LLM SDK available. Install: pip install openai]",
            input_tokens=50, output_tokens=20, total_tokens=70,
            latency_ms=100, cost_usd=0, model=self.model, provider=self.provider_id,
        )

    def _call_litellm(self, messages, temp, max_tok, json_mode, start, ci, co):
        prefix = self.provider_config["litellm_prefix"]
        model_str = f"{prefix}/{self.model}"
        kwargs = {"model": model_str, "messages": messages,
                  "temperature": temp, "max_tokens": max_tok}
        if json_mode:
            kwargs["response_format"] = {"type": "json_object"}
        if self.provider_config.get("api_base"):
            kwargs["api_base"] = self.provider_config["api_base"]

        resp = self._litellm.completion(**kwargs)
        elapsed = (time.perf_counter() - start) * 1000
        u = resp.usage
        return LLMResponse(
            content=resp.choices[0].message.content or "",
            input_tokens=u.prompt_tokens, output_tokens=u.completion_tokens,
            total_tokens=u.total_tokens, latency_ms=elapsed,
            cost_usd=(u.prompt_tokens / 1000 * ci + u.completion_tokens / 1000 * co),
            model=self.model, provider=self.provider_id,
        )

    def _call_anthropic(self, messages, temp, max_tok, start, ci, co):
        sys_msg = next((m["content"] for m in messages if m["role"] == "system"), None)
        user_msgs = [{"role": m["role"], "content": m["content"]} for m in messages if m["role"] != "system"]
        kwargs = {"model": self.model, "max_tokens": max_tok,
                  "temperature": temp, "messages": user_msgs}
        if sys_msg:
            kwargs["system"] = sys_msg
        msg = self._anthropic_client.messages.create(**kwargs)
        elapsed = (time.perf_counter() - start) * 1000
        content = msg.content[0].text if msg.content and msg.content[0].type == "text" else ""
        return LLMResponse(
            content=content,
            input_tokens=msg.usage.input_tokens, output_tokens=msg.usage.output_tokens,
            total_tokens=msg.usage.input_tokens + msg.usage.output_tokens,
            latency_ms=elapsed,
            cost_usd=(msg.usage.input_tokens / 1000 * ci + msg.usage.output_tokens / 1000 * co),
            model=self.model, provider=self.provider_id,
        )

    def _call_openai(self, messages, temp, max_tok, json_mode, start, ci, co):
        kwargs = {"model": self.model, "messages": messages,
                  "temperature": temp, "max_tokens": max_tok}
        if json_mode:
            kwargs["response_format"] = {"type": "json_object"}
        resp = self._openai_client.chat.completions.create(**kwargs)
        elapsed = (time.perf_counter() - start) * 1000
        u = resp.usage
        return LLMResponse(
            content=resp.choices[0].message.content or "",
            input_tokens=u.prompt_tokens if u else 0,
            output_tokens=u.completion_tokens if u else 0,
            total_tokens=u.total_tokens if u else 0, latency_ms=elapsed,
            cost_usd=((u.prompt_tokens if u else 0) / 1000 * ci + (u.completion_tokens if u else 0) / 1000 * co),
            model=self.model, provider=self.provider_id,
        )

    # ── Convenience methods ──────────────────────────

    def generate_answer(self, query, context, system_prompt=None):
        if not system_prompt:
            system_prompt = "Answer accurately using ONLY the provided context. Be concise."
        return self.generate([
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {query}\n\nAnswer:"},
        ], max_tokens=512)

    def extract_entities(self, text):
        return self.generate([
            {"role": "system", "content": 'Extract entities and relationships. Return JSON: {"entities": [{"name": "...", "type": "PERSON|ORG|LOCATION|EVENT|CONCEPT"}], "relations": [{"source": "...", "target": "...", "type": "...", "description": "..."}]}'},
            {"role": "user", "content": text},
        ], max_tokens=2048, json_mode=True)

    def extract_keywords(self, query):
        return self.generate([
            {"role": "system", "content": 'Extract keywords. Return JSON: {"high_level": ["themes"], "low_level": ["entities"]}'},
            {"role": "user", "content": query},
        ], max_tokens=256, json_mode=True)


def get_available_providers() -> List[str]:
    """Return list of provider IDs with valid API keys."""
    available = []
    for pid, cfg in PROVIDERS.items():
        if cfg.get("is_local"):
            available.append(pid)
        elif not cfg.get("api_key_env"):
            available.append(pid)
        elif os.getenv(cfg["api_key_env"]):
            available.append(pid)
    return available


def check_ollama() -> Dict[str, Any]:
    """Check if Ollama is running locally."""
    import urllib.request
    try:
        req = urllib.request.Request("http://localhost:11434/api/tags", method="GET")
        with urllib.request.urlopen(req, timeout=2) as resp:
            data = json.loads(resp.read())
            return {"ok": True, "models": [m["name"] for m in data.get("models", [])]}
    except Exception:
        return {"ok": False, "models": []}