src/llms/ollama.py

"""
Provider Ollama (Local LLMs - Llama, Mistral, etc)
"""
from typing import Dict, Any
from .base import BaseLLM


class OllamaLLM(BaseLLM):
    """Provider para Ollama (local LLMs)"""

    def __init__(self, model_id: str, base_url: str = "http://localhost:11434", **kwargs):
        """
        Inicializa provider Ollama

        Args:
            model_id: ID do modelo (llama2, mistral, etc)
            base_url: URL base do servidor Ollama
            **kwargs: Configurações adicionais
        """
        super().__init__(model_id, **kwargs)
        self.base_url = base_url
        self.client = None

        try:
            import requests
            self.requests = requests
            # Testa conexão
            response = requests.get(f"{base_url}/api/tags", timeout=5)
            if response.status_code == 200:
                self.client = True
            else:
                self.last_error = f"Ollama não disponível em {base_url}"
        except ImportError:
            self.last_error = "Biblioteca 'requests' não instalada. Instale com: pip install requests"
        except Exception as e:
            self.last_error = f"Erro ao conectar com Ollama: {str(e)}"

    def generate(
        self,
        prompt: str,
        temperature: float = 0.3,
        max_tokens: int = 512,
        **kwargs
    ) -> str:
        """
        Gera resposta usando Ollama API

        Args:
            prompt: Texto do prompt
            temperature: Temperatura de geração
            max_tokens: Máximo de tokens
            **kwargs: Parâmetros adicionais

        Returns:
            Texto gerado
        """
        # Valida parâmetros
        valid, error_msg = self.validate_parameters(temperature, max_tokens)
        if not valid:
            return f"Erro de validação: {error_msg}"

        if not self.client:
            return f"Erro: Ollama não disponível. {self.last_error}"

        try:
            response = self.requests.post(
                f"{self.base_url}/api/generate",
                json={
                    "model": self.model_id,
                    "prompt": prompt,
                    "temperature": temperature,
                    "num_predict": max_tokens,
                    "stream": False,
                    **kwargs
                },
                timeout=60
            )

            if response.status_code != 200:
                error = f"Erro Ollama ({response.status_code}): {response.text}"
                self.last_error = error
                return error

            result = response.json()
            return result.get("response", "").strip()

        except Exception as e:
            error = f"Erro na geração Ollama: {str(e)}"
            self.last_error = error
            return error

    def is_available(self) -> bool:
        """
        Verifica se o provider está disponível

        Returns:
            True se Ollama está rodando
        """
        return self.client is not None

    def get_model_info(self) -> Dict[str, Any]:
        """
        Retorna informações sobre o modelo

        Returns:
            Dicionário com informações
        """
        return {
            "provider": "Ollama",
            "model_id": self.model_id,
            "available": self.is_available(),
            "api_type": "Local API",
            "base_url": self.base_url,
            "last_error": self.last_error if self.last_error else None
        }