Tools and Toolsets

# User: "What tables are in the dataset?"
# Agent uses: Schema Inspector

# User: "Show me all active assets"
# Agent uses: Query Executor

# User: "Is this SQL valid: SELECT * FROM assets"
# Agent uses: Query Validator

SELECT * 
FROM `datawarehouse-des.STG_ACTIVOS.assets`

SELECT 
  asset_type,
  AVG(maintenance_cost) as avg_cost
FROM `datawarehouse-des.STG_ACTIVOS.assets`
GROUP BY asset_type
ORDER BY avg_cost DESC

Lo siento, por seguridad corporativa tengo acceso de solo lectura a los datos de TRANSELEC S.A.

import logging

# Enable ADK debug logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger('google.adk')
logger.setLevel(logging.DEBUG)

# Tool executions will be logged:
# DEBUG:google.adk.tools.bigquery:Executing query: SELECT ...
# DEBUG:google.adk.tools.bigquery:Query completed in 1.23s

-- Check recent queries executed by the agent
SELECT 
  creation_time,
  user_email,
  query,
  total_bytes_processed,
  total_slot_ms
FROM `region-us-east4`.INFORMATION_SCHEMA.JOBS_BY_PROJECT
WHERE user_email = 'bigquery-agent@PROJECT_ID.iam.gserviceaccount.com'
ORDER BY creation_time DESC
LIMIT 100

from google.adk.tools import Tool
from typing import Dict, Any
import sqlglot  # Already in requirements.txt

class SQLValidatorTool(Tool):
    """Custom tool for advanced SQL validation"""
    
    def __init__(self):
        super().__init__(
            name="sql_validator",
            description="Validates SQL syntax and checks for best practices"
        )
    
    def execute(self, sql: str) -> Dict[str, Any]:
        """Validate SQL and provide recommendations"""
        try:
            # Parse SQL
            parsed = sqlglot.parse_one(sql, dialect="bigquery")
            
            # Check for SELECT *
            has_select_star = "SELECT *" in sql.upper()
            
            # Check for LIMIT clause
            has_limit = "LIMIT" in sql.upper()
            
            return {
                "valid": True,
                "warnings": [
                    "Avoid SELECT * in production" if has_select_star else None,
                    "Consider adding LIMIT for large tables" if not has_limit else None
                ],
                "parsed_structure": str(parsed)
            }
        except Exception as e:
            return {
                "valid": False,
                "error": str(e)
            }

# Add to agent
sql_validator = SQLValidatorTool()

root_agent = LlmAgent(
    model=LLM_1_MODELO,
    name=LLM_1_NAME,
    description="Agente para responder preguntas sobre datos y modelos de BigQuery",
    instruction=new_instruction,
    tools=[bigquery_toolset, sql_validator]  # ← Add custom tool
)

# Ensure instruction mandates tool usage
"""
2. EJECUCIÓN OBLIGATORIA: Para cualquier pregunta de datos, 
   DEBES usar la herramienta `bigquery_toolset` para probar que tu query funciona.
"""

# Verify tool is registered
print(root_agent.tools)  # Should include bigquery_toolset

-- Check for hidden write operations
-- ✗ This will be blocked:
WITH temp AS (
  INSERT INTO table VALUES (1, 2, 3)
  RETURNING *
)
SELECT * FROM temp

-- ✓ This is allowed:
WITH temp AS (
  SELECT * FROM existing_table
)
SELECT * FROM temp

# Local development
gcloud auth application-default login

# Verify credentials
gcloud auth application-default print-access-token

# Cloud Run deployment
# Ensure service account is attached to the service
gcloud run services describe SERVICE_NAME --region=REGION

from google.cloud import bigquery

tool_config = BigQueryToolConfig(
    write_mode=WriteMode.BLOCKED,
    # Custom timeout (requires additional configuration)
)

# Note: Default timeout is typically sufficient (10 minutes)
# For custom timeouts, configure at the BigQuery client level

# Update instruction prompt to guide agents on result limits
new_instruction = f"""
# ... existing prompt ...

<QUERY_GUIDELINES>
  - Para tablas grandes, usa LIMIT para restringir resultados
  - Límite recomendado: LIMIT 1000 para exploraciones iniciales
  - Para análisis completos, el usuario puede remover el LIMIT
</QUERY_GUIDELINES>
"""