text_2_sql_agentic.py

from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_core.output_parsers.string import StrOutputParser
import sqlite3
from typing import List, Any, Dict
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, END
from urllib.parse import quote_plus
import pyodbc
import os
from langchain_community.callbacks import get_openai_callback
from langchain_core.messages import ToolMessage
from langchain_core.runnables import RunnableLambda, RunnableWithFallbacks
from langgraph.prebuilt import ToolNode
import pandas as pd
import numpy as np
import os
from openai import OpenAI
from langchain_core.tools import tool

from dotenv import load_dotenv
import os

load_dotenv()

## Using the OpenAI client directly.
client = OpenAI(
    api_key=os.environ.get("OPENAI_API_KEY"),
)

class InputState(TypedDict):
    question: str
    database: str
    uri: str 
    schema: Dict[str, Any]
    sql_query: str
    result: List[Any]

class OutputState(TypedDict):
    sql_query: str
    result: List[Any]
    visualization_code: str

def get_schema(state) -> dict:
    """Retrieves the schema of SQL Server database."""
    conn_str = state["uri"]
    conn = pyodbc.connect(conn_str)
    cursor = conn.cursor()
    cursor.execute("SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_TYPE='BASE TABLE';")
    tables = cursor.fetchall()
    schema = {}
    
    for table_name, in tables:
        cursor.execute(f"SELECT COLUMN_NAME FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME='{table_name}'")
        schema[table_name] = {
            "columns": [row[0] for row in cursor.fetchall()],
            "sample_rows": []
        }
        cursor.execute(f"SELECT TOP 3 * FROM {table_name}")
        sample_rows = cursor.fetchall()
        schema[table_name]["sample_rows"] = [list(row) for row in sample_rows]
    
    conn.close()
    return {"schema": schema}

def generate_sql_query(state) -> dict:
    """Generates an SQL Server query from a natural language question."""
    question = state["question"]
    schema = state["schema"]
    
    llm = ChatOpenAI(model="gpt-4o-mini")
    prompt = ChatPromptTemplate.from_messages([
        ("system", """
        You are an SQL expert. Generate an SQL server query based on the database schema and question.
        Always generate syntactically correct SQL Server queries. Avoid including unncessary columns in the query. Only use the columns that are relevant to the question. If the user asks to exclude something from the results, follow the instructions completely. 
        Ensure the output query is executable. Do not put ```sql in the beginning of the query. Just provide raw executable query as a string. Use appropriate column names for the results.
        """),
        ("human", "===Database Schema:\n{schema}\n\n===User Question:\n{question}\n\nSQL Server Query:")
    ])
    
    generate_query_chain = prompt | llm | StrOutputParser()
    response = generate_query_chain.invoke({"question": question, "schema": schema})

    print(f"Original SQL query: {response}")
    
    return {"sql_query": response}

## Need to add query checker for SQL query
def sql_server_query_checker(state) -> dict:
    """Basic SQL Server query checker."""

    query = state["sql_query"]
    question = state["question"]

    query_check_system_prefix = """You are a SQL Server expert with a strong attention to detail.
    You will be provided a SQL Server query and a question. Double check the SQL Server query for common mistakes and verify if it answers the question properly.
    If there are any mistakes, rewrite the query. If there are no mistakes, just reproduce the original query. Always produce the query as raw string which the user can execute directly.
    Do not output anything else other then the query itself. No explanations. Just provide the relevant query."""

    template = f"""
    Check if the SQL query answers the user question properly. Specifically, review if the query satisfies all the information asked by the user and is syntactically correct.

    Question: {question}
    SQL Server Query: {query}
    """

    completion = client.chat.completions.create(
    model="o3-mini", ### Usage of advanced reasoning model
    messages=[
        {"role": "developer", "content": query_check_system_prefix},
        {
            "role": "user",
            "content": template,
        },
    ])

    print(f"Corrected SQL query: {completion.choices[0].message.content}")

    return {"sql_query": completion.choices[0].message.content}

def run_sql_server_query(state) -> dict:
    """Executes a SQL Server query and returns the results with column names embedded."""
    
    conn_str = state["uri"]
    query = state["sql_query"]
    
    try:
        conn = pyodbc.connect(conn_str)
        cursor = conn.cursor()
        cursor.execute(query)
        
        # Fetch column names
        columns = tuple(column[0] for column in cursor.description)
        
        # Fetch results and prepend column names as the first row
        result = [columns] + cursor.fetchall()
        
        return {"result": result}
    
    except Exception as e:
        return {"result": [("Error",)], "error": f"Error executing query: {e}"}
    
    finally:
        conn.close()


### Can also directly use the assistants api from Open AI for writing and executing code.
## More information on this here: https://platform.openai.com/docs/assistants/quickstart
def generate_code_using_openai_call(state) -> dict:
    question = state["question"]
    data = state["result"]
    #error = state["error"]

    system_prefix = """You are an agent designed to generate Python code for creating visualizations based on the data provided.
    Given an input question and data, create a Python code snippet that uses the Plotly library to generate the appropriate visualization. You might get a list of tuples as the data. First flatten the tuple, and then create a DataFrame using the data. Use the DataFrame to create the visualization.
    Ensure the code is syntactically correct and includes necessary imports. Ensure that the code won't throw any errors when executed.
    Do not include the code inside ```python``` tags. Just give the code snippet that the user can run directly using exec(). Therefore, produce the code as raw string.
    Do not include any additional explanations or comments in the code. Do not add any additional logic on the dataframe since the data provided to you already answers the question.
    After importing the necessary libraries, you can start your code directly with "df = pd.DataFrame(result)".
    Always use the right syntax while using functions like df.pivot_table() or df.groupby().sum(). Do not use invalid syntax that will throw errors.
    Use figure size like "plt.figure(figsize=10, 5))" since I have to render this visualization on a web page. Use appropriate colors and labels for the visualization.
    If you do not receive any data, return an empty visualization with a message indicating that the data is not available.
    """

    template = f"""
    Based on the data provided below, write a Python code snippet that would create a appealing visualization based on the data provided. You would be optionally given a chart type in the question. If it is not provided, think of an appropriate chart type based on the data.

    Question: {question}
    Data: {data}
    Visualization Code:
    """

    completion = client.chat.completions.create(
    model="o3-mini",
    messages=[
        {"role": "developer", "content": system_prefix},
        {
            "role": "user",
            "content": template,
        },
    ])


    return {"visualization_code": completion.choices[0].message.content}

def generate_visualization_code(state) -> dict:
    """Generates Python code for visualizing the SQL query results."""
    #print(state)

    question = state["question"]
    data = state["result"]

    llm = ChatOpenAI(model="gpt-4o", temperature=0)

    template = """
    Based on the data provided below, write a Python code snippet that would create a appealing visualization based on the data provided. You would be optionally given a chart type in the question. If it is not provided, think of an appropriate chart type based on the data.

    Question: {question}
    Data: {data}
    Visualization Code:
    """

    system_prefix = """You are an agent designed to generate Python code for creating visualizations based on the data provided.
    Given an input question and data, create a Python code snippet that uses the Plotly library to generate the appropriate visualization. You might get a list of tuples as the data. First flatten the tuple, and then create a DataFrame using the data. Use the DataFrame to create the visualization.
    Ensure the code is syntactically correct and includes necessary imports. Ensure that the code won't throw any errors when executed.
    Do not include the code inside ```python``` tags. Just give the code snippet that the user can run directly using exec(). Therefore, produce the code as raw string.
    Do not include any additional explanations or comments in the code. Do not add any additional logic on the dataframe since the data provided to you already answers the question.
    After importing the necessary libraries, you can start your code directly with "df = pd.DataFrame(result)".
    Always use the right syntax while using functions like df.pivot_table() or df.groupby().sum(). Do not use invalid syntax that will throw errors.
    Use a reasonably big figure size like "plt.figure(figsize=10, 5))" since I have to render this visualization on a web page. Use appropriate colors and labels for the visualization.
    If you do not receive any data, return an empty visualization with a message indicating that the data is not available.
    """

    final_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", system_prefix),
            ("human", template),
        ]
    )

    generate_code_chain = final_prompt | llm | StrOutputParser()

    with get_openai_callback() as cb:
        response = generate_code_chain.invoke({"question": question, "data": data})
        print("GPT -4o callback :::::")
        print(cb)

    return {"visualization_code": response}

### Ideally, it should be like: Generate Code -> Run it and if it fails -> Go to a fallback method that can fix the code
### If it works then continue with the flow
### Same should be the case with SQL query

def run_python_code(state) -> dict:
    visualization_code = state["visualization_code"]

    try:
        exec(visualization_code)
    except Exception as e:
        return {"error": f"Your code failed with the following error: {str(e)}"}

## Also need to add code checker before running it.
## Need to implement handling of errors from the tools. This will take the error returned from any of the generative functions and ask the model to fix it.
def handle_tool_error(state) -> dict:
    error = state.get("error")
    tool_calls = state["messages"][-1].tool_calls
    return {
        "messages": [
            ToolMessage(
                content=f"Error: {repr(error)}\n please fix your mistakes.",
                tool_call_id=tc["id"],
            )
            for tc in tool_calls
        ]
    }

def create_tool_node_with_fallback(tools: list) -> RunnableWithFallbacks[Any, dict]:
    """
    Create a ToolNode with a fallback to handle errors and surface them to the agent.
    """
    return ToolNode(tools).with_fallbacks(
        [RunnableLambda(handle_tool_error)], exception_key="error"
    )


workflow = StateGraph(input=InputState, output=OutputState)
workflow.add_node("get_schema", get_schema)
workflow.add_node("generate_sql", generate_sql_query)
workflow.add_node("check_sql_query", sql_server_query_checker)
workflow.add_node("run_sql", run_sql_server_query)
#workflow.add_node("generate_visualization", generate_visualization_code)
workflow.add_node("generate_visualization", generate_code_using_openai_call)

workflow.add_edge("get_schema", "generate_sql")
workflow.add_edge("generate_sql", "check_sql_query")
workflow.add_edge("check_sql_query", "run_sql")
workflow.add_edge("run_sql", "generate_visualization")
workflow.add_edge("generate_visualization", END)
workflow.set_entry_point("get_schema")


if __name__ == "__main__":
    SERVER_NAME = None
    CONN_STR = f"DRIVER=ODBC Driver 17 for SQL Server;SERVER={SERVER_NAME};DATABASE=GENAIDB;Trusted_Connection=yes;"
    #print(CONN_STR)

    # Create SQLAlchemy Engine for SQL Server
    # CONN_STR = f"mssql+pyodbc://{SQL_USERNAME}:{SQL_PASSWORD}@{SQL_SERVER}/{SQL_DATABASE}?driver={SQL_DRIVER}&trusted_connection=yes"
    #print(CONN_STR)
    ## Define initial state and invoke the graph.
    state = {
        "question": "Show the weekly occurance trend of top 10 defects in October 2024.",
        "database": "GENAIDB",
        "uri": CONN_STR
    }
    
    workflow = workflow.compile()
    result = workflow.invoke(state)
    print(result["result"])
    exec(result["visualization_code"])