Anton
052591281d
- Complete agent system with planning and execution - Bridge API with FastAPI server - WhatsApp-style mobile dashboard - Enhanced web search with query improvement - Security guards and policy checks - Comprehensive documentation
123 lines
3.5 KiB
Python
123 lines
3.5 KiB
Python
"""State models for the agent runtime."""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import asdict, dataclass, field
|
|
from enum import Enum
|
|
from typing import Any, Optional
|
|
|
|
|
|
class StepStatus(str, Enum):
|
|
PENDING = "pending"
|
|
RUNNING = "running"
|
|
DONE = "done"
|
|
FAILED = "failed"
|
|
|
|
|
|
class RunStatus(str, Enum):
|
|
RUNNING = "running"
|
|
DONE = "done"
|
|
FAILED = "failed"
|
|
|
|
|
|
@dataclass
|
|
class PlanStep:
|
|
"""A single step within an agent plan."""
|
|
|
|
step_id: str
|
|
kind: str # "shell" | "write_file" | "read_file" | "finish" | "flaky"
|
|
description: str
|
|
payload: dict[str, Any]
|
|
status: StepStatus = StepStatus.PENDING
|
|
attempts: int = 0 # how many times this step has been attempted (0 = not yet started)
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
d = asdict(self)
|
|
d["status"] = self.status.value
|
|
return d
|
|
|
|
|
|
@dataclass
|
|
class StepResult:
|
|
"""Result of executing a single plan step."""
|
|
|
|
step_id: str
|
|
status: StepStatus
|
|
output_text: Optional[str] = None
|
|
error: Optional[str] = None
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
|
started_at: str = ""
|
|
finished_at: str = ""
|
|
attempt: int = 1 # which attempt produced this result (1-based)
|
|
failure_class: Optional[str] = None # set on FAILED results
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {
|
|
"step_id": self.step_id,
|
|
"status": self.status.value,
|
|
"output_text": self.output_text,
|
|
"error": self.error,
|
|
"metadata": self.metadata,
|
|
"started_at": self.started_at,
|
|
"finished_at": self.finished_at,
|
|
"attempt": self.attempt,
|
|
"failure_class": self.failure_class,
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class Plan:
|
|
"""Ordered list of steps to accomplish a goal."""
|
|
|
|
steps: list[PlanStep] = field(default_factory=list)
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {"steps": [s.to_dict() for s in self.steps]}
|
|
|
|
|
|
@dataclass
|
|
class AgentState:
|
|
"""Full runtime state of a single agent run."""
|
|
|
|
run_id: str
|
|
goal: str
|
|
plan: Plan
|
|
current_step_id: Optional[str]
|
|
history: list[StepResult]
|
|
started_at: str
|
|
updated_at: str
|
|
status: RunStatus
|
|
trace_id: Optional[str] = None
|
|
|
|
# --- History access helpers ---
|
|
|
|
def get_last_output(self) -> Optional[str]:
|
|
"""Return output_text of the most recently completed step, or None."""
|
|
for result in reversed(self.history):
|
|
if result.output_text is not None:
|
|
return result.output_text
|
|
return None
|
|
|
|
def get_step_output(self, step_id: str) -> Optional[str]:
|
|
"""Return output_text for a specific step_id, or None if not found."""
|
|
for result in self.history:
|
|
if result.step_id == step_id:
|
|
return result.output_text
|
|
return None
|
|
|
|
def list_outputs(self) -> list[tuple[str, Optional[str]]]:
|
|
"""Return [(step_id, output_text)] for all history entries."""
|
|
return [(r.step_id, r.output_text) for r in self.history]
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {
|
|
"run_id": self.run_id,
|
|
"trace_id": self.trace_id,
|
|
"goal": self.goal,
|
|
"plan": self.plan.to_dict(),
|
|
"current_step_id": self.current_step_id,
|
|
"history": [r.to_dict() for r in self.history],
|
|
"started_at": self.started_at,
|
|
"updated_at": self.updated_at,
|
|
"status": self.status.value,
|
|
}
|