inference.py

import google.generativeai as genai
from huggingface_utils import HuggingFaceUtils
import time

class InferenceManager:
    def __init__(self, config):
        """Inicializar modelos de IA con validación de configuración"""
        try:
            # Obtener API keys del diccionario anidado
            api_keys = config.get('api_keys', {})
            
            # Configurar parámetros del modelo primero
            self.generation_config = {
                "temperature": 0.9,
                "top_p": 1,
                "top_k": 1,
                "max_output_tokens": 2048,
            }
            
            self.safety_settings = [
                {
                    "category": "HARM_CATEGORY_HARASSMENT",
                    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
                },
                {
                    "category": "HARM_CATEGORY_HATE_SPEECH",
                    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
                },
                {
                    "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
                    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
                },
                {
                    "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
                    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
                },
            ]
            
            # Validar API keys
            self.google_api_key = api_keys.get('GOOGLE_API_KEY', '')
            self.huggingface_token = api_keys.get('HUGGINGFACE_TOKEN', '')
            
            # Configurar HuggingFace primero como fallback
            if self.huggingface_token:
                self.huggingface = HuggingFaceUtils(token=self.huggingface_token)
            else:
                print("Advertencia: No se configuró HuggingFace Token")
                self.huggingface = None
            
            # Intentar configurar Gemini
            self.max_retries = 3
            if self.google_api_key:
                try:
                    genai.configure(api_key=self.google_api_key)
                    self.model = self._init_gemini_model()
                    self.current_model = 'gemini'
                except Exception as e:
                    print(f"Error inicializando Gemini: {e}")
                    if self.huggingface:
                        print("Cambiando a Mixtral como fallback")
                        self.current_model = 'mixtral'
                    else:
                        raise ValueError("No hay modelos disponibles")
            else:
                if self.huggingface:
                    print("Usando Mixtral como modelo principal")
                    self.current_model = 'mixtral'
                else:
                    raise ValueError("No hay modelos disponibles")
            
        except Exception as e:
            print(f"Error en inicialización: {e}")
            raise
    
    def _init_gemini_model(self, attempt=1):
        """Inicializar modelo Gemini con reintentos"""
        try:
            return genai.GenerativeModel(
                'gemini-pro',
                generation_config=self.generation_config,
                safety_settings=self.safety_settings
            )
        except Exception as e:
            if attempt < self.max_retries:
                time.sleep(2 ** attempt)
                return self._init_gemini_model(attempt + 1)
            raise
    
    def change_model(self, model_name):
        """Cambiar modelo de IA con validación"""
        try:
            if model_name == 'gemini' and hasattr(self, 'model'):
                self.current_model = 'gemini'
                return True
            elif model_name == 'mixtral' and self.huggingface:
                self.current_model = 'mixtral'
                return True
            else:
                print(f"Modelo solicitado '{model_name}' no está disponible.")
                return False
        except Exception as e:
            print(f"Error cambiando modelo IA: {e}")
            return False
    
    def _generate_gemini_response(self, prompt, attempt=1):
        """Generar respuesta de Gemini con reintentos y fallback a Mixtral"""
        try:
            response = self.model.generate_content(prompt)
            return response.text
        except Exception as e:
            if "safety" in str(e).lower():
                return "Lo siento, no puedo responder a eso debido a restricciones de seguridad."
            elif attempt < self.max_retries:
                time.sleep(2 ** attempt)
                return self._generate_gemini_response(prompt, attempt + 1)
            else:
                print(f"Error con Gemini, intentando Mixtral: {e}")
                if self.huggingface:
                    self.current_model = 'mixtral'
                    return self.huggingface.generate_response(prompt)
                return "Lo siento, el servicio no está respondiendo. Por favor, intenta de nuevo en unos momentos."
    
    def _generate_mixtral_response(self, prompt, attempt=1):
        """Generar respuesta de Mixtral con reintentos"""
        try:
            return self.huggingface.generate_response(prompt)
        except Exception as e:
            if attempt < self.max_retries:
                print(f"Error con Mixtral (intento {attempt}), reintentando: {e}")
                time.sleep(2 ** attempt)
                return self._generate_mixtral_response(prompt, attempt + 1)
            raise

    def get_response(self, prompt, context='', history=None):
        """Generar respuesta usando el modelo actual con manejo de errores mejorado"""
        try:
            if not prompt:
                return "No se proporcionó un mensaje"
            
            # Construir el prompt completo
            full_prompt = context + "\n\n" if context else ""
            
            # Agregar historial si existe
            if history:
                full_prompt += "Historial de la conversación:\n"
                for msg in history:
                    role = "Usuario" if msg["role"] == "user" else "Asistente"
                    full_prompt += f"{role}: {msg['content']}\n"
            
            full_prompt += f"\nUsuario: {prompt}"
            
            # Generar respuesta según el modelo
            print(f"Usando modelo: {self.current_model}")
            try:
                if self.current_model == 'gemini':
                    response = self._generate_gemini_response(full_prompt)
                    if not response or len(response.strip()) < 10:
                        raise Exception("Respuesta inválida de Gemini")
                    return response
                elif self.current_model == 'mixtral' and self.huggingface:
                    response = self._generate_mixtral_response(full_prompt)
                    if not response or len(response.strip()) < 10:
                        raise Exception("Respuesta inválida de Mixtral")
                    return response
                else:
                    return "Lo siento, no hay modelos disponibles en este momento."
            except Exception as model_error:
                print(f"Error con el modelo {self.current_model}: {model_error}")
                
                # Si Gemini falla, intentar con Mixtral
                if self.current_model == 'gemini' and self.huggingface:
                    print("Cambiando a Mixtral como fallback")
                    self.current_model = 'mixtral'
                    response = self._generate_mixtral_response(full_prompt)
                    if response and len(response.strip()) >= 10:
                        return response
                # Si Mixtral falla y tenemos Gemini disponible, intentar con Gemini
                elif self.current_model == 'mixtral' and hasattr(self, 'model'):
                    print("Cambiando a Gemini como fallback")
                    self.current_model = 'gemini'
                    response = self._generate_gemini_response(full_prompt)
                    if response and len(response.strip()) >= 10:
                        return response
                
                raise Exception("No se pudo obtener una respuesta válida de ningún modelo")
                
        except Exception as e:
            print(f"Error general en get_response: {e}")
            return "Lo siento, hubo un error inesperado. Por favor, intenta de nuevo."