Quality pipeline

from typing import List, Optional

from agno.agent import Agent
from agno.models.anthropic import Claude
from agno.models.google import Gemini
from pydantic import BaseModel, Field


class Contact(BaseModel):
    name: Optional[str] = None
    email: Optional[str] = None
    company: Optional[str] = None


class FieldDisagreement(BaseModel):
    field: str = Field(..., description="Top-level Contact field name")
    value_a: Optional[str] = None
    value_b: Optional[str] = None
    reason: str = Field(..., description="Why this field needs adjudication")


class DisagreementReport(BaseModel):
    disagreements: List[FieldDisagreement] = Field(default_factory=list)
    needs_adjudication: bool = Field(..., description="True if any field disagrees")


class FinalLabel(BaseModel):
    contact: Contact
    notes: Optional[str] = None


LABELER = "Extract contact info. Use exactly what the text shows. Null if missing."

labeler_a = Agent(model=Gemini(id="gemini-3.5-flash"), instructions=LABELER, output_schema=Contact)
labeler_b = Agent(model=Claude(id="claude-opus-4-7"), instructions=LABELER, output_schema=Contact)

reviewer = Agent(
    model=Claude(id="claude-opus-4-7"),
    instructions=(
        "Compare two labelers' Contact outputs field by field. A field "
        "needs adjudication when both are non-null but differ. Set "
        "needs_adjudication=true if any field does."
    ),
    output_schema=DisagreementReport,
)

adjudicator = Agent(
    model=Claude(id="claude-opus-4-7"),
    instructions=(
        "Re-read the original text and resolve every reported "
        "disagreement. Return a FinalLabel with the correct values."
    ),
    output_schema=FinalLabel,
)


def label_with_quality_review(text: str) -> FinalLabel:
    a = labeler_a.run(text).content
    b = labeler_b.run(text).content

    report = reviewer.run(
        f"Labeler A:\n{a.model_dump_json()}\n\nLabeler B:\n{b.model_dump_json()}"
    ).content

    if not report.needs_adjudication:
        return FinalLabel(contact=a, notes="Labelers agreed.")

    return adjudicator.run(
        f"Original input:\n{text}\n\n"
        f"Labeler A:\n{a.model_dump_json()}\n\n"
        f"Labeler B:\n{b.model_dump_json()}\n\n"
        f"Reviewer report:\n{report.model_dump_json()}"
    ).content