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.message import AgentMessage
from app.agents.schemas.task import AgentTask, CollaborationBudget, InterruptRecord, RecoveryRecord, TaskResult, VerificationStatus
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
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


def turn_to_message(turn: ConversationTurn) -> BaseMessage:
    if turn.role == "user":
        return HumanMessage(content=turn.content)
    return AIMessage(content=turn.content)


class AgentState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]
    user_id: str
    conversation_id: str
    parent_conversation_id: str | None
    thread_id: str | None
    last_message_id: str | None
    message_sequence: int
    agent_id: str | None
    parent_agent_id: str | None
    root_agent_id: str | None
    collaboration_depth: int
    spawned_agent_ids: list[str]

    execution_mode: Literal["direct", "collaboration", "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]]
    message_trace: list[AgentMessage | 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]]
    task_hierarchy: dict[str, list[str]]
    interrupted_tasks: list[InterruptRecord | dict[str, Any]]
    recovery_trace: list[RecoveryRecord | dict[str, Any]]
    recovery_points: list[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: CollaborationBudget | dict[str, Any] | None
    collaboration_budget_history: list[CollaborationBudget | dict[str, Any]]

    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,
        parent_conversation_id=None,
        thread_id=None,
        last_message_id=None,
        message_sequence=0,
        agent_id=AgentRole.MASTER.value,
        parent_agent_id=None,
        root_agent_id=AgentRole.MASTER.value,
        collaboration_depth=0,
        spawned_agent_ids=[],
        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=[],
        message_trace=[],
        pending_tasks=[],
        completed_tasks=[],
        active_tasks=[],
        task_results=[],
        task_hierarchy={},
        interrupted_tasks=[],
        recovery_trace=[],
        recovery_points=[],
        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,
        collaboration_budget_history=[],
        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,
    )