ui/agent/graph/state.py

# ui/agent/graph/state.py
from __future__ import annotations

import operator
from typing import Annotated, Any, TypedDict


class TodoItem(TypedDict):
    id: int
    country: str
    methods: str


class Finding(TypedDict, total=False):
    todo_id: int
    todo_title: str
    todo_country: str
    todo_methods: str
    summary: str


class CandidateCountry(TypedDict):
    iso2: str
    name: str
    pathway_hint: str
    label: str


class AgentState(TypedDict, total=False):
    """Top-level workflow state."""

    user_content: str | list[dict[str, Any]]
    history_messages: list[dict[str, Any]]
    profile_summary: str
    candidate_countries: list[CandidateCountry]
    discovery_summary: str
    todos: list[TodoItem]
    findings: Annotated[list[Finding], operator.add]
    final_answer: str
    consolidator_thinking: str


class ResearchTask(TypedDict):
    """Payload sent to each parallel researcher via Send."""

    todo: TodoItem
    profile_summary: str