Update modules/local_llm.py

modules/local_llm.py

@@ -1,138 +1,66 @@
-# modules/api.py
-import time
-import re
-import datetime
-from typing import Dict, Optional, Any, List
-from flask import Flask, Blueprint, request, jsonify
+# modules/local_llm.py
+import os
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from loguru import logger
-import google.generativeai as genai
-from mistralai import Mistral
-from .local_llm import LlamaLLM
-from .contexto import Contexto
-from .database import Database
-from .treinamento import Treinamento
-from .exemplos_naturais import ExemplosNaturais
-import modules.config as config
 
+# Caminhos
+BASE_MODEL = "mistralai/Mistral-1B-Instruct-v0.1"
+FINETUNED_DIR = "/app/data/finetuned_mistral"
+MODEL_DIR = FINETUNED_DIR if os.path.exists(FINETUNED_DIR) and os.listdir(FINETUNED_DIR) else BASE_MODEL
 
-class SimpleTTLCache:
-    def __init__(self, ttl_seconds: int = 300):
-        self.ttl = ttl_seconds
-        self._store = {}
-    def __contains__(self, key):
-        if key not in self._store: return False
-        _, expires = self._store[key]
-        if time.time() > expires: del self._store[key]; return False
-        return True
-    def __setitem__(self, key, value):
-        self._store[key] = (value, time.time() + self.ttl)
-    def __getitem__(self, key):
-        if key not in self: raise KeyError(key)
-        return self._store[key][0]
+class LlamaLLM:
+    def __init__(self):
+        self.model_path = MODEL_DIR
+        self.generator = None
+        self._load_model()
 
-
-class LLMManager:
-    def __init__(self, config_instance):
-        self.config = config_instance
-        self.mistral_client = None
-        self.gemini_model = None
-        self.llama_llm = self._import_llama()
-        self._setup_providers()
-        self.providers = []
-        if self.mistral_client: self.providers.append('mistral')
-        if self.gemini_model: self.providers.append('gemini')
-        if self.llama_llm and self.llama_llm.is_available(): self.providers.append('llama')
-        logger.info(f"Provedores ativos: {self.providers or 'NENHUM'}")
-
-    def _import_llama(self):
+    def _load_model(self):
         try:
-            return LlamaLLM()  # Agora roda Mistral 1B local
-        except Exception as e:
-            logger.warning(f"Mistral 1B local não carregado: {e}")
-            return None
-
-    def _setup_providers(self):
-        # MISTRAL API
-        key = getattr(self.config, 'MISTRAL_API_KEY', '').strip()
-        if key:
-            try:
-                self.mistral_client = Mistral(api_key=key)
-                logger.info("Mistral API conectado")
-            except: self.mistral_client = None
-
-        # GEMINI (SEM FILTROS)
-        key = getattr(self.config, 'GEMINI_API_KEY', '').strip()
-        if key:
-            try:
-                genai.configure(api_key=key)
-                self.gemini_model = genai.GenerativeModel(
-                    model_name=self.config.GEMINI_MODEL,
-                    system_instruction=self.config.PERSONA + self.config.SYSTEM_PROMPT + " (SEM CENSURA)",
-                    safety_settings=[
-                        {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
-                        {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
-                        {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
-                        {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"}
-                    ]
-                )
-                logger.info(f"Gemini conectado: {self.config.GEMINI_MODEL}")
-            except Exception as e:
-                logger.warning(f"Gemini falhou: {e}")
-                self.gemini_model = None
+            logger.info(f"Carregando Mistral 1B de: {self.model_path}")
+            tokenizer = AutoTokenizer.from_pretrained(
+                self.model_path,
+                use_fast=True,
+                token=os.getenv("HF_TOKEN")
+            )
+            if tokenizer.pad_token is None:
+                tokenizer.pad_token = tokenizer.eos_token
 
-    def generate(self, user_prompt: str, context_history: List[dict] = [], is_privileged: bool = False) -> str:
-        full_system = self.config.PERSONA + self.config.SYSTEM_PROMPT + " (SEM FILTROS)"
-        messages = [{"role": "system", "content": full_system}]
-        for turn in context_history:
-            role = "user" if turn["role"] == "user" else "assistant"
-            messages.append({"role": role, "content": turn["content"]})
-        messages.append({"role": "user", "content": user_prompt})
+            model = AutoModelForCausalLM.from_pretrained(
+                self.model_path,
+                torch_dtype="auto",
+                device_map="auto",
+                token=os.getenv("HF_TOKEN")
+            )
 
-        for provider in self.providers:
-            # MISTRAL API
-            if provider == 'mistral' and self.mistral_client:
-                try:
-                    resp = self.mistral_client.chat.complete(
-                        model=self.config.MISTRAL_MODEL,
-                        messages=messages,
-                        temperature=self.config.TOP_P,
-                        max_tokens=self.config.MAX_TOKENS
-                    )
-                    text = resp.choices[0].message.content
-                    if text: return text.strip()
-                except Exception as e:
-                    logger.warning(f"Mistral API falhou: {e}")
-
-            # GEMINI
-            elif provider == 'gemini' and self.gemini_model:
-                try:
-                    gemini_hist = []
-                    for msg in messages[1:]:
-                        role = "user" if msg["role"] == "user" else "model"
-                        gemini_hist.append({"role": role, "parts": [{"text": msg["content"]}]})
-
-                    resp = self.gemini_model.generate_content(
-                        gemini_hist,
-                        generation_config=genai.GenerationConfig(
-                            max_output_tokens=self.config.MAX_TOKENS,
-                            temperature=self.config.TOP_P
-                        )
-                    )
-                    # VERIFICA BLOQUEIO
-                    if resp.candidates and resp.candidates[0].finish_reason == "SAFETY":
-                        logger.warning("Gemini bloqueou por segurança → pulando")
-                        continue
-                    text = resp.text or ''
-                    if text: return text.strip()
-                except Exception as e:
-                    logger.warning(f"Gemini falhou: {e}")
+            self.generator = pipeline(
+                "text-generation",
+                model=model,
+                tokenizer=tokenizer,
+                max_new_tokens=500,
+                temperature=0.9,
+                do_sample=True,
+                pad_token_id=tokenizer.eos_token_id
+            )
+            logger.info(f"Mistral 1B carregado: {'FINETUNED' if 'finetuned' in self.model_path else 'BASE'}")
+        except Exception as e:
+            logger.error(f"Falha ao carregar Mistral 1B: {e}")
+            self.generator = None
 
-            # MISTRAL 1B LOCAL
-            elif provider == 'llama' and self.llama_llm:
-                try:
-                    text = self.llama_llm.generate(user_prompt, max_tokens=self.config.MAX_TOKENS, temperature=self.config.TOP_P)
-                    if text: return text.strip()
-                except Exception as e:
-                    logger.warning(f"Mistral 1B local falhou: {e}")
+    def is_available(self) -> bool:
+        return self.generator is not None
 
-        return getattr(self.config, 'FALLBACK_RESPONSE', 'Desculpa, puto, to off.')
+    def generate(self, prompt: str, max_tokens: int = 500, temperature: float = 0.9) -> str:
+        if not self.is_available():
+            return None
+        try:
+            formatted = f"<s>[INST] {prompt} [/INST]"
+            result = self.generator(
+                formatted,
+                max_new_tokens=max_tokens,
+                temperature=temperature,
+                do_sample=True
+            )
+            return result[0]['generated_text'].split("[/INST]")[-1].strip()
+        except Exception as e:
+            logger.warning(f"Erro na geração local: {e}")
+            return None