packages/core/src/graph/state.py

# ============================================================================
# agent verf - Verification State
# Version: 0.1.0
# Last Updated: 2026-01-13
#
# The state schema that flows through the LangGraph verification workflow.
#
# State is immutable between nodes - each node receives the current state
# and returns updates to be merged.
#
# Key sections:
#   - Input: What the user submitted
#   - Content: Extracted content from URL
#   - Triage: Routing decisions
#   - Evidence: Gathered sources and findings
#   - Verdict: Final output
#   - Metadata: Processing info (timing, costs, errors)
# ============================================================================

from typing import Optional, List, Literal, Annotated
from datetime import datetime
from uuid import UUID, uuid4
from pydantic import BaseModel, Field
import operator

from src.models.content import Content, ContentType, Platform
from src.models.evidence import Evidence, Source
from src.models.verdict import Verdict, VerdictStatus
from src.utils.llm import VerificationMode


# ============================================================================
# Lens Types
# ============================================================================

LensType = Literal[
    "headline_checker",   # News/factual claims
    "guru_auditor",       # Financial/guru claims
    "narrative_tracker",  # Origin/spread tracking
    "general",            # Catch-all
]


# ============================================================================
# Triage Decision
# ============================================================================

class TriageDecision(BaseModel):
    """
    Output from the triage node.

    Determines which lens handles verification and why.
    """
    selected_lens: LensType = Field(
        ...,
        description="Which lens should handle this verification"
    )
    confidence: float = Field(
        default=0.8,
        ge=0.0,
        le=1.0,
        description="Confidence in lens selection"
    )
    reasoning: str = Field(
        ...,
        description="Why this lens was selected"
    )
    detected_claims: List[str] = Field(
        default_factory=list,
        description="Claims identified in the content"
    )
    keywords_matched: List[str] = Field(
        default_factory=list,
        description="Trigger keywords that were found"
    )


# ============================================================================
# Search Result
# ============================================================================

class SearchResult(BaseModel):
    """
    A search result from evidence gathering.
    """
    title: str
    url: str
    snippet: str
    domain: str
    position: int = Field(description="Ranking position in search results")


# ============================================================================
# Verification State
# ============================================================================

class VerificationState(BaseModel):
    """
    The complete state for a verification workflow.

    This flows through all nodes in the graph.
    Each node receives this state and returns partial updates.

    State sections:
      - request_*: Input from user
      - content_*: Extracted content
      - triage_*: Routing decisions
      - evidence_*: Gathered evidence
      - verdict_*: Final output
      - meta_*: Processing metadata
    """

    # -------------------------------------------------------------------------
    # Request Input
    # -------------------------------------------------------------------------
    request_id: UUID = Field(
        default_factory=uuid4,
        description="Unique ID for this verification request"
    )
    request_url: Optional[str] = Field(
        default=None,
        description="URL submitted for verification"
    )
    request_text: Optional[str] = Field(
        default=None,
        description="Raw text submitted (if no URL)"
    )
    request_user_id: Optional[UUID] = Field(
        default=None,
        description="User who submitted request"
    )
    request_platform: str = Field(
        default="web",
        description="Platform request came from (web, telegram, discord)"
    )
    request_scan_mode: Literal["quick_scan", "deep_dive"] = Field(
        default="quick_scan",
        description="Scan depth requested"
    )
    request_verification_mode: VerificationMode = Field(
        default=VerificationMode.FREE,
        description="LLM provider mode (FREE or VENICE)"
    )

    # -------------------------------------------------------------------------
    # Extracted Content
    # -------------------------------------------------------------------------
    content: Optional[Content] = Field(
        default=None,
        description="Extracted content from URL/input"
    )
    content_extraction_success: bool = Field(
        default=False,
        description="Whether content extraction succeeded"
    )
    content_extraction_method: Optional[str] = Field(
        default=None,
        description="How content was extracted (oembed, ytdlp, etc.)"
    )

    # -------------------------------------------------------------------------
    # Triage Results
    # -------------------------------------------------------------------------
    triage_decision: Optional[TriageDecision] = Field(
        default=None,
        description="Triage node output - which lens to use"
    )
    triage_completed: bool = Field(
        default=False,
        description="Whether triage has run"
    )

    # -------------------------------------------------------------------------
    # Evidence Gathering
    # -------------------------------------------------------------------------
    # Using Annotated with operator.add to accumulate evidence across nodes
    search_results: List[SearchResult] = Field(
        default_factory=list,
        description="Raw search results from Brave/etc."
    )
    evidence: List[Evidence] = Field(
        default_factory=list,
        description="Processed evidence items"
    )
    sources_checked: int = Field(
        default=0,
        description="Number of sources consulted"
    )
    evidence_gathering_completed: bool = Field(
        default=False,
        description="Whether evidence gathering is done"
    )

    # -------------------------------------------------------------------------
    # AI Detection (Optional - skipped in MVP)
    # -------------------------------------------------------------------------
    ai_detection_ran: bool = Field(
        default=False,
        description="Whether AI detection was run"
    )
    ai_detection_result: Optional[dict] = Field(
        default=None,
        description="AI detection results if run"
    )

    # -------------------------------------------------------------------------
    # Verdict
    # -------------------------------------------------------------------------
    verdict: Optional[Verdict] = Field(
        default=None,
        description="Final verdict"
    )
    verdict_completed: bool = Field(
        default=False,
        description="Whether verdict has been generated"
    )

    # -------------------------------------------------------------------------
    # Metadata / Processing
    # -------------------------------------------------------------------------
    meta_started_at: datetime = Field(
        default_factory=datetime.utcnow,
        description="When processing started"
    )
    meta_completed_at: Optional[datetime] = Field(
        default=None,
        description="When processing completed"
    )
    meta_total_cost_usd: float = Field(
        default=0.0,
        description="Total LLM cost for this verification"
    )
    meta_models_used: List[str] = Field(
        default_factory=list,
        description="LLM models used"
    )
    meta_tools_used: List[str] = Field(
        default_factory=list,
        description="Tools/APIs used"
    )

    # -------------------------------------------------------------------------
    # Error Handling
    # -------------------------------------------------------------------------
    errors: List[str] = Field(
        default_factory=list,
        description="Errors encountered during processing"
    )
    current_node: Optional[str] = Field(
        default=None,
        description="Currently executing node"
    )

    # -------------------------------------------------------------------------
    # Helper Methods
    # -------------------------------------------------------------------------

    @property
    def is_complete(self) -> bool:
        """Check if verification is complete."""
        return self.verdict_completed

    @property
    def processing_time_ms(self) -> int:
        """Get total processing time in milliseconds."""
        if not self.meta_completed_at:
            return int((datetime.utcnow() - self.meta_started_at).total_seconds() * 1000)
        return int((self.meta_completed_at - self.meta_started_at).total_seconds() * 1000)

    def add_error(self, error: str) -> "VerificationState":
        """Add an error to the state."""
        return self.model_copy(update={"errors": self.errors + [error]})

    def add_tool_used(self, tool: str) -> "VerificationState":
        """Record a tool being used."""
        if tool not in self.meta_tools_used:
            return self.model_copy(update={"meta_tools_used": self.meta_tools_used + [tool]})
        return self

    def add_cost(self, cost_usd: float) -> "VerificationState":
        """Add to the total cost."""
        return self.model_copy(update={"meta_total_cost_usd": self.meta_total_cost_usd + cost_usd})

    class Config:
        # Allow mutation for state updates
        validate_assignment = True


# ============================================================================
# State Factory
# ============================================================================

def create_initial_state(
    url: Optional[str] = None,
    text: Optional[str] = None,
    user_id: Optional[UUID] = None,
    platform: str = "web",
    scan_mode: Literal["quick_scan", "deep_dive"] = "quick_scan",
    mode: VerificationMode = VerificationMode.FREE,
) -> VerificationState:
    """
    Create initial state for a verification request.

    Args:
        url: URL to verify
        text: Raw text to verify (if no URL)
        user_id: User making the request
        platform: Request source (web, telegram, etc.)
        scan_mode: Depth of verification
        mode: LLM provider mode (FREE or VENICE)

    Returns:
        Initial VerificationState ready for the workflow
    """
    return VerificationState(
        request_url=url,
        request_text=text,
        request_user_id=user_id,
        request_platform=platform,
        request_scan_mode=scan_mode,
        request_verification_mode=mode,
    )