Tools / Modules /Deep_Research.py
Nymbo's picture
Update Modules/Deep_Research.py
1e7569e verified
raw
history blame
26.8 kB
from __future__ import annotations
import os
import re
import tempfile
import time
from collections import OrderedDict, deque
from concurrent.futures import Future, ThreadPoolExecutor, as_completed
from datetime import datetime
from typing import Annotated, Callable, Dict, List, Tuple
from urllib.parse import urlparse
import gradio as gr
import requests
from bs4 import BeautifulSoup
from ddgs import DDGS
from huggingface_hub import InferenceClient
from .Web_Fetch import _fullpage_markdown_from_soup, _http_get_enhanced
from app import _log_call_end, _log_call_start, _search_rate_limiter, _truncate_for_log
from ._docstrings import autodoc
HF_TEXTGEN_TOKEN = os.getenv("HF_READ_TOKEN") or os.getenv("HF_TOKEN")
# Single source of truth for the LLM-facing tool description
TOOL_SUMMARY = (
"Write a summary of what the user wants to research, and "
"run multiple DuckDuckGo searches (up to 50 max results between all queries), fetch pages, and a Research agent will produce a comprehensive research report with sources; "
"returns (Markdown report, newline-separated source links, downloadable report path). "
"Provide the user with one-paragraph summary of the research report and the txt file in this format `![research_report](URL)`."
)
RESEARCHER_SYSTEM_PROMPT = (
"You are Nymbot, a helpful deep research assistant. You will be asked a Query from a user and you will create a long, comprehensive, well-structured research report in response to the user's Query.\n\n"
"You will receive a summary of the user question, the search queries used, and the fetched webpages. Follow the guidance below when writing the report.\n\n"
"<report_format>\n"
"Write a well-formatted report in the structure of a scientific report to a broad audience. The report must be readable and have a nice flow of Markdown headers and paragraphs of text. Do NOT use bullet points or lists which break up the natural flow. The report must be exhaustive for comprehensive topics.\n"
"For any given user query, first determine the major themes or areas that need investigation, then structure these as main sections, and develop detailed subsections that explore various facets of each theme. Each section and subsection requires paragraphs of texts that need to all connect into one narrative flow.\n"
"</report_format>\n\n"
"<document_structure>\n"
"- Always begin with a clear title using a single # header\n"
"- Organize content into major sections using ## headers\n"
"- Further divide into subsections using ### headers\n"
"- Use #### headers sparingly for special subsections\n"
"- Never skip header levels\n"
"- Write multiple paragraphs per section or subsection\n"
"- Each paragraph must contain at least 4-5 sentences, present novel insights and analysis grounded in source material, connect ideas to original query, and build upon previous paragraphs to create a narrative flow\n"
"- Never use lists, instead always use text or tables\n\n"
"Mandatory Section Flow:\n"
"1. Title (# level)\n - Before writing the main report, start with one detailed paragraph summarizing key findings\n"
"2. Main Body Sections (## level)\n - Each major topic gets its own section (## level). There MUST BE at least 5 sections.\n - Use ### subsections for detailed analysis\n - Every section or subsection needs at least one paragraph of narrative before moving to the next section\n - Do NOT have a section titled \"Main Body Sections\" and instead pick informative section names that convey the theme of the section\n"
"3. Conclusion (## level)\n - Synthesis of findings\n - Potential recommendations or next steps\n"
"</document_structure>\n\n"
"<planning_rules>\n"
"- Always break it down into multiple steps\n"
"- Assess the different sources and whether they are useful for any steps needed to answer the query\n"
"- Create the best report that weighs all the evidence from the sources\n"
"- Use the current date supplied in the first user message to contextualize findings\n"
"- Make sure that your final report addresses all parts of the query\n"
"- Communicate a brief high-level plan in the introduction; do not reveal chain-of-thought.\n"
"- When referencing sources during analysis, you should still refer to them by index with brackets and follow <citations>\n"
"- As a final step, review your planned report structure and ensure it completely answers the query.\n"
"</planning_rules>\n\n"
)
FILTERER_SYSTEM_PROMPT = (
"You are Nymbot Filterer, an analyst who selects the most relevant sources for a research task. "
"You will be given a summary of the research topic (and optional search queries) followed by multiple fetched documents. "
"Each document includes its URL and a truncated excerpt. Evaluate how well each source helps answer the research topic. "
"Return only the URLs that should be used for the final research step. Output plain text with exactly one URL per line and no additional commentary, bullets, numbering, or explanations. "
"If no sources are relevant, return an empty string."
)
class SlowHost(Exception):
pass
def _normalize_query(q: str) -> str:
if not q:
return ""
repl = {"“": '"', "”": '"', "‘": "'", "’": "'", "`": "'"}
for key, value in repl.items():
q = q.replace(key, value)
q = re.sub(r"\s+", " ", q)
q = re.sub(r'"\s+"', " ", q)
q = q.strip().strip('"').strip()
return q
def _search_urls_only(query: str, max_results: int) -> list[str]:
if not query or not query.strip() or max_results <= 0:
return []
urls: list[str] = []
try:
_search_rate_limiter.acquire()
with DDGS() as ddgs:
for item in ddgs.text(query, region="wt-wt", safesearch="moderate", max_results=max_results):
url = (item.get("href") or item.get("url") or "").strip()
if url:
urls.append(url)
except Exception:
pass
seen = set()
deduped = []
for url in urls:
if url not in seen:
seen.add(url)
deduped.append(url)
return deduped
def _fetch_page_markdown_fast(url: str, max_chars: int = 3000, timeout: float = 10.0) -> str:
try:
resp = _http_get_enhanced(url, timeout=timeout, skip_rate_limit=True)
resp.raise_for_status()
except requests.exceptions.RequestException as exc:
msg = str(exc)
if "timed out" in msg.lower():
raise SlowHost(msg) from exc
return ""
final_url = str(resp.url)
ctype = resp.headers.get("Content-Type", "")
if "html" not in ctype.lower():
return ""
resp.encoding = resp.encoding or resp.apparent_encoding
html = resp.text
soup = BeautifulSoup(html, "lxml")
md_text = _fullpage_markdown_from_soup(soup, final_url, "")
if max_chars > 0 and len(md_text) > max_chars:
md_text = md_text[:max_chars]
return md_text
def _truncate_join(parts: List[str], max_chars: int) -> Tuple[str, bool]:
out = []
total = 0
truncated = False
for part in parts:
if not part:
continue
if total + len(part) > max_chars:
out.append(part[: max(0, max_chars - total)])
truncated = True
break
out.append(part)
total += len(part)
return ("\n\n".join(out), truncated)
def _build_research_prompt(summary: str, queries: List[str], url_list: List[str], pages_map: Dict[str, str]) -> str:
sources_blocks: List[str] = []
indexed_urls: List[str] = []
for idx, url in enumerate(url_list, start=1):
text = pages_map.get(url, "").strip()
if not text:
continue
indexed_urls.append(f"[{idx}] {url}")
sources_blocks.append(f"[Source {idx}] URL: {url}\n\n{text}")
sources_joined, truncated = _truncate_join(sources_blocks, max_chars=100_000)
prompt_parts: List[str] = []
prompt_parts.append("<user_query_summary>\n" + (summary or "") + "\n</user_query_summary>\n")
populated = [q for q in queries if q and q.strip()]
if populated:
prompt_parts.append("<search_queries>\n" + "\n".join(f"- {q.strip()}" for q in populated) + "\n</search_queries>\n")
if indexed_urls:
prompt_parts.append("<sources_list>\n" + "\n".join(indexed_urls) + "\n</sources_list>\n")
prompt_parts.append("<fetched_documents>\n" + sources_joined + ("\n\n[NOTE] Sources truncated due to context limits." if truncated else "") + "\n</fetched_documents>")
return "\n\n".join(prompt_parts)
def _build_filter_prompt(summary: str, queries: List[str], pages_map: Dict[str, str]) -> str:
populated = [q for q in queries if q and q.strip()]
summary_text = summary or ""
prompt_sections: List[str] = []
prompt_sections.append("<research_topic_summary>\n" + summary_text + "\n</research_topic_summary>")
if populated:
prompt_sections.append("<search_queries>\n" + "\n".join(populated) + "\n</search_queries>")
sources: List[str] = []
for idx, (url, text) in enumerate(pages_map.items(), start=1):
content = text.strip()
if not content:
continue
sources.append(f"[Source {idx}] URL: {url}\n\n{content}")
sources_joined, truncated = _truncate_join(sources, max_chars=60_000)
prompt_sections.append("<candidate_sources>\n" + sources_joined + ("\n\n[NOTE] Sources truncated due to context limits." if truncated else "") + "\n</candidate_sources>")
prompt_sections.append(
"<task>\nIdentify which of the provided URLs should be retained for the final research synthesis. "
"Consider coverage, credibility, and relevance to the research topic. "
"Return ONLY the URLs you choose, with one URL per line and no additional text.\n</task>"
)
return "\n\n".join(prompt_sections)
def _parse_filterer_output(raw: str, allowed_urls: List[str]) -> List[str]:
if not raw:
return []
allowed_set = {url.strip(): idx for idx, url in enumerate(allowed_urls)}
found_indices: set[int] = set()
for line in raw.splitlines():
candidate = line.strip()
if not candidate:
continue
if candidate in allowed_set:
found_indices.add(allowed_set[candidate])
continue
match = re.search(r"https?://[^\s]+", candidate)
if not match:
continue
url = match.group(0).rstrip(".,);]")
if url in allowed_set:
found_indices.add(allowed_set[url])
selected = [allowed_urls[idx] for idx in sorted(found_indices)]
return selected
def _write_report_tmp(text: str) -> str:
tmp_dir = tempfile.mkdtemp(prefix="deep_research_")
path = os.path.join(tmp_dir, "research_report.txt")
with open(path, "w", encoding="utf-8") as file:
file.write(text)
return path
def _fetch_pages_within_budget(urls: List[str], char_limit: int, time_left_fn: Callable[[], float]) -> OrderedDict:
pages: dict[str, str] = {}
if not urls:
return OrderedDict()
queue = deque(urls)
attempts: dict[str, int] = {url: 0 for url in urls}
max_attempts = 2
max_workers = min(12, max(4, len(urls)))
in_flight: dict[Future, str] = {}
delayed: list[tuple[float, str]] = []
def schedule_next(executor: ThreadPoolExecutor) -> None:
while queue and len(in_flight) < max_workers:
url = queue.popleft()
if url in pages:
continue
attempts.setdefault(url, 0)
if attempts[url] >= max_attempts:
continue
attempts[url] += 1
tl = time_left_fn()
if tl <= 0.1:
return
per_timeout = 10.0 if tl > 15 else (5.0 if tl > 8 else 2.0)
future = executor.submit(_fetch_page_markdown_fast, url, char_limit, per_timeout)
in_flight[future] = url
with ThreadPoolExecutor(max_workers=max_workers) as executor:
schedule_next(executor)
while (in_flight or queue or delayed) and time_left_fn() > 0.2:
now = time.time()
if delayed:
ready: list[tuple[float, str]] = []
not_ready: list[tuple[float, str]] = []
for ready_time, delayed_url in delayed:
(ready if ready_time <= now else not_ready).append((ready_time, delayed_url))
delayed = not_ready
for _, delayed_url in ready:
queue.append(delayed_url)
if ready:
schedule_next(executor)
done = [future for future in list(in_flight.keys()) if future.done()]
if not done:
if not queue and delayed:
next_ready = min((t for t, _ in delayed), default=time.time())
sleep_for = max(0.0, next_ready - time.time())
time.sleep(max(0.02, min(0.25, sleep_for)))
else:
time.sleep(0.05)
continue
for future in done:
url = in_flight.pop(future)
try:
md = future.result()
if md and not md.startswith("Unsupported content type") and not md.startswith("An error occurred"):
pages[url] = md
try:
print(f"[FETCH OK] {url} (chars={len(md)})", flush=True)
except Exception:
pass
except SlowHost:
if time_left_fn() > 5.0:
delayed.append((time.time() + 3.0, url))
except Exception:
pass
schedule_next(executor)
ordered = OrderedDict((url, pages[url]) for url in urls if url in pages)
return ordered
@autodoc(
summary=TOOL_SUMMARY,
)
def Deep_Research(
summary: Annotated[str, "Summarization of research topic (one or more sentences)."],
query1: Annotated[str, "DDG Search Query 1"],
max1: Annotated[int, "Max results for Query 1 (1-50)"] = 10,
query2: Annotated[str, "DDG Search Query 2"] = "",
max2: Annotated[int, "Max results for Query 2 (1-50)"] = 10,
query3: Annotated[str, "DDG Search Query 3"] = "",
max3: Annotated[int, "Max results for Query 3 (1-50)"] = 10,
query4: Annotated[str, "DDG Search Query 4"] = "",
max4: Annotated[int, "Max results for Query 4 (1-50)"] = 10,
query5: Annotated[str, "DDG Search Query 5"] = "",
max5: Annotated[int, "Max results for Query 5 (1-50)"] = 10,
) -> tuple[str, str, str]:
_log_call_start(
"Deep_Research",
summary=_truncate_for_log(summary or "", 200),
queries=[q for q in [query1, query2, query3, query4, query5] if q],
)
if not HF_TEXTGEN_TOKEN:
_log_call_end("Deep_Research", "error=missing HF token")
raise gr.Error("Please provide a `HF_READ_TOKEN` to enable Deep Research.")
queries = [
_normalize_query(query1 or ""),
_normalize_query(query2 or ""),
_normalize_query(query3 or ""),
_normalize_query(query4 or ""),
_normalize_query(query5 or ""),
]
reqs = [
max(1, min(50, int(max1))),
max(1, min(50, int(max2))),
max(1, min(50, int(max3))),
max(1, min(50, int(max4))),
max(1, min(50, int(max5))),
]
total_requested = sum(reqs)
if total_requested > 50:
reqs = [10, 10, 10, 10, 10]
start_ts = time.time()
budget_seconds = 55.0
deadline = start_ts + budget_seconds
def time_left() -> float:
return max(0.0, deadline - time.time())
now_dt = datetime.now().astimezone()
date_str = now_dt.strftime("%A, %B %d, %Y %I:%M %p %Z").strip()
if not date_str:
date_str = now_dt.isoformat()
all_urls: list[str] = []
tasks = []
with ThreadPoolExecutor(max_workers=min(5, sum(1 for q in queries if q.strip())) or 1) as executor:
for query, count in zip(queries, reqs):
if not query.strip():
continue
tasks.append(executor.submit(_search_urls_only, query.strip(), count))
for future in as_completed(tasks):
try:
urls = future.result() or []
except Exception:
urls = []
for url in urls:
if url not in all_urls:
all_urls.append(url)
if len(all_urls) >= 50:
break
if time_left() <= 0.5:
break
if len(all_urls) > 50:
all_urls = all_urls[:50]
blacklist = {
"homedepot.com",
"tractorsupply.com",
"mcmaster.com",
"mrchain.com",
"answers.com",
"city-data.com",
"dictionary.cambridge.org",
}
def _domain(url: str) -> str:
try:
return urlparse(url).netloc.lower()
except Exception:
return ""
all_urls = [url for url in all_urls if _domain(url) not in blacklist]
skip_exts = (
".pdf",
".ppt",
".pptx",
".doc",
".docx",
".xls",
".xlsx",
".zip",
".gz",
".tgz",
".bz2",
".7z",
".rar",
)
def _skip_url(url: str) -> bool:
try:
path = urlparse(url).path.lower()
except Exception:
return False
return any(path.endswith(ext) for ext in skip_exts)
all_urls = [url for url in all_urls if not _skip_url(url)]
truncated_pages = OrderedDict()
if all_urls and time_left() > 0.2:
truncated_pages = _fetch_pages_within_budget(all_urls, 3000, time_left)
print(
f"[PIPELINE] Initial fetch complete: candidates={len(all_urls)}, truncated_documents={len(truncated_pages)}, time_left={time_left():.2f}s",
flush=True,
)
def _invoke_chat(messages, provider: str, max_tokens: int, temp: float, top_p: float):
client = InferenceClient(provider=provider, api_key=HF_TEXTGEN_TOKEN)
return client.chat.completions.create(
model="Qwen/Qwen3-235B-A22B-Thinking-2507",
messages=messages,
max_tokens=max_tokens,
temperature=temp,
top_p=top_p,
)
filtered_urls: List[str] = list(truncated_pages.keys())
filter_output = ""
filter_used_fallback = False
filter_success = False
if truncated_pages and time_left() > 3.0:
filter_prompt = _build_filter_prompt(summary or "", [q for q in queries if q.strip()], truncated_pages)
filter_messages = [
{"role": "system", "content": FILTERER_SYSTEM_PROMPT},
{"role": "user", "content": f"The current date is {date_str}. Consider how recent each source is when deciding relevance."},
{"role": "user", "content": filter_prompt},
]
filter_completion = None
try:
print("[FILTER] Attempt 1: provider=cerebras, max_tokens=2048", flush=True)
filter_completion = _invoke_chat(filter_messages, "cerebras", 2048, 0.2, 0.9)
except Exception as exc1:
print(f"[FILTER] Attempt 1 failed: {str(exc1)[:200]}", flush=True)
try:
print("[FILTER] Attempt 2: provider=auto, max_tokens=2048", flush=True)
filter_completion = _invoke_chat(filter_messages, "auto", 2048, 0.2, 0.9)
except Exception as exc2:
print(f"[FILTER] Attempt 2 failed: {str(exc2)[:200]}", flush=True)
if filter_completion and filter_completion.choices:
filter_output = filter_completion.choices[0].message.content or ""
filtered_urls = _parse_filterer_output(filter_output, list(truncated_pages.keys()))
filter_success = bool(filter_output.strip()) and bool(filtered_urls)
if not filtered_urls:
filter_used_fallback = True
fallback_count = min(8, len(truncated_pages))
filtered_urls = list(truncated_pages.keys())[:fallback_count]
max_final_urls = 20
if len(filtered_urls) > max_final_urls:
filter_used_fallback = True
filtered_urls = filtered_urls[:max_final_urls]
if not filter_success:
filter_used_fallback = True
print(
f"[FILTER] Selected URLs={len(filtered_urls)}, fallback={filter_used_fallback}, time_left={time_left():.2f}s",
flush=True,
)
final_pages_fetched = OrderedDict()
if filtered_urls and time_left() > 0.2:
final_pages_fetched = _fetch_pages_within_budget(filtered_urls, 8000, time_left)
merged_pages = OrderedDict()
for url in filtered_urls:
content = final_pages_fetched.get(url) or truncated_pages.get(url) or ""
if content:
merged_pages[url] = content
pages = merged_pages
print(
f"[PIPELINE] Final fetch complete: retained_documents={len(pages)}, time_left={time_left():.2f}s",
flush=True,
)
prompt = _build_research_prompt(summary=summary or "", queries=[q for q in queries if q.strip()], url_list=list(pages.keys()), pages_map=pages)
system_message = {"role": "system", "content": RESEARCHER_SYSTEM_PROMPT}
date_message = {"role": "user", "content": f"The current date is {date_str}. Return only the research report."}
messages = [
system_message,
date_message,
{"role": "user", "content": prompt},
]
try:
prompt_chars = len(prompt)
except Exception:
prompt_chars = -1
print(f"[PIPELINE] Fetch complete: pages={len(pages)}, unique_urls={len(pages.keys())}, prompt_chars={prompt_chars}", flush=True)
print("[PIPELINE] Starting inference (provider=cerebras, model=Qwen/Qwen3-235B-A22B-Thinking-2507)", flush=True)
try:
print("[LLM] Attempt 1: provider=cerebras, max_tokens=32768", flush=True)
completion = _invoke_chat(messages, "cerebras", max_tokens=32768, temp=0.3, top_p=0.95)
except Exception as exc1:
print(f"[LLM] Attempt 1 failed: {str(exc1)[:200]}", flush=True)
try:
prompt2 = _build_research_prompt(
summary=summary or "",
queries=[q for q in queries if q.strip()],
url_list=list(pages.keys())[:30],
pages_map={key: pages[key] for key in list(pages.keys())[:30]},
)
messages = [
system_message,
date_message,
{"role": "user", "content": prompt2},
]
print("[LLM] Attempt 2: provider=cerebras (trimmed), max_tokens=16384", flush=True)
completion = _invoke_chat(messages, "cerebras", max_tokens=16384, temp=0.7, top_p=0.95)
except Exception as exc2:
print(f"[LLM] Attempt 2 failed: {str(exc2)[:200]}", flush=True)
try:
print("[LLM] Attempt 3: provider=auto, max_tokens=8192", flush=True)
completion = _invoke_chat(messages, "auto", max_tokens=8192, temp=0.7, top_p=0.95)
except Exception as exc3:
_log_call_end("Deep_Research", f"error={_truncate_for_log(str(exc3), 260)}")
raise gr.Error(f"Researcher model call failed: {exc3}")
raw = completion.choices[0].message.content or ""
try:
no_think = re.sub(r"<think>[\s\S]*?<\\/think>", "", raw, flags=re.IGNORECASE)
no_think = re.sub(r"<\\/?think>", "", no_think, flags=re.IGNORECASE)
except Exception:
no_think = raw
try:
paragraphs = [p for p in re.split(r"\n\s*\n", no_think) if p.strip()]
keep: List[str] = []
removed = 0
planning_re = re.compile(r"\b(let me|now i(?:'ll| will)?|first,|i will now|i will|i'll|let's|now let me|i need to|now i'll|now i will)\b", re.IGNORECASE)
for paragraph in paragraphs:
if planning_re.search(paragraph):
removed += 1
continue
keep.append(paragraph)
report = "\n\n".join(keep).strip()
if not report:
report = no_think.strip()
except Exception:
report = no_think
removed = 0
report = re.sub(r"\n\s*\n\s*\n+", "\n\n", report)
try:
print(f"[POSTPROCESS] removed_planning_paragraphs={removed}, raw_chars={len(raw)}, final_chars={len(report)}", flush=True)
except Exception:
pass
links_text = "\n".join([f"[{i+1}] {url}" for i, url in enumerate(pages.keys())])
if links_text:
sources_section = "\n\n## Sources\n" + "\n".join([f"[{i+1}] {url}" for i, url in enumerate(pages.keys())])
report = report.rstrip() + sources_section
file_path = _write_report_tmp(report)
elapsed = time.time() - start_ts
print(f"[TIMING] Deep_Research elapsed: {elapsed:.2f}s", flush=True)
_log_call_end("Deep_Research", f"urls={len(pages)} file={os.path.basename(file_path)} duration={elapsed:.2f}s")
return report, links_text, file_path
def build_interface() -> gr.Interface:
return gr.Interface(
fn=Deep_Research,
inputs=[
gr.Textbox(label="Summarization of research topic", lines=3, placeholder="Briefly summarize the research topic or user question"),
gr.Textbox(label="DDG Search Query 1", max_lines=1),
gr.Slider(1, 50, value=10, step=1, label="Max results (Q1)"),
gr.Textbox(label="DDG Search Query 2", value="", max_lines=1),
gr.Slider(1, 50, value=10, step=1, label="Max results (Q2)"),
gr.Textbox(label="DDG Search Query 3", value="", max_lines=1),
gr.Slider(1, 50, value=10, step=1, label="Max results (Q3)"),
gr.Textbox(label="DDG Search Query 4", value="", max_lines=1),
gr.Slider(1, 50, value=10, step=1, label="Max results (Q4)"),
gr.Textbox(label="DDG Search Query 5", value="", max_lines=1),
gr.Slider(1, 50, value=10, step=1, label="Max results (Q5)"),
],
outputs=[
gr.Markdown(label="Research Report"),
gr.Textbox(label="Fetched Links", lines=8),
gr.File(label="Download Research Report", file_count="single"),
],
title="Deep Research",
description=(
"<div style=\"text-align:center\">Perform multi-query web research: search with DuckDuckGo, fetch up to 50 pages in parallel, "
"and generate a comprehensive report using a large LLM via Hugging Face Inference Providers (Cerebras). Requires HF_READ_TOKEN.</div>"
),
api_description=TOOL_SUMMARY,
flagging_mode="never",
show_api=True,
)
__all__ = ["Deep_Research", "build_interface"]