polojuan/agentic-workflows

1import streamlit as st2import sqlite33import pandas as pd4from utils.database import create_tables, create_metadata, get_metaschema5from openai import OpenAI6from dotenv import find_dotenv, load_dotenv7from typing import Tuple8 9# Load environment variables10load_dotenv(find_dotenv())11 12PATH = 'data/rca_data.db'13 14def evaluate_and_refine_sql(15    question: str,16    sql_query: str,17    df: pd.DataFrame,18    schema: str,19    model: str = "gpt-5",20    sql_error: Exception = None21) -> Tuple[str, str]:22    """23    Evaluates SQL query results and refines the query if needed to better answer the user's question.24 25    This function uses an LLM to review whether the SQL query output adequately addresses26    the user's original question. If improvements are needed, it generates a refined SQL query.27 28    Args:29        question (str): The original natural language question from the user.30        sql_query (str): The SQL query that was executed.31        df (pd.DataFrame): The resulting DataFrame from executing the SQL query.32        schema (str): The database schema information for reference.33        model (str, optional): The language model to use for evaluation. Defaults to "gpt-5".34 35    Returns:36        Tuple[str, str]: A tuple containing:37            - feedback (str): Brief evaluation and suggestions for improvement.38            - refined_sql (str): The improved SQL query, or the original if no changes needed.39 40    Note:41        If the LLM response is not valid JSON, the function falls back to returning42        the original SQL query with the raw response as feedback.43    """44    # Get client from session state45    client = st.session_state.get("client") or OpenAI()46 47    prompt = f"""48You are an expert SQL reviewer specializing in query optimization and accuracy validation.49 50## CONTEXT51**User's Original Question:**52{question}53 54**Generated SQL Query:**55```sql56{sql_query}57```58 59**Query Results:**60{df.to_string(index=False)}61 62**Table Schema:**63{schema}64 65**SQL Error:**66{sql_error}67 68## YOUR TASK69Analyze whether the SQL query correctly and completely answers the user's question. Include other columns which may be relevant to the question. When aggregations are used, make sure there are no duplications in the RCA ID.70 71Step 1: Briefly evaluate if the SQL output answers the user's question. 72Step 2: If the SQL could be improved, provide a refined SQL query. If SQL Error is not None, rectify the issues in the refined query.73If the original SQL is already correct, return it unchanged.74 75## OUTPUT FORMAT76Return ONLY a valid JSON object with this exact structure:77{{78    "feedback": "Brief evaluation of the query (what's right or what needs improvement)",79    "refined_sql": "The final SQL query to execute (original or improved version)"80}}81 82Do not include any text outside the JSON object.83"""84 85    response = client.chat.completions.create(86        model=model,87        messages=[{"role": "user", "content": prompt}],88    )89 90    import json91    content = response.choices[0].message.content92 93    # Strip markdown code blocks and fix Python-style booleans94    content_clean = content.strip().removeprefix("```json").removeprefix("```").removesuffix("```").strip()95    content_clean = content_clean.replace(": True", ": true").replace(": False", ": false")96 97    try:98        obj = json.loads(content_clean)99        feedback = str(obj.get("feedback", "")).strip()100        refined_sql = str(obj.get("refined_sql", sql_query)).strip()101        if not refined_sql:102            refined_sql = sql_query103    except Exception as e:104        # Fallback if the model does not return valid JSON:105        # use the raw content as feedback and keep the original SQL106        print(f"❌ JSON parsing error in evaluate_and_refine_sql: {e}")107        print(f"Raw response: {content}")108        feedback = content.strip()109        refined_sql = sql_query110 111    return feedback, refined_sql112 113def database_interpreter(query: str, sql_gen_ref: pd.DataFrame, metadata: str, model: str = "gpt-5") -> Tuple[str, bool]:114    """115    Converts SQL query results into a natural language answer for the user's question.116 117    This function takes the raw database query results and translates them into a clear,118    human-readable response that directly answers the user's original question.119 120    Args:121        query (str): The original natural language question from the user.122        sql_gen_ref (pd.DataFrame): The DataFrame containing the SQL query results.123        metadata (str): The database schema metadata for context.124        model (str, optional): The language model to use for interpretation. Defaults to "gpt-5".125 126    Returns:127        Tuple[str, bool]: A tuple containing:128        output (str): A natural language answer that explains the query results in an easy-to-understand format.129        success (bool): True if the final SQL query was adequate to answer the query, False if refined.130 131    Note:132        The response focuses on clarity and readability, avoiding SQL syntax and technical jargon133        where possible. If results are empty, it explains that no data was found.134    """135    # Get client from session state136    client = st.session_state.get("client") or OpenAI()137    prompt = f"""138You are an expert data analyst translating database query results into clear, actionable insights.139 140## CONTEXT141**User's Question:**142{query}143 144**Database Schema:**145{metadata}146 147**Query Results:**148{sql_gen_ref}149 150## YOUR TASK151Provide a natural language answer that directly addresses the user's question based on the SQL query results. Provide complete but concise answer. Provide specific details and figures from the results where relevant. 152 153If the natural language answer fully provides the answers to question, indicate success as True. If the results are empty or insufficient to answer the question, indicate success as False.154 155## OUTPUT FORMAT156Return ONLY a valid JSON object with this exact structure:157{{158    "output": "Brief answer to the query",159    "success": True/False160}}161Do not include any text outside the JSON object.162"""163    response = client.chat.completions.create(164        model=model,165        messages=[{"role": "user", "content": prompt}],166    )167 168    import json169    content = response.choices[0].message.content170    # Strip markdown code blocks and fix Python-style booleans171    content_clean = content.strip().removeprefix("```json").removeprefix("```").removesuffix("```").strip()172    content_clean = content_clean.replace(": True", ": true").replace(": False", ": false")173 174    try:175        obj = json.loads(content_clean)176        output = str(obj["output"]).strip()177        success = bool(obj["success"])178        print("output ", output)179        print("success ", success)180    except Exception as e:181        # Fallback if the model does not return valid JSON182        print(f"❌ JSON parsing error: {e}")183        print(f"Raw response: {content}")184        output = content.strip()185        success = False186    return output, success187 188def database_agent(query: str, model: str = "gpt-5", return_details: bool = False, max_refine_attempts: int = 5) -> str | dict:189    """190    Processes natural language database queries using a two-stage SQL generation and refinement workflow.191 192    This function converts a natural language question into SQL, executes it against a SQLite database,193    refines the query if needed, and returns a natural language interpretation of the results.194 195    Workflow:196        1. Initializes database tables and retrieves schema metadata197        2. Generates an initial SQL query from the natural language question198        3. Executes and evaluates the query results199        4. Refines the SQL query based on feedback if necessary200        5. Returns a natural language interpretation of the final results201 202    Args:203        query (str): The user's natural language question about the database.204        model (str, optional): The language model to use for SQL generation and interpretation.205        return_details (bool, optional): If True, returns a dict with all intermediate steps. Defaults to False.206 207    Returns:208        str or dict: If return_details is False, returns natural language answer.209                     If return_details is True, returns dict with 'answer', 'sql_v1', 'sql_v2',210                     'feedback', 'results_v1', 'results_v2'.211    """212    # Get client from session state213    client = st.session_state.get("client") or OpenAI()214 215    #Initialize DB and get metadata216    conn = sqlite3.connect(PATH)217    create_tables()218    create_metadata()219    meta_schema = get_metaschema()220 221    prompt = f"""222You are an expert SQLite query generator. Your task is to convert natural language questions into accurate, efficient SQL queries. Only answer relevant questions based on the provided database schema. If the question is unrelated to the database, respond with "SELECT NULL;".223 224## TABLE SCHEMA225{meta_schema}226 227## USER QUESTION228{query}229 230## OUTPUT FORMAT231Return ONLY the SQL query without any explanation, markdown formatting, or code blocks.232Do NOT include ```sql``` tags or any other text - just the raw SQL query.233 234Now generate the SQL query for the user's question above:235"""236    response = client.chat.completions.create(237        model=model,238        messages=[{"role": "user", "content": prompt}],239    )240 241    sql_gen_1 = response.choices[0].message.content.strip()242 243    # Execute the first SQL query to get initial results244    q1 = sql_gen_1.strip().removeprefix("```sql").removesuffix("```").strip()245 246    # Check if query is irrelevant (returns SELECT NULL or similar)247    if q1.upper() in ["SELECT NULL;", "SELECT NULL", "NULL"]:248        conn.close()249        irrelevant_msg = "The query is not related to the investigation database."250        if return_details:251            return {252                'answer': irrelevant_msg,253                'sql_v1': q1,254                'sql_v2': q1,255                'feedback': 'Question is not related to the database schema',256                'results_v1': pd.DataFrame(),257                'results_v2': pd.DataFrame()258            }259        return irrelevant_msg260 261    try:262        sql_gen_orig = pd.read_sql_query(q1, conn)263    except Exception as e:264        # If first query fails, create empty dataframe with error265        sql_gen_orig = pd.DataFrame({"error": [str(e)]})266        print(f"❌ Error executing initial query: {e}")267    268    sql_error = None269    for i in range(max_refine_attempts):270        # Evaluate and refine the SQL based on initial results271        if i == 0:272            refined_sql=sql_gen_1,273            sql_gen_ref=sql_gen_orig274        feedback, refined_sql = evaluate_and_refine_sql(275                question=query,276                sql_query=refined_sql,277                df=sql_gen_ref,278                schema=meta_schema,279                model=model,280                sql_error = sql_error281            )282        # Execute the refined SQL query283        q2 = refined_sql.strip().removeprefix("```sql").removesuffix("```").strip()284 285        try:286            sql_gen_ref = pd.read_sql_query(q2, conn)287        except Exception as e:288            print(f"❌ Error executing refined query: {e}")289            sql_error = e290        output, success = database_interpreter(query, sql_gen_ref, metadata=meta_schema, model=model)291 292        print("Refinement Attempt", i+1)293        print("Success or not: ", success)294        print("📝 Reflect on V1 query:\n" + feedback)295        print("🔁 Write V2 query:\n" + refined_sql)296        print("🔁 Answer:\n" + output)297        print("\n")298 299        if success:300            break301    conn.close()302    303 304    if return_details:305        return {306            'answer': output,307            'sql_v1': q1,308            'sql_v2': q2,309            'feedback': feedback,310            'results_v1': sql_gen_orig,311            'results_v2': sql_gen_ref312        }313    return output314