Cambie el modelo antiguo por el nuevo

mojica_agent.py

@@ -1,109 +1,54 @@
-from memory import Memory as ConversationMemory
+from memory import ConversationMemory
 from config import Config
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch, gc
 import unicodedata
-from typing import Dict, Tuple, Optional, Any
-import re
+from typing import Dict, Tuple
+import re 
 import pandas as pd
 import sqlite3
-from intelligent_question_router import IntelligentQuestionRouter
-
 
 class MojicaAgent:
     def __init__(self, config: Config):
         self.config = config
         self.memory = ConversationMemory()
-        self.essential_columns = [
-            {
-                "name": "Descripcion",
-                "type": "TEXT",
-                "description": "Nombre del producto",
-            },
-            {"name": "Cantidad", "type": "REAL", "description": "Unidades vendidas"},
-            {"name": "Cliente", "type": "TEXT", "description": "Código de cliente"},
-            {
-                "name": "Razon Social",
-                "type": "TEXT",
-                "description": "Nombre completo del cliente",
-            },
-            {"name": "Ciudad", "type": "TEXT", "description": "Ciudad del cliente"},
-            {
-                "name": "Fecha",
-                "type": "TEXT",
-                "description": "Fecha de venta (YYYY-MM-DD)",
-            },
-            {"name": "Neto", "type": "REAL", "description": "Valor neto de la venta"},
-        ]
-        self.schema = self._load_schema()
+        self.schema = self._load_schema() 
         self._safe_initializer_model()
 
-    def _initialize_model(self):
+    def _safe_initializer_model(self):
         def try_load_model():
-            self.tokenizer = AutoTokenizer.from_pretrained(self.config.MODEL_NAME)
-            self.model = AutoModelForCausalLM.from_pretrained(
-                self.config.MODEL_NAME,
-                device_map="auto",
-                torch_dtype="auto",
-                trust_remote_code=True,
-            ).eval()
+            dtype = torch.float16 if "cuda" in self.config.DEVICE else torch.float32
+            tokenizer = AutoTokenizer.from_pretrained(self.config.MODEL_NAME)
+            model = (
+                AutoModelForCausalLM.from_pretrained(
+                    self.config.MODEL_NAME, trust_remote_code=True, torch_dtype=dtype
+                )
+                .to(self.config.DEVICE)
+                .eval()
+            )  # eval porque solo se predice
+            return tokenizer, model
 
         try:
-            try_load_model()
+            self.tokenizer, self.model = try_load_model()
         except torch.cuda.OutOfMemoryError:
+            # Liberar memoria y volver a intentar
             gc.collect()
             torch.cuda.empty_cache()
             torch.cuda.ipc_collect()
-            try_load_model()
 
-    def _load_training_data(self):
-        training_examples = [
-            {"question": "productos más vendidos", "category": "producto"},
-            {"question": "mejor producto", "category": "producto"},
-            {"question": "clientes que más compran", "category": "cliente"},
-            {"question": "clientes inactivos", "category": "cliente"},
-        ]
-        try:
-            self.router.semantic_classifier.train(training_examples)
-        except Exception as e:
-            print(f"Error training semantic classifier: {e}")
-
-    def _validate_result_existing(self, result):
-        # Si es un string de error
-        if isinstance(result, str) and "Error" in result:
-            return False
-
-        # Si es un DataFrame vacío
-        if hasattr(result, "empty") and result.empty:
-            return False
-
-        # Si es una lista vacía
-        if isinstance(result, list) and len(result) == 0:
-            return False
-
-        # En cualquier otro caso, asumimos éxito
-        return True
-
-    def _initialize_database(self):
-        self.conn = sqlite3.connect(self.config.DB_PATH)
-
-        cursor = self.conn.cursor()
-        cursor.execute(f"DROP TABLE IF EXISTS {self.config.TABLE_NAME}")
-        self.conn.commit()
-        df = pd.read_csv(self.config.CSV_PATH, low_memory=False)
+            self.tokenizer, self.model = try_load_model()
 
-        real_cols = [
-            col["name"] for col in self.essential_columns if col["type"] == "REAL"
+    def _load_schema(self) -> Dict: 
+        conn = sqlite3.connect(self.config.DB_PATH)
+        cursor = conn.cursor()
+        cursor.execute(f"PRAGMA table_info({self.config.TABLE_NAME})")
+        columns = [
+            {"name": column[1], "type": column[2]} for column in cursor.fetchall()
         ]
-        for col in real_cols:
-            if col in df.columns:
-                df[col] = pd.to_numeric(df[col], errors="coerce")
-
-        df.to_sql(self.config.TABLE_NAME, self.conn, if_exists="replace", index=False)
-        self.schema = self._get_schema_structured()
-        # Configuracion de pandas:
-        pd.set_option("display.float_format", "{:,.2f}".format)
-
+        schema = {"table_name": self.config.TABLE_NAME, "columns": columns}
+        conn.close()
+        return schema
+        
     def _get_schema_structured(self) -> Dict:
         if self.memory.schema_cache:
             return self.memory.schema_cache
@@ -116,23 +61,99 @@ class MojicaAgent:
         self.memory.schema_cache = schema
         return schema
 
-    def _generate_sql_prompt(self, question: str) -> str:
-        memory_context = self.memory.get_relevant_memory(question)
+    def _build_prompt(self, question: str) -> str:
+        memory_context = self.memory.get_context(question)
         table_name = self.schema["table_name"]
-        # Uso del router
-        try:
-            examples_list = self.router.route_question(question)
-            # Convertir ejemplos a texto para el prompt
-            examples_text = "\n".join(
-                [f"-- P: '{ex['pregunta']}'\n{ex['query']}\n" for ex in examples_list]
-            )
-            question_type = "ROUTED_EXAMPLES"
-        except Exception as e:
-            print(f"Router failed, using manual detection: {e}")
-            # Fallback a detección manual
-            # question_type = self._detect_question_type_manual(question)
-            # examples_text = self.examples.get(question_type, "")
 
+        # 1. Detectar tipo de pregunta
+        question_type = (
+            "PRODUCTOS"
+            if "producto" in question.lower()
+            else "CLIENTES" if "cliente" in question.lower() else "GENERAL"
+        )
+
+        # 2. Ejemplos dinámicos
+        examples = {
+            "PRODUCTOS": (
+                "-- P: 'Top 10 productos más vendidos'\n"
+                'SELECT "Descripcion", SUM("Cantidad") AS total_vendido\n'
+                f'FROM "{table_name}"\n'
+                'WHERE "Descripcion" IS NOT NULL\n'
+                'GROUP BY "Descripcion"\n'
+                "ORDER BY total_vendido DESC\n"
+                "LIMIT 10;\n\n"
+                "-- P: 'Productos con mayor valor neto'\n"
+                'SELECT "Descripcion", SUM("Neto") AS valor_total\n'
+                f'FROM "{table_name}"\n'
+                'WHERE "Descripcion" IS NOT NULL\n'
+                'GROUP BY "Descripcion"\n'
+                "ORDER BY valor_total DESC\n"
+                "LIMIT 5;"
+            ),
+            "CLIENTES": (
+                "-- P: 'Top 5 clientes con mayor valor neto'\n"
+                'SELECT "Cliente", SUM("Neto") AS valor_total\n'
+                f'FROM "{table_name}"\n'
+                "WHERE \"Cliente\" IS NOT NULL AND \"Fecha\" BETWEEN '2025-01-01' AND '2025-12-31'\n"
+                'GROUP BY "Cliente"\n'
+                "ORDER BY valor_total DESC\n"
+                "LIMIT 5;\n\n"
+                "-- P: 'Clientes con más compras en marzo'\n"
+                'SELECT "Cliente", COUNT(*) AS total_compras\n'
+                f'FROM "{table_name}"\n'
+                "WHERE \"Cliente\" IS NOT NULL AND strftime('%m', \"Fecha\") = '03'\n"
+                'GROUP BY "Cliente"\n'
+                "ORDER BY total_compras DESC\n"
+                "LIMIT 10;\n\n"
+                "-- P: 'Clientes de Guadalajara con más compras'\n"
+                'SELECT "Cliente", "Razon Social", COUNT(*) AS total_compras\n'
+                f'FROM "{table_name}"\n'
+                'WHERE "Cliente" IS NOT NULL AND "Ciudad" = \'Guadalajara\'\n'
+                'GROUP BY "Cliente", "Razon Social"\n'
+                "ORDER BY total_compras DESC\n"
+                "LIMIT 10;"
+            ),
+            "GENERAL": (
+                "-- P: 'Ventas totales por mes'\n"
+                'SELECT strftime(\'%m\', "Fecha") AS mes, SUM("Neto") AS ventas\n'
+                f'FROM "{table_name}"\n'
+                "WHERE mes IS NOT NULL\n"
+                "GROUP BY mes\n"
+                "ORDER BY mes;\n\n"
+                "-- P: 'Producto menos vendido en 2025'\n"
+                'SELECT "Descripcion", SUM("Cantidad") AS total_vendido\n'
+                f'FROM "{table_name}"\n'
+                "WHERE \"Descripcion\" IS NOT NULL AND \"Fecha\" BETWEEN '2025-01-01' AND '2025-12-31'\n"
+                'GROUP BY "Descripcion"\n'
+                "ORDER BY total_vendido ASC\n"
+                "LIMIT 1;"
+            ),
+        }
+
+        # 3. Columnas esenciales
+        essential_columns = [
+            {
+                "name": "Descripcion",
+                "type": "TEXT",
+                "description": "Nombre del producto",
+            },
+            {"name": "Cantidad", "type": "REAL", "description": "Unidades vendidas"},
+            {"name": "Cliente", "type": "TEXT", "description": "Código de cliente"},
+            {
+                "name": "Razon Social",
+                "type": "TEXT",
+                "description": "Nombre completo del cliente",
+            },
+            {"name": "Ciudad", "type": "TEXT", "description": "Ciudad del cliente"},
+            {
+                "name": "Fecha",
+                "type": "TEXT",
+                "description": "Fecha de venta (YYYY-MM-DD)",
+            },
+            {"name": "Neto", "type": "REAL", "description": "Valor neto de la venta"},
+        ]
+
+        # 4. Prompt final con nueva regla
         return (
             f"""
     ### TAREA ###
@@ -143,7 +164,7 @@ class MojicaAgent:
             + "\n".join(
                 [
                     f"- {col['name']} ({col['type']}): {col['description']}"
-                    for col in self.essential_columns
+                    for col in essential_columns
                 ]
             )
             + f"""
@@ -151,32 +172,22 @@ class MojicaAgent:
     ### CONTEXTO (Últimas interacciones) ###
     {memory_context if memory_context else "Sin historial relevante"}
     
-    ### EJEMPLOS ###
-    {examples_text}
+    ### EJEMPLOS ({question_type}) ###
+    {examples[question_type]}
     
     ### REGLAS CRÍTICAS ###
     - Usar siempre nombres exactos de columnas
-    - Usar solo las columnas listadas
-    - Prohibido inventar columnas
-    - Para el nombre del cliente, usar SIEMPRE "Razon Social".
-    - Para un mes específico usar: strftime('%m', "Fecha") = 'MM'
-    - Para cantidades usar SUM("Cantidad"), para dinero usar SUM("Neto")
     - Agrupar por la dimensión principal (producto/cliente)
     - Ordenar DESC para 'más/mayor', ASC para 'menos/menor'
-    - Contesta siempre en el idioma en el que se te pregunta no traduzcas.
+    - Usar LIMIT para top N
     - Año actual: 2025
-    - No inventes columnas o tablas que no existan
-    - Para preguntas sobre clientes cero, SIEMPRE usar la subconsulta NOT IN con las últimas 4 semanas.
-    - Si se menciona una ciudad, incluir el filtro AND "Ciudad" LIKE '%...%'
-    - Usa LIMIT cuando se te pida un numero finito de datos
+    - Siempre terminar con un LIMIT = 1 en caso que se indique lo contrario
     - Para 'más vendido' usar SUM("Cantidad"), para 'mayor valor' usar SUM("Neto")
     - Usar "Razon Social" cuando pregunten por el nombre del cliente
     - Usar "Ciudad" para filtrar o agrupar por ubicación
     - Queda estrictamente prohibido usar acentos
     - **Siempre excluir valores nulos con 'IS NOT NULL' en las columnas usadas en WHERE, GROUP BY u ORDER BY**
-    - Para preguntas sobre ciudad SIEMPRE incluir "Ciudad" en la query
-    - Para busquedas por Descripcion siempre usar LIKE
-    - Mandar solo la cantidad de rows que el usuario pide.
+        
     ### PREGUNTA ACTUAL ###
     \"\"\"{question}\"\"\"
     
@@ -184,27 +195,7 @@ class MojicaAgent:
     """
         )
 
-    def _generate_analysis_prompt(self, question: str, result: Any) -> str:
-        return f"""
-        Basado EXCLUSIVAMENTE en estos datos: {result}
-
-        Responde esta pregunta: {question}
-
-        Reglas estrictas:
-        - Nunca inventes numeros
-        - Usa solo datos proporcionados
-        - Maximo una oracion
-        """
-
-    def _clean_analysis_output(self, ouput: str) -> Optional[str]:
-        pattern = r"Respuesta:([\s\S]+)"
-        match = re.search(pattern, ouput)
-        if match:
-            return match.group(1).strip()
-        else:
-            return "Sin análisis"
-
-    def _clean_sql_output(self, output: str) -> Optional[str]:
+    def _clean_sql_output(self, output: str) -> str:
         # Encuentra todas las posibles queries completas que terminen en ;
         sql_matches = list(
             re.finditer(
@@ -245,63 +236,38 @@ class MojicaAgent:
         # 2. Agregar LIMIT si no existe
         # ────────────────────────────────
         # Buscar si ya hay un LIMIT en la query
-        # if not re.search(r"\bLIMIT\s+\d+", sql, re.IGNORECASE):
-        #     # Insertar antes del último punto y coma
-        #     sql = sql[:-1] + " LIMIT 1;"  # puedes cambiar 100 por el valor default que quieras
-
-        validate_sql = self._validate_and_correct_sql(sql)
-        return validate_sql
-
-    def _validate_and_correct_sql(self, sql: str) -> str:
-        cur = self.conn.cursor()
-        cur.execute(f'PRAGMA table_info("{self.config.TABLE_NAME}")')
-        real_columns = [row[1] for row in cur.fetchall()]
-        column_lower_map = {col.lower(): col for col in real_columns}
-        aliases = {
-            "city": "Ciudad",
-            "client": "Cliente",
-            "razon_social": "Razon Social",
-            "razón social": "Razon Social",
-            "Sales": "sells",
-            '"Date"': "Fecha",
-            "mojica_Clientes": "sells",
-            "value_total": "valor_total",
-            "strstrftime": "strftime",
-        }
-        alias_map = {k.lower(): v for k, v in aliases.items()}
-
-        pattern = r"\b\w+\b"
-
-        def replace_column(m):
-            candidate = m.group(0)  # Palabra encontrada
-            key = candidate.lower()
-            # ¿Es una columna?
-            corrected = column_lower_map.get(key)
-            if corrected:
-                return corrected
+        if not re.search(r"\bLIMIT\s+\d+", sql, re.IGNORECASE):
+            # Insertar antes del último punto y coma
+            sql = (
+                sql[:-1] + " LIMIT 1;"
+            )  # puedes cambiar 100 por el valor default que quieras
 
-            # ¿Es una alias?
-            corrected = alias_map.get(key)
-            if corrected is not None:
-                return corrected
-            return candidate  # si no encuentra nada, lo deja igual
+        return sql
 
-        return re.sub(pattern, replace_column, sql).replace("\\", "")
-
-    def _execute_sql(self, sql: str) -> Any:
+    def _execute_sql(self, sql: str):
+        conn = sqlite3.connect(self.config.DB_PATH)
         try:
-            return pd.read_sql_query(sql, self.conn)
+            result = pd.read_sql_query(sql, conn)
+            conn.close()
+            return result
         except Exception as e:
-            return f"Error: {str(e)}"
+            return f"Error de ejecución: {str(e)}"
+        finally: 
+            conn.close()
+
+    def consult(self, question: str) -> Tuple[str, any]:
+        prompt = self._build_prompt(question)
 
-    def consult(self, question: str) -> Tuple[str, Any, str]:
-        sql_prompt = self._generate_sql_prompt(question)
         tokenized_input = self.tokenizer(
-            sql_prompt,
+            prompt,
             return_tensors="pt",
             truncation=True,
             max_length=self.config.MAX_TOKENS,
         ).to(self.config.DEVICE)
+
+        # Desactiva el cálculo de gradientes -> Siempre poner cuando se haga prediccion
+        #    - Reduce consumo de memoria
+        #    - Acelera inferencia
         with torch.no_grad():
             tokenized_output_model = self.model.generate(
                 **tokenized_input,
@@ -313,33 +279,17 @@ class MojicaAgent:
                 do_sample=True,
                 pad_token_id=self.tokenizer.eos_token_id,
             )
+
         output_model = self.tokenizer.decode(
             tokenized_output_model[0], skip_special_tokens=True
         )
-        sql = self._clean_sql_output(output_model)
 
-        # * Ejecución de SQL y generación de analisis
-        result = self._execute_sql(sql)
-        # * INICIO DE ANALISIS (COMENTADO)
-        # Analisis 
-        # analysis_prompt = self._generate_analysis_prompt(question, result)
-        # analyzed_token_input = self.tokenizer(
-        #     analysis_prompt,
-        #     return_tensors="pt",
-        #     truncation=True,
-        #     max_length=self.config.MAX_TOKENS,
-        # ).to(self.config.DEVICE)
-        # with torch.no_grad():
-        #     tokenized_analysis_output_model = self.model.generate(
-        #         **analyzed_token_input,
-        #         max_new_tokens=self.config.MAX_NEW_TOKENS,
-        #         temperature=0.65,
-        #     )
-        # analysis = self.tokenizer.decode(
-        #     tokenized_analysis_output_model[0], skip_special_tokens=True
-        # )
-        # analysis = self._clean_analysis_output(analysis)
-        # analysis <- LE quite ese parametro
-        # * FIN DE ANALISIS (COMENTADO)
-        self.memory.add_interaction(question=question, answer=result, sql=sql)
-        return sql, result
+        sql_query = self._clean_sql_output(output_model)
+
+        if not sql_query:
+            return "Error: No se pudo generar SQL válido" + "\n" + output_model, None
+
+        result = self._execute_sql(sql_query)
+        self.memory.add_interaction(question, sql_query, result)
+
+        return sql_query, result