File size: 4,709 Bytes
fb15b03 cc1ea0b fb15b03 |
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 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 |
# deep_research.py
import gradio as gr
from dotenv import load_dotenv
from clarifier_agent import clarifier_agent
from research_manager import ResearchManagerAgent
from agents import Runner
from collections import defaultdict
from datetime import datetime
import time
import logging
load_dotenv(override=True)
# --- Rate Limiter ---
class RateLimiter:
# Rate limit to 2 requests per minute, 10 requests per day
def __init__(self, max_requests=2, time_window=60, daily_quota=10):
self.max_requests = max_requests
self.time_window = time_window # seconds
self.request_history = defaultdict(list)
self.daily_quota = daily_quota
self.daily_counts = defaultdict(lambda: {'date': self._today(), 'count': 0})
def _today(self):
return datetime.utcnow().strftime('%Y-%m-%d')
def is_rate_limited(self, user_id):
now = time.time()
self.request_history[user_id] = [
t for t in self.request_history[user_id] if now - t < self.time_window
]
if len(self.request_history[user_id]) >= self.max_requests:
return True
self.request_history[user_id].append(now)
return False
def is_quota_exceeded(self, user_id):
today = self._today()
user_quota = self.daily_counts[user_id]
if user_quota['date'] != today:
user_quota['date'] = today
user_quota['count'] = 0
if user_quota['count'] >= self.daily_quota:
return True
user_quota['count'] += 1
self.daily_counts[user_id] = user_quota
return False
rate_limiter = RateLimiter()
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)
async def get_user_id(request: gr.Request = None):
user_id = "default_user"
if request is not None:
try:
forwarded = request.headers.get("X-Forwarded-For")
if forwarded:
user_id = forwarded.split(",")[0].strip()
else:
user_id = getattr(request.client, 'host', 'default_user')
except Exception:
pass
logger.debug(f"[RateLimiter] user_id={user_id}")
return user_id
# Step 1 — Generate clarifying questions
async def get_clarifying_questions(query, request: gr.Request = None):
user_id = await get_user_id(request)
if rate_limiter.is_rate_limited(user_id):
return ["Rate limit exceeded. Please wait a minute."], "", "", ""
if rate_limiter.is_quota_exceeded(user_id):
return ["Daily quota exceeded. Try again tomorrow."], "", "", ""
result = await Runner.run(clarifier_agent, input=query)
return result.final_output.questions
# Step 2 — Run full research pipeline via coordinator agent (handoff style)
async def run_with_handoff(query, q1, q2, q3, a1, a2, a3, request: gr.Request = None):
user_id = await get_user_id(request)
if rate_limiter.is_rate_limited(user_id):
yield "Rate limit exceeded. Please wait a minute."
return
if rate_limiter.is_quota_exceeded(user_id):
yield "You have reached your daily quota. Try again tomorrow."
return
questions = [q1, q2, q3]
answers = [a1, a2, a3]
async for chunk in ResearchManagerAgent().run(
query,
questions,
answers,
):
yield chunk
with gr.Blocks(theme=gr.themes.Default(primary_hue="sky")) as ui:
gr.Markdown("# 🔍 Deep Research Agent (Clarify ➡️ Research ➡️ Email)")
query = gr.Textbox(label="🔎 What would you like to research?")
get_questions_btn = gr.Button("Generate Clarifying Questions", variant="primary")
clar_q1 = gr.Textbox(label="Clarifying Question 1", interactive=False)
clar_q2 = gr.Textbox(label="Clarifying Question 2", interactive=False)
clar_q3 = gr.Textbox(label="Clarifying Question 3", interactive=False)
answer_1 = gr.Textbox(label="Your Answer to Q1")
answer_2 = gr.Textbox(label="Your Answer to Q2")
answer_3 = gr.Textbox(label="Your Answer to Q3")
submit_answers_btn = gr.Button("Submit & Run Full Research")
report = gr.Markdown(label="📄 Research Report")
# Step 1
query.submit(
fn=get_clarifying_questions,
inputs=query,
outputs=[clar_q1, clar_q2, clar_q3]
).then(lambda *_: "", outputs=report)
get_questions_btn.click(
fn=get_clarifying_questions,
inputs=query,
outputs=[clar_q1, clar_q2, clar_q3]
).then(lambda *_: "", outputs=report)
# Step 2
submit_answers_btn.click(
fn=run_with_handoff,
inputs=[query, clar_q1, clar_q2, clar_q3, answer_1, answer_2, answer_3],
outputs=report
)
ui.launch(inbrowser=True)
|