Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,9 +1,13 @@
 import streamlit as st
 import requests
 import os
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
 
 # Get Lambda URL from environment variable
-LAMBDA_URL = os.getenv("LAMBDA_URL")
+LAMBDA_URL = os.getenv("LAMBDA_URL") or "http://localhost:8000/chats/messages"
 if not LAMBDA_URL:
     st.error("LAMBDA_URL environment variable is not set. Please check your .env file.")
     st.stop()
@@ -12,6 +16,17 @@ st.set_page_config(page_title="Ask Redshift", page_icon="🤖")
 
 st.title("🧠 Ask your Redshift Database Anything...")
 
+# Initialize conversation_id in session state if not exists
+if "conversation_id" not in st.session_state:
+    st.session_state.conversation_id = None
+
+# Add system prompt input
+system_prompt = st.text_area(
+    "System Prompt (Optional)",
+    help="Customize the AI's behavior and instructions. Leave empty to use default behavior.",
+    height=100
+)
+
 question = st.text_input(
     "Enter your question about your data (e.g., 'What is the average cost per lead for March?')"
 )
@@ -19,13 +34,25 @@ question = st.text_input(
 if st.button("Run Query"):
     if question:
         with st.spinner("Querying..."):
+            # Prepare payload
+            payload = {"question": question}
+            if st.session_state.conversation_id:
+                payload["conversation_id"] = st.session_state.conversation_id
+            if system_prompt:
+                payload["system_prompt"] = system_prompt
+
             # Call the API
             response = requests.post(
                 LAMBDA_URL,
-                json={"question": question},
+                json=payload,
             )
             response.raise_for_status()
             response = response.json()
+
+            # Store conversation_id for future requests
+            if "conversation_id" in response:
+                st.session_state.conversation_id = response["conversation_id"]
+
         st.success(response["answer"])
     else:
         st.error("Please enter a question.")