initial commits

Dockerfile

@@ -17,4 +17,4 @@ EXPOSE 8501
 
 HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health
 
-ENTRYPOINT ["streamlit", "run", "src/streamlit_app.py", "--server.port=8501", "--server.address=0.0.0.0"]
+ENTRYPOINT ["streamlit", "run", "src/chat_ui.py", "--server.port=8501", "--server.address=0.0.0.0"]

src/chat_ui.py

@@ -0,0 +1,315 @@
+
+import streamlit as st
+import requests
+import pandas as pd
+import base64
+import logging
+import os
+
+
+# ===============================
+# Config: Dynamic API_URL based on environment
+# ===============================
+
+# Read API base URL and endpoint from environment variables
+api_base_url = os.getenv("API_BASE_URL", "http://127.0.0.1:8000")
+api_endpoint = os.getenv("API_ENDPOINT", "/api/process-text")
+# Ensure no double slashes
+API_URL = api_base_url.rstrip("/") + "/" + api_endpoint.lstrip("/")
+
+st.set_page_config(page_title="AI Database Assistant", layout="wide")
+
+# Show current API URL in sidebar for debugging
+st.sidebar.info(f"**API URL:** {API_URL}")
+
+# Setup logger
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+
+st.title("💬 AI Database Assistant")
+st.write("Ask questions in plain English to generate and run SQL queries.")
+# st.markdown(
+#     "<h1 style='text-align: center; color: white; background: rgba(0,0,0,0.6); "
+#     "padding: 12px; border-radius: 10px;'>🤖 AI Database Assistant</h1>",
+#     unsafe_allow_html=True,
+# )
+# st.write("Ask questions in plain English to generate and run SQL queries.")
+
+# ===============================
+# Sidebar Settings
+# ===============================
+st.sidebar.header("⚙️ Settings")
+
+model = st.sidebar.selectbox("Select Model", ["gpt-3.5-turbo", "gpt-4", "gpt-4o-mini"])
+agent = st.sidebar.selectbox("Select Agent", ["default", "sql-agent", "custom-agent"])
+theme = st.sidebar.selectbox("Theme", ["Light", "Dark", "Custom"])
+
+# ===============================
+# Theme CSS
+# ===============================
+def inject_css(theme_choice: str):
+    if theme_choice == "Light":
+        css = """
+        body, .stApp { background: #f8fafc !important; }
+        .chat-container { max-width: 900px; margin: auto; }
+        .user-bubble {
+            background: linear-gradient(90deg, #e3f2fd, #ffffff);
+            color: #222;
+            float: right;
+            clear: both;
+        }
+        .ai-bubble {
+            background: linear-gradient(90deg, #f3e5f5, #ffffff);
+            color: #222;
+            float: left;
+            clear: both;
+        }
+        .error-bubble {
+            background: #ffebee; color: #b71c1c; font-weight: bold;
+            float: left; clear: both;
+        }
+        section[data-testid="stSidebar"] { background: #e3f2fd !important; }
+        """
+    elif theme_choice == "Dark":
+        css = """
+        body, .stApp { background: #181825 !important; }
+        .chat-container { max-width: 900px; margin: auto; }
+        .user-bubble {
+            background: linear-gradient(135deg, #3a0ca3, #4361ee);
+            color: #f8fafc;
+            float: right;
+            clear: both;
+        }
+        .ai-bubble {
+            background: linear-gradient(135deg, #14213d, #1d3557);
+            color: #f8fafc;
+            float: left;
+            clear: both;
+        }
+        .error-bubble {
+            background: #ffb4ab; color: #b71c1c; font-weight: bold;
+            float: left; clear: both;
+        }
+        section[data-testid="stSidebar"] { background: #232946 !important; }
+        """
+    else:  # Custom theme
+        css = """
+        body, .stApp { background: #fffbe7 !important; }
+        .chat-container { max-width: 900px; margin: auto; }
+        .user-bubble {
+            background: linear-gradient(90deg, #ffe082, #fffbe7);
+            color: #6d4c00;
+            float: right;
+            clear: both;
+        }
+        .ai-bubble {
+            background: linear-gradient(90deg, #b2dfdb, #fffbe7);
+            color: #004d40;
+            float: left;
+            clear: both;
+        }
+        .error-bubble {
+            background: #ffe0b2; color: #b71c1c; font-weight: bold;
+            float: left; clear: both;
+        }
+        section[data-testid="stSidebar"] { background: #ffe082 !important; }
+        """
+
+    # Common bubble styling
+    css += """
+    .user-bubble, .ai-bubble, .error-bubble {
+        border-radius: 14px;
+        padding: 10px 14px;
+        margin: 6px 0;
+        display: inline-block;
+        max-width: 75%;
+        word-wrap: break-word;
+        box-shadow: 0 2px 8px rgba(0,0,0,0.25);
+    }
+    """
+    st.markdown(f"<style>{css}</style>", unsafe_allow_html=True)
+
+inject_css(theme)
+
+# ===============================
+# Avatar Setup
+# ===============================
+# Default base64 encoded avatars
+DEFAULT_USER_AVATAR = """
+data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHZpZXdCb3g9IjAgMCAyNCAyNCI+PHBhdGggZD0iTTEyIDJDNi40OCAyIDIgNi40OCAyIDEyczQuNDggMTAgMTAgMTAgMTAtNC40OCAxMC0xMFMxNy41MiAyIDEyIDJ6bTAgM2MyLjY3IDAgOC0xLjMzIDgtM3Y0LjE1YzAgMS42My0zLjMzIDIuODUtOCAyLjg1cy04LTEuMjItOC0yLjg1VjJjMCAxLjY3IDUuMzMgMyA4IDN6bTAgMTFjLTIuNjcgMC04LTEuMzMtOC0zdjQuMTVjMCAxLjYzIDMuMzMgMi44NSA4IDIuODVzOC0xLjIyIDgtMi44NVYxM2MwIDEuNjctNS4zMyAzLTggM3oiIGZpbGw9IiM0Q0FGNTAiLz48L3N2Zz4=
+"""
+
+DEFAULT_AI_AVATAR = """
+data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHZpZXdCb3g9IjAgMCAyNCAyNCI+PHBhdGggZD0iTTEyIDJDNi40OCAyIDIgNi40OCAyIDEyczQuNDggMTAgMTAgMTAgMTAtNC40OCAxMC0xMFMxNy41MiAyIDEyIDJ6bS0xIDE2LjkydjEuMDhjMCAuNTUuNDUgMSAxIDFoMmMuNTUgMCAxLS40NSAxLTF2LTEuMDhjLS42My4xMy0xLjI5LjItMS45Ny4yLS43IDAtMS4zNC0uMDctMS45Ny0uMnptNS43OC0yLjc0QzE1LjI0IDE3LjM0IDEzLjY3IDE4IDEyIDE4cy0zLjI0LS42Ni00Ljc4LTEuODJDNC40NCAxNC4zNCAzIDEyLjA0IDMgOS41IDMgNS45MSA1LjkxIDMgOS41IDNoNUM2LjY1IDMgNCA0LjY1IDQgNi43MWMwIDIuMDYgMS42NyAzLjc0IDMuNzMgMy43NC4zNiAwIC43MS0uMDUgMS4wNS0uMTUgMS4xIDEuMTEgMi41NyAxLjc5IDQuMjIgMS43OSAxLjY0IDAgMy4xMi0uNjkgNC4yMi0xLjc5LjMzLjEuNjkuMTUgMS4wNS4xNSAyLjA2IDAgMy43My0xLjY4IDMuNzMtMy43NCAwLTIuMDYtMS42Ny0zLjcxLTMuNzMtMy43MWgtNUM3LjkxIDMgNSA1LjkxIDUgOS41YzAgMi41NCAxLjQ0IDQuODQgMy43OCA2LjY4eiIgZmlsbD0iIzI5NzlGRiIvPjwvc3ZnPg==
+"""
+
+# ===============================
+# Session State
+# ===============================
+if "messages" not in st.session_state:
+    st.session_state["messages"] = []
+
+# ===============================
+# Helper: Render Chat History
+# ===============================
+def render_chat():
+    st.markdown("<div class='chat-container'>", unsafe_allow_html=True)
+    for msg in st.session_state["messages"]:
+        if msg["role"] == "user":
+            st.markdown(f'<div class="user-bubble">{msg["content"]}</div>', unsafe_allow_html=True)
+        elif msg["role"] == "assistant":
+            bubble_class = "error-bubble" if msg.get("is_error") else "ai-bubble"
+            st.markdown(f'<div class="{bubble_class}">{msg["content"]}</div>', unsafe_allow_html=True)
+
+            if msg.get("data"):
+                df = pd.DataFrame(msg["data"])
+                st.dataframe(df, use_container_width=True)
+
+                # Add download button for results
+                csv = df.to_csv(index=False).encode("utf-8")
+                st.download_button("📥 Download CSV", csv, "results.csv", "text/csv")
+
+            if msg.get("chart"):
+                img_data = base64.b64decode(msg["chart"])
+                st.image(img_data, use_column_width=True)
+    st.markdown("</div>", unsafe_allow_html=True)
+
+
+    def render_chat():
+        # Spinner CSS (only inject once)
+        st.markdown('''
+            <style>
+            .spinner {
+                display: inline-block;
+                width: 1.3em;
+                height: 1.3em;
+                border: 3px solid #e0e0e0;
+                border-top: 3px solid #2193b0;
+                border-radius: 50%;
+                animation: spin 0.8s linear infinite;
+                margin-right: 0.7em;
+            }
+            @keyframes spin {
+                0% { transform: rotate(0deg); }
+                100% { transform: rotate(360deg); }
+            }
+            </style>
+        ''', unsafe_allow_html=True)
+        st.markdown("<div class='chat-container'>", unsafe_allow_html=True)
+        user_avatar = DEFAULT_USER_AVATAR
+        ai_avatar = DEFAULT_AI_AVATAR
+        for msg in st.session_state["messages"]:
+            if msg["role"] == "user":
+                st.markdown(
+                    f'''<div style="display: flex; align-items: flex-start; justify-content: flex-end; margin-bottom: 0.5em;">
+                            <div style="margin-right: 0.5em;">
+                                <img src="{user_avatar}" alt="User" style="width: 2.3rem; height: 2.3rem; border-radius: 50%; border: 2px solid #e3f2fd; background: #fff; object-fit: cover;" />
+                            </div>
+                            <div class="user-bubble">{msg["content"]}</div>
+                       </div>''',
+                    unsafe_allow_html=True)
+            elif msg["role"] == "assistant":
+                bubble_class = "error-bubble" if msg.get("is_error") else "ai-bubble"
+                # Show spinner if this is a placeholder 'Thinking...' message
+                if msg.get("is_placeholder"):
+                    st.markdown(
+                        f'''<div style="display: flex; align-items: flex-start; margin-bottom: 0.5em;">
+                                <div style="margin-right: 0.5em;">
+                                    <img src="{ai_avatar_url}" alt="AI" style="width: 2.3rem; height: 2.3rem; border-radius: 50%; border: 2px solid #b2dfdb; background: #fff; object-fit: cover;" />
+                                </div>
+                                <div class="{bubble_class}"><span class="spinner"></span>Thinking...</div>
+                           </div>''',
+                        unsafe_allow_html=True)
+                else:
+                    st.markdown(
+                        f'''<div style="display: flex; align-items: flex-start; margin-bottom: 0.5em;">
+                                <div style="margin-right: 0.5em;">
+                                    <img src="{ai_avatar_url}" alt="AI" style="width: 2.3rem; height: 2.3rem; border-radius: 50%; border: 2px solid #b2dfdb; background: #fff; object-fit: cover;" />
+                                </div>
+                                <div class="{bubble_class}">{msg["content"]}</div>
+                           </div>''',
+                        unsafe_allow_html=True)
+
+                if msg.get("data"):
+                    df = pd.DataFrame(msg["data"])
+                    st.dataframe(df, use_container_width=True)
+
+                    # Add download button for results
+                    csv = df.to_csv(index=False).encode("utf-8")
+                    st.download_button("\U0001F4C5 Download CSV", csv, "results.csv", "text/csv")
+
+                if msg.get("chart"):
+                    img_data = base64.b64decode(msg["chart"])
+                    st.image(img_data, use_column_width=True)
+        st.markdown("</div>", unsafe_allow_html=True)
+
+render_chat()
+
+# ===============================
+# User Input
+# ===============================
+
+# Disable chat input while AI is working
+pending = False
+if st.session_state["messages"]:  # Check if there are any messages
+    if st.session_state["messages"][-1]["role"] == "assistant":
+        pending = st.session_state["messages"][-1].get("is_placeholder", False)
+
+user_query = st.chat_input(
+    "Ask me something about your database...",
+    disabled=pending or False  # Ensure disabled is always boolean
+)
+if user_query and not pending:
+    st.session_state["messages"].append({"role": "user", "content": user_query})
+    st.session_state["messages"].append({"role": "assistant", "content": "🤔 Thinking...", "is_placeholder": True})
+    st.rerun()
+
+# ===============================
+# Handle Pending Assistant Response
+# ===============================
+if (
+    st.session_state["messages"]
+    and st.session_state["messages"][-1].get("is_placeholder")
+    and len(st.session_state["messages"]) >= 2
+    and st.session_state["messages"][-2]["role"] == "user"
+):
+    user_query = st.session_state["messages"][-2]["content"]
+
+    try:
+        with st.spinner("Working..."):
+            payload = {"question": user_query, "model": model, "agent": agent}
+            logging.info(f"Sending API request with: {payload}")
+            response = requests.post(API_URL, json=payload, timeout=60)
+
+            try:
+                result = response.json()
+            except ValueError:
+                result = {"detail": response.text}
+
+        if response.status_code != 200:
+            answer_text = f"❌ Error: {result.get('detail') or result.get('message') or response.text}"
+            rows, chart, is_error = [], None, True
+        elif "message" in result and not ("rows" in result or "chart" in result):
+            answer_text, rows, chart, is_error = result["message"], [], None, False
+        else:
+            rows = result.get("rows", [])
+            chart = result.get("chart", None)
+            if rows:
+                answer_text = f"Here are the top {len(rows)} results I found:"
+            else:
+                answer_text = "I could not find matching records for your query."
+            is_error = False
+
+    except Exception as e:
+        logging.error(f"Exception: {e}")
+        answer_text, rows, chart, is_error = f"⚠️ Exception: {str(e)}", [], None, True
+
+    # Replace placeholder with final assistant message
+    st.session_state["messages"][-1] = {
+        "role": "assistant",
+        "content": answer_text,
+        "data": rows,
+        "chart": chart,
+        "is_error": is_error,
+    }
+
+    st.rerun()