deep_research/deep_research_JD.py

import gradio as gr
from dotenv import load_dotenv
from research_manager_JD import ResearchManager
from planner_agent_JD import get_clarification_questions, refine_query_with_user_input

load_dotenv(override=True)

manager = ResearchManager()

with gr.Blocks(theme=gr.themes.Default(primary_hue="sky")) as ui:
    gr.Markdown("# 🔍 Deep Research Assistant")

    chatbot = gr.Chatbot(label="Deep Research Chat")
    user_input = gr.Textbox(placeholder="Ask a research question...", show_label=False, lines=1)

    query_state = gr.State()         # Stores the original query (if clarification needed)
    clarification_mode = gr.State(value=False)  # Whether awaiting clarification
    message_history = gr.State([])   # Keeps chat history

    # Main chat handler
    async def chat_handler(message, query, awaiting_clarification, history):
        history = history or []

        # Initial user message – run clarification check
        if not query:
            clarification = await get_clarification_questions(message)
            if clarification.lower().strip() == "no clarification needed.":
                response = ""
                async for chunk in manager.run(message):
                    response = chunk
                history.append((message, response))
                return history, None, False, history, gr.update(value="")
            else:
                history.append((message, clarification))
                return history, message, True, history, gr.update(value="")

        # If in clarification mode, treat input as clarification answer
        if awaiting_clarification:
            refined = await refine_query_with_user_input(query, message)
            response = ""
            async for chunk in manager.run(refined):
                response = chunk
            history.append((message, response))
            return history, None, False, history, gr.update(value="")

        # Fallback
        history.append((message, "Hmm, something went wrong."))
        return history, None, False, history, gr.update(value="")

    user_input.submit(
        fn=chat_handler,
        inputs=[user_input, query_state, clarification_mode, message_history],
        outputs=[chatbot, query_state, clarification_mode, message_history, user_input]
    )

ui.launch(inbrowser=True)