models.py

"""
Shared data structures used across agents, verification, and orchestration.
"""

from __future__ import annotations

from datetime import datetime, timezone
from typing import List, Sequence

from pydantic import BaseModel, Field


def utc_now() -> datetime:
    """Return a UTC timestamp for agent outputs."""
    return datetime.now(timezone.utc)


class AgentMessage(BaseModel):
    """Base payload returned by agents with minimal metadata."""

    agent_id: str
    confidence: float
    status: str = "success"
    timestamp: datetime = Field(default_factory=utc_now)


class ResearchOutput(AgentMessage):
    """Structured research response with raw findings."""

    query: str
    findings: List[str] = Field(default_factory=list)
    sources: List[str] = Field(default_factory=list)


class AnalysisOutput(AgentMessage):
    """Analysis distilled from research findings."""

    key_points: List[str] = Field(default_factory=list)
    risks: List[str] = Field(default_factory=list)
    stance: str


class SummaryOutput(AgentMessage):
    """Concise narrative summary with surfaced highlights."""

    summary: str
    highlights: List[str] = Field(default_factory=list)


class VerificationResult(AgentMessage):
    """Cross-check outcome aggregated across verifiers."""

    passed: bool
    score: float
    confidence: float
    issues: List[str] = Field(default_factory=list)
    disagreement: bool
    verifier_results: List[dict] = Field(default_factory=list)
    status: str = "verified"


AgentTrace = List[dict[str, str]]
AgentOutputs = Sequence[AgentMessage]


class TaskContext(BaseModel):
    """Lightweight task context propagated to verifiers."""

    query: str
    agent_trace: AgentTrace = Field(default_factory=list)