Update ai_logic.py

ai_logic.py

@@ -9,6 +9,7 @@ from urllib.parse import urljoin, urlparse
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import FAISS
 from langchain_community.embeddings import HuggingFaceEmbeddings
+from huggingface_hub import InferenceClient
 import asyncio
 
 # --- Configurações ---
@@ -25,24 +26,24 @@ if not HF_TOKEN:
 
 print(f"Token HF encontrado: {HF_TOKEN[:10]}...")
 
-# --- Modelos para teste (versão simplificada e mais robusta) ---
+# --- Modelos para teste (versão com InferenceClient) ---
 MODELS = {}
 
 # Lista de modelos mais estáveis e com maior chance de funcionar
 NEW_MODELS_TO_TEST = [
-    ("GPT-2", "gpt2"),
-    ("DistilGPT-2", "distilgpt2"),
-    ("GPT-2 Medium", "gpt2-medium"),
-    ("Microsoft DialoGPT", "microsoft/DialoGPT-medium"),
-    ("Google T5 Small", "google/flan-t5-small"),
-    ("Google T5 Base", "google/flan-t5-base"),
-    ("Facebook BART", "facebook/bart-base"),
-    ("TinyLlama 1B", "TinyLlama/TinyLlama-1.1B-Chat-v1.0"),
-    ("Phi-3 Mini", "microsoft/Phi-3-mini-4k-instruct"),
+    ("Llama 3.2 1B", "meta-llama/Llama-3.2-1B-Instruct"),
+    ("Llama 3.2 3B", "meta-llama/Llama-3.2-3B-Instruct"),
+    ("Llama 3.1 8B", "meta-llama/Meta-Llama-3.1-8B-Instruct"),
     ("Mistral 7B", "mistralai/Mistral-7B-Instruct-v0.3"),
+    ("Mistral Nemo", "mistralai/Mistral-Nemo-Instruct-2407"),
+    ("Phi-3.5 Mini", "microsoft/Phi-3.5-mini-instruct"),
+    ("Qwen2.5 7B", "Qwen/Qwen2.5-7B-Instruct"),
+    ("Gemma 2 2B", "google/gemma-2-2b-it"),
+    ("CodeLlama 7B", "codellama/CodeLlama-7b-Instruct-hf"),
+    ("Zephyr 7B", "HuggingFaceH4/zephyr-7b-beta"),
 ]
 
-DEFAULT_MODEL = "GPT-2"
+DEFAULT_MODEL = "Llama 3.2 1B"
 
 # --- Gerenciamento de Sessão ---
 user_sessions: Dict[str, Dict[str, List | Dict]] = {}
@@ -226,122 +227,102 @@ def retrieve_context_from_blog(query: str, k: int = 4) -> str:
             print(f"Erro ao buscar contexto: {e}")
     return ""
 
-# --- API Client (Versão Melhorada) ---
-class HuggingFaceAPIClient:
+# --- Inference Client (Versão Melhorada com huggingface_hub) ---
+class HuggingFaceInferenceClient:
     def __init__(self, token: str):
-        self.headers = {"Authorization": f"Bearer {token}"}
-        self.base_url = "https://api-inference.huggingface.co/models/"
-
+        self.token = token
+        self.clients = {}  # Cache de clientes para diferentes modelos
+    
+    def get_client(self, model_name: str) -> InferenceClient:
+        """Obtém ou cria um cliente para o modelo especificado."""
+        if model_name not in self.clients:
+            self.clients[model_name] = InferenceClient(
+                model=model_name,
+                token=self.token
+            )
+        return self.clients[model_name]
+    
     def check_model_status(self, model_name: str) -> Tuple[bool, str]:
-        """Verifica se um modelo está disponível via API."""
-        url = f"{self.base_url}{model_name}"
-        
-        # Teste simples para verificar se o modelo responde
-        test_payload = {
-            "inputs": "Hello",
-            "parameters": {
-                "max_new_tokens": 5,
-                "temperature": 0.1,
-                "return_full_text": False
-            }
-        }
-        
+        """Verifica se um modelo está disponível."""
         try:
             print(f"  Testando {model_name}...")
-            response = requests.post(url, headers=self.headers, json=test_payload, timeout=30)
+            client = self.get_client(model_name)
+            
+            # Teste simples com mensagem básica
+            test_messages = [
+                {"role": "user", "content": "Hello"}
+            ]
+            
+            # Tenta fazer uma requisição de teste
+            response = client.chat_completion(
+                messages=test_messages,
+                max_tokens=5,
+                temperature=0.1
+            )
             
-            if response.status_code == 200:
-                result = response.json()
-                if isinstance(result, list) and len(result) > 0:
-                    return True, "Modelo disponível"
-                elif isinstance(result, dict) and 'generated_text' in result:
-                    return True, "Modelo disponível"
-                else:
-                    return False, f"Resposta inesperada: {result}"
-                    
-            elif response.status_code == 503:
-                return False, "Modelo carregando (503)"
-            elif response.status_code == 401:
-                return False, "Token inválido (401)"
-            elif response.status_code == 400:
-                error_msg = response.json().get('error', 'Erro desconhecido')
-                if 'loading' in error_msg.lower():
-                    return False, "Modelo carregando"
-                return False, f"Erro 400: {error_msg}"
+            if response and hasattr(response, 'choices') and len(response.choices) > 0:
+                return True, "Modelo disponível"
             else:
-                return False, f"HTTP {response.status_code}: {response.text[:100]}"
+                return False, "Resposta inválida do modelo"
                 
-        except requests.exceptions.Timeout:
-            return False, "Timeout"
-        except requests.exceptions.RequestException as e:
-            return False, f"Erro de conexão: {str(e)[:100]}"
         except Exception as e:
-            return False, f"Erro inesperado: {str(e)[:100]}"
+            error_msg = str(e).lower()
+            if 'loading' in error_msg or 'currently loading' in error_msg:
+                return False, "Modelo carregando"
+            elif 'rate limit' in error_msg:
+                return False, "Rate limit atingido"
+            elif 'token' in error_msg or 'unauthorized' in error_msg:
+                return False, "Token inválido"
+            elif 'model not found' in error_msg:
+                return False, "Modelo não encontrado"
+            else:
+                return False, f"Erro: {str(e)[:100]}"
 
-    def query_model(self, model_name: str, messages: List[Dict], max_tokens: int = 500) -> str:
-        """Faz requisição ao modelo."""
-        # Converte mensagens para formato de prompt simples
-        prompt = self._convert_messages_to_prompt(messages)
-        
-        url = f"{self.base_url}{model_name}"
-        payload = {
-            "inputs": prompt,
-            "parameters": {
-                "max_new_tokens": max_tokens,
-                "temperature": 0.7,
-                "do_sample": True,
-                "return_full_text": False
-            }
-        }
-        
+    def query_model(self, model_name: str, messages: List[Dict], max_tokens: int = 500, temperature: float = 0.7) -> str:
+        """Faz requisição ao modelo usando chat completion."""
         try:
-            response = requests.post(url, headers=self.headers, json=payload, timeout=60)
+            client = self.get_client(model_name)
+            
+            # Faz a requisição usando chat completion
+            response = client.chat_completion(
+                messages=messages,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                stream=False
+            )
             
-            if response.status_code == 200:
-                result = response.json()
-                if isinstance(result, list) and len(result) > 0:
-                    return result[0].get('generated_text', '').strip()
-                elif isinstance(result, dict) and 'generated_text' in result:
-                    return result['generated_text'].strip()
-                else:
-                    return f"Formato de resposta inesperado: {result}"
+            # Extrai a resposta
+            if response and hasattr(response, 'choices') and len(response.choices) > 0:
+                content = response.choices[0].message.content
+                return content.strip() if content else "Resposta vazia do modelo"
             else:
-                return f"Erro na requisição: {response.status_code} - {response.text[:200]}"
+                return "Erro: Resposta inválida do modelo"
                 
         except Exception as e:
-            return f"Erro ao consultar modelo: {str(e)}"
-
-    def _convert_messages_to_prompt(self, messages: List[Dict]) -> str:
-        """Converte mensagens para prompt simples."""
-        prompt_parts = []
-        for msg in messages:
-            role = msg['role']
-            content = msg['content']
-            
-            if role == 'system':
-                prompt_parts.append(f"Sistema: {content}")
-            elif role == 'user':
-                prompt_parts.append(f"Usuário: {content}")
-            elif role == 'assistant':
-                prompt_parts.append(f"Assistente: {content}")
-                
-        prompt_parts.append("Assistente:")
-        return "\n\n".join(prompt_parts)
+            error_msg = str(e)
+            if 'loading' in error_msg.lower():
+                return f"Modelo {model_name} está carregando. Tente novamente em alguns minutos."
+            elif 'rate limit' in error_msg.lower():
+                return "Rate limit atingido. Aguarde alguns momentos e tente novamente."
+            elif 'token' in error_msg.lower() or 'unauthorized' in error_msg.lower():
+                return "Erro de autenticação. Verifique seu token HuggingFace."
+            else:
+                return f"Erro ao consultar modelo: {error_msg}"
 
-# --- Função para Testar e Atualizar Modelos (SEMPRE RETORNA TODOS) ---
+# --- Função para Testar e Atualizar Modelos ---
 def test_and_update_models() -> int:
     """Testa modelos e adiciona TODOS à lista MODELS, independente da disponibilidade."""
     print("Testando disponibilidade dos modelos...")
     print(f"Token HF disponível: {'Sim' if HF_TOKEN else 'Não'}")
     print("-" * 60)
     
-    api_client = HuggingFaceAPIClient(HF_TOKEN)
+    inference_client = HuggingFaceInferenceClient(HF_TOKEN)
     model_status = {}  # Para armazenar status de cada modelo
     
     # Testa todos os modelos mas adiciona TODOS à lista MODELS
     for model_label, model_name in NEW_MODELS_TO_TEST:
         try:
-            is_available, message = api_client.check_model_status(model_name)
+            is_available, message = inference_client.check_model_status(model_name)
             
             # Armazena o status para exibição
             model_status[model_label] = {
@@ -364,7 +345,7 @@ def test_and_update_models() -> int:
             }
             
         # Pausa para evitar rate limiting
-        time.sleep(2)
+        time.sleep(3)
     
     # SEMPRE adiciona TODOS os modelos, independente da disponibilidade
     global MODELS
@@ -453,14 +434,13 @@ def responder_como_aldo(session_id: str, pergunta: str, modelo: str = None) -> s
         {"role": "user", "content": mensagem_usuario}
     ]
     
-    # Faz requisição
-    api_client = HuggingFaceAPIClient(HF_TOKEN)
+    # Faz requisição usando InferenceClient
+    inference_client = HuggingFaceInferenceClient(HF_TOKEN)
     model_name = MODELS[modelo]
-    resposta = api_client.query_model(model_name, messages)
+    resposta = inference_client.query_model(model_name, messages)
     
-    # Adiciona informação sobre tentativa com modelo indisponível
-    if "Erro na requisição:" in resposta or "Erro ao consultar modelo:" in resposta:
-        resposta += f"\n\n*Nota: Tentativa feita com {modelo}, mas o modelo pode estar indisponível. Tente outro modelo ou aguarde.*"
+    # Adiciona informação sobre modelo usado
+    resposta += f"\n\n*Resposta gerada pelo modelo: {modelo} ({model_name})*"
     
     # Salva na memória
     add_to_memory(session_id, pergunta, resposta)
@@ -472,6 +452,15 @@ def inicializar_sistema():
     print("Inicializando Chatbot Dr. Aldo...")
     print("=" * 50)
     
+    # Verificar se huggingface_hub está instalado
+    try:
+        from huggingface_hub import InferenceClient
+        print("✓ huggingface_hub disponível")
+    except ImportError:
+        print("⚠ AVISO: huggingface_hub não encontrado!")
+        print("Execute: pip install huggingface_hub")
+        return False, {}
+    
     # Testa modelos (agora sempre retorna todos)
     num_total_models = test_and_update_models()
     
@@ -515,13 +504,15 @@ if __name__ == "__main__":
         print(f"\n3. {clear_memory(session_id)}")
         
         print("\n" + "="*50)
-        print("SISTEMA PRONTO PARA USO! ")
+        print("SISTEMA PRONTO PARA USO!")
         print("="*50)
+        print("✓ Usando huggingface_hub.InferenceClient")
         print("⚠ Lembre-se: nem todos os modelos podem estar funcionais.")
         print("⚠ Teste diferentes modelos se encontrar erros.")
         
     else:
         print("\n" + "="*50)
-        print("ERRO INESPERADO NA INICIALIZAÇÃO")
+        print("ERRO NA INICIALIZAÇÃO")
         print("="*50)
-        print("Verifique as configurações e tente novamente.")
+        print("Instale as dependências necessárias:")
+        print("pip install huggingface_hub")