backend/models/providers.py

import logging
from typing import Dict, Any
from abc import ABC, abstractmethod

import httpx

logger = logging.getLogger("model_providers")


class BaseProvider(ABC):
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.api_key = config.get("api_key")
        self.base_url = config.get("base_url")
        self.models = config.get("models", [])
        self.default_model = config.get("default_model")

    @abstractmethod
    async def generate(self, model: str, prompt: str, **kwargs) -> Dict[str, Any]:
        pass

    def get_cost(self, model: str, tokens_prompt: int, tokens_completion: int) -> float:
        return (tokens_prompt + tokens_completion) * 0.000002


async def _openai_compatible_generate(
    base_url: str,
    api_key: str,
    model: str,
    prompt: str,
    extra_headers: Dict[str, str] = None,
    **kwargs,
) -> Dict[str, Any]:
    """
    Shared implementation for any provider exposing an OpenAI-compatible
    /chat/completions endpoint (Groq, Cerebras, OpenRouter all qualify).
    """
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    if extra_headers:
        headers.update(extra_headers)

    payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": kwargs.get("max_tokens", 1024),
        "temperature": kwargs.get("temperature", 0.7),
    }

    async with httpx.AsyncClient(timeout=60.0) as client:
        response = await client.post(
            f"{base_url}/chat/completions",
            headers=headers,
            json=payload,
        )
        response.raise_for_status()
        data = response.json()

    choice = data["choices"][0]["message"]["content"]
    usage = data.get("usage", {})

    return {
        "text": choice,
        "tokens_prompt": usage.get("prompt_tokens", 0),
        "tokens_completion": usage.get("completion_tokens", 0),
        "cost": 0.0,
    }


class GroqClient(BaseProvider):
    BASE_URL = "https://api.groq.com/openai/v1"

    async def generate(self, model: str, prompt: str, **kwargs) -> Dict[str, Any]:
        return await _openai_compatible_generate(
            self.BASE_URL, self.api_key, model, prompt, **kwargs
        )


class CerebrasClient(BaseProvider):
    BASE_URL = "https://api.cerebras.ai/v1"

    async def generate(self, model: str, prompt: str, **kwargs) -> Dict[str, Any]:
        return await _openai_compatible_generate(
            self.BASE_URL, self.api_key, model, prompt, **kwargs
        )


class OpenRouterClient(BaseProvider):
    BASE_URL = "https://openrouter.ai/api/v1"

    async def generate(self, model: str, prompt: str, **kwargs) -> Dict[str, Any]:
        return await _openai_compatible_generate(
            self.BASE_URL,
            self.api_key,
            model,
            prompt,
            extra_headers={"HTTP-Referer": "https://dolor3v.studio"},
            **kwargs,
        )


class HuggingFaceClient(BaseProvider):
    BASE_URL = "https://api-inference.huggingface.co/models"

    async def generate(self, model: str, prompt: str, **kwargs) -> Dict[str, Any]:
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }
        payload = {
            "inputs": prompt,
            "parameters": {
                "max_new_tokens": kwargs.get("max_tokens", 1024),
                "temperature": kwargs.get("temperature", 0.7),
            },
        }
        async with httpx.AsyncClient(timeout=60.0) as client:
            response = await client.post(
                f"{self.BASE_URL}/{model}",
                headers=headers,
                json=payload,
            )
            response.raise_for_status()
            data = response.json()

        text = ""
        if isinstance(data, list) and data:
            text = data[0].get("generated_text", "")
        elif isinstance(data, dict):
            text = data.get("generated_text", "")

        return {
            "text": text,
            "tokens_prompt": 0,
            "tokens_completion": 0,
            "cost": 0.0,
        }


PROVIDER_CLASSES = {
    "groq": GroqClient,
    "cerebras": CerebrasClient,
    "openrouter": OpenRouterClient,
    "huggingface": HuggingFaceClient,
}