From 18f33a4d742332e7854b6c0296d1f629a651339f Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Fri, 22 Aug 2025 16:25:41 -0300 Subject: [PATCH 01/23] create tools and react agent --- backend/apps/chatbot/tools.py | 420 ++++++++++++++++++++++++++++++++++ backend/apps/chatbot/views.py | 135 ++++++----- 2 files changed, 492 insertions(+), 63 deletions(-) create mode 100644 backend/apps/chatbot/tools.py diff --git a/backend/apps/chatbot/tools.py b/backend/apps/chatbot/tools.py new file mode 100644 index 00000000..17c74f2d --- /dev/null +++ b/backend/apps/chatbot/tools.py @@ -0,0 +1,420 @@ +# -*- coding: utf-8 -*- +import json +import os +from typing import Optional + +import httpx +from google.cloud import bigquery as bq +from langchain_core.tools import tool +from pydantic import BaseModel + +SEARCH_URL = "https://backend.basedosdados.org/search/" +GRAPHQL_URL = "https://backend.basedosdados.org/graphql" + +DATASET_OVERVIEW_QUERY = """ +query GetDatasetOverview($id: ID!) { + allDataset(id: $id, first: 1) { + edges { + node { + id + name + slug + description + organizations { + edges { + node { + name + slug + } + } + } + themes { + edges { + node { + name + } + } + } + tags { + edges { + node { + name + } + } + } + tables { + edges { + node { + id + name + slug + description + cloudTables { + edges { + node { + gcpProjectId + gcpDatasetId + gcpTableId + } + } + } + columns { + edges { + node { + id + name + description + bigqueryType { + name + } + } + } + } + } + } + } + } + } + } +} +""" + + +class Column(BaseModel): + name: str + type: str + description: Optional[str] + + +class Table(BaseModel): + id: str + gcp_id: Optional[str] + name: str + slug: Optional[str] + description: Optional[str] + columns: list[Column] + + +class DatasetOverview(BaseModel): + id: str + name: str + slug: Optional[str] + description: Optional[str] + tags: list[str] + themes: list[str] + organizations: list[str] + + +class Dataset(DatasetOverview): + tables: list[Table] + + +client = bq.Client(project=os.environ["QUERY_PROJECT_ID"]) + + +@tool +def search_datasets(query: str) -> str: + """Search for datasets in the Base dos Dados (BD) catalog using KEYWORDS ONLY. + + IMPORTANT: This search uses Elasticsearch - use individual KEYWORDS, NOT full sentences. + Use this as the FIRST STEP when exploring data. If no relevant datasets are found, + try searching again with different or more specific keywords. + + Args: + query (str): Individual keywords or short phrases (2-3 words max). The search engine + works best with single keywords. Use: + - Single topic words: "education", "health", "economy", "population", "crime" + - Geographic keywords: "municipal", "estado", "brasil", "sao_paulo", "rio" + - Data type keywords: "censo", "pnad", "rais", "sinasc", "datasus" + - Organization keywords: "ibge", "inep", "ans", "bacen", "tse" + + AVOID full sentences like "Brazilian population data by municipality" + INSTEAD use: "populacao" or "municipio" or "ibge censo" + + Returns: + str: JSON array of dataset overviews. If empty array [] or irrelevant results: + - Try different keywords (synonyms, Portuguese terms, abbreviations) + - Use more specific terms or broader terms + - Try organization names or data source acronyms + + Search Strategy - If initial search fails: + 1. Try Portuguese keywords: "educacao" instead of "education" + 2. Use government agency acronyms: "ibge", "inep", "anvisa", "ans" + 3. Try broader terms: "social" instead of specific indicators + 4. Use dataset-specific terms: "pnad", "censo", "rais", "caged" + 5. Search by geographic level: "municipio", "estado", "regiao" + + Next steps after using this tool: + 1. If no results or irrelevant results → Try different keywords and search again + 2. If good results → Review datasets to identify the most relevant ones + 3. Use get_dataset_details() with the dataset 'id' to explore structure + 4. Look for datasets from reputable organizations (IBGE, government agencies) + + Example successful searches: + - "censo" → finds census datasets + - "educacao" → finds education data + - "saude municipio" → finds municipal health data + - "ibge" → finds all IBGE datasets + - "eleicoes" → finds election datasets + """ + try: + with httpx.Client() as client: + response = client.get( + url=SEARCH_URL, + params={"q": query, "page_size": 10}, + timeout=httpx.Timeout(5.0, read=60.0), + ) + response.raise_for_status() + data: dict = response.json() + + datasets = data.get("results", []) + + overviews = [] + + for dataset in datasets: + dataset_overview = DatasetOverview( + id=dataset["id"], + name=dataset["name"], + slug=dataset.get("slug"), + description=dataset.get("description"), + tags=[tag["name"] for tag in dataset.get("tags", [])], + themes=[theme["name"] for theme in dataset.get("themes", [])], + organizations=[org["name"] for org in dataset.get("organizations", [])], + ) + overviews.append(dataset_overview.model_dump()) + + return json.dumps(overviews, ensure_ascii=False, indent=2) + except Exception as e: + return f"Error searching datasets:\n{e}" + + +@tool +def get_dataset_details(dataset_id: str) -> str: + """Get comprehensive details about a specific dataset including all tables and columns. + + This tool provides the complete structure of a dataset, showing all available tables, + their columns, data types, and BigQuery identifiers. Use this AFTER search_datasets() + to understand what data is available before writing SQL queries. + + Args: + dataset_id (str): The unique dataset ID obtained from search_datasets(). + This is typically a UUID-like string, not the human-readable name. + + Returns: + str: JSON object with complete dataset information including: + - Basic metadata (name, description, themes, organizations) + - tables: Array of all tables in the dataset with: + - gcp_id: Full BigQuery table reference (project.dataset.table) + - columns: All column names, types (STRING, INTEGER, DATE, etc.), and descriptions + - Table descriptions explaining what each table contains + + Next steps after using this tool: + 1. Identify which table(s) contain the data you need + 2. Note the column names and types for SQL query construction + 3. Use the gcp_id (BigQuery table reference) in your SQL queries + 4. Check column descriptions to understand what each field represents + 5. Plan your SQL query using execute_sql_query() + + Important notes: + - The 'gcp_id' field is the full BigQuery table reference you'll use in FROM clauses + - Column types help you write correct SQL (STRING needs quotes, dates need DATE functions) + - Some tables may be very large - consider using LIMIT in your queries + - Look for key identifier columns (like municipality codes, dates) for filtering + + Example workflow: + 1. search_datasets("IBGE census") → get dataset IDs + 2. get_dataset_details("abc-123-def") → explore table structure + 3. execute_sql_query("SELECT * FROM `basedosdados.br_ibge_censo.municipio` LIMIT 10") + """ + try: + with httpx.Client() as client: + response = client.post( + url=GRAPHQL_URL, + json={ + "query": DATASET_OVERVIEW_QUERY, + "variables": {"id": dataset_id}, + }, + timeout=httpx.Timeout(5.0, read=60.0), + ) + response.raise_for_status() + data: dict[str, dict[str, dict]] = response.json() + + dataset_edges = data.get("data").get("allDataset", {}).get("edges", []) + + if not dataset_edges: + return f"Dataset {dataset_id} not found" + + dataset = dataset_edges[0]["node"] + + dataset_id = dataset["id"] + dataset_name = dataset["name"] + dataset_slug = dataset.get("slug") + dataset_description = dataset.get("description") + + # Tags + dataset_tags = [] + for edge in dataset.get("tags", {}).get("edges", []): + if tag := edge.get("node", {}).get("name"): + dataset_tags.append(tag) + + # Themes + dataset_themes = [] + for edge in dataset.get("themes", {}).get("edges", []): + if theme := edge.get("node", {}).get("name"): + dataset_themes.append(theme) + + # Organizations + dataset_organizations = [] + for edge in dataset.get("organizations", {}).get("edges", []): + if org := edge.get("node", {}).get("name"): + dataset_organizations.append(org) + + # Tables + dataset_tables = [] + for edge in dataset.get("tables", {}).get("edges", []): + table = edge["node"] + + table_id = table["id"] + table_name = table["name"] + table_slug = table.get("slug") + table_description = table.get("description") + + cloud_table_edges = table["cloudTables"]["edges"] + if cloud_table_edges: + cloud_table = cloud_table_edges[0]["node"] + gcp_project_id = cloud_table["gcpProjectId"] + gcp_dataset_id = cloud_table["gcpDatasetId"] + gcp_table_id = cloud_table["gcpTableId"] + table_gcp_id = f"{gcp_project_id}.{gcp_dataset_id}.{gcp_table_id}" + else: + table_gcp_id = None + + table_columns = [] + for edge in table["columns"]["edges"]: + column = edge["node"] + table_columns.append( + Column( + name=column["name"], + type=column["bigqueryType"]["name"], + description=column.get("description"), + ) + ) + + dataset_tables.append( + Table( + id=table_id, + gcp_id=table_gcp_id, + name=table_name, + slug=table_slug, + description=table_description, + columns=table_columns, + ) + ) + + dataset = Dataset( + id=dataset_id, + name=dataset_name, + slug=dataset_slug, + description=dataset_description, + tags=dataset_tags, + themes=dataset_themes, + organizations=dataset_organizations, + tables=dataset_tables, + ).model_dump() + + return json.dumps(dataset, ensure_ascii=False, indent=2) + except Exception as e: + return f"Error fetching dataset details:\n{e}" + + +@tool +def execute_sql_query(sql_query: str) -> str: + """Execute a SQL query against BigQuery tables from the Base dos Dados catalog. + + This tool runs your SQL query and returns the results. It includes safety checks + to prevent expensive queries (>10GB processing limit). Use this as the FINAL STEP + after identifying the right datasets and understanding their structure. + + Args: + sql_query (str): Standard SQL query using BigQuery syntax. Must reference + tables using their full gcp_id from get_dataset_details(). + + SQL Best Practices: + - Use backticks around table names: `project.dataset.table` + - Add LIMIT clause for exploration: LIMIT 10, LIMIT 100 + - Filter early with WHERE clauses to reduce data processing + - Select only needed columns, avoid SELECT * + - Use appropriate data types in comparisons + + Returns: + str: Query results as JSON array of objects, where each object represents a row. + Special cases: + - Empty results: Returns "[]" + - Query too large: Returns error with processing size details + - SQL errors: Returns error message with details for debugging + + Safety Features: + - Dry run validation before execution + - 10GB processing limit to prevent expensive queries + - Automatic query optimization suggestions if limit exceeded + + Next steps after using this tool: + 1. If results are empty, check your filters and table names + 2. If you get size limit errors, add more specific WHERE conditions + 3. For large result sets, consider using LIMIT or aggregation + 4. Analyze results and refine query if needed + 5. Format or visualize results for final presentation + + Common SQL patterns for Base dos Dados: + ```sql + -- Basic exploration + SELECT * FROM `basedosdados.br_ibge_censo_2010.municipio` LIMIT 10 + + -- Filtered by geography + SELECT * FROM `table` WHERE sigla_uf = 'SP' LIMIT 100 + + -- Aggregation by region + SELECT sigla_uf, COUNT(*) FROM `table` GROUP BY sigla_uf + + -- Time series data + SELECT ano, SUM(valor) FROM `table` + WHERE ano BETWEEN 2010 AND 2020 GROUP BY ano ORDER BY ano + ``` + + Troubleshooting: + - "Table not found" → Check gcp_id from get_dataset_details() + - "Column not found" → Verify column names match dataset structure + - "Query too large" → Add WHERE filters or select fewer columns + - "Syntax error" → Check BigQuery SQL syntax and table backticks + """ + try: + job_config = bq.QueryJobConfig(dry_run=True, use_query_cache=False) + query_job = client.query(sql_query, job_config=job_config) + + limit_bytes = 20e9 # 10GB + + if query_job.total_bytes_processed > limit_bytes: + return json.dumps( + { + "status": "error", + "error_type": "query_limit_exceeded", + "total_processed_bytes": query_job.total_bytes_processed, + "limit_bytes": limit_bytes, + "message": ( + "Query aborted: Data processed exceeds the 10GB per-query limit. " + "Consider optimizing by adding filters, selecting fewer columns, " + "or using a LIMIT clause before retrying." + ), + }, + indent=2, + ) + + rows = client.query(sql_query).result() + + results = [dict(row) for row in rows] + + if results: + return json.dumps(results, ensure_ascii=False, default=str) + return "[]" + except Exception as e: + return f"SQL query execution failed:\n{e}" diff --git a/backend/apps/chatbot/views.py b/backend/apps/chatbot/views.py index 2bf3772f..0d919fe8 100644 --- a/backend/apps/chatbot/views.py +++ b/backend/apps/chatbot/views.py @@ -9,9 +9,8 @@ from django.http import StreamingHttpResponse from graphql_jwt.shortcuts import get_user_by_token from langchain.chat_models import init_chat_model -from langchain_google_vertexai import VertexAIEmbeddings -from langchain_postgres import PGVector from langgraph.checkpoint.postgres import PostgresSaver +from langgraph.prebuilt import create_react_agent from loguru import logger from rest_framework import exceptions, status from rest_framework.parsers import JSONParser @@ -22,7 +21,6 @@ from rest_framework.views import APIView from rest_framework_simplejwt.tokens import RefreshToken -from backend.apps.chatbot.context_provider import PostgresContextProvider from backend.apps.chatbot.feedback_sender import LangSmithFeedbackSender from backend.apps.chatbot.models import Feedback, MessagePair, Thread from backend.apps.chatbot.serializers import ( @@ -33,9 +31,7 @@ ThreadSerializer, UserMessageSerializer, ) -from backend.apps.chatbot.utils.stream import process_chunk -from chatbot.assistants import SQLAssistant, format_sql_agent_response -from chatbot.formatters import SQLPromptFormatter +from backend.apps.chatbot.tools import execute_sql_query, get_dataset_details, search_datasets ModelSerializer = TypeVar("ModelSerializer", bound=Serializer) @@ -302,42 +298,42 @@ def _get_feedback_sender() -> LangSmithFeedbackSender: return LangSmithFeedbackSender() -@cache -def _get_context_provider(connection: str) -> PostgresContextProvider: - """Provide a configured `PostgresContextProvider` context provider. +# @cache +# def _get_context_provider(connection: str) -> PostgresContextProvider: +# """Provide a configured `PostgresContextProvider` context provider. - Args: - connection (str): A vector database connection for providing few-shot examples. +# Args: +# connection (str): A vector database connection for providing few-shot examples. - Returns: - PostgresContextProvider: An instance of `PostgresContextProvider`. - """ - bq_billing_project = os.environ["BILLING_PROJECT_ID"] - bq_query_project = os.environ["QUERY_PROJECT_ID"] +# Returns: +# PostgresContextProvider: An instance of `PostgresContextProvider`. +# """ +# bq_billing_project = os.environ["BILLING_PROJECT_ID"] +# bq_query_project = os.environ["QUERY_PROJECT_ID"] - embedding_model = os.getenv("EMBEDDING_MODEL") - pgvector_collection = os.getenv("PGVECTOR_COLLECTION") +# embedding_model = os.getenv("EMBEDDING_MODEL") +# pgvector_collection = os.getenv("PGVECTOR_COLLECTION") - if embedding_model and pgvector_collection: - embeddings = VertexAIEmbeddings(embedding_model) +# if embedding_model and pgvector_collection: +# embeddings = VertexAIEmbeddings(embedding_model) - vector_store = PGVector( - embeddings=embeddings, - connection=connection, - collection_name=pgvector_collection, - use_jsonb=True, - ) - else: - vector_store = None - - context_provider = PostgresContextProvider( - billing_project=bq_billing_project, - query_project=bq_query_project, - metadata_vector_store=vector_store, - top_k=5, - ) +# vector_store = PGVector( +# embeddings=embeddings, +# connection=connection, +# collection_name=pgvector_collection, +# use_jsonb=True, +# ) +# else: +# vector_store = None - return context_provider +# context_provider = PostgresContextProvider( +# billing_project=bq_billing_project, +# query_project=bq_query_project, +# metadata_vector_store=vector_store, +# top_k=5, +# ) + +# return context_provider @contextmanager @@ -355,23 +351,29 @@ def _get_sql_assistant(): conn = f"postgresql://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}" - context_provider = _get_context_provider(conn) + # context_provider = _get_context_provider(conn) - prompt_formatter = SQLPromptFormatter(vector_store=None) + # prompt_formatter = SQLPromptFormatter(vector_store=None) with PostgresSaver.from_conn_string(conn) as checkpointer: checkpointer.setup() model = init_chat_model(MODEL_URI, temperature=0) - assistant = SQLAssistant( - model=model, - context_provider=context_provider, - prompt_formatter=prompt_formatter, - checkpointer=checkpointer, - ) + tools = [search_datasets, get_dataset_details, execute_sql_query] - yield assistant + agent = create_react_agent(model=model, tools=tools, checkpointer=checkpointer) + + # assistant = SQLAssistant( + # model=model, + # context_provider=context_provider, + # prompt_formatter=prompt_formatter, + # checkpointer=checkpointer, + # ) + + # yield assistant + + yield agent def _stream_sql_assistant_response( @@ -388,36 +390,40 @@ def _stream_sql_assistant_response( Iterator[str]: JSON string containing the streaming status and the current step data. """ steps = [] - last_chunk = None + # last_chunk = None try: logger.info("Calling SQLAssistant...") with _get_sql_assistant() as assistant: - for mode, chunk in assistant.stream( - message=message, - config=config, - stream_mode=["updates", "values"], - rewrite_query=True, - ): - last_chunk = chunk + agent_response = assistant.invoke( + input={"messages": [{"role": "user", "content": message}]}, config=config + ) + # for mode, chunk in assistant.stream( + # message=message, + # config=config, + # stream_mode=["updates", "values"], + # rewrite_query=True, + # ): + # last_chunk = chunk - # Skip "values" chunks during streaming. We only need the final one at the end, - # as it contains the SQL Agent's final state. - if mode == "values": - continue + # # Skip "values" chunks during streaming. We only need the final one at the end, + # # as it contains the SQL Agent's final state. + # if mode == "values": + # continue - step = process_chunk(chunk) + # step = process_chunk(chunk) - if step is None: - continue + # if step is None: + # continue - steps.append(step.model_dump()) + # steps.append(step.model_dump()) - yield json.dumps({"status": "running", "data": step.model_dump_json()}) + "\n\n" + # yield json.dumps({"status": "running", "data": step.model_dump_json()}) + "\n\n" # The last chunk represents the SQLAgent's final state, # so we format it as the final response for the user. - response = format_sql_agent_response(last_chunk) + # response = format_sql_agent_response(last_chunk) + response = {} response["error_message"] = None except Exception: logger.exception(f"Error responding message {config['run_id']}:") @@ -434,6 +440,9 @@ def _stream_sql_assistant_response( response["id"] = config["run_id"] + response["content"] = agent_response["messages"][-1].content + response["sql_queries"] = None + message_pair = MessagePair.objects.create( id=response["id"], thread=thread, From 4bc38892c3c90f40bd1f9c7c5be207ced9bcc503 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Fri, 22 Aug 2025 18:18:00 -0300 Subject: [PATCH 02/23] commit save: new tools and system prompt --- backend/apps/chatbot/tools.py | 224 +++++++++++++++++++++++++++++++++- backend/apps/chatbot/views.py | 141 ++++++++++++++++++++- 2 files changed, 360 insertions(+), 5 deletions(-) diff --git a/backend/apps/chatbot/tools.py b/backend/apps/chatbot/tools.py index 17c74f2d..a18e2759 100644 --- a/backend/apps/chatbot/tools.py +++ b/backend/apps/chatbot/tools.py @@ -391,13 +391,13 @@ def execute_sql_query(sql_query: str) -> str: job_config = bq.QueryJobConfig(dry_run=True, use_query_cache=False) query_job = client.query(sql_query, job_config=job_config) - limit_bytes = 20e9 # 10GB + limit_bytes = 10e9 # 10GB if query_job.total_bytes_processed > limit_bytes: return json.dumps( { "status": "error", - "error_type": "query_limit_exceeded", + "error_type": "query_too_large", "total_processed_bytes": query_job.total_bytes_processed, "limit_bytes": limit_bytes, "message": ( @@ -418,3 +418,223 @@ def execute_sql_query(sql_query: str) -> str: return "[]" except Exception as e: return f"SQL query execution failed:\n{e}" + + +@tool +def inspect_column_values(table_gcp_id: str, column_name: str, limit: int = 100) -> str: + """Inspect the actual values present in a specific column to understand data patterns and encoding. + + This tool helps you understand what values actually exist in a column before writing filters + or joins. Essential for avoiding "value not found" errors and discovering encoded values + that need dictionary lookups. + + Args: + table_gcp_id (str): Full BigQuery table reference from get_table_details() (project.dataset.table) + column_name (str): Exact column name as shown in get_table_details() + limit (int, optional): Maximum number of distinct values to return (default: 100) + Use smaller limits (10-20) for initial exploration, larger for comprehensive views + + Returns: + str: JSON array of distinct values found in the column, sorted by frequency. + - Shows actual data patterns (encoded IDs, abbreviations, etc.) + - Reveals if values are numeric codes that need dictionary decoding + - Helps identify filtering options and data quality issues + + Use this tool when: + - Planning WHERE clauses to see available filter values + - Column values look like codes (1, 2, 3 or 'A', 'B', 'C') that might need decoding + - Getting "no results" from queries and need to verify correct values + - Understanding data patterns before analysis + - Checking for NULL values or data quality issues + + Next steps after using this tool: + 1. If values look like codes → Use search_dictionary_table() to find meanings + 2. If values are clear → Use them directly in WHERE clauses + 3. If unexpected values → Investigate data quality or documentation + 4. Plan your main analysis query with correct filter values + + Example workflow: + 1. get_table_details("table-id") → see columns + 2. inspect_column_values("project.dataset.table", "tipo_escola") → see actual values + 3. If returns [1,2,3,4] → search dictionary for meanings + 4. execute_sql_query with proper filters based on real values + + Common patterns in Brazilian data: + - Estado codes: 11, 12, 13 (need IBGE state dictionary) + - Municipality codes: 3550308 (São Paulo city IBGE code) + - Category codes: 1,2,3,4 (often need dataset dictionary) + - Yes/No: 0/1 or S/N depending on dataset + """ # noqa: E501 + sql_query = f"SELECT DISTINCT({column_name}) FROM `{table_gcp_id}` LIMIT {limit}" + + try: + # Check query size first + job_config = bq.QueryJobConfig(dry_run=True, use_query_cache=False) + query_job = client.query(sql_query, job_config=job_config) + + limit_bytes = 5e9 # 5GB limit for inspection queries + if query_job.total_bytes_processed and query_job.total_bytes_processed > limit_bytes: + bytes_processed_gb = query_job.total_bytes_processed / 1e9 + return json.dumps( + { + "status": "error", + "error_type": "query_too_large", + "total_processed_bytes": query_job.total_bytes_processed, + "limit_bytes": limit_bytes, + "message": ( + f"Column inspection would process {bytes_processed_gb:.1f} GB. " + "Try a smaller limit or add WHERE filters to reduce data size." + ), + }, + indent=2, + ) + + rows = client.query(sql_query).result() + results = [dict(row) for row in rows] + if results: + return json.dumps(results, ensure_ascii=False, default=str) + return "[]" + except Exception as e: + return f"Failed to inspect colum values:\n{e}" + + +@tool +def search_dictionary_table(dataset_gcp_id: str, table_name: str, column_name: str) -> str: + """Look up the meaning of encoded values using the dataset's dictionary table. + + Many Brazilian datasets use numeric codes or abbreviations for storage efficiency. + The 'dicionario' table provides mappings from encoded values (chave) to human-readable + meanings (valor). Use this after inspect_column_values() reveals coded values. + + Args: + dataset_gcp_id (str): Dataset portion of BigQuery reference (project.dataset_name) + - Extract from table gcp_id: if table is "basedosdados.br_ibge_censo.municipio" + - Use dataset_gcp_id: "basedosdados.br_ibge_censo" + table_name (str): Just the table name (not full gcp_id) + - Extract from gcp_id: if "basedosdados.br_ibge_censo.municipio" + - Use table_name: "municipio" + column_name (str): Exact column name that contains encoded values + + Returns: + str: JSON array of dictionary mappings with structure: + - chave: The encoded value found in your data (1, 2, 'A', 'B', etc.) + - valor: Human-readable meaning/description + + If no mappings found, returns guidance on alternative approaches. + + Next steps after using this tool: + 1. Match the returned 'chave' values with codes in your data + 2. Use JOIN with dictionary table for readable results: + ```sql + SELECT t.*, d.valor as decoded_column + FROM `project.dataset.table` t + LEFT JOIN `project.dataset.dicionario` d + ON t.encoded_column = d.chave + AND d.id_tabela = 'table_name' + AND d.nome_coluna = 'column_name' + ``` + 3. Or use the mappings to understand what codes mean for filtering + + Example usage: + # After inspect_column_values shows tipo_escola has values [1,2,3,4]: + search_dictionary_table("basedosdados.br_inep_censo_escolar", "escola", "tipo_escola") + + # Result might show: 1='Pública', 2='Privada', 3='Federal', 4='Municipal' + + # Then use in queries: + SELECT tipo_escola, COUNT(*) FROM `basedosdados.br_inep_censo_escolar.escola` + WHERE tipo_escola IN (1, 2) -- Now you know 1=Pública, 2=Privada + GROUP BY tipo_escola + + Common encoded column types in Brazilian data: + - Geographic codes: estado, municipio (IBGE codes) + - Category codes: tipo_*, categoria_*, situacao_* + - Status codes: ativo, situacao, condicao + - Classification codes: nivel_*, grau_*, classe_* + + If no dictionary found: + - Dataset might not have encoded values + - Values might be in separate reference tables + - Try common directory datasets like br_bd_diretorios_brasil + """ + try: + # Build the specific query for this table/column combination + dict_table_id = f"{dataset_gcp_id}.dicionario" + + search_query = f""" + SELECT chave, valor + FROM `{dict_table_id}` + WHERE id_tabela = '{table_name}' + AND nome_coluna = '{column_name}' + ORDER BY chave + """ + + # Check if query is reasonable size + job_config = bq.QueryJobConfig(dry_run=True, use_query_cache=False) + query_job = client.query(search_query, job_config=job_config) + + # Dictionary queries should be small, but let's be safe + if query_job.total_bytes_processed > 1e9: # 1GB limit + return json.dumps( + { + "status": "error", + "message": ( + "Dictionary table is unexpectedly large. " + "This might not be the right approach." + ), + "suggestion": ( + "Check if this dataset actually uses " + "encoded values or try a different table." + ), + } + ) + + rows = client.query(search_query).result() + results = [dict(row) for row in rows] + + if results: + return json.dumps(results, ensure_ascii=False, default=str) + else: + # Try to see if dictionary table exists and what it contains + fallback_query = f""" + SELECT DISTINCT id_tabela, nome_coluna + FROM `{dict_table_id}` + WHERE id_tabela LIKE '%{table_name}%' OR nome_coluna LIKE '%{column_name}%' + LIMIT 20 + """ + + try: + fallback_rows = client.query(fallback_query).result() + similar_entries = [dict(row) for row in fallback_rows] + + return json.dumps( + { + "status": "no_exact_match", + "message": ( + f"No dictionary entries found for table='{table_name}', " + f"column='{column_name}'" + ), + "similar_entries": similar_entries, + "suggestion": ( + "Check if table_name or column_name needs " + "adjustment based on similar entries above." + ), + }, + indent=2, + ) + + except Exception: + return json.dumps( + { + "status": "no_dictionary", + "message": f"No dictionary table found in {dataset_gcp_id}", + "suggestion": ( + "This dataset might not use encoded values, " + "or dictionary might be in a different location." + ), + }, + indent=2, + ) + + except Exception as e: + return f"Failed to search dictionary table:\n{e}" diff --git a/backend/apps/chatbot/views.py b/backend/apps/chatbot/views.py index 0d919fe8..f8855c54 100644 --- a/backend/apps/chatbot/views.py +++ b/backend/apps/chatbot/views.py @@ -31,7 +31,13 @@ ThreadSerializer, UserMessageSerializer, ) -from backend.apps.chatbot.tools import execute_sql_query, get_dataset_details, search_datasets +from backend.apps.chatbot.tools import ( + execute_sql_query, + get_dataset_details, + inspect_column_values, + search_datasets, + search_dictionary_table, +) ModelSerializer = TypeVar("ModelSerializer", bound=Serializer) @@ -360,9 +366,17 @@ def _get_sql_assistant(): model = init_chat_model(MODEL_URI, temperature=0) - tools = [search_datasets, get_dataset_details, execute_sql_query] + tools = [ + search_datasets, + get_dataset_details, + execute_sql_query, + inspect_column_values, + search_dictionary_table, + ] - agent = create_react_agent(model=model, tools=tools, checkpointer=checkpointer) + agent = create_react_agent( + model=model, tools=tools, prompt=SYSTEM_PROMPT, checkpointer=checkpointer + ) # assistant = SQLAssistant( # model=model, @@ -529,3 +543,124 @@ def _validate(request: Request, model_serializer: Type[ModelSerializer]) -> Mode raise exceptions.ValidationError(serializer.errors) return serializer + + +SYSTEM_PROMPT = """# Base dos Dados Research Assistant System Prompt + +You are a specialized research assistant expert in Brazilian data analysis using the Base dos Dados (BD) platform. Your role is to help users find, explore, and analyze Brazilian public data through BigQuery using a systematic approach. + +## Your Expertise +- **Brazilian Data Landscape**: Deep knowledge of Brazilian government agencies, data sources, and common data structures +- **Base dos Dados Platform**: Expert user of BD's catalog and BigQuery integration +- **Research Methodology**: Systematic approach to data discovery and analysis +- **SQL & BigQuery**: Proficient in writing efficient, cost-effective queries + +## Core Principles + +### 1. Systematic Data Discovery +Always follow this workflow: +1. **Search** → Use keywords to find relevant datasets +2. **Explore** → Review dataset overviews and available tables +3. **Investigate** → Get detailed column information for specific tables +4. **Query** → Execute targeted SQL queries with proper filters + +### 2. Search Strategy +- Use **individual keywords** only (Elasticsearch backend) +- Try **Portuguese terms** first ("educacao" not "education") +- Use **agency acronyms** (IBGE, INEP, ANVISA, ANS) +- **Retry with different keywords** if initial search fails +- **Never give up** after one failed search - always try alternatives + +### 3. Brazilian Data Context +Understand common data patterns: +- **Geographic**: sigla_uf (state codes), id_municipio (municipality codes) +- **Temporal**: ano (year), mes (month), data (date) +- **Identifiers**: id_*, codigo_*, sigla_* +- **Values**: *_valor, *_quantidade, *_taxa, *_indice, *_percentual + +### 4. Query Optimization +- Always start with **LIMIT clauses** for exploration +- Use **WHERE filters** early to reduce processing costs +- Select **specific columns** rather than SELECT * +- Be mindful of the **20GB query limit** + +## Key Brazilian Data Sources +- **IBGE**: Census, demographic, economic surveys (censo, pnad, pof) +- **INEP**: Education data (censo escolar, enem, prova brasil) +- **MS/DataSUS**: Health data (sinasc, sim, sinan) +- **MTE**: Employment data (rais, caged) +- **TSE**: Electoral data (eleicoes, candidatos) +- **BACEN**: Financial data (taxa selic, cambio) + +## Response Guidelines + +### When Searching Fails +- **Never give up after one search** +- Try multiple keyword variations: + - Portuguese vs English terms + - Synonyms and related terms + - Broader or narrower concepts + - Agency names and acronyms +- Explain your search strategy to the user +- Ask for clarification if search terms are unclear + +### When Presenting Results +- **Summarize findings clearly** before showing raw data +- **Explain data context**: What agency published it, what it represents +- **Highlight key insights** from query results +- **Suggest follow-up analyses** or related datasets + +### When Queries Fail +- **Analyze error messages** and suggest fixes +- **Recommend query optimizations** (filters, limits, column selection) +- **Explain BigQuery concepts** when helpful +- **Try alternative approaches** if initial query doesn't work + +## Communication Style +- **Friendly and helpful**: Make complex data accessible +- **Educational**: Explain Brazilian data landscape and sources +- **Systematic**: Show your step-by-step approach +- **Persistent**: Don't give up easily on difficult requests +- **Context-aware**: Understand Brazilian geography, politics, and institutions + +## Tool Usage Best Practices + +### search_datasets +- Use single keywords or short phrases +- Try Portuguese terms first +- Be persistent - retry with alternatives if no results +- Focus on topic, geography, or agency keywords + +### get_dataset_details +- Use to understand dataset scope and available tables +- Review themes, organizations, and tags for context +- Identify the most relevant tables before diving deeper + +### get_table_details +- Get complete column information and BigQuery references +- Note the gcp_id for SQL queries +- Understand column types and descriptions + +### execute_sql_query +- Always start with LIMIT for exploration +- Use proper BigQuery syntax with backticks +- Filter early with WHERE clauses +- Handle errors gracefully and suggest improvements + +## Error Recovery +- **Search failures**: Try different keywords, explain your strategy +- **Dataset not found**: Verify IDs, suggest alternatives +- **Query errors**: Analyze error messages, suggest fixes +- **Large query warnings**: Recommend optimization strategies +- **Empty results**: Check filters, suggest broader queries + +## Success Metrics +Your success is measured by: +- Finding relevant datasets for user queries +- Providing accurate, well-contextualized data analysis +- Teaching users about Brazilian data landscape +- Writing efficient, cost-effective BigQuery queries +- Persistence in overcoming search and query challenges + +Remember: You are not just a tool executor, but a knowledgeable research partner who understands Brazilian data, institutions, and research needs. Be proactive, educational, and persistent in helping users discover insights from Brazil's rich public data ecosystem. +""" # noqa: E501 From 81dfc286893469efb54b9dc13e066eb771bb432c Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Mon, 25 Aug 2025 11:58:47 -0300 Subject: [PATCH 03/23] refactor tools --- backend/apps/chatbot/tools.py | 337 +++++++++++++++++----------------- backend/apps/chatbot/views.py | 4 +- 2 files changed, 172 insertions(+), 169 deletions(-) diff --git a/backend/apps/chatbot/tools.py b/backend/apps/chatbot/tools.py index a18e2759..98445079 100644 --- a/backend/apps/chatbot/tools.py +++ b/backend/apps/chatbot/tools.py @@ -1,9 +1,10 @@ # -*- coding: utf-8 -*- import json import os -from typing import Optional +from typing import Any, Literal, Optional import httpx +from google.api_core import exceptions from google.cloud import bigquery as bq from langchain_core.tools import tool from pydantic import BaseModel @@ -11,7 +12,7 @@ SEARCH_URL = "https://backend.basedosdados.org/search/" GRAPHQL_URL = "https://backend.basedosdados.org/graphql" -DATASET_OVERVIEW_QUERY = """ +DATASET_DETAILS_QUERY = """ query GetDatasetOverview($id: ID!) { allDataset(id: $id, first: 1) { edges { @@ -109,54 +110,59 @@ class Dataset(DatasetOverview): tables: list[Table] +class ToolOutput(BaseModel): + status: Literal["success", "error"] + results: Optional[Any] = None + error_details: Optional[dict[str, Any]] = None + + client = bq.Client(project=os.environ["QUERY_PROJECT_ID"]) @tool def search_datasets(query: str) -> str: - """Search for datasets in the Base dos Dados (BD) catalog using KEYWORDS ONLY. + """Search for datasets in the Base dos Dados (BD) database using KEYWORDS ONLY. IMPORTANT: This search uses Elasticsearch - use individual KEYWORDS, NOT full sentences. Use this as the FIRST STEP when exploring data. If no relevant datasets are found, try searching again with different or more specific keywords. Args: - query (str): Individual keywords or short phrases (2-3 words max). The search engine - works best with single keywords. Use: - - Single topic words: "education", "health", "economy", "population", "crime" - - Geographic keywords: "municipal", "estado", "brasil", "sao_paulo", "rio" - - Data type keywords: "censo", "pnad", "rais", "sinasc", "datasus" - - Organization keywords: "ibge", "inep", "ans", "bacen", "tse" + query (str): Individual keywords (2-3 words max). The search engine works best with single keywords. Use: + - Single topic words: "educação", "saúde", "economia", "emprego", "população", "criminalidade" + - Geographic keywords: "municipal", "estado", "brasil", "sao paulo", "rio" + - Dataset names: "censo", "pnad", "rais", "caged" + - Organization acronyms: "ibge", "inep", "anvisa", "tse" AVOID full sentences like "Brazilian population data by municipality" - INSTEAD use: "populacao" or "municipio" or "ibge censo" + INSTEAD use: "censo" or "ibge" or "populacao" or "municipio" Returns: - str: JSON array of dataset overviews. If empty array [] or irrelevant results: - - Try different keywords (synonyms, Portuguese terms, abbreviations) + str: JSON array of dataset overviews. If empty array `[]` or irrelevant results: + - Try different keywords (synonyms, abbreviations) - Use more specific terms or broader terms - - Try organization names or data source acronyms + - Try organization names or data source acronyms (caged, ibge, rais, censo, etc.) - Search Strategy - If initial search fails: - 1. Try Portuguese keywords: "educacao" instead of "education" - 2. Use government agency acronyms: "ibge", "inep", "anvisa", "ans" - 3. Try broader terms: "social" instead of specific indicators - 4. Use dataset-specific terms: "pnad", "censo", "rais", "caged" - 5. Search by geographic level: "municipio", "estado", "regiao" + Search Strategy: + 1. Use portuguese keywords: "educacao" instead of "education" + 2. Use government agency acronyms: "ibge", "inep", "anvisa" + 3. Use dataset-specific terms: "censo", "rais", "pnad", "caged" + 4. Search by geographic level: "municipio", "estado", "regiao" Next steps after using this tool: - 1. If no results or irrelevant results → Try different keywords and search again - 2. If good results → Review datasets to identify the most relevant ones - 3. Use get_dataset_details() with the dataset 'id' to explore structure - 4. Look for datasets from reputable organizations (IBGE, government agencies) + 1. If no results or irrelevant results: Try different keywords and search again + 2. If good results: Review datasets to identify the most relevant ones + 3. Use `get_dataset_details()` with the dataset `id` to explore structure + 4. Look for datasets from reputable organizations (IBGE and other government agencies) Example successful searches: - "censo" → finds census datasets - "educacao" → finds education data + - "emprego" → finds labour data - "saude municipio" → finds municipal health data - "ibge" → finds all IBGE datasets - "eleicoes" → finds election datasets - """ + """ # noqa: E501 try: with httpx.Client() as client: response = client.get( @@ -183,9 +189,13 @@ def search_datasets(query: str) -> str: ) overviews.append(dataset_overview.model_dump()) - return json.dumps(overviews, ensure_ascii=False, indent=2) + tool_output = ToolOutput(status="success", results=overviews).model_dump(exclude_none=True) except Exception as e: - return f"Error searching datasets:\n{e}" + tool_output = ToolOutput( + status="error", error_details={"message": f"Error searching datasets:\n{e}"} + ).model_dump(exclude_none=True) + + return json.dumps(tool_output, ensure_ascii=False, indent=2) @tool @@ -193,45 +203,45 @@ def get_dataset_details(dataset_id: str) -> str: """Get comprehensive details about a specific dataset including all tables and columns. This tool provides the complete structure of a dataset, showing all available tables, - their columns, data types, and BigQuery identifiers. Use this AFTER search_datasets() + their columns, data types, and BigQuery identifiers. Use this AFTER `search_datasets()` to understand what data is available before writing SQL queries. Args: - dataset_id (str): The unique dataset ID obtained from search_datasets(). + dataset_id (str): The unique dataset ID obtained from `search_datasets()`. This is typically a UUID-like string, not the human-readable name. Returns: str: JSON object with complete dataset information including: - - Basic metadata (name, description, themes, organizations) + - Basic metadata (name, description, tags, themes, organizations) - tables: Array of all tables in the dataset with: - - gcp_id: Full BigQuery table reference (project.dataset.table) + - gcp_id: Full BigQuery table reference (`project.dataset.table`) - columns: All column names, types (STRING, INTEGER, DATE, etc.), and descriptions - Table descriptions explaining what each table contains Next steps after using this tool: 1. Identify which table(s) contain the data you need - 2. Note the column names and types for SQL query construction - 3. Use the gcp_id (BigQuery table reference) in your SQL queries - 4. Check column descriptions to understand what each field represents - 5. Plan your SQL query using execute_sql_query() - - Important notes: - - The 'gcp_id' field is the full BigQuery table reference you'll use in FROM clauses - - Column types help you write correct SQL (STRING needs quotes, dates need DATE functions) + 2. Check column descriptions to understand what each field represents + 3. Note the column names and types for SQL query construction + 4. Use the `gcp_id` (BigQuery table reference) in your SQL queries + 5. Plan your SQL query using `execute_sql_query()` + + IMPORTANT: + - The `gcp_id` field is the full BigQuery table reference you'll use in FROM clauses + - Column types help you write correct SQL (STRING needs quotes, dates need DATE functions, etc.) - Some tables may be very large - consider using LIMIT in your queries - - Look for key identifier columns (like municipality codes, dates) for filtering + - Look for key identifier columns (like municipality codes, dates, etc.) for filtering Example workflow: - 1. search_datasets("IBGE census") → get dataset IDs - 2. get_dataset_details("abc-123-def") → explore table structure - 3. execute_sql_query("SELECT * FROM `basedosdados.br_ibge_censo.municipio` LIMIT 10") - """ + 1. `search_datasets("IBGE")` → get dataset IDs + 2. `get_dataset_details("abc-123-def")` → explore tables structure + 3. `execute_sql_query("SELECT * FROM `basedosdados.br_ibge_censo.municipio` LIMIT 10")` + """ # noqa: E501 try: with httpx.Client() as client: response = client.post( url=GRAPHQL_URL, json={ - "query": DATASET_OVERVIEW_QUERY, + "query": DATASET_DETAILS_QUERY, "variables": {"id": dataset_id}, }, timeout=httpx.Timeout(5.0, read=60.0), @@ -322,111 +332,122 @@ def get_dataset_details(dataset_id: str) -> str: tables=dataset_tables, ).model_dump() - return json.dumps(dataset, ensure_ascii=False, indent=2) + tool_output = ToolOutput(status="success", results=dataset).model_dump(exclude_none=True) except Exception as e: - return f"Error fetching dataset details:\n{e}" + tool_output = ToolOutput( + status="error", error_details={"message": f"Error fetching dataset details:\n{e}"} + ).model_dump(exclude_none=True) + + return json.dumps(tool_output, ensure_ascii=False, indent=2) @tool -def execute_sql_query(sql_query: str) -> str: - """Execute a SQL query against BigQuery tables from the Base dos Dados catalog. +def execute_bigquery_sql(sql_query: str) -> str: + """Execute a SQL query against BigQuery tables from the Base dos Dados database. This tool runs your SQL query and returns the results. It includes safety checks - to prevent expensive queries (>10GB processing limit). Use this as the FINAL STEP + to prevent expensive queries (>20GB processing limit). Use this as the FINAL STEP after identifying the right datasets and understanding their structure. Args: sql_query (str): Standard SQL query using BigQuery syntax. Must reference - tables using their full gcp_id from get_dataset_details(). + tables using their full `gcp_id` from get_dataset_details(). SQL Best Practices: - - Use backticks around table names: `project.dataset.table` + - Use fully qualified table names: `project.dataset.table` + - Select only needed columns, avoid SELECT * - Add LIMIT clause for exploration: LIMIT 10, LIMIT 100 - Filter early with WHERE clauses to reduce data processing - - Select only needed columns, avoid SELECT * + - Order results by relevant columns to present the most significant results first. + - Never use DDL (e.g, `CREATE`, `ALTER`, `DROP`) or DML (e.g., `INSERT`, `UPDATE`, `DELETE`) commands - Use appropriate data types in comparisons Returns: - str: Query results as JSON array of objects, where each object represents a row. - Special cases: + str: Query results as JSON array of objects, where each object represents a row. Special cases: - Empty results: Returns "[]" - Query too large: Returns error with processing size details - SQL errors: Returns error message with details for debugging Safety Features: - Dry run validation before execution - - 10GB processing limit to prevent expensive queries - - Automatic query optimization suggestions if limit exceeded + - 20GB processing limit to prevent expensive queries Next steps after using this tool: 1. If results are empty, check your filters and table names 2. If you get size limit errors, add more specific WHERE conditions 3. For large result sets, consider using LIMIT or aggregation 4. Analyze results and refine query if needed - 5. Format or visualize results for final presentation + 5. Format results for final presentation Common SQL patterns for Base dos Dados: ```sql -- Basic exploration - SELECT * FROM `basedosdados.br_ibge_censo_2010.municipio` LIMIT 10 + SELECT * FROM `basedosdados.br_inep_ideb.municipio` LIMIT 10 -- Filtered by geography - SELECT * FROM `table` WHERE sigla_uf = 'SP' LIMIT 100 + SELECT * FROM `basedosdados.br_inep_ideb.municipio` + WHERE sigla_uf = 'SP' LIMIT 100 -- Aggregation by region - SELECT sigla_uf, COUNT(*) FROM `table` GROUP BY sigla_uf + SELECT sigla_uf, COUNT(*) + FROM `basedosdados.br_inep_ideb.municipio` + GROUP BY sigla_uf -- Time series data - SELECT ano, SUM(valor) FROM `table` - WHERE ano BETWEEN 2010 AND 2020 GROUP BY ano ORDER BY ano + SELECT ano, AVG(ideb) FROM `basedosdados.br_inep_ideb.municipio` + WHERE ano BETWEEN 2010 AND 2020 + GROUP BY ano + ORDER BY ano ``` Troubleshooting: - - "Table not found" → Check gcp_id from get_dataset_details() - - "Column not found" → Verify column names match dataset structure + - "Not found" → Check `gcp_id` from `get_dataset_details()` + - "Unrecognized name" → Verify column names match table structure from `get_dataset_details()` - "Query too large" → Add WHERE filters or select fewer columns - - "Syntax error" → Check BigQuery SQL syntax and table backticks - """ + - "Syntax error" → Check GoogleSQL syntax + """ # noqa: E501 try: job_config = bq.QueryJobConfig(dry_run=True, use_query_cache=False) query_job = client.query(sql_query, job_config=job_config) - limit_bytes = 10e9 # 10GB + limit_bytes = 20e9 # 20GB limit for queries if query_job.total_bytes_processed > limit_bytes: - return json.dumps( - { - "status": "error", - "error_type": "query_too_large", - "total_processed_bytes": query_job.total_bytes_processed, + tool_output = ToolOutput( + status="error", + error_details={ + "type": "QueryTooLarge", "limit_bytes": limit_bytes, + "total_processed_bytes": query_job.total_bytes_processed, "message": ( - "Query aborted: Data processed exceeds the 10GB per-query limit. " + "Query aborted: Data processed exceeds the 20GB per-query limit. " "Consider optimizing by adding filters, selecting fewer columns, " "or using a LIMIT clause before retrying." ), }, - indent=2, ) + return tool_output.model_dump_json(indent=2, exclude_none=True) rows = client.query(sql_query).result() results = [dict(row) for row in rows] - if results: - return json.dumps(results, ensure_ascii=False, default=str) - return "[]" + tool_output = ToolOutput(status="success", results=results).model_dump(exclude_none=True) except Exception as e: - return f"SQL query execution failed:\n{e}" + tool_output = ToolOutput( + status="error", error_details={"message": f"SQL query execution failed:\n{e}"} + ).model_dump(exclude_none=True) + + return json.dumps(tool_output, ensure_ascii=False) @tool def inspect_column_values(table_gcp_id: str, column_name: str, limit: int = 100) -> str: - """Inspect the actual values present in a specific column to understand data patterns and encoding. + """Inspect the actual values in a specific column to understand data content and patterns. - This tool helps you understand what values actually exist in a column before writing filters - or joins. Essential for avoiding "value not found" errors and discovering encoded values - that need dictionary lookups. + Use this tool as a FALLBACK when `search_dictionary_table()` fails or returns no results. + Always try `search_dictionary_table()` FIRST before using this tool, as dictionary tables + provide more meaningful information about coded values. Args: table_gcp_id (str): Full BigQuery table reference from get_table_details() (project.dataset.table) @@ -465,46 +486,51 @@ def inspect_column_values(table_gcp_id: str, column_name: str, limit: int = 100) - Category codes: 1,2,3,4 (often need dataset dictionary) - Yes/No: 0/1 or S/N depending on dataset """ # noqa: E501 - sql_query = f"SELECT DISTINCT({column_name}) FROM `{table_gcp_id}` LIMIT {limit}" + sql_query = f"SELECT DISTINCT {column_name} FROM {table_gcp_id} LIMIT {limit}" try: - # Check query size first job_config = bq.QueryJobConfig(dry_run=True, use_query_cache=False) query_job = client.query(sql_query, job_config=job_config) limit_bytes = 5e9 # 5GB limit for inspection queries + if query_job.total_bytes_processed and query_job.total_bytes_processed > limit_bytes: bytes_processed_gb = query_job.total_bytes_processed / 1e9 - return json.dumps( - { + return ToolOutput( + status="error", + error_details={ "status": "error", - "error_type": "query_too_large", - "total_processed_bytes": query_job.total_bytes_processed, + "type": "QueryTooLarge", "limit_bytes": limit_bytes, + "total_processed_bytes": query_job.total_bytes_processed, "message": ( - f"Column inspection would process {bytes_processed_gb:.1f} GB. " + f"Column inspection exceeds the 5GB per-query " + f"limit for inspection({bytes_processed_gb:.1f} GB). " "Try a smaller limit or add WHERE filters to reduce data size." ), }, - indent=2, - ) + ).model_dump_json(indent=2, exclude_none=True) rows = client.query(sql_query).result() - results = [dict(row) for row in rows] - if results: - return json.dumps(results, ensure_ascii=False, default=str) - return "[]" + + results = [row.get(column_name) for row in rows] + + tool_output = ToolOutput(status="success", results=results).model_dump(exclude_none=True) except Exception as e: - return f"Failed to inspect colum values:\n{e}" + tool_output = ToolOutput( + status="error", error_details={"message": f"Failed to inspect colum values:\n{e}"} + ).model_dump(exclude_none=True) + + return json.dumps(tool_output, ensure_ascii=False) @tool def search_dictionary_table(dataset_gcp_id: str, table_name: str, column_name: str) -> str: - """Look up the meaning of encoded values using the dataset's dictionary table. + """Look up the meaning of coded values in a column using the dataset's dictionary table - Many Brazilian datasets use numeric codes or abbreviations for storage efficiency. - The 'dicionario' table provides mappings from encoded values (chave) to human-readable - meanings (valor). Use this after inspect_column_values() reveals coded values. + **ALWAYS USE THIS TOOL FIRST** when you need to understand column values before `inspect_column_values()`. + Many Base dos Dados datasets use coded values or abbreviations for storage efficiency. + This tool provides the authoritative meanings of these codes. Args: dataset_gcp_id (str): Dataset portion of BigQuery reference (project.dataset_name) @@ -556,85 +582,62 @@ def search_dictionary_table(dataset_gcp_id: str, table_name: str, column_name: s - Dataset might not have encoded values - Values might be in separate reference tables - Try common directory datasets like br_bd_diretorios_brasil - """ - try: - # Build the specific query for this table/column combination - dict_table_id = f"{dataset_gcp_id}.dicionario" + """ # noqa: E501 - search_query = f""" - SELECT chave, valor - FROM `{dict_table_id}` + # Build the specific query for this table/column combination + dict_table_id = f"{dataset_gcp_id}.dicionario" + + search_query = f""" + SELECT nome_coluna, chave, valor + FROM {dict_table_id} WHERE id_tabela = '{table_name}' - AND nome_coluna = '{column_name}' - ORDER BY chave - """ + """ + + if column_name is not None: + search_query += f"AND nome_coluna = '{column_name}'" - # Check if query is reasonable size + search_query += "ORDER BY nome_coluna, chave" + + try: job_config = bq.QueryJobConfig(dry_run=True, use_query_cache=False) query_job = client.query(search_query, job_config=job_config) - # Dictionary queries should be small, but let's be safe - if query_job.total_bytes_processed > 1e9: # 1GB limit - return json.dumps( - { - "status": "error", + limit_bytes = 1e9 # 1GB limit for dictionary queries + + if query_job.total_bytes_processed > limit_bytes: + return ToolOutput( + status="error", + error_details={ + "type": "QueryTooLarge", + "limit_bytes": limit_bytes, + "total_processed_bytes": query_job.total_bytes_processed, "message": ( "Dictionary table is unexpectedly large. " - "This might not be the right approach." - ), - "suggestion": ( + "This might not be the right approach. " "Check if this dataset actually uses " "encoded values or try a different table." ), - } - ) + }, + ).model_dump_json(indent=2, exclude_none=True) rows = client.query(search_query).result() results = [dict(row) for row in rows] - - if results: - return json.dumps(results, ensure_ascii=False, default=str) - else: - # Try to see if dictionary table exists and what it contains - fallback_query = f""" - SELECT DISTINCT id_tabela, nome_coluna - FROM `{dict_table_id}` - WHERE id_tabela LIKE '%{table_name}%' OR nome_coluna LIKE '%{column_name}%' - LIMIT 20 - """ - - try: - fallback_rows = client.query(fallback_query).result() - similar_entries = [dict(row) for row in fallback_rows] - - return json.dumps( - { - "status": "no_exact_match", - "message": ( - f"No dictionary entries found for table='{table_name}', " - f"column='{column_name}'" - ), - "similar_entries": similar_entries, - "suggestion": ( - "Check if table_name or column_name needs " - "adjustment based on similar entries above." - ), - }, - indent=2, - ) - - except Exception: - return json.dumps( - { - "status": "no_dictionary", - "message": f"No dictionary table found in {dataset_gcp_id}", - "suggestion": ( - "This dataset might not use encoded values, " - "or dictionary might be in a different location." - ), - }, - indent=2, - ) - + tool_output = ToolOutput(status="success", results=results).model_dump(exclude_none=True) + except exceptions.NotFound: + return ToolOutput( + status="error", + error_details={ + "type": "TableNotFound", + "message": ( + f"Dictionary table not found for dataset {dataset_gcp_id}. " + "This indicates this dataset does not contain a dicionary table. " + "Consider using the `inspect_column_values` tool to inspect column values." + ), + }, + ).model_dump_json(indent=2, exclude_none=True) except Exception as e: - return f"Failed to search dictionary table:\n{e}" + tool_output = ToolOutput( + status="error", error_details={"message": f"Failed to search dictionary table:\n{e}"} + ).model_dump(exclude_none=True) + + return json.dumps(tool_output, ensure_ascii=False) diff --git a/backend/apps/chatbot/views.py b/backend/apps/chatbot/views.py index f8855c54..d3003d17 100644 --- a/backend/apps/chatbot/views.py +++ b/backend/apps/chatbot/views.py @@ -32,7 +32,7 @@ UserMessageSerializer, ) from backend.apps.chatbot.tools import ( - execute_sql_query, + execute_bigquery_sql, get_dataset_details, inspect_column_values, search_datasets, @@ -369,7 +369,7 @@ def _get_sql_assistant(): tools = [ search_datasets, get_dataset_details, - execute_sql_query, + execute_bigquery_sql, inspect_column_values, search_dictionary_table, ] From c9bd70701f1fedca6d0b6f9ebbac8e043ccd429c Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Mon, 25 Aug 2025 20:29:14 -0300 Subject: [PATCH 04/23] refactor tools: simplify docstrings and fix bugs --- backend/apps/chatbot/tools.py | 376 +++++++++++----------------------- 1 file changed, 115 insertions(+), 261 deletions(-) diff --git a/backend/apps/chatbot/tools.py b/backend/apps/chatbot/tools.py index 98445079..c2f6569a 100644 --- a/backend/apps/chatbot/tools.py +++ b/backend/apps/chatbot/tools.py @@ -121,47 +121,20 @@ class ToolOutput(BaseModel): @tool def search_datasets(query: str) -> str: - """Search for datasets in the Base dos Dados (BD) database using KEYWORDS ONLY. + """Search for datasets in Base dos Dados using keywords. - IMPORTANT: This search uses Elasticsearch - use individual KEYWORDS, NOT full sentences. - Use this as the FIRST STEP when exploring data. If no relevant datasets are found, - try searching again with different or more specific keywords. + CRITICAL: Use individual KEYWORDS only, not full sentences. The search engine uses Elasticsearch. Args: - query (str): Individual keywords (2-3 words max). The search engine works best with single keywords. Use: - - Single topic words: "educação", "saúde", "economia", "emprego", "população", "criminalidade" - - Geographic keywords: "municipal", "estado", "brasil", "sao paulo", "rio" - - Dataset names: "censo", "pnad", "rais", "caged" - - Organization acronyms: "ibge", "inep", "anvisa", "tse" - - AVOID full sentences like "Brazilian population data by municipality" - INSTEAD use: "censo" or "ibge" or "populacao" or "municipio" + query (str): 2-3 keywords maximum. Use Portuguese terms, organization acronyms, or dataset acronyms. + Good Examples: "censo", "educacao", "ibge", "inep", "rais", "saude" + Avoid: "Brazilian population data by municipality" Returns: - str: JSON array of dataset overviews. If empty array `[]` or irrelevant results: - - Try different keywords (synonyms, abbreviations) - - Use more specific terms or broader terms - - Try organization names or data source acronyms (caged, ibge, rais, censo, etc.) - - Search Strategy: - 1. Use portuguese keywords: "educacao" instead of "education" - 2. Use government agency acronyms: "ibge", "inep", "anvisa" - 3. Use dataset-specific terms: "censo", "rais", "pnad", "caged" - 4. Search by geographic level: "municipio", "estado", "regiao" - - Next steps after using this tool: - 1. If no results or irrelevant results: Try different keywords and search again - 2. If good results: Review datasets to identify the most relevant ones - 3. Use `get_dataset_details()` with the dataset `id` to explore structure - 4. Look for datasets from reputable organizations (IBGE and other government agencies) - - Example successful searches: - - "censo" → finds census datasets - - "educacao" → finds education data - - "emprego" → finds labour data - - "saude municipio" → finds municipal health data - - "ibge" → finds all IBGE datasets - - "eleicoes" → finds election datasets + str: JSON array of datasets. If empty/irrelevant results, try different keywords. + + Strategy: Start with broad terms like "censo", "ibge", "inep", "rais", then get specific if needed. + Next step: Use `get_dataset_details()` with returned dataset IDs. """ # noqa: E501 try: with httpx.Client() as client: @@ -202,39 +175,21 @@ def search_datasets(query: str) -> str: def get_dataset_details(dataset_id: str) -> str: """Get comprehensive details about a specific dataset including all tables and columns. - This tool provides the complete structure of a dataset, showing all available tables, - their columns, data types, and BigQuery identifiers. Use this AFTER `search_datasets()` - to understand what data is available before writing SQL queries. + Use AFTER `search_datasets()` to understand data structure before writing queries. Args: - dataset_id (str): The unique dataset ID obtained from `search_datasets()`. + dataset_id (str): Dataset ID obtained from `search_datasets()`. This is typically a UUID-like string, not the human-readable name. Returns: - str: JSON object with complete dataset information including: + str: JSON object with complete dataset information, including: - Basic metadata (name, description, tags, themes, organizations) - tables: Array of all tables in the dataset with: - gcp_id: Full BigQuery table reference (`project.dataset.table`) - - columns: All column names, types (STRING, INTEGER, DATE, etc.), and descriptions - - Table descriptions explaining what each table contains - - Next steps after using this tool: - 1. Identify which table(s) contain the data you need - 2. Check column descriptions to understand what each field represents - 3. Note the column names and types for SQL query construction - 4. Use the `gcp_id` (BigQuery table reference) in your SQL queries - 5. Plan your SQL query using `execute_sql_query()` - - IMPORTANT: - - The `gcp_id` field is the full BigQuery table reference you'll use in FROM clauses - - Column types help you write correct SQL (STRING needs quotes, dates need DATE functions, etc.) - - Some tables may be very large - consider using LIMIT in your queries - - Look for key identifier columns (like municipality codes, dates, etc.) for filtering - - Example workflow: - 1. `search_datasets("IBGE")` → get dataset IDs - 2. `get_dataset_details("abc-123-def")` → explore tables structure - 3. `execute_sql_query("SELECT * FROM `basedosdados.br_ibge_censo.municipio` LIMIT 10")` + - columns: All column names, types, and descriptions + - table descriptions explaining what each table contains + + Next step: Use `execute_bigquery_sql()` to execute queries. """ # noqa: E501 try: with httpx.Client() as client: @@ -345,80 +300,39 @@ def get_dataset_details(dataset_id: str) -> str: def execute_bigquery_sql(sql_query: str) -> str: """Execute a SQL query against BigQuery tables from the Base dos Dados database. - This tool runs your SQL query and returns the results. It includes safety checks - to prevent expensive queries (>20GB processing limit). Use this as the FINAL STEP - after identifying the right datasets and understanding their structure. + Use AFTER identifying the right datasets and understanding tables structure. + It includes a 20GB processing limit for safety. Args: - sql_query (str): Standard SQL query using BigQuery syntax. Must reference - tables using their full `gcp_id` from get_dataset_details(). - - SQL Best Practices: - - Use fully qualified table names: `project.dataset.table` - - Select only needed columns, avoid SELECT * - - Add LIMIT clause for exploration: LIMIT 10, LIMIT 100 - - Filter early with WHERE clauses to reduce data processing - - Order results by relevant columns to present the most significant results first. - - Never use DDL (e.g, `CREATE`, `ALTER`, `DROP`) or DML (e.g., `INSERT`, `UPDATE`, `DELETE`) commands - - Use appropriate data types in comparisons + sql_query (str): Standard GoogleSQL query. Must reference + tables using their full `gcp_id` from `get_dataset_details()`. + + Best practices: + - Use fully qualified names: `project.dataset.table` + - Select only needed columns, avoid `SELECT *` + - Add `LIMIT` for exploration + - Filter early with `WHERE` clauses + - Order by relevant columns + - Never use DDL/DML commands + - Use appropriate data types in comparisons Returns: - str: Query results as JSON array of objects, where each object represents a row. Special cases: - - Empty results: Returns "[]" - - Query too large: Returns error with processing size details - - SQL errors: Returns error message with details for debugging - - Safety Features: - - Dry run validation before execution - - 20GB processing limit to prevent expensive queries - - Next steps after using this tool: - 1. If results are empty, check your filters and table names - 2. If you get size limit errors, add more specific WHERE conditions - 3. For large result sets, consider using LIMIT or aggregation - 4. Analyze results and refine query if needed - 5. Format results for final presentation - - Common SQL patterns for Base dos Dados: - ```sql - -- Basic exploration - SELECT * FROM `basedosdados.br_inep_ideb.municipio` LIMIT 10 - - -- Filtered by geography - SELECT * FROM `basedosdados.br_inep_ideb.municipio` - WHERE sigla_uf = 'SP' LIMIT 100 - - -- Aggregation by region - SELECT sigla_uf, COUNT(*) - FROM `basedosdados.br_inep_ideb.municipio` - GROUP BY sigla_uf - - -- Time series data - SELECT ano, AVG(ideb) FROM `basedosdados.br_inep_ideb.municipio` - WHERE ano BETWEEN 2010 AND 2020 - GROUP BY ano - ORDER BY ano - ``` - - Troubleshooting: - - "Not found" → Check `gcp_id` from `get_dataset_details()` - - "Unrecognized name" → Verify column names match table structure from `get_dataset_details()` - - "Query too large" → Add WHERE filters or select fewer columns - - "Syntax error" → Check GoogleSQL syntax + str: Query results as JSON array. Empty results return "[]". """ # noqa: E501 try: job_config = bq.QueryJobConfig(dry_run=True, use_query_cache=False) query_job = client.query(sql_query, job_config=job_config) limit_bytes = 20e9 # 20GB limit for queries + total_bytes = query_job.total_bytes_processed - if query_job.total_bytes_processed > limit_bytes: + if total_bytes and total_bytes > limit_bytes: tool_output = ToolOutput( status="error", error_details={ "type": "QueryTooLarge", "limit_bytes": limit_bytes, - "total_processed_bytes": query_job.total_bytes_processed, + "total_processed_bytes": total_bytes, "message": ( "Query aborted: Data processed exceeds the 20GB per-query limit. " "Consider optimizing by adding filters, selecting fewer columns, " @@ -438,154 +352,40 @@ def execute_bigquery_sql(sql_query: str) -> str: status="error", error_details={"message": f"SQL query execution failed:\n{e}"} ).model_dump(exclude_none=True) - return json.dumps(tool_output, ensure_ascii=False) + return json.dumps(tool_output, ensure_ascii=False, default=str) @tool -def inspect_column_values(table_gcp_id: str, column_name: str, limit: int = 100) -> str: - """Inspect the actual values in a specific column to understand data content and patterns. +def decode_table_values(table_gcp_id: str, column_name: Optional[str] = None) -> str: + """Decode coded values from a table. - Use this tool as a FALLBACK when `search_dictionary_table()` fails or returns no results. - Always try `search_dictionary_table()` FIRST before using this tool, as dictionary tables - provide more meaningful information about coded values. + Use when column values appear to be codes (e.g., 1,2,3 or A,B,C). + Many datasets use codes for storage efficiency. This tool provides + the authoritative meanings of these codes. Args: - table_gcp_id (str): Full BigQuery table reference from get_table_details() (project.dataset.table) - column_name (str): Exact column name as shown in get_table_details() - limit (int, optional): Maximum number of distinct values to return (default: 100) - Use smaller limits (10-20) for initial exploration, larger for comprehensive views + table_gcp_id (str): Full BigQuery table reference + column_name (Optional[str], optional): Column with coded values. If `None`, + all columns will be used. Defaults to `None`. Returns: - str: JSON array of distinct values found in the column, sorted by frequency. - - Shows actual data patterns (encoded IDs, abbreviations, etc.) - - Reveals if values are numeric codes that need dictionary decoding - - Helps identify filtering options and data quality issues - - Use this tool when: - - Planning WHERE clauses to see available filter values - - Column values look like codes (1, 2, 3 or 'A', 'B', 'C') that might need decoding - - Getting "no results" from queries and need to verify correct values - - Understanding data patterns before analysis - - Checking for NULL values or data quality issues - - Next steps after using this tool: - 1. If values look like codes → Use search_dictionary_table() to find meanings - 2. If values are clear → Use them directly in WHERE clauses - 3. If unexpected values → Investigate data quality or documentation - 4. Plan your main analysis query with correct filter values - - Example workflow: - 1. get_table_details("table-id") → see columns - 2. inspect_column_values("project.dataset.table", "tipo_escola") → see actual values - 3. If returns [1,2,3,4] → search dictionary for meanings - 4. execute_sql_query with proper filters based on real values - - Common patterns in Brazilian data: - - Estado codes: 11, 12, 13 (need IBGE state dictionary) - - Municipality codes: 3550308 (São Paulo city IBGE code) - - Category codes: 1,2,3,4 (often need dataset dictionary) - - Yes/No: 0/1 or S/N depending on dataset + str: JSON array with chave (code) and valor (meaning) mappings. """ # noqa: E501 - sql_query = f"SELECT DISTINCT {column_name} FROM {table_gcp_id} LIMIT {limit}" - try: - job_config = bq.QueryJobConfig(dry_run=True, use_query_cache=False) - query_job = client.query(sql_query, job_config=job_config) - - limit_bytes = 5e9 # 5GB limit for inspection queries - - if query_job.total_bytes_processed and query_job.total_bytes_processed > limit_bytes: - bytes_processed_gb = query_job.total_bytes_processed / 1e9 - return ToolOutput( - status="error", - error_details={ - "status": "error", - "type": "QueryTooLarge", - "limit_bytes": limit_bytes, - "total_processed_bytes": query_job.total_bytes_processed, - "message": ( - f"Column inspection exceeds the 5GB per-query " - f"limit for inspection({bytes_processed_gb:.1f} GB). " - "Try a smaller limit or add WHERE filters to reduce data size." - ), - }, - ).model_dump_json(indent=2, exclude_none=True) - - rows = client.query(sql_query).result() - - results = [row.get(column_name) for row in rows] - - tool_output = ToolOutput(status="success", results=results).model_dump(exclude_none=True) - except Exception as e: - tool_output = ToolOutput( - status="error", error_details={"message": f"Failed to inspect colum values:\n{e}"} - ).model_dump(exclude_none=True) - - return json.dumps(tool_output, ensure_ascii=False) - - -@tool -def search_dictionary_table(dataset_gcp_id: str, table_name: str, column_name: str) -> str: - """Look up the meaning of coded values in a column using the dataset's dictionary table - - **ALWAYS USE THIS TOOL FIRST** when you need to understand column values before `inspect_column_values()`. - Many Base dos Dados datasets use coded values or abbreviations for storage efficiency. - This tool provides the authoritative meanings of these codes. - - Args: - dataset_gcp_id (str): Dataset portion of BigQuery reference (project.dataset_name) - - Extract from table gcp_id: if table is "basedosdados.br_ibge_censo.municipio" - - Use dataset_gcp_id: "basedosdados.br_ibge_censo" - table_name (str): Just the table name (not full gcp_id) - - Extract from gcp_id: if "basedosdados.br_ibge_censo.municipio" - - Use table_name: "municipio" - column_name (str): Exact column name that contains encoded values - - Returns: - str: JSON array of dictionary mappings with structure: - - chave: The encoded value found in your data (1, 2, 'A', 'B', etc.) - - valor: Human-readable meaning/description - - If no mappings found, returns guidance on alternative approaches. - - Next steps after using this tool: - 1. Match the returned 'chave' values with codes in your data - 2. Use JOIN with dictionary table for readable results: - ```sql - SELECT t.*, d.valor as decoded_column - FROM `project.dataset.table` t - LEFT JOIN `project.dataset.dicionario` d - ON t.encoded_column = d.chave - AND d.id_tabela = 'table_name' - AND d.nome_coluna = 'column_name' - ``` - 3. Or use the mappings to understand what codes mean for filtering - - Example usage: - # After inspect_column_values shows tipo_escola has values [1,2,3,4]: - search_dictionary_table("basedosdados.br_inep_censo_escolar", "escola", "tipo_escola") - - # Result might show: 1='Pública', 2='Privada', 3='Federal', 4='Municipal' - - # Then use in queries: - SELECT tipo_escola, COUNT(*) FROM `basedosdados.br_inep_censo_escolar.escola` - WHERE tipo_escola IN (1, 2) -- Now you know 1=Pública, 2=Privada - GROUP BY tipo_escola - - Common encoded column types in Brazilian data: - - Geographic codes: estado, municipio (IBGE codes) - - Category codes: tipo_*, categoria_*, situacao_* - - Status codes: ativo, situacao, condicao - - Classification codes: nivel_*, grau_*, classe_* - - If no dictionary found: - - Dataset might not have encoded values - - Values might be in separate reference tables - - Try common directory datasets like br_bd_diretorios_brasil - """ # noqa: E501 + project_name, dataset_name, table_name = table_gcp_id.split(".") + except ValueError: + return ToolOutput( + status="error", + error_details={ + "message": ( + f"{table_gcp_id} is not a valid table reference. " + "Please, provide a valid table reference in the format `project.dataset.table`" + ) + }, + ).model_dump_json(indent=2, exclude_none=True) - # Build the specific query for this table/column combination - dict_table_id = f"{dataset_gcp_id}.dicionario" + dataset_id = f"{project_name}.{dataset_name}" + dict_table_id = f"{dataset_id}.dicionario" search_query = f""" SELECT nome_coluna, chave, valor @@ -603,14 +403,15 @@ def search_dictionary_table(dataset_gcp_id: str, table_name: str, column_name: s query_job = client.query(search_query, job_config=job_config) limit_bytes = 1e9 # 1GB limit for dictionary queries + total_bytes = query_job.total_bytes_processed - if query_job.total_bytes_processed > limit_bytes: + if total_bytes and total_bytes > limit_bytes: return ToolOutput( status="error", error_details={ "type": "QueryTooLarge", "limit_bytes": limit_bytes, - "total_processed_bytes": query_job.total_bytes_processed, + "total_processed_bytes": total_bytes, "message": ( "Dictionary table is unexpectedly large. " "This might not be the right approach. " @@ -629,7 +430,7 @@ def search_dictionary_table(dataset_gcp_id: str, table_name: str, column_name: s error_details={ "type": "TableNotFound", "message": ( - f"Dictionary table not found for dataset {dataset_gcp_id}. " + f"Dictionary table not found for dataset {dataset_id}. " "This indicates this dataset does not contain a dicionary table. " "Consider using the `inspect_column_values` tool to inspect column values." ), @@ -637,7 +438,60 @@ def search_dictionary_table(dataset_gcp_id: str, table_name: str, column_name: s ).model_dump_json(indent=2, exclude_none=True) except Exception as e: tool_output = ToolOutput( - status="error", error_details={"message": f"Failed to search dictionary table:\n{e}"} + status="error", error_details={"message": f"Failed to decode table values:\n{e}"} + ).model_dump(exclude_none=True) + + return json.dumps(tool_output, ensure_ascii=False, default=str) + + +@tool +def inspect_column_values(table_gcp_id: str, column_name: str, limit: int = 100) -> str: + """Show actual distinct values in a column. + + FALLBACK tool to use when `search_dictionary_table()` fails. + Useful for understanding data patterns and planning `WHERE` clauses. + + Args: + table_gcp_id (str): Full BigQuery table reference + column_name (str): Column name from `get_dataset_details()` + limit (int, optional): Max distinct values to return (default: 100) + + Returns: + str: JSON array of distinct values. + """ # noqa: E501 + sql_query = f"SELECT DISTINCT {column_name} FROM {table_gcp_id} LIMIT {limit}" + + try: + job_config = bq.QueryJobConfig(dry_run=True, use_query_cache=False) + query_job = client.query(sql_query, job_config=job_config) + + limit_bytes = 5e9 # 5GB limit for inspection queries + total_bytes = query_job.total_bytes_processed + + if total_bytes and total_bytes > limit_bytes: + return ToolOutput( + status="error", + error_details={ + "status": "error", + "type": "QueryTooLarge", + "limit_bytes": limit_bytes, + "total_processed_bytes": total_bytes, + "message": ( + f"Column inspection exceeds the 5GB per-query " + f"limit for inspection({total_bytes/1e9:.1f} GB). " + "Try a smaller limit or add WHERE filters to reduce data size." + ), + }, + ).model_dump_json(indent=2, exclude_none=True) + + rows = client.query(sql_query).result() + + results = [row.get(column_name) for row in rows] + + tool_output = ToolOutput(status="success", results=results).model_dump(exclude_none=True) + except Exception as e: + tool_output = ToolOutput( + status="error", error_details={"message": f"Failed to inspect colum values:\n{e}"} ).model_dump(exclude_none=True) - return json.dumps(tool_output, ensure_ascii=False) + return json.dumps(tool_output, ensure_ascii=False, default=str) From d4134926752a3d9f564272098facae225656d88c Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Tue, 26 Aug 2025 08:59:22 -0300 Subject: [PATCH 05/23] add validator to `ToolOutput` model --- backend/apps/chatbot/tools.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/backend/apps/chatbot/tools.py b/backend/apps/chatbot/tools.py index c2f6569a..f83a44a3 100644 --- a/backend/apps/chatbot/tools.py +++ b/backend/apps/chatbot/tools.py @@ -1,13 +1,13 @@ # -*- coding: utf-8 -*- import json import os -from typing import Any, Literal, Optional +from typing import Any, Literal, Optional, Self import httpx from google.api_core import exceptions from google.cloud import bigquery as bq from langchain_core.tools import tool -from pydantic import BaseModel +from pydantic import BaseModel, model_validator SEARCH_URL = "https://backend.basedosdados.org/search/" GRAPHQL_URL = "https://backend.basedosdados.org/graphql" @@ -115,6 +115,12 @@ class ToolOutput(BaseModel): results: Optional[Any] = None error_details: Optional[dict[str, Any]] = None + @model_validator(mode="after") + def check_passwords_match(self) -> Self: + if (self.results is None) ^ (self.error_details is None): + return self + raise ValueError("Only one of 'results' or 'error_details' should be set") + client = bq.Client(project=os.environ["QUERY_PROJECT_ID"]) From b22be9b8d516b470f35c05b50307058e552fa71b Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Tue, 26 Aug 2025 13:09:34 -0300 Subject: [PATCH 06/23] refactor and bugfix in tools module --- backend/apps/chatbot/tools.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/backend/apps/chatbot/tools.py b/backend/apps/chatbot/tools.py index f83a44a3..24428d85 100644 --- a/backend/apps/chatbot/tools.py +++ b/backend/apps/chatbot/tools.py @@ -1,6 +1,7 @@ # -*- coding: utf-8 -*- import json import os +from functools import cache from typing import Any, Literal, Optional, Self import httpx @@ -122,7 +123,9 @@ def check_passwords_match(self) -> Self: raise ValueError("Only one of 'results' or 'error_details' should be set") -client = bq.Client(project=os.environ["QUERY_PROJECT_ID"]) +@cache +def get_bigquery_client(): + return bq.Client(project=os.environ["QUERY_PROJECT_ID"]) @tool @@ -210,7 +213,7 @@ def get_dataset_details(dataset_id: str) -> str: response.raise_for_status() data: dict[str, dict[str, dict]] = response.json() - dataset_edges = data.get("data").get("allDataset", {}).get("edges", []) + dataset_edges = data.get("data", {}).get("allDataset", {}).get("edges", []) if not dataset_edges: return f"Dataset {dataset_id} not found" @@ -325,6 +328,8 @@ def execute_bigquery_sql(sql_query: str) -> str: Returns: str: Query results as JSON array. Empty results return "[]". """ # noqa: E501 + client = get_bigquery_client() + try: job_config = bq.QueryJobConfig(dry_run=True, use_query_cache=False) query_job = client.query(sql_query, job_config=job_config) @@ -377,6 +382,8 @@ def decode_table_values(table_gcp_id: str, column_name: Optional[str] = None) -> Returns: str: JSON array with chave (code) and valor (meaning) mappings. """ # noqa: E501 + client = get_bigquery_client() + try: project_name, dataset_name, table_name = table_gcp_id.split(".") except ValueError: @@ -465,6 +472,8 @@ def inspect_column_values(table_gcp_id: str, column_name: str, limit: int = 100) Returns: str: JSON array of distinct values. """ # noqa: E501 + client = get_bigquery_client() + sql_query = f"SELECT DISTINCT {column_name} FROM {table_gcp_id} LIMIT {limit}" try: From ec754c5810103b8892529d8ad9c988d808973d1a Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Tue, 26 Aug 2025 13:11:39 -0300 Subject: [PATCH 07/23] create `prompts` module --- backend/apps/chatbot/prompts.py | 101 ++++++++++++++++++++++++++++++++ 1 file changed, 101 insertions(+) create mode 100644 backend/apps/chatbot/prompts.py diff --git a/backend/apps/chatbot/prompts.py b/backend/apps/chatbot/prompts.py new file mode 100644 index 00000000..172a4213 --- /dev/null +++ b/backend/apps/chatbot/prompts.py @@ -0,0 +1,101 @@ +# -*- coding: utf-8 -*- +SQL_AGENT_SYSTEM_PROMPT = """# Persona: Base dos Dados Research Assistant +You are a specialized AI research assistant, an expert in the Base dos Dados (BD) platform and the landscape of Brazilian public data. Your mission is to be a knowledgeable, systematic, and persistent research partner. You don't just execute tools; you guide users through the complexities of Brazil's data ecosystem, explaining the context behind the data and teaching best practices along the way. + +--- + +# Core Directives (Mandatory) + +**1. Embody the Persona**: Act as Base dos Dados Research Asisstant, the expert guide. Be proactive, educational, and systematic in every interaction. +**2. Adhere to the Workflow**: Your reasoning process **MUST** follow this strict cycle for every user request. Do not deviate. + - **Thought**: First, state your goal and reasoning. Formulate a hypothesis, select the appropriate tool, and explain your chosen parameters. This inner monologue is crucial for transparency. + - **Action**: Execute a single tool call based on your thought process (`search_datasets`, `get_dataset_details`, `decode_table_values`, `inspect_column_values`, or `execute_bigquery_sql`). +**3. Search Protocol is Mandatory**: **Always** begin your investigation with the single-keyword search strategy outlined below. This is the most critical rule for success. +**4. Explain Everything**: Never just show data. Summarize your findings, explain the data's source and context, highlight key insights, and suggest logical next steps. + +--- + +# Knowledge Base: Brazilian Data Essentials + +### Key Data Sources +- **Instituto Brasileiro de Geografia e Estatística (IBGE)**: Census, demographics, economic surveys (`censo`, `pnad`, `pof`). +- **Instituto Nacional de Estudos e Pesquisas Educacionais Anísio Teixeira (INEP)**: Education data (`ideb`, `censo escolar`, `enem`). +- **Ministério da Saúde (MS)**: Health data (`pns`, `sinasc`, `sinan`, `sim`). +- **Ministério da Economia (ME)**: Employment & economic data (`rais`, `caged`). +- **Tribunal Superior Eleitoral (TSE)**: Electoral data (`eleicoes`, `filiados`). +- **Banco Central do Brasil (BCB)**: Financial data (`taxa selic`, `cambio`, `ipca`). + +### Common Data Patterns +- **Geographic**: `sigla_uf` (state), `id_municipio` (municipality), `regiao`. +- **Temporal**: `ano` (year), `mes` (month), `data` (date), `semestre`, `trimestre`. +- **Identifiers**: `id_*`, `codigo_*`, `sigla_*`. +- **Coded Values**: Many columns use codes for efficiency (e.g., `id_municipio`). **Always** prioritize `decode_table_values` to understand them. Use `inspect_column_values` as a fallback for exploration. + +--- + +# Search Protocol + +You **MUST** follow this tiered search funnel. Do not skip steps. Justify your keyword choices in your **Thought** process. + +### Tier 1: High-Confidence Single Keywords (Always Try First) +*Start every search with a single, high-probability keyword, tried in this specific order.* +1. **Dataset Name**: If the user's query mentions a known dataset name (`censo`, `rais`, `enem`, `sinasc`), use it directly. +2. **Organization Acronym**: If a government organization is relevant (`ibge`, `inep`, `ms`, `tse`, `bcb`), use its acronym. +3. **Core Theme (Portuguese)**: Use a broad, common theme in Portuguese (`educacao`, `saude`, `economia`, `emprego`, `eleicoes`). + + +**User:** Como foi o desempenho em matemática dos alunos no brasil nos últimos anos? +**Thought:** The user is asking about student performance. The organization `inep` might be a good data source. I will start by searching with the keyword "inep". If that fails, I will try searching for the theme "educacao". +**Action:** `search_datasets(inep)` + + +### Tier 2: Alternative Single Keywords (If Tier 1 Fails) +*If and only if Tier 1 yields no relevant results, document the failure and proceed to these options.* +- **Synonyms**: Try a Portuguese synonym (`ensino` for `educacao`, `trabalho` for `emprego`). +- **Broader Concepts**: Use a more general term (`social`, `demografia`, `infraestrutura`). +- **English Equivalents**: As a last resort for single keywords, try English (`health`, `education`). + +### Tier 3: Multi-Keyword Search (Last Resort) +*Only use 2-3 keywords if all single-keyword searches have failed. This is an exception, not the rule.* +- **Theme + Agency**: `saude ms`, `educacao inep` +- **Dataset + Geography**: `censo municipio`, `rais estado` + +--- + +# BigQuery SQL Protocol + +- **Reference Full IDs**: Always use the full table ID: `project.dataset.table`. +- **Select Specific Columns**: Never use `SELECT *`. Explicitly list the columns you need. +- **Limit for Exploration**: When first inspecting a table, **always** use a `LIMIT` clause. You can query without a `LIMIT` clause later. +- **Filter Early and Often**: Use `WHERE` clauses on partitioned or clustered columns (usually `ano`) to drastically reduce query cost. +- **Default to Most Recent Data**: If the user does not specify a time range, your default behavior **MUST** be to query for the most recent data. Find the latest year or date in the relevant column (e.g., `ano`) and use it to filter the query. You **MUST** state that you queried the most recent data available. +- **Order for Insights**: Use `ORDER BY` to present data logically. +- **No DDL/DML:** NEVER run DDL/DML commands (`CREATE`, `ALTER`, `DROP`, `INSERT`, `UPDATE`, `DELETE`) + +--- + +# User Communication Protocol + +### Response Structure +1. **Summary of Findings**: Start with a clear, concise summary of the answer. +2. **Context**: Explain what the data represents. Mention the source organization (e.g., "Data from IBGE's 2010 Census..."), the time period, and the geographic level. +3. **Data/Results**: Present the data clearly. Use Markdown tables for structured results, bullet points for lists, etc. Display null/empty values as "N/A" for clarity. +4. **Key Insights**: Highlight 1-3 important points or patterns from the results. +5. **Suggested Next Steps**: Propose a relevant follow-up question, a related dataset to explore, or a way to refine the current analysis. + +### Handling Failures +- **Search Fails**: Explain your keyword strategy, state why it failed (e.g., "The search for 'cnes' returned no datasets"), and describe your next attempt based on the Search Protocol. +- **Query Errors**: Analyze the BigQuery error message. Suggest a specific fix (e.g., "The query is too large. I will add a `WHERE` clause to filter by year to reduce the data processed."). +- **Empty Results**: Hypothesize why the result is empty. Check your filters, the data's time range, or if you are filtering on a coded value incorrectly. Suggest a modified query. + +--- + +# Final Reminder + +**Before executing any action, ensure you are compliant:** +- **NEVER** use `SELECT *`. +- **NEVER** query a table without a `LIMIT` clause during initial exploration. +- **NEVER** run Data Definition/Manipulation Language (`CREATE`, `ALTER`, `DROP`, `INSERT`, `UPDATE`, `DELETE`). Your access is strictly read-only. +- **NEVER** start with a multi-keyword search. The Search Protocol is mandatory. +- **NEVER** present raw data without a summary and context first. +- **NEVER** give up after one failed attempt. Show persistence and a systematic problem-solving approach.""" # noqa: E501 From 53802f58c552054438c683916574464d91ea0344 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Tue, 26 Aug 2025 13:12:04 -0300 Subject: [PATCH 08/23] update chatbot views --- backend/apps/chatbot/views.py | 143 +++++----------------------------- 1 file changed, 18 insertions(+), 125 deletions(-) diff --git a/backend/apps/chatbot/views.py b/backend/apps/chatbot/views.py index d3003d17..6d9ae766 100644 --- a/backend/apps/chatbot/views.py +++ b/backend/apps/chatbot/views.py @@ -23,6 +23,7 @@ from backend.apps.chatbot.feedback_sender import LangSmithFeedbackSender from backend.apps.chatbot.models import Feedback, MessagePair, Thread +from backend.apps.chatbot.prompts import SQL_AGENT_SYSTEM_PROMPT from backend.apps.chatbot.serializers import ( FeedbackCreateSerializer, FeedbackSerializer, @@ -32,11 +33,11 @@ UserMessageSerializer, ) from backend.apps.chatbot.tools import ( + decode_table_values, execute_bigquery_sql, get_dataset_details, inspect_column_values, search_datasets, - search_dictionary_table, ) ModelSerializer = TypeVar("ModelSerializer", bound=Serializer) @@ -370,12 +371,22 @@ def _get_sql_assistant(): search_datasets, get_dataset_details, execute_bigquery_sql, + decode_table_values, inspect_column_values, - search_dictionary_table, ] + # agent = ReActAgent( + # model=model, + # tools=tools, + # system_message=SQL_AGENT_SYSTEM_PROMPT, + # checkpointer=checkpointer, + # ) + agent = create_react_agent( - model=model, tools=tools, prompt=SYSTEM_PROMPT, checkpointer=checkpointer + model=model, + tools=tools, + prompt=SQL_AGENT_SYSTEM_PROMPT, + checkpointer=checkpointer, ) # assistant = SQLAssistant( @@ -410,7 +421,10 @@ def _stream_sql_assistant_response( logger.info("Calling SQLAssistant...") with _get_sql_assistant() as assistant: agent_response = assistant.invoke( - input={"messages": [{"role": "user", "content": message}]}, config=config + # question=message, + # config=config + input={"messages": [{"role": "user", "content": message}]}, + config=config, ) # for mode, chunk in assistant.stream( # message=message, @@ -543,124 +557,3 @@ def _validate(request: Request, model_serializer: Type[ModelSerializer]) -> Mode raise exceptions.ValidationError(serializer.errors) return serializer - - -SYSTEM_PROMPT = """# Base dos Dados Research Assistant System Prompt - -You are a specialized research assistant expert in Brazilian data analysis using the Base dos Dados (BD) platform. Your role is to help users find, explore, and analyze Brazilian public data through BigQuery using a systematic approach. - -## Your Expertise -- **Brazilian Data Landscape**: Deep knowledge of Brazilian government agencies, data sources, and common data structures -- **Base dos Dados Platform**: Expert user of BD's catalog and BigQuery integration -- **Research Methodology**: Systematic approach to data discovery and analysis -- **SQL & BigQuery**: Proficient in writing efficient, cost-effective queries - -## Core Principles - -### 1. Systematic Data Discovery -Always follow this workflow: -1. **Search** → Use keywords to find relevant datasets -2. **Explore** → Review dataset overviews and available tables -3. **Investigate** → Get detailed column information for specific tables -4. **Query** → Execute targeted SQL queries with proper filters - -### 2. Search Strategy -- Use **individual keywords** only (Elasticsearch backend) -- Try **Portuguese terms** first ("educacao" not "education") -- Use **agency acronyms** (IBGE, INEP, ANVISA, ANS) -- **Retry with different keywords** if initial search fails -- **Never give up** after one failed search - always try alternatives - -### 3. Brazilian Data Context -Understand common data patterns: -- **Geographic**: sigla_uf (state codes), id_municipio (municipality codes) -- **Temporal**: ano (year), mes (month), data (date) -- **Identifiers**: id_*, codigo_*, sigla_* -- **Values**: *_valor, *_quantidade, *_taxa, *_indice, *_percentual - -### 4. Query Optimization -- Always start with **LIMIT clauses** for exploration -- Use **WHERE filters** early to reduce processing costs -- Select **specific columns** rather than SELECT * -- Be mindful of the **20GB query limit** - -## Key Brazilian Data Sources -- **IBGE**: Census, demographic, economic surveys (censo, pnad, pof) -- **INEP**: Education data (censo escolar, enem, prova brasil) -- **MS/DataSUS**: Health data (sinasc, sim, sinan) -- **MTE**: Employment data (rais, caged) -- **TSE**: Electoral data (eleicoes, candidatos) -- **BACEN**: Financial data (taxa selic, cambio) - -## Response Guidelines - -### When Searching Fails -- **Never give up after one search** -- Try multiple keyword variations: - - Portuguese vs English terms - - Synonyms and related terms - - Broader or narrower concepts - - Agency names and acronyms -- Explain your search strategy to the user -- Ask for clarification if search terms are unclear - -### When Presenting Results -- **Summarize findings clearly** before showing raw data -- **Explain data context**: What agency published it, what it represents -- **Highlight key insights** from query results -- **Suggest follow-up analyses** or related datasets - -### When Queries Fail -- **Analyze error messages** and suggest fixes -- **Recommend query optimizations** (filters, limits, column selection) -- **Explain BigQuery concepts** when helpful -- **Try alternative approaches** if initial query doesn't work - -## Communication Style -- **Friendly and helpful**: Make complex data accessible -- **Educational**: Explain Brazilian data landscape and sources -- **Systematic**: Show your step-by-step approach -- **Persistent**: Don't give up easily on difficult requests -- **Context-aware**: Understand Brazilian geography, politics, and institutions - -## Tool Usage Best Practices - -### search_datasets -- Use single keywords or short phrases -- Try Portuguese terms first -- Be persistent - retry with alternatives if no results -- Focus on topic, geography, or agency keywords - -### get_dataset_details -- Use to understand dataset scope and available tables -- Review themes, organizations, and tags for context -- Identify the most relevant tables before diving deeper - -### get_table_details -- Get complete column information and BigQuery references -- Note the gcp_id for SQL queries -- Understand column types and descriptions - -### execute_sql_query -- Always start with LIMIT for exploration -- Use proper BigQuery syntax with backticks -- Filter early with WHERE clauses -- Handle errors gracefully and suggest improvements - -## Error Recovery -- **Search failures**: Try different keywords, explain your strategy -- **Dataset not found**: Verify IDs, suggest alternatives -- **Query errors**: Analyze error messages, suggest fixes -- **Large query warnings**: Recommend optimization strategies -- **Empty results**: Check filters, suggest broader queries - -## Success Metrics -Your success is measured by: -- Finding relevant datasets for user queries -- Providing accurate, well-contextualized data analysis -- Teaching users about Brazilian data landscape -- Writing efficient, cost-effective BigQuery queries -- Persistence in overcoming search and query challenges - -Remember: You are not just a tool executor, but a knowledgeable research partner who understands Brazilian data, institutions, and research needs. Be proactive, educational, and persistent in helping users discover insights from Brazil's rich public data ecosystem. -""" # noqa: E501 From 25d204b5545e94a99b1fe15ecfa0d6695e67ce58 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Wed, 27 Aug 2025 17:58:24 -0300 Subject: [PATCH 09/23] create `agent` module --- backend/apps/chatbot/agent/__init__.py | 0 backend/apps/chatbot/{ => agent}/prompts.py | 0 backend/apps/chatbot/{ => agent}/tools.py | 3 +-- 3 files changed, 1 insertion(+), 2 deletions(-) create mode 100644 backend/apps/chatbot/agent/__init__.py rename backend/apps/chatbot/{ => agent}/prompts.py (100%) rename backend/apps/chatbot/{ => agent}/tools.py (99%) diff --git a/backend/apps/chatbot/agent/__init__.py b/backend/apps/chatbot/agent/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/backend/apps/chatbot/prompts.py b/backend/apps/chatbot/agent/prompts.py similarity index 100% rename from backend/apps/chatbot/prompts.py rename to backend/apps/chatbot/agent/prompts.py diff --git a/backend/apps/chatbot/tools.py b/backend/apps/chatbot/agent/tools.py similarity index 99% rename from backend/apps/chatbot/tools.py rename to backend/apps/chatbot/agent/tools.py index 24428d85..a17bf516 100644 --- a/backend/apps/chatbot/tools.py +++ b/backend/apps/chatbot/agent/tools.py @@ -492,8 +492,7 @@ def inspect_column_values(table_gcp_id: str, column_name: str, limit: int = 100) "limit_bytes": limit_bytes, "total_processed_bytes": total_bytes, "message": ( - f"Column inspection exceeds the 5GB per-query " - f"limit for inspection({total_bytes/1e9:.1f} GB). " + "Column inspection exceeds the 5GB per-query limit for inspection. " "Try a smaller limit or add WHERE filters to reduce data size." ), }, From 8184636c34d74e400e60469c0cd6a6dc6e7bee32 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Wed, 27 Aug 2025 18:06:02 -0300 Subject: [PATCH 10/23] refactored views and message streaming --- .../0007_rename_steps_messagepair_events.py | 18 ++ backend/apps/chatbot/models.py | 2 +- backend/apps/chatbot/utils/stream.py | 276 +++++------------- backend/apps/chatbot/views.py | 187 +++--------- 4 files changed, 133 insertions(+), 350 deletions(-) create mode 100644 backend/apps/chatbot/migrations/0007_rename_steps_messagepair_events.py diff --git a/backend/apps/chatbot/migrations/0007_rename_steps_messagepair_events.py b/backend/apps/chatbot/migrations/0007_rename_steps_messagepair_events.py new file mode 100644 index 00000000..e5b8c698 --- /dev/null +++ b/backend/apps/chatbot/migrations/0007_rename_steps_messagepair_events.py @@ -0,0 +1,18 @@ +# -*- coding: utf-8 -*- +# Generated by Django 4.2.10 on 2025-08-27 20:15 + +from django.db import migrations + + +class Migration(migrations.Migration): + dependencies = [ + ("chatbot", "0006_messagepair_error_message_messagepair_steps_and_more"), + ] + + operations = [ + migrations.RenameField( + model_name="messagepair", + old_name="steps", + new_name="events", + ), + ] diff --git a/backend/apps/chatbot/models.py b/backend/apps/chatbot/models.py index a2898d40..70735bee 100644 --- a/backend/apps/chatbot/models.py +++ b/backend/apps/chatbot/models.py @@ -26,7 +26,7 @@ class MessagePair(models.Model): generated_queries = models.JSONField(null=True, blank=True) generated_chart = models.JSONField(null=True, blank=True) created_at = models.DateTimeField(auto_now_add=True) - steps = models.JSONField(null=True) + events = models.JSONField(null=True) class Meta: constraints = [ diff --git a/backend/apps/chatbot/utils/stream.py b/backend/apps/chatbot/utils/stream.py index e768f1a8..3eb6366d 100644 --- a/backend/apps/chatbot/utils/stream.py +++ b/backend/apps/chatbot/utils/stream.py @@ -1,231 +1,89 @@ # -*- coding: utf-8 -*- -import re -from collections.abc import Callable -from typing import TypeAlias +from typing import Any, Literal, Optional -import sqlparse from langchain_core.messages import AIMessage, ToolMessage -from pydantic import BaseModel, Field +from pydantic import BaseModel -# SQLAgent nodes -LIST_DATASETS = "list_datasets" -CALL_SELECT_DATASETS = "call_select_datasets" -TABLES_INFO = "tables_info" -QUERY_AGENT = "query_agent" -GET_ANSWER = "get_answer" -SQL_TOOLS = "tools" +class ToolCall(BaseModel): + id: str + name: str + args: dict[str, Any] -class StepContent(BaseModel): - title: str | None = Field(default=None, description="A title to be displayed in the UI") - body: str = Field(description="Content to be displayed in the UI") +class ToolOutput(BaseModel): + status: Literal["error", "success"] + tool_call_id: str + tool_name: str + output: str -class Step(BaseModel): - label: str = Field( - description="Short description of the action taken by the agent at this step" - ) - content: list[StepContent] = Field( - description="Detailed outputs generated by the agent for this step" - ) +EventType = Literal[ + "tool_call", + "tool_result", + "final_answer", + "error", +] -Handler: TypeAlias = Callable[[dict], Step | None] +class EventData(BaseModel): + message: Optional[str] = None + tool_calls: Optional[list[ToolCall]] = None + tool_outputs: Optional[list[ToolOutput]] = None + error_details: Optional[dict[str, Any]] = None -# ============================== Helper Functions ============================== -def _format_sql(sql: str, reindent: bool = True, keyword_case: str = "upper") -> str: - return sqlparse.format(sql=sql, reindent=reindent, keyword_case=keyword_case) +class StreamEvent(BaseModel): + type: EventType + data: EventData -def _format_datasets_info(text: str) -> str: - matches = re.findall( - pattern=r"^# (.*?)\s+### Description:\s+(.*?)### Tables:", - string=text, - flags=re.MULTILINE | re.DOTALL, - ) + def to_sse(self) -> str: + return self.model_dump_json() + "\n\n" - parts = [] - for title, description in matches: - parts.append(f"##### {title.strip()}\n\n{description.strip()}") +def process_chunk(chunk: dict[str, Any]) -> StreamEvent: + """Process a streaming chunk from an agent workflow into a standardized StreamEvent. - return "\n\n".join(parts) - - -def _format_tables_info(text: str) -> str: - matches = re.findall( - pattern=r"^# (.*?)\s+### Description:\s+(.*?)### Schema:", - string=text, - flags=re.MULTILINE | re.DOTALL, - ) - - parts = [] - - for title, description in matches: - parts.append(f"##### {title}\n\n{description}") - - return "\n\n".join(parts) - - -# ========================== SQLAgent Chunk Handlers ========================== -def _handle_list_datasets(chunk: dict) -> Step: - message: ToolMessage = chunk[LIST_DATASETS]["messages"][0] - - label = "Selecionando conjuntos de dados..." - - content = StepContent( - title="Buscando Conjuntos de Dados:", body=_format_datasets_info(message.content) - ) - - return Step(label=label, content=[content]) - - -def _handle_call_select_datasets(chunk: dict) -> Step: - message: AIMessage = chunk[CALL_SELECT_DATASETS]["messages"][-1] - names: list[str] = [] - - for tool_call in message.tool_calls: - dataset_names: str = tool_call["args"].get("dataset_names") - if dataset_names: - names.extend(dataset_names.split(",")) - - bolded_names = [f"**{name.strip()}**" for name in names] - - label = "Selecionando conjuntos de dados..." - - if not bolded_names: - body = "" - - body = "Conjuntos de dados selecionandos:\n- " + "\n- ".join(bolded_names) - - content = StepContent(title="Selecionando Conjuntos de Dados:", body=body) - - return Step(label=label, content=[content]) - - -def _handle_tables_info(chunk: dict) -> Step: - messages: list[ToolMessage] = chunk[TABLES_INFO]["messages"] - - formatted_tables_info = [] - - for message in messages: - if message.status == "error": - formatted_tables_info.append( - ":red[**Erro ao selecionar tabelas. Tentando novamente...**]" - ) - else: - formatted_tables_info.append(_format_tables_info(message.content)) - - label = "Selecionando tabelas..." - - content = StepContent( - title="Selecionando Tabelas:", body="\n\n---\n\n".join(formatted_tables_info) - ) - - return Step(label=label, content=[content]) - - -def _handle_query_agent(chunk: dict) -> Step | None: - message: AIMessage = chunk[QUERY_AGENT]["messages"][0] - - if not message.tool_calls: - return None - - content_parts = [] - - for tool_call in message.tool_calls: - tool_name = tool_call["name"] - - if tool_name == "sql_query_check": - query = tool_call["args"].get("query") - label = "Verificando consulta..." - if query: - body = f"```sql\n{_format_sql(query)}\n```" - else: - body = "red[**Erro na chamada da ferramenta**]" - part = StepContent(title="Verificando Consulta:", body=body) - content_parts.append(part) - - elif tool_name == "sql_query_exec": - query = tool_call["args"].get("query") - label = "Executando consulta..." - if query: - body = f"```sql\n{_format_sql(query)}\n```" - else: - body = "red[**Erro na chamada da ferramenta**]" - part = StepContent(title="Executando Consulta:", body=body) - content_parts.append(part) - - if len(content_parts) > 1: - label = "Consultando banco de dados..." - - return Step(label=label, content=content_parts) - - -def _handle_sql_tools(chunk: dict) -> Step | None: - updates = chunk[SQL_TOOLS] - - # single tool call - if isinstance(updates, dict): - messages: list[ToolMessage] = updates["messages"] - # multiple parallel tool calls - elif isinstance(updates, list): - messages: list[ToolMessage] = [ - update["messages"][0] for update in updates if "messages" in update - ] - - content_parts = [] - - for message in messages: - if message.status == "success": - continue - - if message.name == "sql_query_check": - label = "Verificando consulta..." - part = StepContent(body=":red[**Erro na verificação da consulta**]") - elif message.name == "sql_query_exec": - label = "Executando consulta..." - part = StepContent(body=":red[**Erro na execução da consulta**]") - - content_parts.append(part) - - if content_parts: - if len(content_parts) > 1: - label = "Consultando banco de dados..." - return Step(label=label, content=content_parts) - - return None - - -def _handle_get_answer(chunk: dict) -> Step: - final_answer = chunk[GET_ANSWER]["final_answer"] - label = "Gerando resposta..." - content = StepContent(title="Resposta Final:", body=final_answer) - return Step(label=label, content=[content]) - - -# ========= Mapping of SQLAgent chunks to their corresponding handlers ========= -SQL_AGENT_HANDLERS: dict[str, Handler] = { - LIST_DATASETS: _handle_list_datasets, - CALL_SELECT_DATASETS: _handle_call_select_datasets, - TABLES_INFO: _handle_tables_info, - QUERY_AGENT: _handle_query_agent, - SQL_TOOLS: _handle_sql_tools, - GET_ANSWER: _handle_get_answer, -} - - -def process_chunk(chunk: dict) -> Step | None: - """Process a single chunk from the stream by dispatching it to the correct handler. + Converts LangChain agent and tool execution chunks into structured events for + consistent handling in streaming applications. Args: - chunk (dict): A chunk from the stream. + chunk (dict[str, Any]): Raw chunk dictionary containing either: + - "agent" key with AIMessage for tool calls or final answers + - "tools" key with ToolMessage list for tool execution results Returns: - Step | None: The processed results if a handler was found. `None` otherwise. + StreamEvent with appropriate type: + - "tool_call" for agent messages with tool invocations + - "final_answer" for agent messages without tool calls + - "tool_result" for tool execution outputs """ - for key, handler in SQL_AGENT_HANDLERS.items(): - if key in chunk: - return handler(chunk) - return None + if "agent" in chunk: + message: AIMessage = chunk["agent"]["messages"][0] + + if message.tool_calls: + tool_calls = [ + ToolCall(id=tool_call["id"], name=tool_call["name"], args=tool_call["args"]) + for tool_call in message.tool_calls + ] + event_type = "tool_call" + event_data = EventData(message=message.content, tool_calls=tool_calls) + else: + event_type = "final_answer" + event_data = EventData(message=message.content) + + return StreamEvent(type=event_type, data=event_data) + elif "tools" in chunk: + messages: list[ToolMessage] = chunk["tools"]["messages"] + + tool_outputs = [ + ToolOutput( + status=message.status, + tool_call_id=message.tool_call_id, + tool_name=message.name, + output=message.content, + ) + for message in messages + ] + + return StreamEvent(type="tool_result", data=EventData(tool_outputs=tool_outputs)) diff --git a/backend/apps/chatbot/views.py b/backend/apps/chatbot/views.py index 6d9ae766..bce1761b 100644 --- a/backend/apps/chatbot/views.py +++ b/backend/apps/chatbot/views.py @@ -1,5 +1,4 @@ # -*- coding: utf-8 -*- -import json import os import uuid from contextlib import contextmanager @@ -21,9 +20,16 @@ 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.tools import ( + decode_table_values, + execute_bigquery_sql, + get_dataset_details, + inspect_column_values, + search_datasets, +) from backend.apps.chatbot.feedback_sender import LangSmithFeedbackSender from backend.apps.chatbot.models import Feedback, MessagePair, Thread -from backend.apps.chatbot.prompts import SQL_AGENT_SYSTEM_PROMPT from backend.apps.chatbot.serializers import ( FeedbackCreateSerializer, FeedbackSerializer, @@ -32,13 +38,7 @@ ThreadSerializer, UserMessageSerializer, ) -from backend.apps.chatbot.tools import ( - decode_table_values, - execute_bigquery_sql, - get_dataset_details, - inspect_column_values, - search_datasets, -) +from backend.apps.chatbot.utils.stream import EventData, StreamEvent, process_chunk ModelSerializer = TypeVar("ModelSerializer", bound=Serializer) @@ -176,7 +176,7 @@ def delete(self, request: Request, thread_id: uuid.UUID) -> Response: try: thread.deleted = True thread.save() - with _get_sql_assistant() as assistant: + with _get_sql_agent() as assistant: assistant.clear_thread(str(thread_id)) return Response({"detail": "Thread deleted successfully"}) except Exception: @@ -305,46 +305,8 @@ def _get_feedback_sender() -> LangSmithFeedbackSender: return LangSmithFeedbackSender() -# @cache -# def _get_context_provider(connection: str) -> PostgresContextProvider: -# """Provide a configured `PostgresContextProvider` context provider. - -# Args: -# connection (str): A vector database connection for providing few-shot examples. - -# Returns: -# PostgresContextProvider: An instance of `PostgresContextProvider`. -# """ -# bq_billing_project = os.environ["BILLING_PROJECT_ID"] -# bq_query_project = os.environ["QUERY_PROJECT_ID"] - -# embedding_model = os.getenv("EMBEDDING_MODEL") -# pgvector_collection = os.getenv("PGVECTOR_COLLECTION") - -# if embedding_model and pgvector_collection: -# embeddings = VertexAIEmbeddings(embedding_model) - -# vector_store = PGVector( -# embeddings=embeddings, -# connection=connection, -# collection_name=pgvector_collection, -# use_jsonb=True, -# ) -# else: -# vector_store = None - -# context_provider = PostgresContextProvider( -# billing_project=bq_billing_project, -# query_project=bq_query_project, -# metadata_vector_store=vector_store, -# top_k=5, -# ) - -# return context_provider - - @contextmanager -def _get_sql_assistant(): +def _get_sql_agent(): """Provide a configured `SQLAssistant`. Yields: @@ -358,47 +320,27 @@ def _get_sql_assistant(): conn = f"postgresql://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}" - # context_provider = _get_context_provider(conn) - - # prompt_formatter = SQLPromptFormatter(vector_store=None) - with PostgresSaver.from_conn_string(conn) as checkpointer: checkpointer.setup() model = init_chat_model(MODEL_URI, temperature=0) tools = [ - search_datasets, - get_dataset_details, - execute_bigquery_sql, decode_table_values, + execute_bigquery_sql, + get_dataset_details, inspect_column_values, + search_datasets, ] - # agent = ReActAgent( - # model=model, - # tools=tools, - # system_message=SQL_AGENT_SYSTEM_PROMPT, - # checkpointer=checkpointer, - # ) - - agent = create_react_agent( + sql_agent = create_react_agent( model=model, tools=tools, prompt=SQL_AGENT_SYSTEM_PROMPT, checkpointer=checkpointer, ) - # assistant = SQLAssistant( - # model=model, - # context_provider=context_provider, - # prompt_formatter=prompt_formatter, - # checkpointer=checkpointer, - # ) - - # yield assistant - - yield agent + yield sql_agent def _stream_sql_assistant_response( @@ -414,88 +356,53 @@ def _stream_sql_assistant_response( Yields: Iterator[str]: JSON string containing the streaming status and the current step data. """ - steps = [] - # last_chunk = None + events = [] + message_pair = {} try: logger.info("Calling SQLAssistant...") - with _get_sql_assistant() as assistant: - agent_response = assistant.invoke( - # question=message, - # config=config + with _get_sql_agent() as agent: + for chunk in agent.stream( input={"messages": [{"role": "user", "content": message}]}, + stream_mode="updates", config=config, - ) - # for mode, chunk in assistant.stream( - # message=message, - # config=config, - # stream_mode=["updates", "values"], - # rewrite_query=True, - # ): - # last_chunk = chunk - - # # Skip "values" chunks during streaming. We only need the final one at the end, - # # as it contains the SQL Agent's final state. - # if mode == "values": - # continue - - # step = process_chunk(chunk) - - # if step is None: - # continue + ): + event = process_chunk(chunk) + events.append(event.model_dump()) - # steps.append(step.model_dump()) + yield event.to_sse() - # yield json.dumps({"status": "running", "data": step.model_dump_json()}) + "\n\n" - - # The last chunk represents the SQLAgent's final state, - # so we format it as the final response for the user. - # response = format_sql_agent_response(last_chunk) - response = {} - response["error_message"] = None + # The last event always contains the agent's final answer, + # so we use it to save the message pair in the database + message_pair["content"] = event.data.message + message_pair["error_message"] = None + logger.success("SQLAssistant called successfully") except Exception: logger.exception(f"Error responding message {config['run_id']}:") - response = { - "error_message": ( - "Ops, algo deu errado! Ocorreu um erro inesperado. Por favor, tente novamente. " - "Se o problema persistir, avise-nos. Obrigado pela paciência!" - ), + error_message = ( + "Ops, algo deu errado! Ocorreu um erro inesperado. Por favor, tente novamente. " + "Se o problema persistir, avise-nos. Obrigado pela paciência!" + ) + + yield StreamEvent( + type="error", data=EventData(error_details={"message": error_message}) + ).to_sse() + + message_pair = { "content": None, "sql_queries": None, + "error_message": error_message, } - logger.success("SQLAssistant called successfully") - - response["id"] = config["run_id"] - - response["content"] = agent_response["messages"][-1].content - response["sql_queries"] = None - - message_pair = MessagePair.objects.create( - id=response["id"], + MessagePair.objects.create( + id=config["run_id"], thread=thread, model_uri=MODEL_URI, user_message=message, - assistant_message=response["content"], - error_message=response["error_message"], - generated_queries=response["sql_queries"], - steps=steps, - ) - - yield ( - json.dumps( - { - "status": "complete", - "data": { - "id": message_pair.id, - "user_message": message_pair.user_message, - "assistant_message": message_pair.assistant_message, - "error_message": message_pair.error_message, - "generated_queries": message_pair.generated_queries, - }, - } - ) - + "\n\n" + assistant_message=message_pair["content"], + error_message=message_pair["error_message"], + generated_queries=None, + events=events, ) From 0ef708a74386e9132ec1d92a76b08f65b29ca870 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Thu, 28 Aug 2025 11:35:43 -0300 Subject: [PATCH 11/23] refactor views and message streaming --- backend/apps/chatbot/utils/stream.py | 18 +++++++++++------- backend/apps/chatbot/views.py | 26 +++++++++++--------------- 2 files changed, 22 insertions(+), 22 deletions(-) diff --git a/backend/apps/chatbot/utils/stream.py b/backend/apps/chatbot/utils/stream.py index 3eb6366d..dbeeef8c 100644 --- a/backend/apps/chatbot/utils/stream.py +++ b/backend/apps/chatbot/utils/stream.py @@ -2,7 +2,7 @@ from typing import Any, Literal, Optional from langchain_core.messages import AIMessage, ToolMessage -from pydantic import BaseModel +from pydantic import UUID4, BaseModel class ToolCall(BaseModel): @@ -20,14 +20,16 @@ class ToolOutput(BaseModel): EventType = Literal[ "tool_call", - "tool_result", + "tool_output", "final_answer", "error", + "complete", ] class EventData(BaseModel): - message: Optional[str] = None + run_id: Optional[UUID4] = None + content: Optional[str] = None tool_calls: Optional[list[ToolCall]] = None tool_outputs: Optional[list[ToolOutput]] = None error_details: Optional[dict[str, Any]] = None @@ -55,8 +57,10 @@ def process_chunk(chunk: dict[str, Any]) -> StreamEvent: Returns: StreamEvent with appropriate type: - "tool_call" for agent messages with tool invocations + - "tool_output" for tool execution outputs - "final_answer" for agent messages without tool calls - - "tool_result" for tool execution outputs + - "error" if an error occurs + - "complete" when streaming has finished """ if "agent" in chunk: message: AIMessage = chunk["agent"]["messages"][0] @@ -67,10 +71,10 @@ def process_chunk(chunk: dict[str, Any]) -> StreamEvent: for tool_call in message.tool_calls ] event_type = "tool_call" - event_data = EventData(message=message.content, tool_calls=tool_calls) + event_data = EventData(content=message.content, tool_calls=tool_calls) else: event_type = "final_answer" - event_data = EventData(message=message.content) + event_data = EventData(content=message.content) return StreamEvent(type=event_type, data=event_data) elif "tools" in chunk: @@ -86,4 +90,4 @@ def process_chunk(chunk: dict[str, Any]) -> StreamEvent: for message in messages ] - return StreamEvent(type="tool_result", data=EventData(tool_outputs=tool_outputs)) + return StreamEvent(type="tool_output", data=EventData(tool_outputs=tool_outputs)) diff --git a/backend/apps/chatbot/views.py b/backend/apps/chatbot/views.py index bce1761b..7c629f0c 100644 --- a/backend/apps/chatbot/views.py +++ b/backend/apps/chatbot/views.py @@ -357,10 +357,9 @@ def _stream_sql_assistant_response( Iterator[str]: JSON string containing the streaming status and the current step data. """ events = [] - message_pair = {} try: - logger.info("Calling SQLAssistant...") + logger.info("Calling SQL Agent...") with _get_sql_agent() as agent: for chunk in agent.stream( input={"messages": [{"role": "user", "content": message}]}, @@ -374,11 +373,12 @@ def _stream_sql_assistant_response( # The last event always contains the agent's final answer, # so we use it to save the message pair in the database - message_pair["content"] = event.data.message - message_pair["error_message"] = None - logger.success("SQLAssistant called successfully") + assistant_message = event.data.content + error_message = None + logger.success("SQL Agent called successfully. Saving message pair...") except Exception: logger.exception(f"Error responding message {config['run_id']}:") + assistant_message = None error_message = ( "Ops, algo deu errado! Ocorreu um erro inesperado. Por favor, tente novamente. " "Se o problema persistir, avise-nos. Obrigado pela paciência!" @@ -388,22 +388,18 @@ def _stream_sql_assistant_response( type="error", data=EventData(error_details={"message": error_message}) ).to_sse() - message_pair = { - "content": None, - "sql_queries": None, - "error_message": error_message, - } - - MessagePair.objects.create( + message_pair = MessagePair.objects.create( id=config["run_id"], thread=thread, model_uri=MODEL_URI, user_message=message, - assistant_message=message_pair["content"], - error_message=message_pair["error_message"], - generated_queries=None, + assistant_message=assistant_message, + error_message=error_message, events=events, ) + logger.success(f"Message pair {message_pair.id} saved successfully") + + yield StreamEvent(type="complete", data=EventData(run_id=message_pair.id)).to_sse() def _get_thread_by_id(thread_id: uuid.UUID) -> Thread: From d7881ae5c8b155192ea60d4aafd260cebb437694 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Fri, 29 Aug 2025 15:53:02 -0300 Subject: [PATCH 12/23] update type annotation and docstring --- backend/apps/chatbot/utils/stream.py | 22 +++++++++------------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/backend/apps/chatbot/utils/stream.py b/backend/apps/chatbot/utils/stream.py index dbeeef8c..3997ae42 100644 --- a/backend/apps/chatbot/utils/stream.py +++ b/backend/apps/chatbot/utils/stream.py @@ -43,24 +43,19 @@ def to_sse(self) -> str: return self.model_dump_json() + "\n\n" -def process_chunk(chunk: dict[str, Any]) -> StreamEvent: - """Process a streaming chunk from an agent workflow into a standardized StreamEvent. - - Converts LangChain agent and tool execution chunks into structured events for - consistent handling in streaming applications. +def process_chunk(chunk: dict[str, Any]) -> StreamEvent | None: + """Process a streaming chunk from a react agent workflow into a standardized StreamEvent. Args: - chunk (dict[str, Any]): Raw chunk dictionary containing either: - - "agent" key with AIMessage for tool calls or final answers - - "tools" key with ToolMessage list for tool execution results + chunk (dict[str, Any]): Raw chunk from agent workflow. + Only processes "agent" and "tools" nodes. Returns: - StreamEvent with appropriate type: - - "tool_call" for agent messages with tool invocations - - "tool_output" for tool execution outputs + StreamEvent | None: Structured event or None if the chunk is ignored: + - "tool_call" for agent messages with tool calls + - "tool_output" for tool execution results - "final_answer" for agent messages without tool calls - - "error" if an error occurs - - "complete" when streaming has finished + - None for ignored chunks """ if "agent" in chunk: message: AIMessage = chunk["agent"]["messages"][0] @@ -91,3 +86,4 @@ def process_chunk(chunk: dict[str, Any]) -> StreamEvent: ] return StreamEvent(type="tool_output", data=EventData(tool_outputs=tool_outputs)) + return None From be3acb51109ece4e07fa16153f869493a67cbd9a Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Fri, 29 Aug 2025 15:53:23 -0300 Subject: [PATCH 13/23] create `get_tools` helper method --- backend/apps/chatbot/agent/tools.py | 29 ++++++++++++++++++++++++++--- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/backend/apps/chatbot/agent/tools.py b/backend/apps/chatbot/agent/tools.py index a17bf516..972ede99 100644 --- a/backend/apps/chatbot/agent/tools.py +++ b/backend/apps/chatbot/agent/tools.py @@ -5,9 +5,9 @@ from typing import Any, Literal, Optional, Self import httpx -from google.api_core import exceptions +from google.api_core import exceptions as google_api_exceptions from google.cloud import bigquery as bq -from langchain_core.tools import tool +from langchain_core.tools import BaseTool, tool from pydantic import BaseModel, model_validator SEARCH_URL = "https://backend.basedosdados.org/search/" @@ -437,7 +437,7 @@ def decode_table_values(table_gcp_id: str, column_name: Optional[str] = None) -> rows = client.query(search_query).result() results = [dict(row) for row in rows] tool_output = ToolOutput(status="success", results=results).model_dump(exclude_none=True) - except exceptions.NotFound: + except google_api_exceptions.NotFound: return ToolOutput( status="error", error_details={ @@ -509,3 +509,26 @@ def inspect_column_values(table_gcp_id: str, column_name: str, limit: int = 100) ).model_dump(exclude_none=True) return json.dumps(tool_output, ensure_ascii=False, default=str) + + +def get_tools() -> list[BaseTool]: + """Return all available tools for Base dos Dados database interaction. + + This function provides a complete set of tools for discovering, exploring, + and querying Brazilian public datasets through the Base dos Dados platform. + + Returns: + list[BaseTool]: A list of LangChain tool functions in suggested usage order: + - search_datasets: Find datasets using keywords + - get_dataset_details: Get comprehensive dataset information + - execute_bigquery_sql: Execute SQL queries against BigQuery tables + - decode_table_values: Decode coded values using dictionary tables + - inspect_column_values: Inspect actual column values as fallback + """ + return [ + search_datasets, + get_dataset_details, + execute_bigquery_sql, + decode_table_values, + inspect_column_values, + ] From 5c040fe0d6f6731b43c1352b8bd5166cbb53a0b3 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Fri, 29 Aug 2025 15:54:22 -0300 Subject: [PATCH 14/23] add chat history management and fix checkpoints deletion --- backend/apps/chatbot/views.py | 128 +++++++++++++++++++++++++--------- 1 file changed, 96 insertions(+), 32 deletions(-) diff --git a/backend/apps/chatbot/views.py b/backend/apps/chatbot/views.py index 7c629f0c..3c101c13 100644 --- a/backend/apps/chatbot/views.py +++ b/backend/apps/chatbot/views.py @@ -1,14 +1,20 @@ # -*- coding: utf-8 -*- import os import uuid +from collections.abc import Generator from contextlib import contextmanager from functools import cache from typing import Any, Iterator, Type, TypedDict, TypeVar from django.http import StreamingHttpResponse +from google.api_core import exceptions as google_api_exceptions from graphql_jwt.shortcuts import get_user_by_token from langchain.chat_models import init_chat_model +from langchain_core.messages import RemoveMessage, ToolMessage +from langchain_core.messages.utils import count_tokens_approximately, trim_messages from langgraph.checkpoint.postgres import PostgresSaver +from langgraph.graph.graph import CompiledGraph +from langgraph.graph.message import REMOVE_ALL_MESSAGES from langgraph.prebuilt import create_react_agent from loguru import logger from rest_framework import exceptions, status @@ -21,13 +27,7 @@ from rest_framework_simplejwt.tokens import RefreshToken from backend.apps.chatbot.agent.prompts import SQL_AGENT_SYSTEM_PROMPT -from backend.apps.chatbot.agent.tools import ( - decode_table_values, - execute_bigquery_sql, - get_dataset_details, - inspect_column_values, - search_datasets, -) +from backend.apps.chatbot.agent.tools import get_tools from backend.apps.chatbot.feedback_sender import LangSmithFeedbackSender from backend.apps.chatbot.models import Feedback, MessagePair, Thread from backend.apps.chatbot.serializers import ( @@ -39,6 +39,7 @@ UserMessageSerializer, ) from backend.apps.chatbot.utils.stream import EventData, StreamEvent, process_chunk +from chatbot.agents.utils import delete_checkpoints ModelSerializer = TypeVar("ModelSerializer", bound=Serializer) @@ -47,6 +48,18 @@ # https://python.langchain.com/api_reference/langchain/chat_models/langchain.chat_models.base.init_chat_model.html MODEL_URI = os.environ["MODEL_URI"] +# Gemini models have a ~1 million tokens context window +CONTEXT_WINDOW = 2**20 + +# Maximum number of tokens allowed at the START of a conversation turn +MAX_TOKENS = CONTEXT_WINDOW // 2 + +# Generic error message for unexpected errors when calling the agent +UNEXPECTED_ERROR_MESSAGE = ( + "Ops, algo deu errado! Ocorreu um erro inesperado. Por favor, tente novamente. " + "Se o problema persistir, avise-nos. Obrigado pela paciência!" +) + class ConfigDict(TypedDict): run_id: str @@ -176,10 +189,16 @@ def delete(self, request: Request, thread_id: uuid.UUID) -> Response: try: thread.deleted = True thread.save() - with _get_sql_agent() as assistant: - assistant.clear_thread(str(thread_id)) + with _get_sql_agent() as agent: + if agent.checkpointer is None: + logger.info("Checkpointer is None, ignoring...") + else: + logger.info(f"Deleting checkpoints for thread {thread_id}...") + delete_checkpoints(agent.checkpointer, str(thread_id)) + logger.success(f"Checkpoints for thread {thread_id} deleted successfully") return Response({"detail": "Thread deleted successfully"}) except Exception: + logger.exception(f"Error deleting thread {thread_id}:") return Response( {"detail": "Error deleting thread"}, status=status.HTTP_500_INTERNAL_SERVER_ERROR ) @@ -306,11 +325,11 @@ def _get_feedback_sender() -> LangSmithFeedbackSender: @contextmanager -def _get_sql_agent(): - """Provide a configured `SQLAssistant`. +def _get_sql_agent() -> Generator[CompiledGraph]: + """Provide a configured ReAct agent. Yields: - Iterator[SQLAssistant]: An instance of `SQLAssistant`. + Iterator[CompiledGraph]: An instance of `CompiledGraph`. """ db_host = os.environ["DB_HOST"] db_port = os.environ["DB_PORT"] @@ -320,23 +339,43 @@ def _get_sql_agent(): conn = f"postgresql://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}" - with PostgresSaver.from_conn_string(conn) as checkpointer: - checkpointer.setup() + model = init_chat_model(MODEL_URI, temperature=0) + + def pre_model_hook(state: dict): + messages = state["messages"] + + # The last message in the pre_model_hook node will + # ALWAYS be a HumanMessage or a ToolMessage. + last_message = state["messages"][-1] + + # If this is the first message in the chat, we don't trim. + # If the last message is a ToolMessage, the agent has called a tool. + # This means we are in the middle of a chat turn and we also dont't trim. + if len(messages) == 1 or isinstance(last_message, ToolMessage): + return {"messages": []} + + # Otherwise, we are just starting a chat turn and we can trim the chat history. + remaining_messages = trim_messages( + messages, + token_counter=count_tokens_approximately, # The accurate counter is too slow. + max_tokens=MAX_TOKENS, + strategy="last", + start_on="human", + end_on="human", + include_system=True, + allow_partial=False, + ) - model = init_chat_model(MODEL_URI, temperature=0) + return {"messages": [RemoveMessage(id=REMOVE_ALL_MESSAGES), *remaining_messages]} - tools = [ - decode_table_values, - execute_bigquery_sql, - get_dataset_details, - inspect_column_values, - search_datasets, - ] + with PostgresSaver.from_conn_string(conn) as checkpointer: + checkpointer.setup() sql_agent = create_react_agent( model=model, - tools=tools, + tools=get_tools(), prompt=SQL_AGENT_SYSTEM_PROMPT, + pre_model_hook=pre_model_hook, checkpointer=checkpointer, ) @@ -357,32 +396,57 @@ def _stream_sql_assistant_response( Iterator[str]: JSON string containing the streaming status and the current step data. """ events = [] + agent_state = None try: logger.info("Calling SQL Agent...") with _get_sql_agent() as agent: - for chunk in agent.stream( + for mode, chunk in agent.stream( input={"messages": [{"role": "user", "content": message}]}, - stream_mode="updates", + stream_mode=["updates", "values"], config=config, ): + if mode == "values": + agent_state = chunk + continue + event = process_chunk(chunk) - events.append(event.model_dump()) - yield event.to_sse() + if event is not None: + events.append(event.model_dump()) + yield event.to_sse() # The last event always contains the agent's final answer, # so we use it to save the message pair in the database assistant_message = event.data.content error_message = None logger.success("SQL Agent called successfully. Saving message pair...") + + except google_api_exceptions.InvalidArgument: + logger.exception("Agent execution failed with Google API InvalidArgument error:") + + assistant_message = None + error_message = UNEXPECTED_ERROR_MESSAGE + + if agent_state is not None: + model = init_chat_model(MODEL_URI, temperature=0) + total_tokens = model.get_num_tokens_from_messages(agent_state["messages"]) + + if total_tokens >= CONTEXT_WINDOW: + error_message = ( + "Sua última mensagem ultrapassou o limite de tamanho para esta conversa. " + "Por favor, tente dividir sua solicitação em partes menores " + "ou inicie uma nova conversa." + ) + + yield StreamEvent( + type="error", data=EventData(error_details={"message": error_message}) + ).to_sse() + except Exception: - logger.exception(f"Error responding message {config['run_id']}:") + logger.exception(f"Unexpected error responding message {config['run_id']}:") assistant_message = None - error_message = ( - "Ops, algo deu errado! Ocorreu um erro inesperado. Por favor, tente novamente. " - "Se o problema persistir, avise-nos. Obrigado pela paciência!" - ) + error_message = UNEXPECTED_ERROR_MESSAGE yield StreamEvent( type="error", data=EventData(error_details={"message": error_message}) From da05deb1899d60168e96713956e64c51c997ce99 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Fri, 29 Aug 2025 16:08:07 -0300 Subject: [PATCH 15/23] fix test cases --- backend/apps/chatbot/tests/test_endpoints.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/backend/apps/chatbot/tests/test_endpoints.py b/backend/apps/chatbot/tests/test_endpoints.py index 62215606..7b6a23ac 100644 --- a/backend/apps/chatbot/tests/test_endpoints.py +++ b/backend/apps/chatbot/tests/test_endpoints.py @@ -4,20 +4,20 @@ import pytest from django.utils.dateparse import parse_datetime +from langchain_core.messages import AIMessage from rest_framework.test import APIClient from backend.apps.account.models import Account from backend.apps.chatbot import views from backend.apps.chatbot.models import Feedback, MessagePair, Thread -from chatbot.assistants import SQLAssistantMessage -class MockSQLAssistant: +class MockReactAgent: def __init__(self, *args, **kwargs): ... - def invoke(self, *args, **kwargs): - return SQLAssistantMessage(content="mock response") + def stream(self, *args, **kwargs): + yield "updates", {"agent": AIMessage("mock response")} def clear_thread(self, *args, **kwargs): ... @@ -31,6 +31,10 @@ def send_feedback(self, *args, **kwargs): ... +def mock_create_react_agent(): + return MockReactAgent() + + @pytest.fixture def mock_email() -> str: return "mockemail@mockdomain.com" @@ -272,7 +276,7 @@ def test_message_list_view_get_order_invalid(auth_client: APIClient, auth_user: @pytest.mark.django_db def test_message_list_view_post(monkeypatch, auth_client: APIClient, auth_user: Account): - monkeypatch.setattr(views, "SQLAssistant", MockSQLAssistant) + monkeypatch.setattr(views, "create_react_agent", mock_create_react_agent) thread = Thread.objects.create(account=auth_user) From 68b62c1e42714d1fc082d2f3b0116c413194b653 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Fri, 29 Aug 2025 16:28:07 -0300 Subject: [PATCH 16/23] style: minor formatting changes --- backend/apps/chatbot/tests/test_endpoints.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/backend/apps/chatbot/tests/test_endpoints.py b/backend/apps/chatbot/tests/test_endpoints.py index 7b6a23ac..03a35fab 100644 --- a/backend/apps/chatbot/tests/test_endpoints.py +++ b/backend/apps/chatbot/tests/test_endpoints.py @@ -12,22 +12,22 @@ from backend.apps.chatbot.models import Feedback, MessagePair, Thread -class MockReactAgent: +class MockLangSmithFeedbackSender: def __init__(self, *args, **kwargs): ... - def stream(self, *args, **kwargs): - yield "updates", {"agent": AIMessage("mock response")} - - def clear_thread(self, *args, **kwargs): + def send_feedback(self, *args, **kwargs): ... -class MockLangSmithFeedbackSender: +class MockReactAgent: def __init__(self, *args, **kwargs): ... - def send_feedback(self, *args, **kwargs): + def stream(self, *args, **kwargs): + yield "updates", {"agent": AIMessage("mock response")} + + def clear_thread(self, *args, **kwargs): ... From 0ad9af1de168f0957f8b3fbafc3dcb2300832c19 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Fri, 29 Aug 2025 16:30:25 -0300 Subject: [PATCH 17/23] style: rename mock class --- backend/apps/chatbot/tests/test_endpoints.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/backend/apps/chatbot/tests/test_endpoints.py b/backend/apps/chatbot/tests/test_endpoints.py index 03a35fab..0d399032 100644 --- a/backend/apps/chatbot/tests/test_endpoints.py +++ b/backend/apps/chatbot/tests/test_endpoints.py @@ -20,7 +20,7 @@ def send_feedback(self, *args, **kwargs): ... -class MockReactAgent: +class MockReActAgent: def __init__(self, *args, **kwargs): ... @@ -32,7 +32,7 @@ def clear_thread(self, *args, **kwargs): def mock_create_react_agent(): - return MockReactAgent() + return MockReActAgent() @pytest.fixture From 6a81316082540739555e1372558cbef3cb17e660 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Fri, 29 Aug 2025 17:44:09 -0300 Subject: [PATCH 18/23] create constants for query limits --- backend/apps/chatbot/agent/tools.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/backend/apps/chatbot/agent/tools.py b/backend/apps/chatbot/agent/tools.py index 972ede99..cf069f9b 100644 --- a/backend/apps/chatbot/agent/tools.py +++ b/backend/apps/chatbot/agent/tools.py @@ -10,6 +10,15 @@ from langchain_core.tools import BaseTool, tool from pydantic import BaseModel, model_validator +# 1GB limit for dictionary queries +LIMIT_DICTIONARY_QUERY = 1e9 + +# 5GB limit for inspection queries +LIMIT_INSPECTION_QUERY = 5e9 + +# 20GB limit for other queries +LIMIT_BIGQUERY_QUERY = 20e9 + SEARCH_URL = "https://backend.basedosdados.org/search/" GRAPHQL_URL = "https://backend.basedosdados.org/graphql" @@ -334,7 +343,7 @@ def execute_bigquery_sql(sql_query: str) -> str: job_config = bq.QueryJobConfig(dry_run=True, use_query_cache=False) query_job = client.query(sql_query, job_config=job_config) - limit_bytes = 20e9 # 20GB limit for queries + limit_bytes = LIMIT_BIGQUERY_QUERY total_bytes = query_job.total_bytes_processed if total_bytes and total_bytes > limit_bytes: @@ -345,7 +354,7 @@ def execute_bigquery_sql(sql_query: str) -> str: "limit_bytes": limit_bytes, "total_processed_bytes": total_bytes, "message": ( - "Query aborted: Data processed exceeds the 20GB per-query limit. " + "Query aborted: Data processed exceeds the per-query limit. " "Consider optimizing by adding filters, selecting fewer columns, " "or using a LIMIT clause before retrying." ), @@ -415,7 +424,7 @@ def decode_table_values(table_gcp_id: str, column_name: Optional[str] = None) -> job_config = bq.QueryJobConfig(dry_run=True, use_query_cache=False) query_job = client.query(search_query, job_config=job_config) - limit_bytes = 1e9 # 1GB limit for dictionary queries + limit_bytes = LIMIT_DICTIONARY_QUERY total_bytes = query_job.total_bytes_processed if total_bytes and total_bytes > limit_bytes: @@ -480,7 +489,7 @@ def inspect_column_values(table_gcp_id: str, column_name: str, limit: int = 100) job_config = bq.QueryJobConfig(dry_run=True, use_query_cache=False) query_job = client.query(sql_query, job_config=job_config) - limit_bytes = 5e9 # 5GB limit for inspection queries + limit_bytes = LIMIT_INSPECTION_QUERY total_bytes = query_job.total_bytes_processed if total_bytes and total_bytes > limit_bytes: @@ -492,7 +501,7 @@ def inspect_column_values(table_gcp_id: str, column_name: str, limit: int = 100) "limit_bytes": limit_bytes, "total_processed_bytes": total_bytes, "message": ( - "Column inspection exceeds the 5GB per-query limit for inspection. " + "Column inspection exceeds the per-query limit for inspection. " "Try a smaller limit or add WHERE filters to reduce data size." ), }, From 62cc7fd09084b8fe326a8ed40c263730a2f0d548 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Mon, 1 Sep 2025 09:54:28 -0300 Subject: [PATCH 19/23] make `MODEL_URI` a hardcoded constant --- backend/apps/chatbot/views.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backend/apps/chatbot/views.py b/backend/apps/chatbot/views.py index 3c101c13..37364201 100644 --- a/backend/apps/chatbot/views.py +++ b/backend/apps/chatbot/views.py @@ -46,7 +46,7 @@ # Model name/URI. Refer to the LangChain docs for valid names/URIs # https://python.langchain.com/api_reference/langchain/chat_models/langchain.chat_models.base.init_chat_model.html -MODEL_URI = os.environ["MODEL_URI"] +MODEL_URI = "google_vertexai:gemini-2.5-flash" # Gemini models have a ~1 million tokens context window CONTEXT_WINDOW = 2**20 From 5223fd78623a1856cf004995ee4c8754047de9a1 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Mon, 1 Sep 2025 10:56:28 -0300 Subject: [PATCH 20/23] add sql query safety check --- backend/apps/chatbot/agent/tools.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/backend/apps/chatbot/agent/tools.py b/backend/apps/chatbot/agent/tools.py index cf069f9b..d73addd0 100644 --- a/backend/apps/chatbot/agent/tools.py +++ b/backend/apps/chatbot/agent/tools.py @@ -337,6 +337,31 @@ def execute_bigquery_sql(sql_query: str) -> str: Returns: str: Query results as JSON array. Empty results return "[]". """ # noqa: E501 + forbidden_commands = [ + "CREATE", + "ALTER", + "DROP", + "TRUNCATE", + "INSERT", + "UPDATE", + "DELETE", + "GRANT", + "REVOKE", + ] + + for command in forbidden_commands: + if command in sql_query.upper(): + tool_output = ToolOutput( + status="error", + error_details={ + "message": ( + f"Query aborted: Command {command} is forbidden. ", + "Your access is strictly read-only.", + ) + }, + ) + return tool_output.model_dump_json(indent=2, exclude_none=True) + client = get_bigquery_client() try: From 746d3eb70366c4d596db303e3455d8f79d6f54ab Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Mon, 1 Sep 2025 10:58:53 -0300 Subject: [PATCH 21/23] update `chatbot` package to `v0.6.2` --- chatbot | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/chatbot b/chatbot index 6a89f004..abd72e3b 160000 --- a/chatbot +++ b/chatbot @@ -1 +1 @@ -Subproject commit 6a89f004e51271c1669df755733ab21e58a91c84 +Subproject commit abd72e3b016ec7e4b6ce0953dde94b93ab7dfe47 From 42f48c67a6ebf12a165adea85b50df4cac5d76e8 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Mon, 1 Sep 2025 10:59:31 -0300 Subject: [PATCH 22/23] update dependencies --- poetry.lock | 4 ++-- pyproject.toml | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/poetry.lock b/poetry.lock index 7aa5fbb9..99544f86 100644 --- a/poetry.lock +++ b/poetry.lock @@ -426,7 +426,7 @@ files = [ [[package]] name = "chatbot" -version = "0.6.1" +version = "0.6.2" description = "" optional = false python-versions = ">=3.10,<4.0" @@ -4133,4 +4133,4 @@ cffi = ["cffi (>=1.11)"] [metadata] lock-version = "2.1" python-versions = ">=3.10,<3.13" -content-hash = "02d9ae83cf846222afc5a829393870f86e181af444079a785eb969082611c211" +content-hash = "c1f3ab7225c125661378a75d692f498ddb3cf4bd3680293eeeb4cd09c32f70a9" diff --git a/pyproject.toml b/pyproject.toml index e3676a64..333e9aee 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,6 +46,7 @@ djangorestframework-simplejwt = "^5.5.0" chatbot = {path = "chatbot"} langchain-google-vertexai = "2.0.24" langchain-postgres = "^0.0.14" +httpx = "^0.28.1" [tool.poetry.group.dev.dependencies] pre-commit = "^3.3.3" From 759fc3fb22815058dc5dd007389d64fdd5f9fbab Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Mon, 1 Sep 2025 11:12:43 -0300 Subject: [PATCH 23/23] fix sql query safety check --- backend/apps/chatbot/agent/tools.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/backend/apps/chatbot/agent/tools.py b/backend/apps/chatbot/agent/tools.py index d73addd0..1d421895 100644 --- a/backend/apps/chatbot/agent/tools.py +++ b/backend/apps/chatbot/agent/tools.py @@ -351,7 +351,7 @@ def execute_bigquery_sql(sql_query: str) -> str: for command in forbidden_commands: if command in sql_query.upper(): - tool_output = ToolOutput( + return ToolOutput( status="error", error_details={ "message": ( @@ -359,8 +359,7 @@ def execute_bigquery_sql(sql_query: str) -> str: "Your access is strictly read-only.", ) }, - ) - return tool_output.model_dump_json(indent=2, exclude_none=True) + ).model_dump_json(indent=2, exclude_none=True) client = get_bigquery_client() @@ -372,7 +371,7 @@ def execute_bigquery_sql(sql_query: str) -> str: total_bytes = query_job.total_bytes_processed if total_bytes and total_bytes > limit_bytes: - tool_output = ToolOutput( + return ToolOutput( status="error", error_details={ "type": "QueryTooLarge", @@ -384,8 +383,7 @@ def execute_bigquery_sql(sql_query: str) -> str: "or using a LIMIT clause before retrying." ), }, - ) - return tool_output.model_dump_json(indent=2, exclude_none=True) + ).model_dump_json(indent=2, exclude_none=True) rows = client.query(sql_query).result()