RAGGER/ollama_client.py

import httpx
import warnings
from typing import Optional
from config import settings
from logger import get_logger

warnings.filterwarnings("ignore", message="Unverified HTTPS request")

log = get_logger("ollama_client")

PROVIDER_BASE_URLS = {
    "openrouter": "https://openrouter.ai/api/v1",
    "openai": "https://api.openai.com/v1",
    "groq": "https://api.groq.com/openai/v1",
    "opencode": "https://opencode.ai/zen/v1",
}


class LLMClient:
    def __init__(self):
        self.provider = settings.LLM_PROVIDER
        self.timeout = httpx.Timeout(300.0, connect=30.0)
        log.info(f"Initializing LLMClient: provider={self.provider}")
        
        if self.provider == "ollama":
            self.host = settings.OLLAMA_HOST
            log.debug(f"Ollama host: {self.host}")
            self.client = httpx.Client(timeout=self.timeout, verify=False)
        else:
            base_url = settings.LLM_BASE_URL or PROVIDER_BASE_URLS.get(self.provider, "")
            api_key = settings.LLM_API_KEY or ""
            log.debug(f"External API: {self.provider} -> {base_url}")
            self.client = httpx.Client(
                timeout=self.timeout,
                base_url=base_url,
                headers={
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json"
                }
            )
            if self.provider == "openrouter":
                self.client.headers["HTTP-Referer"] = "https://github.com/ragger"
                self.client.headers["X-Title"] = "RAGGER"
    
    def embed(self, text: str) -> list[float]:
        model = settings.OLLAMA_EMBED_MODEL
        log.debug(f"Embedding: model={model}, text_len={len(text)}")
        
        url = f"{settings.OLLAMA_HOST.rstrip('/')}/api/embeddings"
        log.debug(f"POST {url}")
        
        try:
            response = self.client.post(
                url,
                json={"model": model, "prompt": text}
            )
            response.raise_for_status()
            embedding = response.json()["embedding"]
            log.debug(f"Embedding success: {len(embedding)} dims")
            return embedding
        except httpx.HTTPStatusError as e:
            log.error(f"HTTP error: {e.response.status_code} - {e.response.text}")
            raise
        except Exception as e:
            log.error(f"Embedding error: {e}")
            raise
    
    def generate(
        self,
        prompt: str,
        system: Optional[str] = None,
        context: Optional[str] = None
    ) -> str:
        log.debug(f"Generate: provider={self.provider}, prompt_len={len(prompt)}")
        
        if self.provider == "ollama":
            return self._generate_ollama(prompt, system, context)
        else:
            return self._generate_openai(prompt, system, context)
    
    def _generate_ollama(
        self,
        prompt: str,
        system: Optional[str] = None,
        context: Optional[str] = None
    ) -> str:
        model = settings.OLLAMA_MODEL
        log.debug(f"Ollama generate: model={model}")
        
        full_prompt = ""
        if system:
            full_prompt += f"System: {system}\n\n"
        if context:
            full_prompt += f"Contesto:\n{context}\n\n"
        full_prompt += prompt
        
        log.debug(f"Full prompt length: {len(full_prompt)} chars")
        
        url = f"{settings.OLLAMA_HOST.rstrip('/')}/api/generate"
        log.debug(f"POST {url}")
        
        try:
            response = self.client.post(
                url,
                json={"model": model, "prompt": full_prompt, "stream": False}
            )
            response.raise_for_status()
            result = response.json()["response"]
            log.debug(f"Response length: {len(result)} chars")
            log.info(f"Generated response ({len(result)} chars)")
            return result
        except httpx.HTTPStatusError as e:
            log.error(f"HTTP error: {e.response.status_code} - {e.response.text}")
            raise
        except Exception as e:
            log.error(f"Generation error: {e}")
            raise
    
    def _generate_openai(
        self,
        prompt: str,
        system: Optional[str] = None,
        context: Optional[str] = None
    ) -> str:
        model = settings.LLM_MODEL or "openrouter/auto"
        log.debug(f"OpenAI generate: model={model}")
        
        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        
        user_content = ""
        if context:
            user_content += f"Contesto:\n{context}\n\n"
        user_content += prompt
        messages.append({"role": "user", "content": user_content})
        
        log.debug(f"Messages: {len(messages)} items")
        
        try:
            response = self.client.post(
                "/chat/completions",
                json={"model": model, "messages": messages}
            )
            response.raise_for_status()
            result = response.json()["choices"][0]["message"]["content"]
            log.info(f"Generated response ({len(result)} chars)")
            return result
        except httpx.HTTPStatusError as e:
            log.error(f"HTTP error: {e.response.status_code} - {e.response.text}")
            raise
        except Exception as e:
            log.error(f"Generation error: {e}")
            raise
    
    def is_available(self) -> bool:
        log.debug("Checking availability...")
        try:
            if self.provider == "ollama":
                url = f"{settings.OLLAMA_HOST.rstrip('/')}/api/tags"
                log.debug(f"GET {url}")
                response = self.client.get(url, timeout=5.0)
                available = response.status_code == 200
                log.debug(f"Ollama available: {available}")
                return available
            else:
                log.debug("External API - assuming available")
                return True
        except Exception as e:
            log.warning(f"Availability check failed: {e}")
            return False
    
    def list_models(self) -> list[str]:
        if self.provider != "ollama":
            return []
        
        try:
            url = f"{settings.OLLAMA_HOST.rstrip('/')}/api/tags"
            response = self.client.get(url, timeout=5.0)
            response.raise_for_status()
            models = response.json().get("models", [])
            return [m["name"] for m in models]
        except Exception as e:
            log.error(f"Error listing models: {e}")
            return []


llm = LLMClient()
ollama = llm