| import json |
| import time |
| from collections.abc import Callable, Iterator |
|
|
| from _core.llm import LLMClient |
| from _core.models import estimate_cost |
| from _core.tracer import Step |
|
|
| PLAN_PROMPT = ( |
| "Break the research topic into 2-4 specific sub-questions that together cover it. " |
| "Output ONLY a JSON array of strings." |
| ) |
| SYNTH_PROMPT = ( |
| "Write a concise research brief answering the topic. Cite sources inline as [n] " |
| "using the provided numbered source list." |
| ) |
|
|
| SearchFn = Callable[[str], list[dict]] |
|
|
|
|
| class DeepResearchAgent: |
| def __init__(self, llm: LLMClient, search_fn: SearchFn, max_subquestions: int = 4): |
| self.llm = llm |
| self.search_fn = search_fn |
| self.max_subquestions = max_subquestions |
|
|
| def _ask(self, system: str, user: str): |
| start = time.monotonic() |
| resp = self.llm.chat( |
| [{"role": "system", "content": system}, {"role": "user", "content": user}] |
| ) |
| latency = int((time.monotonic() - start) * 1000) |
| cost = estimate_cost(self.llm.model, resp.prompt_tokens, resp.completion_tokens) |
| return resp, latency, cost |
|
|
| def run(self, topic: str) -> Iterator[Step]: |
| resp, lat, cost = self._ask(PLAN_PROMPT, topic) |
| try: |
| parsed = json.loads(resp.content or "[]") |
| except json.JSONDecodeError: |
| parsed = [] |
| subqs = parsed[: self.max_subquestions] if isinstance(parsed, list) else [] |
| if not subqs: |
| subqs = [topic] |
| yield Step( |
| kind="thought", |
| content="Sub-questions:\n" + "\n".join(f"- {q}" for q in subqs), |
| tokens=resp.prompt_tokens + resp.completion_tokens, |
| cost_usd=cost, |
| latency_ms=lat, |
| ) |
|
|
| sources: list[dict] = [] |
| for q in subqs: |
| yield Step(kind="action", content=f"search: {q}") |
| results = self.search_fn(q) |
| ids = [] |
| for r in results: |
| sources.append(r) |
| ids.append(len(sources)) |
| yield Step( |
| kind="observation", |
| content="\n".join( |
| f"[{i}] {sources[i - 1].get('title', '')}: {sources[i - 1].get('body', '')}" |
| for i in ids |
| ) |
| or "No results.", |
| ) |
|
|
| source_list = "\n".join( |
| f"[{i + 1}] {s.get('title', '')} — {s.get('url', '')}" for i, s in enumerate(sources) |
| ) |
| fresp, flat, fcost = self._ask( |
| SYNTH_PROMPT, f"Topic: {topic}\nNumbered sources:\n{source_list}" |
| ) |
| yield Step( |
| kind="final", |
| content=(fresp.content or "") + "\n\n### Sources\n" + source_list, |
| tokens=fresp.prompt_tokens + fresp.completion_tokens, |
| cost_usd=fcost, |
| latency_ms=flat, |
| ) |
|
|
|
|
| def make_search_fn(tavily_key: str | None = None, max_results: int = 3) -> SearchFn: |
| if tavily_key: |
| import requests |
|
|
| def search(query: str) -> list[dict]: |
| resp = requests.post( |
| "https://api.tavily.com/search", |
| json={"api_key": tavily_key, "query": query, "max_results": max_results}, |
| timeout=30, |
| ) |
| resp.raise_for_status() |
| data = resp.json() |
| return [ |
| {"title": r.get("title", ""), "body": r.get("content", ""), "url": r.get("url", "")} |
| for r in data.get("results", []) |
| ] |
|
|
| return search |
|
|
| def search(query: str) -> list[dict]: |
| try: |
| from ddgs import DDGS |
| except ImportError: |
| from duckduckgo_search import DDGS |
| out: list[dict] = [] |
| with DDGS() as ddgs: |
| for r in ddgs.text(query, max_results=max_results): |
| out.append( |
| { |
| "title": r.get("title", ""), |
| "body": r.get("body", ""), |
| "url": r.get("href", ""), |
| } |
| ) |
| return out |
|
|
| return search |
|
|
