fix: fix UI and add some extra-features

Author: hanshal101

.env

@@ -1,4 +1,4 @@
-GROQ_API_KEY=""
+GROQ_API_KEY="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"
 GROQ_MODEL="llama-3.3-70b-specdec"
-ANALYZING_MODEL="deepseek-r1-distill-llama-70b-specdec"
+ANALYZING_MODEL="deepseek-r1-distill-llama-70b"
 OLLAMA_MODEL="llama3.2"

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 from performer.performer import graph
 from agentstate.agent_state import AgentState
+from utils.sql_utils import extract_sql_queries
 
 st.title("PostgreSQL Database Optimization Assistant")
 st.markdown("""
@@ -29,28 +30,83 @@ def run_analysis():
             if "analysis" in event:
                 st.session_state.analysis_history.append(event["analysis"])
 
+def extract_thinking_content(analysis):
+    """Extracts the content between <think> and </think> tags."""
+    start_tag = "<think>"
+    end_tag = "</think>"
+    start_index = analysis.find(start_tag)
+    end_index = analysis.find(end_tag)
+    if start_index != -1 and end_index != -1:
+        return analysis[start_index + len(start_tag):end_index].strip()
+    return None
+
+def execute_sql_queries(edited_queries):
+    """Extract and execute SQL queries from the final analysis."""
+    current_state = graph.get_state(thread)
+    schema = current_state.values.get("schema", "")
+
+    # Update state with edited queries
+    graph.update_state(thread, {"execute_query": edited_queries})
+
+    # Execute each query
+    for query in edited_queries.split(";"):
+        query = query.strip()
+        if query:  # Skip empty queries
+            st.write(f"Executing: {query}")
+            try:
+                # Simulate execution (replace with actual SQL execution logic if needed)
+                st.write("Query executed successfully.")
+            except Exception as e:
+                st.error(f"Failed to execute query: {query}. Error: {str(e)}")
+
 if st.button("Analyze"):
     if query and schema:
         run_analysis()
     else:
         st.error("Please provide both query and schema")
 
 if st.session_state.analysis_history:
+    latest_analysis = st.session_state.analysis_history[-1]
+    start_tag = "<think>"
+    end_tag = "</think>"
+    if start_tag in latest_analysis and end_tag in latest_analysis:
+        start_index = latest_analysis.index(start_tag) + len(start_tag)
+        end_index = latest_analysis.index(end_tag)
+        thinking_content = latest_analysis[start_index:end_index].strip()
+    else:
+        thinking_content = "No detailed reasoning available for this analysis."
+
+    with st.expander("Thinking Mode: View Detailed Reasoning", expanded=False):
+        st.markdown(thinking_content)
+
     st.subheader("Analysis History")
     for i, analysis in enumerate(st.session_state.analysis_history, 1):
         st.write(f"Iteration {i}:")
-        st.code(analysis)
+        st.markdown(analysis)
 
     st.subheader("Feedback")
     col1, col2 = st.columns(2)
 
     with col1:
         if st.button("Yes - Accept Analysis"):
             graph.update_state(thread, {"execute": True})
-            st.success("Analysis accepted! Proceeding to execution...")
-            st.query_params = {"status": "completed"}
-            st.rerun()
-    
+            st.success("Analysis accepted! Proceeding to SQL execution...")
+            
+            # Extract SQL queries from the latest analysis
+            current_state = graph.get_state(thread)
+            analysis = current_state.values.get("analysis", "")
+            sql_queries = extract_sql_queries(analysis)
+
+            if not sql_queries:
+                st.warning("No SQL queries found in the analysis.")
+            else:
+                # Display SQL queries in an editable text area
+                edited_queries = st.text_area("Edit SQL Queries:", value=sql_queries), height=200)
+
+                # Add an "Execute" button to confirm and execute the edited queries
+                if st.button("Execute Edited Queries"):
+                    execute_sql_queries(edited_queries)
+
     with col2:
         if st.button("No - Revise Analysis"):
             st.session_state.show_feedback = True

performance_test.py

@@ -0,0 +1,93 @@
+import sys
+from io import StringIO
+import logging
+from performer.performer import graph  # Import your LangGraph setup
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(message)s",
+    handlers=[logging.StreamHandler(sys.stdout)]
+)
+
+def get_mock_schema():
+    """Returns a mock schema for testing."""
+    return """
+     table_name  | column_name 
+    -------------+-------------
+     customers   | id
+     customers   | name
+     customers   | email
+     customers   | address
+     orders      | id
+     orders      | customer_id
+     orders      | order_date
+     orders      | total
+     products    | id
+     products    | name
+     products    | description
+     products    | price
+     products    | stock_quantity
+    """
+
+def test_workflow():
+    """Simulates the workflow of the PostgreSQL optimization assistant with detailed logging."""
+    logging.info("Starting test workflow...")
+    
+    mock_inputs = [
+        "no",  
+        "Just give me details about index creation.",  
+        "yes" 
+    ]
+    mock_query = "Identify and give me solutions to optimize my postgres database"
+    mock_schema = get_mock_schema()
+
+    logging.info(f"Mock query: {mock_query}")
+    logging.info(f"Mock schema:\n{mock_schema}")
+
+    original_stdin = sys.stdin
+    original_stdout = sys.stdout
+
+    try:
+        sys.stdin = StringIO("\n".join(mock_inputs))
+
+        captured_output = StringIO()
+        sys.stdout = captured_output
+
+        thread = {"configurable": {"thread_id": "performance_optimization_test"}}
+        logging.info(f"Initialized thread with ID: {thread['configurable']['thread_id']}")
+
+        logging.info("Starting analysis phase...")
+        for event in graph.stream(
+            {"query": mock_query, "schema": mock_schema},
+            thread,
+            stream_mode="values"
+        ):
+            if "analysis" in event:
+                logging.info("Analysis generated:")
+                logging.info(event["analysis"])
+                print("\n**Analysis**")
+                print(event["analysis"])
+
+        current_state = graph.get_state(thread)
+        logging.info(f"Current state after analysis: {current_state.values}")
+        
+        if current_state.values.get("execute"):
+            logging.info("User approved analysis. Proceeding to SQL execution...")
+            sql_queries = current_state.values.get("execute_query", [])
+            logging.info(f"Extracted SQL queries: {sql_queries}")
+            
+            print("\n**Executing SQL Commands**")
+            for cmd in sql_queries:
+                logging.info(f"Executing SQL command: {cmd}")
+                print(f"Executing: {cmd}")
+
+    finally:
+        sys.stdin = original_stdin
+        sys.stdout = original_stdout
+
+    logging.info("Captured output from the workflow:")
+    logging.info(captured_output.getvalue())
+
+if __name__ == "__main__":
+    test_workflow()

performer/__pycache__/performer.cpython-312.pyc

@@ -1,7 +1,7 @@
 from agentstate.agent_state import AgentState
 from langgraph.checkpoint.memory import MemorySaver
 from langchain_core.messages import SystemMessage, HumanMessage, RemoveMessage
-from llm.llm import llm # Change according to where you want to test with
+from llm.llm import llm, analyze_llm, ollama_llm
 from langgraph.graph import START, END, StateGraph
 from sql.sql_agent import SQLAgent
 from langchain_core.tools import tool
@@ -14,13 +14,13 @@
 db_config = {
     "user": "postgres",
     "password": "postgres",
-    "host": "localhost",
+    "host": "0.0.0.0",
     "port": "5432",
     "database": "ecommerce_db"
 }
 
 sql_agent = SQLAgent(db_config=db_config,name="SQLAgent")
-get_dB_schema = sql_agent.get_schema() # Use this schema for dynamic databases
+get_dB_schema = sql_agent.get_schema()
 
 
 # @tool
@@ -64,7 +64,7 @@ def analyze_database(state: AgentState):
     ]
 
     response = llm.invoke(message)
-    # sql_commands = extract_sql_commands(response.content)  # New helper function
+    # sql_commands = extract_sql_commands(response.content)
     return {
         "analysis": response.content,
         "feedback": fdb,

performer/performer.py

@@ -62,7 +62,7 @@ def get_schema_info(db_config: dict):
             cursor.execute("""
                 SELECT table_name, column_name 
                 FROM information_schema.columns 
-                WHERE table_schema = 'public'
+                WHERE table_schema = 'ecommerce'
                 ORDER BY table_name, ordinal_position;
             """)
             rows = cursor.fetchall()