File size: 2,976 Bytes
c049b12
 
 
 
 
 
 
f3da3e8
 
 
 
 
c049b12
fb1de84
f3da3e8
 
fb1de84
 
f3da3e8
 
 
 
 
 
 
fb1de84
 
f3da3e8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c049b12
 
 
f3da3e8
 
 
fb1de84
f3da3e8
 
c049b12
fb1de84
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
import gradio as gr
from dotenv import load_dotenv

from research_manager import ResearchManager

load_dotenv()

WELCOME_MESSAGE = (
    "Welcome to **Deep Research**!\n"
    "Send me a research topic and I will ask clarifying questions.\n"
    "Answer them here to receive a detailed report."
)


def _format_clarifications(questions: list[str]) -> str:
    return "\n".join(f"{idx + 1}. {q}" for idx, q in enumerate(questions))


def _build_clarification_block(questions: list[str], answers: str) -> str:
    lines = [line.strip() for line in answers.split("\n")]
    while len(lines) < len(questions):
        lines.append("")
    return "\n".join(
        f"Q{idx + 1}: {q}\nA{idx + 1}: {lines[idx]}" for idx, q in enumerate(questions)
    )


async def respond(
    message: str, history: list[tuple[str | None, str | None]], state: dict
):
    if not state:
        state = {"stage": "awaiting_query", "query": "", "questions": []}

    history.append((message, None))

    if state["stage"] == "awaiting_query":
        state["query"] = message
        manager = ResearchManager()
        questions = await manager.get_clarifying_questions(message)
        state["questions"] = questions
        if questions:
            state["stage"] = "awaiting_answers"
            q_text = _format_clarifications(questions)
            history.append(
                (
                    None,
                    f"Please answer the following questions, one per line:\n{q_text}",
                )
            )
            yield history, state
        else:
            state["stage"] = "running"
            history.append((None, ""))
            async for chunk in manager.run(message, ""):
                history[-1] = (None, (history[-1][1] or "") + chunk)
                yield history, state
            state["stage"] = "awaiting_query"
            yield history, state
    elif state["stage"] == "awaiting_answers":
        answers_block = _build_clarification_block(state["questions"], message)
        manager = ResearchManager()
        state["stage"] = "running"
        history.append((None, ""))
        async for chunk in manager.run(state["query"], answers_block):
            history[-1] = (None, (history[-1][1] or "") + chunk)
            yield history, state
        state["stage"] = "awaiting_query"
        yield history, state
    else:
        history.append((None, "Please wait for the current task to finish."))
        yield history, state


def reset():
    return [(None, WELCOME_MESSAGE)], {
        "stage": "awaiting_query",
        "query": "",
        "questions": [],
    }


with gr.Blocks(theme=gr.themes.Default(primary_hue="yellow")) as ui:
    chatbot = gr.Chatbot()
    state = gr.State({})
    msg = gr.Textbox(placeholder="Type your message and press Enter")

    ui.load(fn=reset, outputs=[chatbot, state])
    msg.submit(respond, inputs=[msg, chatbot, state], outputs=[chatbot, state])

if __name__ == "__main__":
    ui.launch()