Lab — Orchestration Workshop

import asyncio
import os
from dataclasses import dataclass, field
from typing import Literal

from dotenv import load_dotenv
from pydantic import BaseModel, Field
from pydantic_ai import Agent, ModelRetry, RunContext
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.openai import OpenAIProvider

load_dotenv()

PROXY_URL = "https://litellm.6640.ucf.spencerlyon.com"


def get_model(model_name: str) -> OpenAIChatModel:
    """Create a model connection through our LiteLLM proxy."""
    return OpenAIChatModel(
        model_name,
        provider=OpenAIProvider(
            base_url=PROXY_URL,
            api_key=os.environ["CAP6640_API_KEY"],
        ),
    )

import logfire

logfire.configure(send_to_logfire=False)
logfire.instrument_pydantic_ai()

# Seeded fact base — topic keyword → list of bullet-point facts with "source".
FACT_DB = {
    "rag": [
        ("Lewis et al. 2020 introduced RAG at NeurIPS, combining DPR retrieval with BART generation.", "Lewis 2020"),
        ("Modern RAG uses hybrid search (BM25 + dense) and a cross-encoder reranker — 'retrieve wide, rerank narrow'.", "HF course Ch. 5"),
        ("Chunking strategy accounts for ~80% of RAG retrieval failures in production.", "industry survey 2025"),
        ("RAGAS-style metrics score faithfulness, answer relevance, and context precision/recall.", "RAGAS docs"),
    ],
    "agents": [
        ("An agent is an LLM with access to tools, running in a loop.", "Anthropic 2024"),
        ("Orchestrator-workers is 90% faster than single-agent for parallelizable research.", "Anthropic 2025"),
        ("Claude Code uses a single main thread with subagents depth ≤ 1 for tightly-coupled coding tasks.", "MinusX 2025"),
        ("Cemri et al. 2025 found multi-agent systems correct only ~25% of the time across 5 frameworks.", "Cemri 2025"),
    ],
    "evaluation": [
        ("BLEU measures n-gram precision; ROUGE measures n-gram recall; BERTScore measures semantic similarity.", "J&M Ch. 12"),
        ("pydantic-evals ships LLMJudge for rubric-based evaluation with score vs. assertion modes.", "PydanticAI docs"),
        ("Goodhart's Law: when a measure becomes a target, it ceases to be a good measure.", "Strathern 1997"),
    ],
}


def lookup_facts(topic_keyword: str) -> list[tuple[str, str]]:
    """Return seeded facts for a topic keyword, or [] if unknown."""
    return FACT_DB.get(topic_keyword.lower().strip(), [])


# Quick sanity check
print(f"'rag' facts: {len(lookup_facts('rag'))}")
print(f"'agents' facts: {len(lookup_facts('agents'))}")
print(f"'unicorns' facts: {len(lookup_facts('unicorns'))}")

@dataclass
class BriefingDeps:
    """External state for the briefing workflow."""
    fact_db: dict = field(default_factory=lambda: FACT_DB)
    allowed_destinations: tuple[str, ...] = ("team-channel", "archive")
    max_revisions: int = 3
    quality_threshold: float = 0.75

class ResearchFinding(BaseModel):
    """One evidence-grounded answer to a sub-question."""
    sub_question: str
    answer: str = Field(description="2-3 sentence answer grounded in the retrieved facts.")
    citations: list[str] = Field(default_factory=list)


research_subagent = Agent(
    get_model("claude-haiku-4-5"),
    deps_type=BriefingDeps,
    output_type=ResearchFinding,
    system_prompt=(
        "You are a research specialist. Answer ONE focused sub-question by calling "
        "search_facts for the relevant topic keyword. Ground every claim in a cited fact. "
        "If no facts are returned, say so honestly."
    ),
)


@research_subagent.tool
def search_facts(ctx: RunContext[BriefingDeps], topic_keyword: str) -> str:
    """Look up evidence facts for a topic keyword.

    topic_keyword should be a single lowercase word like 'rag', 'agents', 'evaluation'.
    """
    facts = ctx.deps.fact_db.get(topic_keyword.lower().strip(), [])
    if not facts:
        raise ModelRetry(
            f"No facts found for '{topic_keyword}'. "
            f"Available keywords: {list(ctx.deps.fact_db.keys())}"
        )
    return "\n".join(f"- {fact} (source: {src})" for fact, src in facts)

class ResearchPlan(BaseModel):
    """Decomposition of a topic into sub-questions."""
    topic: str
    sub_questions: list[str] = Field(
        min_length=2,
        max_length=5,
        description="3-5 focused questions that together cover the topic.",
    )


planner = Agent(
    get_model("claude-haiku-4-5"),
    output_type=ResearchPlan,
    system_prompt=(
        "You plan research on a topic by decomposing it into 3-5 focused sub-questions. "
        "Each sub-question should be answerable by looking up one topic keyword from this list: "
        "rag, agents, evaluation. Keep sub-questions specific and non-overlapping."
    ),
)

class Briefing(BaseModel):
    """The output artifact of the workflow."""
    topic: str
    summary: str = Field(description="3-5 sentence executive summary.")
    key_findings: list[str] = Field(description="Bullet-point key findings with citations inline.")
    open_questions: list[str] = Field(description="What the briefing does NOT answer.")


synthesizer = Agent(
    get_model("claude-haiku-4-5"),
    output_type=Briefing,
    system_prompt=(
        "You are a technical writer. Combine a set of research findings into a clean briefing "
        "with a short summary, 3-5 key findings (each citing a source), and a list of open questions. "
        "Do NOT introduce claims not present in the findings."
    ),
)

async def research_one(sub_q: str, deps: BriefingDeps) -> ResearchFinding:
    """Run the research subagent on one sub-question."""
    result = await research_subagent.run(sub_q, deps=deps)
    return result.output


async def generate_briefing_v1(topic: str, deps: BriefingDeps) -> Briefing:
    """Version 1: plan → research in parallel → synthesize. No evaluator yet."""
    # Step 1: plan
    plan_result = await planner.run(f"Topic: {topic}")
    plan = plan_result.output
    print(f"Plan: {len(plan.sub_questions)} sub-questions")

    # Step 2: parallel research (orchestrator-workers pattern)
    findings = await asyncio.gather(
        *[research_one(sq, deps) for sq in plan.sub_questions]
    )

    # Step 3: synthesize
    findings_text = "\n\n".join(
        f"Q: {f.sub_question}\nA: {f.answer}\nCitations: {', '.join(f.citations)}"
        for f in findings
    )
    synth_result = await synthesizer.run(
        f"Topic: {topic}\n\nFindings:\n{findings_text}"
    )
    return synth_result.output

deps = BriefingDeps()

briefing = await generate_briefing_v1("the current state of RAG and agents", deps)

print(f"\nTopic: {briefing.topic}")
print(f"\nSummary:\n  {briefing.summary}")
print(f"\nKey findings:")
for f in briefing.key_findings:
    print(f"  - {f}")
print(f"\nOpen questions:")
for q in briefing.open_questions:
    print(f"  - {q}")