Refactored & Langsmith Prompt

app.py

@@ -5,16 +5,19 @@ import matplotlib.pyplot as plt
 from transformers import HfEngine, ReactCodeAgent
 from transformers.agents import Tool
 from langsmith import traceable
+from langchain import hub
+
+
 # Height of the Tabs Text Area
 TAB_LINES = 8
-# Load Token
-md_token = os.getenv('MD_TOKEN')
-os.environ['HF_TOKEN'] = os.getenv('HF_TOKEN')
 
-print('Connecting to DB...')
-# Connect to DB
+
+#----------CONNECT TO DATABASE----------
+md_token = os.getenv('MD_TOKEN')
 conn = duckdb.connect(f"md:my_db?motherduck_token={md_token}", read_only=True)
+#---------------------------------------
 
+#-------LOAD HUGGINGFACE MODEL-------
 models = ["Qwen/Qwen2.5-72B-Instruct","meta-llama/Meta-Llama-3-70B-Instruct",
           "meta-llama/Llama-3.1-70B-Instruct"]
 
@@ -31,7 +34,14 @@ for model in models:
 
 if not model_loaded:
     gr.Warning(f"❌ None of the model form {models} are available. {e}")
+#---------------------------------------
+
+#-----LOAD PROMPT FROM LANCHAIN HUB-----
+prompt = hub.pull("viz-prompt")
+#-------------------------------------
+
 
+#--------------ALL UTILS----------------
 def get_schemas():
     schemas = conn.execute("""
     SELECT DISTINCT schema_name
@@ -65,92 +75,6 @@ def get_table_schema(table):
     return ddl_create, full_path
 
 
-def get_visualization(question, tool, schema, table_name):
-    agent = ReactCodeAgent(tools=[tool], llm_engine=llm_engine, add_base_tools=True,
-                           additional_authorized_imports=['matplotlib.pyplot',
-                                                 'pandas', 'plotly.express',
-                                                 'seaborn'], max_iterations=10)
-    results = agent.run(
-
-    task= f'''
-
-    Here are the steps you should follow while writing code for Visualization:
-    1. You have access to the database with the `sql_engine` tool, which allows you to run DuckDB SQL queries and return results as a df.
-    2. Query the database using `sql_engine`, print the first 5 rows to inspect the data.
-    3. Select the most appropriate chart type for the data:
-       - Use bar charts for categorical comparisons, line charts for trends over time, scatter plots for relationships between variables, pie charts for proportions, histograms for distribution, and box plots for data spread and outliers.
-    4. Analyze the data and choose the best visualization type to answer the query.
-    5. Always include a plot in your answer.
-    6. Use Seaborn for the plots.
-    7. In the end, return a dictionary containing the final figure (`fig` key), the generated SQL (`sql` key), and the data as a dataframe (`data` key) using the `final_answer` tool, e.g. `final_answer(answer={{"fig": 'fig.png', "sql": sql, "data": data}})`.
-
-    Example:
-
-    ```python
-    # Input query
-    query_description = 'Average tip amount based on the ride time length in minutes.'
-
-    # SQL Query to get ride time length and average tip amount
-    query = """
-        SELECT 
-            EXTRACT(EPOCH FROM (tpep_dropoff_datetime - tpep_pickup_datetime)) / 60 AS ride_time_length,
-            AVG(tip_amount) AS avg_tip_amount
-        FROM 
-            sample_data.nyc.taxi
-        GROUP BY 
-            EXTRACT(EPOCH FROM (tpep_dropoff_datetime - tpep_pickup_datetime)) / 60
-    """
-
-    # Execute the query using the sql_engine tool
-    df = sql_engine(query=query)
-
-    # Print the result to observe the data
-    print(df)
-
-    # Create a line plot using seaborn
-    import seaborn as sns
-    import matplotlib.pyplot as plt
-
-    plt.figure(figsize=(10,6))
-    sns.lineplot(x="ride_time_length", y="avg_tip_amount", data=df)
-
-    # Set the title and labels
-    plt.title("Average Tip Amount vs Ride Time Length")
-    plt.xlabel("Ride Time Length (minutes)")
-    plt.ylabel("Average Tip Amount")
-
-    # Print the plot to observe the results
-    print("Plot created")
-
-    # Since we are required to return a fig, sql, and data, let's store the plot in a variable
-    fig = plt.gcf()
-
-    # Store the query in a variable
-    sql = query
-
-    # Store the dataframe in a variable
-    data = df
-
-    # Return the final answer
-    final_answer(answer={{"fig": fig, "sql": sql, "data": data}})
-    ```
-
-
-    Here is the query you should generate a plot for: '{question}'.
-    Here is the schema: '{schema}' and here is the table name: '{table_name}
-
-      
-    '''
-    )
-
-    return results
-
-
-@traceable()
-def query_response(input_prompt, generated_sql):
-  return generated_sql
-
-
 class SQLExecutorTool(Tool):
     name = "sql_engine"
     inputs = {
@@ -168,6 +92,27 @@ class SQLExecutorTool(Tool):
     
 tool = SQLExecutorTool()
 
+def process_outputs(output) :
+    if 'data' in output:
+        output['data'] = "<DataFrame is hidden>"
+    if 'fig' in output:
+        output['fig'] = "<Figure is hidden>"
+    return output
+
+@traceable(process_outputs=process_outputs)
+def get_visualization(question, schema, table_name):
+    agent = ReactCodeAgent(tools=[tool], llm_engine=llm_engine, add_base_tools=True,
+                           additional_authorized_imports=['matplotlib.pyplot',
+                                                 'pandas', 'plotly.express',
+                                                 'seaborn'], max_iterations=10)
+    results = agent.run(
+    task= prompt.format(question=question, schema=schema, table_name=table_name)
+    )
+
+    return results
+#---------------------------------------
+
+
 def main(table, text_query):
     # Empty Fig
     fig, ax = plt.subplots()
@@ -175,14 +120,12 @@ def main(table, text_query):
     
     schema, table_name = get_table_schema(table)
     
-    
     try:
-        output = get_visualization(question=text_query, tool=tool, schema=schema, table_name=table_name)
+        output = get_visualization(question=text_query, schema=schema, table_name=table_name)
         fig = output.get('fig', None)
         generated_sql = output.get('sql', None)
         data = output.get('data', None)
-        input_prompt = text_query + '\n' + table_name + '\n' + schema
-        _ = query_response(input_prompt, generated_sql)
+
     except Exception as e:
         gr.Warning(f"❌ Unable to generate the visualization. {e}")
         
@@ -246,4 +189,8 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple", secondary_hue="indigo"
         generate_query_button.click(main, inputs=[tables_dropdown, query_input], outputs=[result_plot, generated_sql, data])
 
 if __name__ == "__main__":
-    demo.launch(debug=True)
+    demo.launch(debug=True)
+    
+    
+    
+    

requirements.txt

@@ -5,4 +5,5 @@ huggingface_hub
 accelerate==0.34.2
 transformers==4.44.2
 duckdb==1.1.1
-langsmith==0.1.135
+langsmith==0.1.135
+langchain==0.3.4