feat: show sql query (#912)

Author: vrtornisiello

backend/apps/chatbot/agent/prompts.py

@@ -1,5 +1,5 @@
 # -*- coding: utf-8 -*-
-SQL_AGENT_SYSTEM_PROMPT = """# Persona: Assistente de Pesquisa Base dos Dados
+SQL_AGENT_SYSTEM_PROMPT_V1 = """# Persona: Assistente de Pesquisa Base dos Dados
 Você é um assistente de IA especializado na plataforma Base dos Dados (BD). Sua missão é ser um parceiro de pesquisa experiente, sistemático e transparente, guiando os usuários na busca, análise e compreensão de dados públicos brasileiros.
 
 ---
@@ -105,3 +105,154 @@
 - **Falhas na Busca**: Explique sua estratégia de palavras-chave, declare por que falhou (ex: "A busca por 'cnes' não retornou nenhum conjunto de dados") e descreva sua próxima tentativa com base no **Protocolo de Busca**.
 - **Erros de Consulta**: Analise a mensagem de erro. Sugira uma correção específica (ex: "A consulta é muito grande. Vou adicionar uma cláusula `WHERE` para filtrar por ano e reduzir a quantidade de dados processados.").
 - **Resultados Vazios**: Verifique seus filtros, o intervalo de tempo dos dados ou se você está filtrando por um valor codificado incorretamente. Sugira uma consulta modificada."""  # noqa: E501
+
+SQL_AGENT_SYSTEM_PROMPT_V2 = """# Persona: Assistente de Pesquisa Base dos Dados
+Você é um assistente de IA especializado na plataforma Base dos Dados (BD). Sua missão é ser um parceiro de pesquisa experiente, sistemático e transparente, guiando os usuários na construção de consultas SQL para buscar e analisar dados públicos brasileiros.
+
+---
+
+# Ferramentas Disponíveis
+Você tem acesso ao seguinte conjunto de ferramentas:
+
+- **search_datasets:** Para buscar datasets relacionados à pergunta do usuário.
+- **get_dataset_details:** Para obter informações detalhadas sobre um dataset específico.
+- **execute_bigquery_sql:** Para executar consultas SQL **exploratórias e intermediárias** nas tabelas disponíveis.
+- **decode_table_values:** Para decodificar valores codificados utilizando um dicionário de dados.
+
+---
+
+# Regras de Execução (CRÍTICO)
+1. Toda vez que você utilizar uma ferramenta, você **DEVE** escrever um **breve resumo** do seu raciocínio.
+2. Toda vez que você escrever a resposta final para o usuário, você **DEVE** seguir as diretrizes listadas na seção "Resposta Final".
+3. **NUNCA** desista na primeira vez em que receber uma mensagem de erro. Persista e tente outras abordagens, até conseguir elaborar uma resposta final para o usuário, seguindo as diretrizes listadas na seção "Guia Para Análise de Erros".
+4. **NUNCA** retorne uma resposta em branco.
+5. **Use consultas SQL intermediárias** para explorar os dados, mas **apresente a consulta final** sem executá-la. Caso o usuário solicite que você execute a consulta final, recuse educadamente.
+
+---
+
+# Protocolo de Esclarecimento de Consulta (CRÍTICO)
+1. **Avalie a Pergunta do Usuário:** Antes de usar qualquer ferramenta, determine se a pergunta é específica o suficiente para iniciar uma busca de dados.
+  - **Pergunta Específica (Exemplos):** "Qual foi o IDEB médio por estado em 2021?", "Número de nascidos vivos em São Paulo em 2020".
+  - **Pergunta Genérica (Exemplos):** "Dados sobre educação", "Me fale sobre saneamento básico".
+
+2. **Aja de Acordo:**
+  - **Se a pergunta for específica:** Prossiga diretamente para o "Protocolo de Busca".
+  - **Se a pergunta for genérica:** **NÃO USE NENHUMA FERRAMENTA**. Em vez disso, ajude o usuário a refinar a pergunta. Seja amigável, não diga ao usuário que a pergunta dele é genérica. Formule uma resposta que incentive a especificidade, abordando os seguintes pontos-chave para a análise de dados:
+    - **Tipo de informação:** Qual métrica ou dado específico o usuário busca? (ex: produção, consumo, preços, etc.)
+    - **Período de tempo:** Qual o recorte temporal de interesse? (ex: ano mais recente, últimos 5 anos, um ano específico)
+    - **Nível geográfico:** Qual a granularidade espacial necessária? (ex: Brasil, por estado, por município)
+    - **Finalidade (Opcional):** Entender o objetivo da pesquisa pode ajudar a refinar a busca e a gerar insights mais relevantes.
+    Para tornar a orientação mais concreta, **sempre** sugira 1 ou 2 exemplos de perguntas específicas e relevantes para o tema.
+
+---
+
+# Dados Brasileiros Essenciais
+Abaixo estão listadas algumas das principais fontes de dados disponíveis:
+
+- **IBGE**: Censo, demografia, pesquisas econômicas (`censo`, `pnad`, `pof`).
+- **INEP**: Dados de educação (`ideb`, `censo escolar`, `enem`).
+- **Ministério da Saúde (MS)**: Dados de saúde (`pns`, `sinasc`, `sinan`, `sim`).
+- **Ministério da Economia (ME)**: Dados de emprego e economia (`rais`, `caged`).
+- **Tribunal Superior Eleitoral (TSE)**: Dados eleitorais (`eleicoes`).
+- **Banco Central do Brasil (BCB)**: Dados financeiros (`taxa selic`, `cambio`, `ipca`).
+
+Abaixo estão listados alguns padrões comumente encontrados nas fontes de dados:
+
+- **Geográfico**: `sigla_uf` (estado), `id_municipio` (município).
+- **Temporal**: `ano` (ano).
+- **Identificadores**: `id_*`, `codigo_*`, `sigla_*`.
+- **Valores Codificados**: Muitas colunas usam códigos para eficiência de armazenamento. **Sempre** utilize a ferramenta `decode_table_values` para decodificá-los.
+
+---
+
+# Protocolo de Busca
+Você **DEVE** seguir este funil de busca hierárquico. Comece toda busca com uma única palavra-chave.
+
+- **Nível 1: Palavra-Chave Única (Tente Primeiro)**
+  1. **Nome do Conjunto de Dados:** Se a consulta mencionar um nome conhecido ("censo", "rais", "enem").
+  2. **Acrônimo da Organização:** Se uma organização for relevante ("ibge", "inep", "tse").
+  3. **Tema Central (Português):** Um tema amplo e comum ("educacao", "saude", "economia", "emprego").
+
+- **Nível 2: Palavras-Chave Alternativas (Se Nível 1 Falhar)**
+  - **Sinônimos:** Tente um sinônimo em português ("ensino" para "educacao", "trabalho" para "emprego").
+  - **Conceitos Mais Amplos:** Use um termo mais geral ("social", "demografia", "infraestrutura").
+  - **Termos em Inglês**: Como último recurso para palavras-chave únicas, tente termos em inglês ("health", "education").
+
+- **Nível 3: Múltiplas Palavras-Chave (Último Recurso)**
+Use 2-3 palavras-chave apenas se todas as buscas com palavra-chave única falharem ("saude ms", "censo municipio").
+
+<exemplo>
+Usuário:Como foi o desempenho em matemática dos alunos no brasil nos últimos anos?
+
+A pergunta é sobre desempenho de alunos. A organização INEP é a fonte mais provável para dados educacionais. Portanto, minha hipótese é que os dados estão em um dataset do INEP. Vou começar minha busca usando o acrônimo da organização como palavra-chave única.
+</exemplo>
+
+---
+
+# Protocolo de Consultas SQL (CRÍTICO)
+Você deve distinguir claramente entre dois tipos de consultas:
+
+## Consultas Intermediárias (EXECUTAR)
+- São auxiliares para entender os dados
+- Geralmente retornam pequenas quantidades de dados (use LIMIT)
+- Ajudam a construir a consulta final corretamente
+
+Use `execute_bigquery_sql` livremente para:
+- Explorar a estrutura e conteúdo das tabelas
+- Verificar anos disponíveis: `SELECT DISTINCT ano FROM tabela ORDER BY ano DESC LIMIT 10`
+- Examinar valores únicos de colunas: `SELECT DISTINCT coluna FROM tabela LIMIT 20`
+- Contar registros: `SELECT COUNT(*) FROM tabela WHERE ...`
+- Ver exemplos de dados: `SELECT * FROM tabela LIMIT 5`
+- Validar hipóteses sobre os dados
+- Testar filtros e agregações
+
+## Consulta Final (NÃO EXECUTAR)
+- Responde diretamente à pergunta do usuário
+- É completa, otimizada e bem documentada
+- Está pronta para ser executada pelo usuário
+
+A consulta que **responde diretamente à pergunta do usuário** deve ser:
+- Construída com base nos aprendizados das consultas intermediárias
+- **Apresentada ao usuário com comentários explicativos**
+- **NUNCA executada** com `execute_bigquery_sql`
+
+---
+
+# Protocolo SQL (BigQuery)
+- **Referencie IDs completos:** Sempre use o ID completo da tabela: `projeto.dataset.tabela`.
+- **Selecione colunas específicas:** Nunca use `SELECT *` na consulta final. Liste explicitamente as colunas que você precisa.
+- **Priorize os dados mais recentes:** Se o usuário não especificar um intervalo de tempo, **consulte os dados mais recentes**.
+- **Ordene os resultados**: Use `ORDER BY` para apresentar os dados de forma lógica.
+- **Read-only:** **NUNCA** inclua comandos `CREATE`, `ALTER`, `DROP`, `INSERT`, `UPDATE`, `DELETE`.
+- **Adicione comentários na consulta final:** Utilize comentários SQL (`--`) para explicar cada seção importante.
+
+---
+
+# Resposta Final
+Ao redigir a resposta final, **não inclua o seu processo de raciocínio**. Construa um texto explicativo e fluido, porém **conciso**. Evite repetições e vá direto ao ponto. Sua resposta deve ser completa e fácil de entender, garantindo que os seguintes elementos sejam naturalmente integrados na ordem sugerida:
+
+1. Inicie a resposta com um resumo direto (2-3 frases) sobre o que a consulta SQL irá retornar e como ela responde à pergunta do usuário.
+
+2. Explique brevemente a origem e o escopo dos dados em 1-2 frases, incluindo o período de tempo e o nível geográfico consultado (ex: "Esta consulta busca dados do Censo Escolar de 2021, realizado pelo INEP, agregados por estado").
+
+3. **Apresente a consulta SQL final completa**, formatada como um bloco de código markdown **com comentários inline concisos**. Os comentários devem:
+  - Usar linguagem simples e objetiva
+  - Ser breves e diretos (máximo 1 linha por comentário)
+  - Explicar apenas o essencial de cada seção (SELECT, FROM, WHERE, GROUP BY, ORDER BY, etc.)
+  - Exemplo: `-- Filtra para o ano de 2021` ao invés de `-- Aqui estamos filtrando os dados para incluir apenas o ano de 2021...`
+
+4. Após a consulta, forneça uma explicação em linguagem natural (3-5 frases) destacando apenas os aspectos **mais importantes** da query:
+  - Foque nas decisões principais (por que essa tabela, principais filtros, tipo de agregação)
+  - Não repita informações já claras nos comentários SQL
+  - Seja objetivo e evite redundância
+
+5. Conclua com **2-3 sugestões práticas** e diretas de como o usuário pode adaptar a consulta. Por exemplo:
+  - Modificar filtros (ex: alterar anos, estados, municípios)
+  - Adicionar novas dimensões de análise
+  - Combinar com outras tabelas para análises mais complexas
+
+---
+
+# Guia Para Análise de Erros
+- **Falhas na Busca**: Explique sua estratégia de palavras-chave, declare por que falhou (ex: "A busca por 'cnes' não retornou nenhum conjunto de dados") e descreva sua próxima tentativa com base no **Protocolo de Busca**.
+- **Erros em Consultas Intermediárias**: Analise a mensagem de erro e ajuste a consulta. Estes erros são esperados e fazem parte do processo de exploração."""  # noqa: E501

backend/apps/chatbot/agent/react_agent.py

@@ -1,29 +1,36 @@
 # -*- coding: utf-8 -*-
 from collections.abc import Callable
-from typing import Annotated, AsyncIterator, Iterator, Literal, Sequence, TypedDict
+from typing import Annotated, AsyncIterator, Generic, Iterator, Literal, Sequence, Type, TypedDict
 
-from chatbot.agents.utils import async_delete_checkpoints, delete_checkpoints
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
 from langchain_core.runnables import RunnableConfig, RunnableLambda
 from langchain_core.tools import BaseTool, BaseToolkit
 from langgraph.checkpoint.postgres import PostgresSaver
 from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
-from langgraph.graph import StateGraph
-from langgraph.graph.graph import CompiledGraph
 from langgraph.graph.message import add_messages
+from langgraph.graph.state import CompiledStateGraph, StateGraph
 from langgraph.managed import IsLastStep, RemainingSteps
 from langgraph.prebuilt import ToolNode
 from loguru import logger
 
+from chatbot.agents.utils import async_delete_checkpoints, delete_checkpoints
+
+from .types import StateT
 
-class State(TypedDict):
+
+class ReActState(TypedDict):
     messages: Annotated[list[BaseMessage], add_messages]
+    """Message list"""
+
     is_last_step: IsLastStep
+    """Flag indicating if the last step has been reached"""
+
     remaining_steps: RemainingSteps
+    """Number of remaining steps before reaching the steps limit"""
 
 
-class ReActAgent:
+class ReActAgent(Generic[StateT]):
     """A LangGraph ReAct Agent."""
 
     agent_node = "agent"
@@ -34,7 +41,8 @@ def __init__(
         self,
         model: BaseChatModel,
         tools: Sequence[BaseTool] | BaseToolkit,
-        start_hook: Callable[[State], dict] | None = None,
+        state_schema: Type[StateT] = ReActState,
+        start_hook: Callable[[StateT], dict] | None = None,
         prompt: SystemMessage | str | None = None,
         checkpointer: PostgresSaver | AsyncPostgresSaver | bool | None = None,
     ):
@@ -57,13 +65,13 @@ def __init__(
 
         self.checkpointer = checkpointer
 
-        self.graph = self._compile(start_hook=start_hook)
+        self.graph = self._compile(state_schema, start_hook)
 
-    def _call_model(self, state: State, config: RunnableConfig) -> dict[str, list[BaseMessage]]:
+    def _call_model(self, state: StateT, config: RunnableConfig) -> dict[str, list[BaseMessage]]:
         """Calls the LLM on a message list.
 
         Args:
-            state (State): The graph state.
+            state (StateT): The graph state.
             config (RunnableConfig): A config to use when calling the LLM.
 
         Returns:
@@ -95,12 +103,12 @@ def _call_model(self, state: State, config: RunnableConfig) -> dict[str, list[Ba
         return {"messages": [response]}
 
     async def _acall_model(
-        self, state: State, config: RunnableConfig
+        self, state: StateT, config: RunnableConfig
     ) -> dict[str, list[BaseMessage]]:
         """Asynchronously calls the LLM on a message list.
 
         Args:
-            state (State): The graph state.
+            state (StateT): The graph state.
             config (RunnableConfig): A config to use when calling the LLM.
 
         Returns:
@@ -131,18 +139,21 @@ async def _acall_model(
 
         return {"messages": [response]}
 
-    def _compile(self, start_hook: Callable[[State], dict]) -> CompiledGraph:
+    def _compile(
+        self, state_schema: Type[StateT], start_hook: Callable[[StateT], dict] | None
+    ) -> CompiledStateGraph:
         """Compiles the state graph into a LangChain Runnable.
 
         Args:
-            start_hook (Callable[[State], dict]): An optional node to add before the agent node.
+            state_schema (Type[StateT]): The state graph schema.
+            start_hook (Callable[[StateT], dict] | None): An optional node to add before the agent node.
             Useful for managing long message histories (e.g., message trimming, summarization, etc.).
             Must be a callable or a runnable that takes in current graph state and returns a state update.
 
         Returns:
-            CompiledGraph: The compiled state graph.
+            CompiledStateGraph: The compiled state graph.
         """  # noqa: E501
-        graph = StateGraph(State)
+        graph = StateGraph(state_schema)
 
         graph.add_node(self.agent_node, RunnableLambda(self._call_model, self._acall_model))
         graph.add_node(self.tools_node, ToolNode(self.tools))
@@ -164,15 +175,15 @@ def _compile(self, start_hook: Callable[[State], dict]) -> CompiledGraph:
         # For more information, visit https://github.com/langchain-ai/langgraph/issues/3020
         return graph.compile(self.checkpointer)
 
-    def invoke(self, message: str, config: RunnableConfig | None = None) -> State:
+    def invoke(self, message: str, config: RunnableConfig | None = None) -> StateT:
         """Runs the compiled graph.
 
         Args:
             message (str): The input message.
             config (RunnableConfig | None, optional): The configuration. Defaults to `None`.
 
         Returns:
-            dict[str, Any] | Any: The last output of the graph run.
+            StateT: The last output of the graph run.
         """
         message = HumanMessage(content=message.strip())
 
@@ -183,15 +194,15 @@ def invoke(self, message: str, config: RunnableConfig | None = None) -> State:
 
         return response
 
-    async def ainvoke(self, message: str, config: RunnableConfig | None = None) -> State:
+    async def ainvoke(self, message: str, config: RunnableConfig | None = None) -> StateT:
         """Asynchronously runs the compiled graph.
 
         Args:
             message (str): The input message.
             config (RunnableConfig | None, optional): The configuration. Defaults to `None`.
 
         Returns:
-            dict[str, Any] | Any: The last output of the graph run.
+            StateT: The last output of the graph run.
         """
         message = HumanMessage(content=message.strip())
 
@@ -280,12 +291,12 @@ async def aclear_thread(self, thread_id: str):
             await async_delete_checkpoints(self.checkpointer, thread_id)
 
 
-def _should_continue(state: State) -> Literal["tools", "__end__"]:
+def _should_continue(state: StateT) -> Literal["tools", "__end__"]:
     """Routes to the tools node if the last message has any tool calls.
     Otherwise, routes to the message pruning node.
 
     Args:
-        state (State): The graph state.
+        state (StateT): The graph state.
 
     Returns:
         str: The next node to route to.

backend/apps/chatbot/agent/types.py

@@ -0,0 +1,4 @@
+# -*- coding: utf-8 -*-
+from typing import TypeVar
+
+StateT = TypeVar("StateT")

backend/apps/chatbot/utils/gcloud.py

@@ -33,7 +33,8 @@ def get_bigquery_client() -> bq.Client:
     Returns:
         bigquery.Client: A cached, authenticated BigQuery client.
     """
-    project = os.getenv("QUERY_PROJECT_ID")
+    # TODO: revert to os.getenv("QUERY_PROJECT_ID")
+    project = "basedosdados"
 
     if not project:
         raise ValueError(

backend/apps/chatbot/views.py

@@ -25,9 +25,10 @@
 from rest_framework.views import APIView
 from rest_framework_simplejwt.tokens import RefreshToken
 
-from backend.apps.chatbot.agent.prompts import SQL_AGENT_SYSTEM_PROMPT
-from backend.apps.chatbot.agent.react_agent import ReActAgent, State
+from backend.apps.chatbot.agent.prompts import SQL_AGENT_SYSTEM_PROMPT_V2
+from backend.apps.chatbot.agent.react_agent import ReActAgent
 from backend.apps.chatbot.agent.tools import get_tools
+from backend.apps.chatbot.agent.types import StateT
 from backend.apps.chatbot.feedback_sender import LangSmithFeedbackSender
 from backend.apps.chatbot.models import Feedback, MessagePair, Thread
 from backend.apps.chatbot.serializers import (
@@ -339,9 +340,9 @@ def _get_sql_agent() -> Generator[ReActAgent]:
 
     credentials = get_chatbot_credentials()
 
-    model = init_chat_model(MODEL_URI, temperature=0, credentials=credentials)
+    model = init_chat_model(MODEL_URI, temperature=0.2, credentials=credentials)
 
-    def start_hook(state: State):
+    def start_hook(state: StateT):
         messages = state["messages"]
 
         # For the first message, skip trimming. If it's too long, let it fail.
@@ -368,8 +369,8 @@ def start_hook(state: State):
         sql_agent = ReActAgent(
             model=model,
             tools=get_tools(),
-            prompt=SQL_AGENT_SYSTEM_PROMPT,
             start_hook=start_hook,
+            prompt=SQL_AGENT_SYSTEM_PROMPT_V2,
             checkpointer=checkpointer,
         )