|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import asyncio |
| 4 | + |
| 5 | +from rich.console import Console |
| 6 | + |
| 7 | +from agents import Runner, gen_trace_id, trace |
| 8 | + |
| 9 | +from .agents.contact_finder_agent import ContactInfo, contact_finder_agent |
| 10 | +from .agents.info_agent import LeadInfo, info_agent |
| 11 | +from .printer import Printer |
| 12 | + |
| 13 | + |
| 14 | +class SalesLeadResearcher: |
| 15 | + def __init__(self) -> None: |
| 16 | + self.console = Console() |
| 17 | + self.printer = Printer(self.console) |
| 18 | + |
| 19 | + async def run(self, targets: list[tuple[str, str | None]]) -> list[LeadInfo]: |
| 20 | + trace_id = gen_trace_id() |
| 21 | + with trace("Sales lead research", trace_id=trace_id): |
| 22 | + self.printer.update_item( |
| 23 | + "trace_id", |
| 24 | + f"View trace: https://platform.openai.com/traces/trace?trace_id={trace_id}", |
| 25 | + is_done=True, |
| 26 | + hide_checkmark=True, |
| 27 | + ) |
| 28 | + tasks = [ |
| 29 | + asyncio.create_task(self._process(company, person)) for company, person in targets |
| 30 | + ] |
| 31 | + results: list[LeadInfo] = [] |
| 32 | + num_done = 0 |
| 33 | + for task in asyncio.as_completed(tasks): |
| 34 | + result = await task |
| 35 | + results.append(result) |
| 36 | + num_done += 1 |
| 37 | + self.printer.update_item("progress", f"Processed {num_done}/{len(tasks)} leads") |
| 38 | + self.printer.mark_item_done("progress") |
| 39 | + self.printer.end() |
| 40 | + return results |
| 41 | + |
| 42 | + async def _process(self, company: str, person: str | None) -> LeadInfo: |
| 43 | + if not person: |
| 44 | + result = await Runner.run(contact_finder_agent, company) |
| 45 | + info = result.final_output_as(ContactInfo) |
| 46 | + person = info.full_name |
| 47 | + input_data = f"Person: {person}\nCompany: {company}" |
| 48 | + result = await Runner.run(info_agent, input_data) |
| 49 | + return result.final_output_as(LeadInfo) |
0 commit comments