Update app.py

app.py

@@ -2,9 +2,15 @@ import os
 import streamlit as st
 import pandas as pd
 import sqlite3
+import numpy as np  # New import
 from langchain import OpenAI, LLMChain, PromptTemplate
 import sqlparse
 import logging
+# Removed unused import: from sql_metadata import Parser
+from sklearn.linear_model import LinearRegression  # New import
+from sklearn.model_selection import train_test_split  # New import
+from sklearn.metrics import mean_squared_error, r2_score  # New import
+
 
 # Initialize conversation history
 if 'history' not in st.session_state:
@@ -66,24 +72,58 @@ def process_input():
                 st.session_state.history.append({"role": "assistant", "content": assistant_response})
             else:
                 columns = ', '.join(valid_columns)
-                generated_sql = sql_generation_chain.run({
+                response = sql_generation_chain.run({
                     'question': user_prompt,
                     'table_name': table_name,
                     'columns': columns
                 })
 
-                # Debug: Display generated SQL query for inspection
-                st.write(f"Generated SQL Query:\n{generated_sql}")
-
-                # Attempt to execute SQL query and handle exceptions
-                try:
-                    result = pd.read_sql_query(generated_sql, conn)
-                    assistant_response = f"Generated SQL Query:\n{generated_sql}"
-                    st.session_state.history.append({"role": "assistant", "content": assistant_response})
-                    st.session_state.history.append({"role": "assistant", "content": result})
-                except Exception as e:
-                    logging.error(f"An error occurred during SQL execution: {e}")
-                    assistant_response = f"Error executing SQL query: {e}"
+                # Extract code from response
+                code = extract_code(response)
+                if code:
+                    # Determine if the code is SQL or Python
+                    if code.strip().lower().startswith('select'):
+                        # It's a SQL query
+                        st.write(f"Generated SQL Query:\n{code}")
+                        try:
+                            result = pd.read_sql_query(code, conn)
+                            assistant_response = f"Generated SQL Query:\n{code}"
+                            st.session_state.history.append({"role": "assistant", "content": assistant_response})
+                            st.session_state.history.append({"role": "assistant", "content": result})
+                        except Exception as e:
+                            logging.error(f"An error occurred during SQL execution: {e}")
+                            assistant_response = f"Error executing SQL query: {e}"
+                            st.session_state.history.append({"role": "assistant", "content": assistant_response})
+                    else:
+                        # It's Python code
+                        st.write(f"Generated Python Code:\n{code}")
+                        try:
+                            # Prepare the local namespace
+                            local_vars = {
+                                'pd': pd,
+                                'np': np,
+                                'data': data.copy(),
+                                'result': None,
+                                'LinearRegression': LinearRegression,
+                                'train_test_split': train_test_split,
+                                'mean_squared_error': mean_squared_error,
+                                'r2_score': r2_score
+                            }
+                            exec(code, {}, local_vars)
+                            result = local_vars.get('result')
+                            if result is not None:
+                                assistant_response = "Result:"
+                                st.session_state.history.append({"role": "assistant", "content": assistant_response})
+                                st.session_state.history.append({"role": "assistant", "content": result})
+                            else:
+                                assistant_response = "Code executed successfully."
+                                st.session_state.history.append({"role": "assistant", "content": assistant_response})
+                        except Exception as e:
+                            logging.error(f"An error occurred during code execution: {e}")
+                            assistant_response = f"Error executing code: {e}"
+                            st.session_state.history.append({"role": "assistant", "content": assistant_response})
+                else:
+                    assistant_response = response.strip()
                     st.session_state.history.append({"role": "assistant", "content": assistant_response})
 
         except Exception as e:
@@ -94,16 +134,31 @@ def process_input():
         # Reset the user_input in session state
         st.session_state['user_input'] = ''
 
+def extract_code(response):
+    """Extracts code enclosed between <CODE> and </CODE> tags."""
+    import re
+    pattern = r"<CODE>(.*?)</CODE>"
+    match = re.search(pattern, response, re.DOTALL)
+    if match:
+        return match.group(1).strip()
+    else:
+        return None
+
+# Display the conversation history
 # Display the conversation history
 for message in st.session_state.history:
     if message['role'] == 'user':
         st.markdown(f"**User:** {message['content']}")
     elif message['role'] == 'assistant':
-        if isinstance(message['content'], pd.DataFrame):
-            st.markdown("**Assistant:** Query Results:")
-            st.dataframe(message['content'])
+        content = message['content']
+        if isinstance(content, pd.DataFrame):
+            st.markdown("**Assistant:** Here are the results:")
+            st.dataframe(content)
+        elif isinstance(content, (int, float, str, list, dict)):
+            st.markdown(f"**Assistant:** {content}")
         else:
-            st.markdown(f"**Assistant:** {message['content']}")
+            st.markdown(f"**Assistant:** {content}")
+
 
 # Place the input field at the bottom with the callback
 st.text_input("Enter your message:", key='user_input', on_change=process_input)