wehub-resource-sync
266 lines
10 KiB
Python
266 lines
10 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import time
|
||
from collections.abc import Sequence
|
||
from datetime import datetime, timezone
|
||
|
||
from pydantic import BaseModel
|
||
from rich.console import Console
|
||
|
||
from agents import Runner, RunResult, RunResultStreaming, custom_span, gen_trace_id, trace
|
||
from examples.web_search_utils import extract_url_citations, extract_web_search_source_urls
|
||
|
||
from .agents.financials_agent import financials_agent
|
||
from .agents.planner_agent import FinancialSearchItem, FinancialSearchPlan, planner_agent
|
||
from .agents.risk_agent import risk_agent
|
||
from .agents.search_agent import FinancialSearchSummary, search_agent
|
||
from .agents.verifier_agent import VerificationResult, verifier_agent
|
||
from .agents.writer_agent import REVISION_PROMPT, FinancialReportData, writer_agent
|
||
from .printer import Printer
|
||
|
||
|
||
class FinancialSource(BaseModel):
|
||
title: str
|
||
url: str
|
||
|
||
|
||
class FinancialSearchEvidence(BaseModel):
|
||
query: str
|
||
reason: str
|
||
summary: str
|
||
sources: list[FinancialSource]
|
||
retrieved_at: str
|
||
|
||
|
||
def _extract_financial_sources(items: Sequence[object]) -> list[FinancialSource]:
|
||
sources: list[FinancialSource] = []
|
||
seen: set[str] = set()
|
||
|
||
for citation in extract_url_citations(items):
|
||
if citation.url in seen:
|
||
continue
|
||
seen.add(citation.url)
|
||
sources.append(FinancialSource(title=citation.title, url=citation.url))
|
||
|
||
for url in extract_web_search_source_urls(items):
|
||
if url in seen:
|
||
continue
|
||
seen.add(url)
|
||
sources.append(FinancialSource(title=url, url=url))
|
||
|
||
return sources
|
||
|
||
|
||
async def _summary_extractor(run_result: RunResult | RunResultStreaming) -> str:
|
||
"""Custom output extractor for sub‑agents that return an AnalysisSummary."""
|
||
# The financial/risk analyst agents emit an AnalysisSummary with a `summary` field.
|
||
# We want the tool call to return just that summary text so the writer can drop it inline.
|
||
return str(run_result.final_output.summary)
|
||
|
||
|
||
class FinancialResearchManager:
|
||
"""
|
||
Orchestrates the full flow: planning, searching, sub‑analysis, writing, and verification.
|
||
"""
|
||
|
||
def __init__(self) -> None:
|
||
self.console = Console()
|
||
self.printer = Printer(self.console)
|
||
self.research_cutoff = datetime.now(timezone.utc).date().isoformat()
|
||
|
||
async def run(self, query: str) -> None:
|
||
trace_id = gen_trace_id()
|
||
try:
|
||
with trace("Financial research trace", trace_id=trace_id):
|
||
self.printer.update_item(
|
||
"trace_id",
|
||
f"View trace: https://platform.openai.com/traces/trace?trace_id={trace_id}",
|
||
is_done=True,
|
||
hide_checkmark=True,
|
||
)
|
||
self.printer.update_item("start", "Starting financial research...", is_done=True)
|
||
search_plan = await self._plan_searches(query)
|
||
search_results = await self._perform_searches(search_plan)
|
||
report, verification = await self._produce_verified_report(query, search_results)
|
||
|
||
final_report = f"Report summary\n\n{report.short_summary}"
|
||
self.printer.update_item("final_report", final_report, is_done=True)
|
||
finally:
|
||
self.printer.end()
|
||
|
||
# Print to stdout
|
||
print("\n\n=====REPORT=====\n\n")
|
||
print(f"Report:\n{report.markdown_report}")
|
||
print("\n\n=====FOLLOW UP QUESTIONS=====\n\n")
|
||
print("\n".join(report.follow_up_questions))
|
||
print("\n\n=====VERIFICATION=====\n\n")
|
||
print(verification)
|
||
|
||
async def _produce_verified_report(
|
||
self,
|
||
query: str,
|
||
search_results: Sequence[FinancialSearchEvidence],
|
||
) -> tuple[FinancialReportData, VerificationResult]:
|
||
report = await self._write_report(query, search_results)
|
||
verification = await self._verify_report(query, report, search_results)
|
||
if verification.verified:
|
||
return report, verification
|
||
|
||
report = await self._revise_report(query, report, search_results, verification)
|
||
verification = await self._verify_report(query, report, search_results)
|
||
if not verification.verified:
|
||
raise RuntimeError(
|
||
"Financial report failed evidence verification after one revision: "
|
||
f"{verification.model_dump_json()}"
|
||
)
|
||
return report, verification
|
||
|
||
async def _plan_searches(self, query: str) -> FinancialSearchPlan:
|
||
self.printer.update_item("planning", "Planning searches...")
|
||
result = await Runner.run(planner_agent, f"Query: {query}")
|
||
self.printer.update_item(
|
||
"planning",
|
||
f"Will perform {len(result.final_output.searches)} searches",
|
||
is_done=True,
|
||
)
|
||
return result.final_output_as(FinancialSearchPlan)
|
||
|
||
async def _perform_searches(
|
||
self, search_plan: FinancialSearchPlan
|
||
) -> Sequence[FinancialSearchEvidence]:
|
||
with custom_span("Search the web"):
|
||
self.printer.update_item("searching", "Searching...")
|
||
tasks = [asyncio.create_task(self._search(item)) for item in search_plan.searches]
|
||
results: list[FinancialSearchEvidence] = []
|
||
num_completed = 0
|
||
num_succeeded = 0
|
||
num_failed = 0
|
||
for task in asyncio.as_completed(tasks):
|
||
result = await task
|
||
if result is not None:
|
||
results.append(result)
|
||
num_succeeded += 1
|
||
else:
|
||
num_failed += 1
|
||
num_completed += 1
|
||
status = f"Searching... {num_completed}/{len(tasks)} finished"
|
||
if num_failed:
|
||
status += f" ({num_succeeded} succeeded, {num_failed} failed)"
|
||
self.printer.update_item(
|
||
"searching",
|
||
status,
|
||
)
|
||
summary = f"Searches finished: {num_succeeded}/{len(tasks)} succeeded"
|
||
if num_failed:
|
||
summary += f", {num_failed} failed"
|
||
self.printer.update_item("searching", summary, is_done=True)
|
||
return results
|
||
|
||
async def _search(self, item: FinancialSearchItem) -> FinancialSearchEvidence | None:
|
||
input_data = f"Search term: {item.query}\nReason: {item.reason}"
|
||
try:
|
||
result = await Runner.run(search_agent, input_data)
|
||
search_summary = result.final_output_as(FinancialSearchSummary)
|
||
sources = _extract_financial_sources(result.new_items)
|
||
if not sources:
|
||
return None
|
||
return FinancialSearchEvidence(
|
||
query=item.query,
|
||
reason=item.reason,
|
||
summary=search_summary.summary,
|
||
sources=sources,
|
||
retrieved_at=self.research_cutoff,
|
||
)
|
||
except Exception:
|
||
return None
|
||
|
||
async def _write_report(
|
||
self,
|
||
query: str,
|
||
search_results: Sequence[FinancialSearchEvidence],
|
||
) -> FinancialReportData:
|
||
# Expose the specialist analysts as tools so the writer can invoke them inline
|
||
# and still produce the final FinancialReportData output.
|
||
fundamentals_tool = financials_agent.as_tool(
|
||
tool_name="fundamentals_analysis",
|
||
tool_description="Use to get a short write‑up of key financial metrics",
|
||
custom_output_extractor=_summary_extractor,
|
||
)
|
||
risk_tool = risk_agent.as_tool(
|
||
tool_name="risk_analysis",
|
||
tool_description="Use to get a short write‑up of potential red flags",
|
||
custom_output_extractor=_summary_extractor,
|
||
)
|
||
writer_with_tools = writer_agent.clone(tools=[fundamentals_tool, risk_tool])
|
||
self.printer.update_item("writing", "Thinking about report...")
|
||
input_data = self._report_input(query, search_results)
|
||
result = Runner.run_streamed(writer_with_tools, input_data)
|
||
update_messages = [
|
||
"Planning report structure...",
|
||
"Writing sections...",
|
||
"Finalizing report...",
|
||
]
|
||
last_update = time.time()
|
||
next_message = 0
|
||
async for _ in result.stream_events():
|
||
if time.time() - last_update > 5 and next_message < len(update_messages):
|
||
self.printer.update_item("writing", update_messages[next_message])
|
||
next_message += 1
|
||
last_update = time.time()
|
||
self.printer.mark_item_done("writing")
|
||
return result.final_output_as(FinancialReportData)
|
||
|
||
async def _revise_report(
|
||
self,
|
||
query: str,
|
||
report: FinancialReportData,
|
||
search_results: Sequence[FinancialSearchEvidence],
|
||
verification: VerificationResult,
|
||
) -> FinancialReportData:
|
||
self.printer.update_item("revising", "Revising report from verification feedback...")
|
||
revision_agent = writer_agent.clone(instructions=REVISION_PROMPT)
|
||
input_data = (
|
||
f"{self._report_input(query, search_results)}\n"
|
||
f"Existing report:\n{report.model_dump_json()}\n"
|
||
f"Verification feedback:\n{verification.model_dump_json()}"
|
||
)
|
||
result = await Runner.run(revision_agent, input_data)
|
||
self.printer.mark_item_done("revising")
|
||
return result.final_output_as(FinancialReportData)
|
||
|
||
async def _verify_report(
|
||
self,
|
||
query: str,
|
||
report: FinancialReportData,
|
||
search_results: Sequence[FinancialSearchEvidence],
|
||
) -> VerificationResult:
|
||
self.printer.update_item("verifying", "Verifying report...")
|
||
input_data = json.dumps(
|
||
{
|
||
"original_query": query,
|
||
"research_cutoff": self.research_cutoff,
|
||
"report": report.model_dump(mode="json"),
|
||
"evidence": [item.model_dump(mode="json") for item in search_results],
|
||
},
|
||
ensure_ascii=False,
|
||
)
|
||
result = await Runner.run(verifier_agent, input_data)
|
||
self.printer.mark_item_done("verifying")
|
||
return result.final_output_as(VerificationResult)
|
||
|
||
def _report_input(
|
||
self,
|
||
query: str,
|
||
search_results: Sequence[FinancialSearchEvidence],
|
||
) -> str:
|
||
return json.dumps(
|
||
{
|
||
"original_query": query,
|
||
"research_cutoff": self.research_cutoff,
|
||
"evidence": [item.model_dump(mode="json") for item in search_results],
|
||
},
|
||
ensure_ascii=False,
|
||
)
|