Files
JARVIS/backend/app/agents/state.py

250 lines
7.7 KiB
Python

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
AgentPhase = Literal[
"phase_0_bootstrap",
"phase_1_routing",
"phase_2_controlled_collaboration",
"phase_3_dynamic_collaboration",
"phase_4_visibility_and_verification",
]
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]]
isolation_mode: str
isolation_id: str | None
isolation_workspace_path: str | None
isolation_parent_conversation_id: str | None
isolation_metadata: dict[str, Any]
input_tokens: int
output_tokens: int
estimated_cost: float | None
budget_warning: bool
cost_by_agent: dict[str, dict[str, Any]]
cost_thresholds: dict[str, Any]
budget_state: CollaborationBudget | dict[str, Any] | None
collaboration_budget_history: list[CollaborationBudget | dict[str, Any]]
current_phase: AgentPhase
phase_history: list[dict[str, Any]]
current_checkpoint: str | None
checkpoint_history: list[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=[],
isolation_mode="none",
isolation_id=None,
isolation_workspace_path=None,
isolation_parent_conversation_id=None,
isolation_metadata={},
input_tokens=0,
output_tokens=0,
estimated_cost=None,
budget_warning=False,
cost_by_agent={},
cost_thresholds={},
budget_state=None,
collaboration_budget_history=[],
current_phase="phase_0_bootstrap",
phase_history=[
{
"phase": "phase_0_bootstrap",
"reason": "initial_state_created",
}
],
current_checkpoint="bootstrap.initialized",
checkpoint_history=[
{
"checkpoint": "bootstrap.initialized",
"phase": "phase_0_bootstrap",
"reason": "initial_state_created",
}
],
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,
)