JARVIS/backend/app/agents/state.py

from dataclasses import dataclass
from enum import Enum
from typing import Annotated, Any, Literal, TypedDict

from app.agents.schemas.event import AgentEvent
from app.agents.schemas.task import AgentTask, TaskResult, VerificationStatus
from langchain_core.messages import BaseMessage
from langgraph.graph.message import add_messages


class AgentRole(str, Enum):
    MASTER = "master"
    SCHEDULE_PLANNER = "schedule_planner"
    EXECUTOR = "executor"
    LIBRARIAN = "librarian"
    ANALYST = "analyst"


@dataclass
class ConversationTurn:
    role: str  # "user" | "assistant"
    content: str
    agent: AgentRole | None = None
    model: str | None = None


class AgentState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]
    user_id: str
    conversation_id: str

    execution_mode: Literal["direct", "delegated", "verified"]
    current_agent: str | None
    next_step: str | None
    active_agents: list[AgentRole]
    current_sub_commander: str | None
    active_sub_commanders: list[str]
    sub_commander_trace: list[dict[str, Any]]
    agent_trace: list[str]
    event_trace: list[AgentEvent | dict[str, Any]]

    pending_tasks: list[dict[str, Any]]
    completed_tasks: list[dict[str, Any]]
    active_tasks: list[AgentTask | dict[str, Any]]
    task_results: list[TaskResult | dict[str, Any]]
    tool_calls: list[dict[str, Any]]
    last_tool_result: str | None
    action_results: list[dict[str, Any]]
    created_entities: list[dict[str, Any]]
    tool_outcomes: list[dict[str, Any]]
    task_result_summary: dict[str, Any] | None
    verifier_hints: dict[str, Any] | None

    verification_status: VerificationStatus | None
    verification_summary: str | None
    verification_evidence: list[dict[str, Any]]
    budget_state: dict[str, Any] | None

    tool_strategy_used: str | None
    tool_round_count: int
    max_tool_rounds: int
    retry_count: int
    max_retries: int
    iteration_count: int
    max_iterations: int
    routing_hops: int
    max_routing_hops: int
    terminated_due_to_loop_guard: bool
    retrieval_trace: list[dict[str, Any]]
    stop_reason: str | None

    clarification_needed: bool
    clarification_question: str | None
    fallback_parse_error: str | None
    should_respond: bool

    knowledge_context: str | None
    graph_context: str | None
    schedule_context_summary: str | None
    plan: str | None
    plan_steps: list[dict[str, Any]]
    analysis_report: str | None
    final_response: str | None

    memory_context: str | None
    current_datetime_context: str | None
    current_datetime_reference: dict[str, str] | None

    turn_context: dict[str, Any] | None
    routing_decision: dict[str, Any] | None
    continuity_state: dict[str, Any] | None
    pending_action: dict[str, Any] | None
    last_completed_action: dict[str, Any] | None
    clarification_context: dict[str, Any] | None

    user_llm_config: dict[str, Any] | None
    provider_capabilities: dict[str, Any] | None


def initial_state(user_id: str, conversation_id: str) -> AgentState:
    return AgentState(
        messages=[],
        user_id=user_id,
        conversation_id=conversation_id,
        execution_mode="direct",
        current_agent=AgentRole.MASTER.value,
        next_step=None,
        active_agents=[AgentRole.MASTER],
        current_sub_commander=None,
        active_sub_commanders=[],
        sub_commander_trace=[],
        agent_trace=[AgentRole.MASTER.value],
        event_trace=[],
        pending_tasks=[],
        completed_tasks=[],
        active_tasks=[],
        task_results=[],
        tool_calls=[],
        last_tool_result=None,
        action_results=[],
        created_entities=[],
        tool_outcomes=[],
        task_result_summary=None,
        verifier_hints=None,
        verification_status=None,
        verification_summary=None,
        verification_evidence=[],
        budget_state=None,
        tool_strategy_used=None,
        tool_round_count=0,
        max_tool_rounds=2,
        retry_count=0,
        max_retries=1,
        iteration_count=0,
        max_iterations=3,
        routing_hops=0,
        max_routing_hops=2,
        terminated_due_to_loop_guard=False,
        retrieval_trace=[],
        stop_reason=None,
        clarification_needed=False,
        clarification_question=None,
        fallback_parse_error=None,
        should_respond=True,
        knowledge_context=None,
        graph_context=None,
        schedule_context_summary=None,
        plan=None,
        plan_steps=[],
        analysis_report=None,
        final_response=None,
        memory_context=None,
        current_datetime_context=None,
        current_datetime_reference=None,
        turn_context=None,
        routing_decision=None,
        continuity_state=None,
        pending_action=None,
        last_completed_action=None,
        clarification_context=None,
        user_llm_config=None,
        provider_capabilities=None,
    )