src/llms/huggingface.py

"""
Provider HuggingFace usando Inference API
"""
from typing import Dict, Any
from huggingface_hub import InferenceClient
from .base import BaseLLM


class HuggingFaceLLM(BaseLLM):
    """Provider para HuggingFace Inference API"""

    def __init__(self, model_id: str, api_token: str, **kwargs):
        """
        Inicializa provider HuggingFace

        Args:
            model_id: ID do modelo no Hub
            api_token: Token de API do HuggingFace
            **kwargs: Configurações adicionais
        """
        super().__init__(model_id, **kwargs)
        self.api_token = api_token
        self.client = None

        if api_token:
            try:
                self.client = InferenceClient(token=api_token)
            except Exception as e:
                self.last_error = f"Erro ao inicializar InferenceClient: {str(e)}"

    def generate(
        self,
        prompt: str,
        temperature: float = 0.3,
        max_tokens: int = 512,
        **kwargs
    ) -> str:
        """
        Gera resposta usando HuggingFace Inference 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: Cliente HuggingFace não inicializado. {self.last_error}"

        try:
            response = self.client.text_generation(
                prompt,
                model=self.model_id,
                temperature=temperature,
                max_new_tokens=max_tokens,
                return_full_text=False,
                **kwargs
            )

            return response.strip() if response else "Sem resposta do modelo"

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

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

        Returns:
            True se cliente foi inicializado
        """
        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": "HuggingFace",
            "model_id": self.model_id,
            "available": self.is_available(),
            "api_type": "Inference API",
            "last_error": self.last_error if self.last_error else None
        }