Update app.py

app.py

@@ -1,142 +1,172 @@
-# app.py (The "Glass Box" UI)
-
-import streamlit as st
+import os
+import gradio as gr
 import google.generativeai as genai
 from tavily import TavilyClient
-from rag_agent import ResearchAgent, AITools # We now import the agent class
+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
 
-# --- Page Configuration ---
-st.set_page_config(page_title="DeepSearch Agent", layout="wide", initial_sidebar_state="expanded")
+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
 
-# --- 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)
+# --- 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>")
 
-# --- 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 = ""
+    agent_state = gr.State("INITIAL") 
+    initial_topic_state = gr.State("")
 
-# --- 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")
+    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)
 
-    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 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)
+        }
 
-# --- 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)
+    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
 
-# Display chat history
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
+        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.")
 
-# Main chat input logic
-if st.session_state.agent_state != "INITIAL":
-    if prompt := st.chat_input("Enter your research topic or answer the questions..."):
+    def chat_step(user_input, history, current_agent_state, topic_state):
+        history = history or []
+        history.append((user_input, None))
         
-        # Add user message to chat
-        st.session_state.messages.append({"role": "user", "content": prompt})
-        with st.chat_message("user"):
-            st.markdown(prompt)
+        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...")
 
-        # 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()
+        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...")
             
-            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"
+            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)
                 
-                # 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
-                )
+                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)
                 
-                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)
+                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?")
 
-                    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
+    # --- 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_chunk':
-                        # This replaces the streamed content with the final, formatted chunk
-                        full_report += content
-                        with report_container: st.markdown(full_report)
+    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
+    )
 
-                # 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
+app.launch(debug=True)