Gonalb/multi_agentic_sql_generator

1import openai2from config import PROJECT_ID, DATASET_ID3from utils.bigquery_utils import get_bigquery_schema_info4 5def sql_generation_agent(state):6    """Generates a SQL query based on the natural language query and sample data."""7    natural_language_query = state["sql_query"]8    relevant_tables = state.get("relevant_tables", [])9    sample_data = state.get("sample_data", {})10    client = state["client"]11    12    if client is None:13        return {"generated_sql": "-- Error: Failed to connect to BigQuery."}14    15    schema_info = get_bigquery_schema_info(client, PROJECT_ID, DATASET_ID)16    17    # Format the schema for the prompt18    schema_text = ""19    for table_name, columns in schema_info.items():20        if f"{DATASET_ID}.{table_name}" in relevant_tables:21            schema_text += f"- **{DATASET_ID}.{table_name}** ({', '.join(columns)})\n"22    23    # Format sample data for the prompt24    sample_data_text = ""25    for table, rows in sample_data.items():26        if isinstance(rows, list) and rows:27            sample_data_text += f"\n**Sample data from {table}:**\n"28            # Get column names from the first row29            columns = list(rows[0].keys())30            sample_data_text += "| " + " | ".join(columns) + " |\n"31            sample_data_text += "| " + " | ".join(["---"] * len(columns)) + " |\n"32            33            # Add row data34            for row in rows:35                sample_data_text += "| " + " | ".join([str(row.get(col, "")) for col in columns]) + " |\n"36    37    prompt = f"""38    Generate a BigQuery SQL query to answer the following question:39    40    **Question:** "{natural_language_query}"41    42    **Relevant Tables Schema:**43    {schema_text}44    45    **Sample Data:**46    {sample_data_text}47    48    **Rules:**49    - Use only the provided tables with their full dataset.table_name format (e.g., {DATASET_ID}.users).50    - Ensure correct column names as shown in the schema.51    - Use appropriate joins based on the relationships visible in the sample data.52    - Use BigQuery SQL syntax.53    - Return ONLY the SQL query without any explanations or markdown formatting.54    """55    56    response = openai.chat.completions.create(57        model="gpt-4o-mini",58        messages=[{"role": "user", "content": prompt}],59        temperature=0.060    )61    62    generated_sql = response.choices[0].message.content.strip()63    64    # Remove markdown code block formatting if present65    if generated_sql.startswith("```sql"):66        generated_sql = generated_sql.replace("```sql", "").replace("```", "").strip()67    68    return {"generated_sql": generated_sql}