src/csv_agent.py

from database_manager import DatabaseManager
from prompts import *
from typing import Optional, Dict, Any
from langchain_mistralai import ChatMistralAI
from langchain_core.tools import tool
from langchain_core.messages import HumanMessage, SystemMessage
from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from dotenv import load_dotenv
import os
import json
import re

load_dotenv()

class CSVAgent:
    """Agente para análise de dados CSV usando LangChain e Mistral AI"""

    def __init__(self,
                 database_manager: DatabaseManager,
                 general_llm: Optional[ChatMistralAI] = None,
                 sql_llm: Optional[ChatMistralAI] = None):
        """
        Inicializa o agente CSV

        Args:
            database_manager: Instância do DatabaseManager
            general_llm: LLM para uso geral (opcional)
            sql_llm: LLM para geração de SQL (opcional)
        """
        self.db_manager = database_manager

        # Configura LLMs padrão se não fornecidos
        if general_llm is None:
            self.llm = ChatMistralAI(
                model="mistral-medium-latest",
                api_key=os.getenv('MISTRAL_API_KEY'),
                temperature=0.1
            )
        else:
            self.llm = general_llm

        if sql_llm is None:
            self.sql_llm = ChatMistralAI(
                model="codestral-latest",
                api_key=os.getenv('MISTRAL_API_KEY'),
                temperature=0.1
            )
        else:
            self.sql_llm = sql_llm

        # Configura ferramentas e agente
        self._setup_tools_and_agent()

    def _setup_tools_and_agent(self):
        """Configura ferramentas e agente do LangChain"""

        # Cria a ferramenta de consulta ao banco de dados
        @tool
        def consultar_banco_dados(consulta_linguagem_natural: str) -> str:
            """
            Executa consultas SQL nos dados CSV carregados.
            Use esta ferramenta quando o usuário pedir dados específicos, contagens, estatísticas,
            filtragem, análise ou qualquer operação que envolva os dados das tabelas.

            Args:
                consulta_linguagem_natural: A pergunta do usuário sobre os dados

            Returns:
                String JSON com os resultados da consulta
            """
            try:
                print(f"Executando consulta no banco para: {consulta_linguagem_natural}")

                # Gera consulta SQL usando Codestral
                sql_query = self._generate_sql_query(consulta_linguagem_natural)
                print(f"SQL gerado: {sql_query}")

                # Executa consulta
                results = self.db_manager.execute_sql_query(sql_query)
                print(f"Consulta retornou {len(results)} resultados")

                return json.dumps({
                    "sucesso": True,
                    "consulta_sql": sql_query,
                    "resultados": results,
                    "total_resultados": len(results),
                    "tabelas_disponiveis": self.db_manager.get_tables_list()
                }, indent=2, default=str, ensure_ascii=False)

            except Exception as e:
                print(f"Erro na consulta ao banco: {str(e)}")
                return json.dumps({
                    "sucesso": False,
                    "erro": str(e),
                    "consulta_sql": None,
                    "resultados": [],
                    "total_resultados": 0,
                    "tabelas_disponiveis": self.db_manager.get_tables_list()
                }, ensure_ascii=False)

        # Armazena a ferramenta
        self.database_tool = consultar_banco_dados

        # Cria prompt do sistema
        self.system_prompt = get_system_prompt(
            database_schema=self.db_manager.get_database_schema(),
            tables_list=', '.join(self.db_manager.get_tables_list())
        )
        # Cria lista de ferramentas
        self.tools = [self.database_tool]

        # Cria agente
        try:
            prompt = ChatPromptTemplate.from_messages([
                ("system", self.system_prompt),
                ("human", "{input}"),
                MessagesPlaceholder("agent_scratchpad")
            ])

            self.agent = create_tool_calling_agent(self.llm, self.tools, prompt)
            self.agent_executor = AgentExecutor(
                agent=self.agent,
                tools=self.tools,
                verbose=True,
                handle_parsing_errors=True,
                max_iterations=3,
                return_intermediate_steps=True
            )

        except Exception as e:
            print(f"Aviso: Não foi possível criar o executor do agente: {e}")
            self.agent_executor = None

    def _generate_sql_query(self, user_query: str) -> str:
        """Gera consulta SQL a partir de linguagem natural usando LLM"""
        sql_prompt = get_sql_prompt(
            database_schema=self.db_manager.get_database_schema(),
            tables_list=', '.join(self.db_manager.get_tables_list()),
            user_query=user_query
        )

        try:
            response = self.sql_llm.invoke(sql_prompt)

            # Limpa a consulta SQL
            sql_query = response.content
            sql_query = re.sub(r'```sql\n?', '', sql_query)
            sql_query = re.sub(r'```\n?', '', sql_query)
            sql_query = sql_query.strip()

            # Remove ponto e vírgula final se presente
            if sql_query.endswith(';'):
                sql_query = sql_query[:-1]

            return sql_query

        except Exception as e:
            raise Exception(f"Erro ao gerar SQL: {str(e)}")

    def query(self, user_input: str) -> Dict[str, Any]:
        """
        Método principal para processar consultas do usuário

        Args:
            user_input: Consulta em linguagem natural do usuário

        Returns:
            Dicionário contendo resposta e metadados
        """
        try:
            if self.agent_executor:
                print(f"Processando consulta com agente Mistral: {user_input}")

                # Executa com agente
                result = self.agent_executor.invoke({"input": user_input})

                # Captura a consulta SQL dos intermediate_steps
                sql_query = None
                function_called = False

                for step in result.get("intermediate_steps", []):
                    if len(step) > 1 and "consultar_banco_dados" in str(step[0]):
                        function_called = True
                        try:
                            resposta = json.loads(step[1])
                            sql_query = resposta.get("consulta_sql")
                            break
                        except:
                            pass
                return {
                    "resposta": result["output"],
                    "tipo_consulta": "executor_agent",
                    "funcao_chamada": function_called,
                    "consulta_sql": sql_query,
                    "sucesso": True,
                    "modelo_usado": self.llm.model,
                    "tabelas_disponiveis": self.db_manager.get_tables_list(),
                    "resultado_agente_bruto": result
                }

        except Exception as e:
            error_msg = f"Encontrei um erro: {str(e)}"
            print(f"Erro no processamento da consulta: {str(e)}")

            return {
                "resposta": error_msg,
                "tipo_consulta": "erro",
                "funcao_chamada": False,
                "sucesso": False,
                "erro": str(e),
                "tabelas_disponiveis": self.db_manager.get_tables_list()
            }