Update app.py

app.py

@@ -1,172 +1,142 @@
-import os
-import gradio as gr
+# app.py (The "Glass Box" UI)
+
+import streamlit as st
 import google.generativeai as genai
 from tavily import TavilyClient
-from sentence_transformers import SentenceTransformer, CrossEncoder
-
-from research_agent.config import AgentConfig
-from research_agent.agent import get_clarifying_questions, research_and_plan, write_report_stream
-
-# --- CSS for a professional, ChatGPT-inspired look ---
-CSS = """
-body, .gradio-container { font-family: 'Inter', sans-serif; background-color: #343541; color: #ECECEC; }
-.gradio-container { max-width: 800px !important; margin: auto !important; padding-top: 2rem !important;}
-h1 { text-align: center; font-weight: 700; font-size: 2.5em; color: white; }
-.sub-header { text-align: center; color: #C5C5D2; margin-bottom: 2rem; font-size: 1.1em; }
-.accordion { background-color: #40414F; border: 1px solid #565869 !important; border-radius: 8px !important; }
-.accordion .gr-button { background-color: #4B4C5A; color: white; }
-#chatbot { box-shadow: none !important; border: none !important; background-color: transparent !important; }
-.message-bubble { background: #40414F !important; border: 1px solid #565869 !important; color: #ECECEC !important;}
-.message-bubble.user { background: #343541 !important; border: none !important; }
-footer { display: none !important; }
-.gr-box.gradio-container { padding: 0 !important; }
-.gr-form { background-color: transparent !important; border: none !important; box-shadow: none !important; }
-.gradio-container .gr-form .gr-button { display: none; } /* Hide the default submit button */
-#chat-input-container { position: relative; }
-#chat-input-container textarea { background-color: #40414F; color: white; border: 1px solid #565869 !important; }
-#submit-button { position: absolute; right: 10px; top: 50%; transform: translateY(-50%); background: #2563EB; color: white; border-radius: 4px; padding: 4px 8px; }
-"""
-
-# --- Model Initialization ---
-config = AgentConfig()
-writer_model, planner_model, embedding_model, reranker, tavily_client = None, None, None, None, None
-IS_PROCESSING = False # Add a lock to prevent concurrent runs
+from rag_agent import ResearchAgent, AITools # We now import the agent class
 
-def initialize_models(google_key, tavily_key):
-    global writer_model, planner_model, embedding_model, reranker, tavily_client, IS_PROCESSING
-    if not google_key or not tavily_key:
-        raise gr.Error("API keys are required.")
-    try:
-        genai.configure(api_key=google_key)
-        tavily_client = TavilyClient(api_key=tavily_key)
-        writer_model = genai.GenerativeModel(config.WRITER_MODEL)
-        planner_model = genai.GenerativeModel(config.WRITER_MODEL)
-        embedding_model = SentenceTransformer('all-MiniLM-L6-v2', device='cpu')
-        reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', device='cpu')
-    except Exception as e:
-        raise gr.Error(f"Failed to initialize models. Error: {str(e)}")
-    IS_PROCESSING = False # Ensure lock is free on initialization
+# --- Page Configuration ---
+st.set_page_config(page_title="DeepSearch Agent", layout="wide", initial_sidebar_state="expanded")
 
-# --- Gradio Application Logic ---
-with gr.Blocks(css=CSS, theme=gr.themes.Base()) as app:
-    gr.Markdown("<h1>Mini DeepSearch Agent</h1>")
-    gr.Markdown("<p class='sub-header'>Your AI partner for in-depth research and analysis.</p>")
+# --- CSS for a professional look ---
+st.markdown("""
+<style>
+    /* Main container and text styling */
+    .stApp {
+        background-color: #0F172A; /* Slate 900 */
+        color: #E2E8F0; /* Slate 200 */
+    }
+    h1, h2, h3 {
+        color: #F8FAFC; /* Slate 50 */
+    }
+    /* Expander for logs */
+    .streamlit-expanderHeader {
+        background-color: #1E293B; /* Slate 800 */
+        color: #E2E8F0;
+    }
+    /* Chat input styling */
+    .stTextInput > div > div > input {
+        background-color: #1E293B;
+        color: #E2E8F0;
+    }
+    /* Button styling */
+    .stButton > button {
+        background-color: #2563EB; /* Blue 600 */
+        color: white;
+        border: none;
+    }
+    .stButton > button:hover {
+        background-color: #1D4ED8; /* Blue 700 */
+    }
+</style>
+""", unsafe_allow_html=True)
 
-    agent_state = gr.State("INITIAL") 
-    initial_topic_state = gr.State("")
+# --- Session State Management ---
+if "agent" not in st.session_state:
+    st.session_state.agent = None
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+if "agent_state" not in st.session_state:
+    st.session_state.agent_state = "INITIAL" # INITIAL, CLARIFYING, GENERATING
+if "initial_topic" not in st.session_state:
+    st.session_state.initial_topic = ""
 
-    with gr.Accordion("API & Settings", open=True, elem_classes="accordion") as settings_accordion:
-        with gr.Row():
-            google_api_key_input = gr.Textbox(label="Google API Key", type="password", placeholder="Enter Google AI API Key", scale=2)
-            tavily_api_key_input = gr.Textbox(label="Tavily API Key", type="password", placeholder="Enter Tavily Search API Key", scale=2)
-            init_button = gr.Button("Initialize Agent", scale=1)
-    
-    chatbot = gr.Chatbot(
-        elem_id="chatbot",
-        bubble_full_width=False,
-        height=500,
-        visible=False
-    )
-    
-    with gr.Row(elem_id="chat-input-container"):
-        chat_input = gr.Textbox(placeholder="What would you like to research?", interactive=False, visible=False, show_label=False, scale=8)
-        submit_button = gr.Button("Submit", elem_id="submit-button", visible=False, scale=1)
+# --- Sidebar for API Keys ---
+with st.sidebar:
+    st.header("🔑 API Configuration")
+    google_key = st.text_input("Google Gemini API Key", type="password", key="google_api_key")
+    tavily_key = st.text_input("Tavily API Key", type="password", key="tavily_api_key")
 
-    def handle_initialization(google_key, tavily_key):
-        initialize_models(google_key, tavily_key)
-        return {
-            chatbot: gr.update(visible=True, value=[(None, "Agent initialized. Please enter your research topic to begin.")]),
-            chat_input: gr.update(interactive=True, visible=True),
-            submit_button: gr.update(visible=True),
-            settings_accordion: gr.update(open=False)
-        }
+    if st.button("Initialize Agent"):
+        if google_key and tavily_key:
+            with st.spinner("Initializing Agent's RAG models... This may take a moment."):
+                tools = AITools(api_keys={'google': google_key, 'tavily': tavily_key})
+                st.session_state.agent = ResearchAgent(tools_instance=tools)
+            st.success("Agent Initialized!")
+            st.session_state.agent_state = "READY"
+        else:
+            st.error("Please provide all API keys.")
 
-    def chat_step_wrapper(user_input, history, current_agent_state, topic_state):
-        """A wrapper to manage the processing lock."""
-        global IS_PROCESSING
-        if IS_PROCESSING:
-            print("Ignoring duplicate request while processing.")
-            if False: # This makes the function a generator
-                yield
-            return
+# --- Main Application UI ---
+st.title("Mini DeepSearch Agent")
+st.markdown("<p>Your AI partner for in-depth research and analysis.</p>", unsafe_allow_html=True)
 
-        IS_PROCESSING = True
-        try:
-            # Yield all updates from the actual agent logic
-            for update in chat_step(user_input, history, current_agent_state, topic_state):
-                yield update
-        except Exception as e:
-            error_message = f"An error occurred: {str(e)}"
-            history.append((None, error_message))
-            yield history, "INITIAL", "", gr.update(interactive=True, placeholder="Let's try again. What's the topic?")
-        finally:
-            # Release the lock once the generator is exhausted or an error occurs
-            IS_PROCESSING = False
-            print("Processing finished. Lock released.")
+# Display chat history
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
 
-    def chat_step(user_input, history, current_agent_state, topic_state):
-        history = history or []
-        history.append((user_input, None))
+# Main chat input logic
+if st.session_state.agent_state != "INITIAL":
+    if prompt := st.chat_input("Enter your research topic or answer the questions..."):
         
-        if current_agent_state == "INITIAL":
-            yield history, "CLARIFYING", user_input, gr.update(interactive=False, placeholder="Thinking...")
-            questions = get_clarifying_questions(planner_model, user_input)
-            history[-1] = (user_input, "To give you the best report, could you answer these questions for me?\n\n" + questions)
-            yield history, "CLARIFYING", user_input, gr.update(interactive=True, placeholder="Provide your answers to the questions above...")
+        # Add user message to chat
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        with st.chat_message("user"):
+            st.markdown(prompt)
 
-        elif current_agent_state == "CLARIFYING":
-            thinking_message = "Got it. Generating your full research report. This will take a moment..."
-            history[-1] = (user_input, thinking_message)
-            yield history, "GENERATING", topic_state, gr.update(interactive=False, placeholder="Generating...")
+        # Assistant's response
+        with st.chat_message("assistant"):
+            # Create containers for the two-column layout
+            log_container = st.expander("Agent's Thought Process", expanded=True)
+            report_container = st.container()
             
-            try:
-                plan = research_and_plan(config, planner_model, tavily_client, topic_state, user_input)
-                report_generator = write_report_stream(config, writer_model, tavily_client, embedding_model, reranker, plan)
+            full_report = ""
+            log_messages = []
+            
+            # --- Logic for different agent states ---
+            if st.session_state.agent_state == "READY":
+                st.session_state.initial_topic = prompt
+                st.session_state.agent_state = "CLARIFYING"
+                
+                # First call to the agent to get clarifying questions
+                for msg_type, content in st.session_state.agent.run(prompt):
+                    if msg_type == 'log':
+                        log_messages.append(content)
+                        with log_container: st.write("\n".join(log_messages))
+                    elif msg_type == 'clarification':
+                        st.session_state.messages.append({"role": "assistant", "content": content})
+                        report_container.markdown(content)
+
+            elif st.session_state.agent_state == "CLARIFYING":
+                st.session_state.agent_state = "GENERATING"
                 
-                stream_content = ""
-                for update in report_generator:
-                    stream_content = update
-                    history[-1] = (user_input, stream_content)
-                    yield history, "GENERATING", topic_state, gr.update(interactive=False)
+                # Second call to the agent to continue the run with refinements
+                report_generator = st.session_state.agent.continue_run(
+                    st.session_state.initial_topic, 
+                    user_refinements=prompt
+                )
                 
-                yield history, "INITIAL", "", gr.update(interactive=True, placeholder="Research complete. What's the next topic?")
-            
-            except Exception as e:
-                error_message = f"An error occurred: {str(e)}"
-                history.append((None, error_message))
-                yield history, "INITIAL", "", gr.update(interactive=True, placeholder="Let's try again. What's the topic?")
+                for msg_type, content in report_generator:
+                    if msg_type == 'log':
+                        log_messages.append(content)
+                        with log_container: st.write("\n".join(log_messages))
+                    
+                    elif msg_type == 'report_start':
+                        full_report = content
+                        with report_container: st.markdown(full_report)
 
-    # --- Event Listeners ---
-    # This section is rewritten to prevent duplicate triggers.
-    
-    init_button.click(
-        fn=handle_initialization,
-        inputs=[google_api_key_input, tavily_api_key_input],
-        outputs=[chatbot, chat_input, submit_button, settings_accordion]
-    )
-    
-    # We define a single submission event and trigger it from both the button and the textbox.
-    # It now calls the wrapper function to handle the processing lock.
-    submit_event = submit_button.click(
-        fn=chat_step_wrapper,
-        inputs=[chat_input, chatbot, agent_state, initial_topic_state],
-        outputs=[chatbot, agent_state, initial_topic_state, chat_input],
-    ).then(
-        fn=lambda: gr.update(value=""),
-        inputs=None,
-        outputs=[chat_input],
-        queue=False
-    )
+                    elif msg_type == 'report_stream':
+                        # This part handles the live-writing effect
+                        full_report = content # The content is the full report so far
+                        with report_container: st.markdown(full_report + "▌") # Add a blinking cursor
 
-    chat_input.submit(
-        fn=chat_step_wrapper,
-        inputs=[chat_input, chatbot, agent_state, initial_topic_state],
-        outputs=[chatbot, agent_state, initial_topic_state, chat_input],
-    ).then(
-        fn=lambda: gr.update(value=""),
-        inputs=None,
-        outputs=[chat_input],
-        queue=False
-    )
+                    elif msg_type == 'report_chunk':
+                        # This replaces the streamed content with the final, formatted chunk
+                        full_report += content
+                        with report_container: st.markdown(full_report)
 
-app.launch(debug=True) 
+                # Final update after the loop
+                with report_container: st.markdown(full_report)
+                st.session_state.messages.append({"role": "assistant", "content": full_report})
+                st.session_state.agent_state = "READY" # Reset for the next query