pip install perplexityai pydantic

import json
import argparse
from typing import List, Optional
from pydantic import BaseModel
from perplexity import Perplexity


class ReportSource(BaseModel):
    title: str
    url: str
    relevance: str


class ReportSection(BaseModel):
    heading: str
    content: str
    confidence: float
    sources: List[ReportSource]


class ResearchReport(BaseModel):
    title: str
    summary: str
    sections: List[ReportSection]
    conclusion: str
    overall_confidence: float
    total_sources: int


def run_research(topic: str, model: Optional[str] = None) -> ResearchReport:
    """Conduct deep research on a topic and return a structured report."""
    client = Perplexity()

    params = {
        "preset": "deep-research",
        "input": (
            f"Conduct thorough research on the following topic and produce a "
            f"detailed report with multiple sections, cited sources, and "
            f"confidence scores for each section.\n\nTopic: {topic}"
        ),
        "response_format": {
            "type": "json_schema",
            "json_schema": {
                "name": "research_report",
                "schema": ResearchReport.model_json_schema(),
            },
        },
    }

    if model:
        params["model"] = model

    response = client.responses.create(**params)
    return ResearchReport.model_validate_json(response.output_text)


def format_report(report: ResearchReport) -> str:
    """Format a ResearchReport into human-readable text."""
    lines = [f"{'=' * 60}", f"RESEARCH REPORT: {report.title}", f"{'=' * 60}", ""]
    lines += [f"SUMMARY:", report.summary, ""]

    for i, section in enumerate(report.sections, 1):
        lines.append(f"--- Section {i}: {section.heading} ---")
        lines.append(f"Confidence: {section.confidence:.0%}\n")
        lines.append(section.content)
        if section.sources:
            lines.append("\nSources:")
            for src in section.sources:
                lines.append(f"  - {src.title} ({src.relevance})")
                lines.append(f"    {src.url}")
        lines.append("")

    lines += [f"{'=' * 60}", "CONCLUSION:", report.conclusion, ""]
    lines += [f"Overall Confidence: {report.overall_confidence:.0%}"]
    lines += [f"Total Sources: {report.total_sources}", f"{'=' * 60}"]
    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(description="Agent Research Assistant")
    parser.add_argument("topic", help="The research topic")
    parser.add_argument("--model", help="Override the default model", default=None)
    parser.add_argument("--json", action="store_true", help="Output raw JSON")
    args = parser.parse_args()

    print(f"Researching: {args.topic}")
    print("This may take a moment (deep research uses multi-step reasoning)...\n")

    report = run_research(args.topic, model=args.model)

    if args.json:
        print(json.dumps(report.model_dump(), indent=2))
    else:
        print(format_report(report))


if __name__ == "__main__":
    main()