> ## Documentation Index
> Fetch the complete documentation index at: https://arkos.mintlify.site/llms.txt
> Use this file to discover all available pages before exploring further.

# State Module Overview

> YAML-configurable state machine for conversation flow management

# State Module Overview

The State Module implements a YAML-configurable state machine that manages conversation flow and agent behavior in ARKOS.

## Core Concepts

<Info>
  The State Module provides a declarative way to define conversation flows using YAML configuration files, with automatic transition handling based on context.
</Info>

### What is a State?

A state represents a distinct phase in the conversation or agent processing pipeline. Each state:

* Has a unique name
* Defines specific behavior when active
* Specifies possible transitions to other states
* Can be marked as terminal (ending the flow)

## Architecture

```mermaid theme={null}
graph TB
    subgraph "State Module"
        SH[StateHandler]
        SH --> SR[State Registry]
        SH --> YAML[YAML Config]
    end

    subgraph "State Types"
        SR --> AS[AgentState Enum]
        AS --> GU[GREET_USER]
        AS --> FP[FETCH_PRODUCT]
        AS --> SR2[SUMMARIZE_RESULT]
        AS --> DONE[DONE]
    end

    subgraph "State Base"
        S[State Class]
        S --> RUN[run method]
        S --> CHECK[check_transition_ready]
    end
```

## State Class

The base `State` class that all states extend:

```python theme={null}
from state_module.state import State

class State:
    def __init__(self, name: str, config: Dict[str, Any]):
        self.name = name
        self.is_terminal: bool = False
        self.transition = config.get("transition", {})

    def check_transition_ready(self, context: Dict[str, Any]) -> bool:
        """
        Override to define when this state can transition.
        """
        raise NotImplementedError

    def run(self, context: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """
        Override to define state behavior.
        """
        raise NotImplementedError
```

## AgentState Enum

Predefined states available in the system:

```python theme={null}
from enum import Enum

class AgentState(Enum):
    """Central registry of all possible agent states."""

    GREET_USER = "greet_user"
    FETCH_PRODUCT = "fetch_product"
    SUMMARIZE_RESULT = "summarize_result"
    DONE = "done"
```

## StateHandler

Manages the state machine lifecycle and transitions:

```python theme={null}
from state_module.state_handler import StateHandler

# Initialize with YAML configuration
handler = StateHandler(yaml_path="state_module/state_graph.yaml")

# Get initial state
initial_state = handler.get_initial_state()

# Get a specific state
state = handler.get_state("agent_reply")

# Get possible transitions from current state
transitions = handler.get_transitions(current_state, messages)
# Returns: {"tt": ["state1", "state2"], "td": ["description1", "description2"]}
```

## YAML Configuration

### Basic State Graph

```yaml theme={null}
# state_module/state_graph.yaml
 ask_user:
    description: "Waiting for user input"
    type: user
    transition:
      next: [agent_reply]

  agent_reply:
    description: "AI reasoning and response generation"
    type: agent
    transition:
      next: [ask_user, use_tool]

  use_tool:
    description: "Execute a specific tool"
    type: tool
    transition:
      next: [agent_reply, ask_user]

```

### Transition Format

Transitions use a list of names corresponding to the next state

```yaml theme={null}
transition: 
  next: [agent_reply, ask_user, use_tool]
```

## State Execution Flow

```mermaid theme={null}
sequenceDiagram
    participant Agent
    participant StateHandler
    participant State
    participant LLM

    Agent->>StateHandler: get_initial_state()
    StateHandler-->>Agent: current_state

    loop Until Terminal State
        Agent->>State: run(context, agent)
        State-->>Agent: update (AIMessage/ToolMessage)

        Agent->>State: check_transition_ready(messages)
        State-->>Agent: ready (bool)

        alt Single Transition
            Agent->>StateHandler: get_state(next_name)
        else Multiple Transitions
            Agent->>LLM: choose_transition(options)
            LLM-->>Agent: selected_state
            Agent->>StateHandler: get_state(selected)
        end
    end
```

## Integration with Agent

The Agent uses StateHandler for flow control:

```python theme={null}
class Agent:
    def __init__(self, ..., flow: StateHandler, ...):
        self.flow = flow
        self.current_state = self.flow.get_initial_state()

    async def step(self, messages, user_id=None):
        while not self.current_state.is_terminal:
            # Run current state
            context = self.get_context()
            update = await self.current_state.run(context, self)

            if update:
                self.add_context([update])

            # Check for transition
            if self.current_state.check_transition_ready(messages):
                transition_dict = self.flow.get_transitions(
                    self.current_state, messages
                )
                transition_names = transition_dict["tt"]

                if len(transition_names) == 1:
                    next_state_name = transition_names[0]
                else:
                    # LLM chooses between options
                    next_state_name = await self.choose_transition(
                        transition_dict, messages
                    )

                self.current_state = self.flow.get_state(next_state_name)
```

## Transition Selection

When multiple transitions are possible, the LLM chooses:

```python theme={null}
async def choose_transition(self, transitions_dict, messages):
    """Use LLM to choose the best next state."""

    # Create tuples of (state_name, description)
    transition_tuples = list(zip(
        transitions_dict["tt"],
        transitions_dict["td"]
    ))

    prompt = f"""Given the conversation context and these options:
    {transition_tuples}
    Choose the most appropriate next state."""

    # Create Pydantic model for structured output
    NextStates = self.create_next_state_class(transition_tuples)

    # Get LLM decision
    output = await self.call_llm(
        context=[SystemMessage(content=prompt)] + messages,
        json_schema=NextStates.model_json_schema()
    )

    return json.loads(output.content)["next_state"]
```

## Custom State Types

### Creating a Custom State

### (NOTE: all states must be registered with @register\_state and follow state\_\[name].py convention)

```python theme={null}
from state_module.state import State
from model_module.ArkModelNew import AIMessage

@register_state
class GreetingState(State):
    def __init__(self, name: str, config: dict):
        super().__init__(name, config)
        self.greeting = config.get("greeting", "Hello!")

    def check_transition_ready(self, context) -> bool:
        # Always ready to transition after greeting
        return True

    def run(self, context, agent=None):
        return AIMessage(content=self.greeting)
```

### Tool State Example

```python theme={null}
from model_module.ArkModelNew import SystemMessage
from tool_module.tool_call import AuthRequiredError

from state_module.state import State
from state_module.state_registry import register_state


@register_state
class StateTool(State):
    type = "tool"

    def __init__(self, name: str, config: dict):
        super().__init__(name, config)
        self.is_terminal = False

    def check_transition_ready(self, context):
        return True

    async def choose_tool(self, context, agent):
        """
        Chooses tool to use based on the context and server
        """
        prompt = "Based on the user request, choose the best tool. Return JSON with tool_name field."
        instructions = context + [SystemMessage(content=prompt)]

        # Get available tools and build schema
        tool_option_class = await agent.create_tool_option_class()
        json_schema = {
            "type": "json_schema",
            "json_schema": {
                "name": "tool_choice",
                "schema": tool_option_class.model_json_schema(),
            },
        }

        # Get tool choice from LLM
        output = await agent.call_llm(instructions, json_schema)
        parsed = json.loads(output.content)
        tool_name = parsed["tool_name"]

        # Get tool spec
        server_name = agent.tool_manager._tool_registry[tool_name]
        all_tools = await agent.tool_manager.list_all_tools()
        tool_spec = all_tools[server_name][tool_name]

        # Get tool arguments
        args_prompt = f"Fill in arguments for tool '{tool_name}'. Return JSON matching the schema."
        args_context = context + [SystemMessage(content=args_prompt)]
        args_schema = {
            "type": "json_schema",
            "json_schema": {
                "name": "tool_args",
                "schema": tool_spec.get("inputSchema", {"type": "object", "properties": {}}),
            },
        }

        args_output = await agent.call_llm(args_context, args_schema)
        tool_args = json.loads(args_output.content)

        return {"tool_name": tool_name, "tool_args": tool_args}

    async def execute_tool(self, tool_call, agent):
        """
        Parses and fills args for chosen tool for tool call execution
        """
        tool_name = tool_call["tool_name"]
        tool_args = tool_call["tool_args"]

        tool_result = await agent.tool_manager.call_tool(
            tool_name=tool_name,
            arguments=tool_args,
            user_id=agent.current_user_id,
        )

        return tool_result

    async def run(self, context, agent=None):
        try:
            tool_arg_dict = await self.choose_tool(context=context, agent=agent)
            tool_result = await self.execute_tool(tool_call=tool_arg_dict, agent=agent)
            return SystemMessage(content=str(tool_result))

        except AuthRequiredError as e:
            # Return friendly message with connect link
            return SystemMessage(
                content=f"To complete this request, please connect your {e.service_info.get('name', e.service)}:\n\n"
                        f" {e.connect_url}\n\n"
                        f"After connecting, try your request again."
            )
```

## State Types Reference

| Type       | Purpose               | Terminal |
| ---------- | --------------------- | -------- |
| `agent`    | Generate AI response  | No       |
| `user`     | Wait for user input   | No       |
| `tool`     | Execute tool call     | No       |
| `decision` | Make routing decision | No       |
| `terminal` | End conversation      | Yes      |

## Best Practices

### State Design Guidelines

1. **Single Responsibility**: Each state should have one clear purpose
2. **Clear Transitions**: Define unambiguous transition conditions
3. **Terminal States**: Ensure all paths eventually reach a terminal
4. **Descriptive Names**: Use clear, action-oriented state names

### YAML Configuration Tips

```yaml theme={null}
# Good: Clear state names and descriptions
states:
  gather_requirements:
    type: user
    transition:
      next: [validate_input, ask_clarification]
     

# Avoid: Ambiguous transitions
states:
  process:
    type: unknown
    transition:
      next: [next1, next2]
```

## Debugging States

### Logging State Transitions

```python theme={null}
# Agent logs state changes
print(f"agent.py CURR STATE: {self.current_state}")
print(f"agent.py IS TERMINAL?: {self.current_state.is_terminal}")
```

### Common Issues

<AccordionGroup>
  <Accordion title="State not transitioning">
    Check that `check_transition_ready()` returns `True`:

    ```python theme={null}
    ready = state.check_transition_ready(messages)
    print(f"Transition ready: {ready}")
    ```
  </Accordion>

  <Accordion title="Wrong state selected">
    Improve transition descriptions for better LLM selection:

    ```yaml theme={null}
    transition:
      next: [help, search]
      t
    ```
  </Accordion>

  <Accordion title="Infinite loop">
    Ensure there's always a path to a terminal state and MAX\_ITER limit is respected.
  </Accordion>
</AccordionGroup>

## Example: Complete State Graph

```yaml theme={null}

initial: agent_reply

states:
  # === Core States ===
  ask_user:
    description: "Waiting for user input"
    type: user
    transition:
      next: [agent_reply]

  agent_reply:
    description: "AI reasoning and response generation"
    type: agent
    transition:
      next: [agent_reply, ask_user, use_tool, autonomous_plan, learn_skill]

  use_tool:
    description: "Execute a specific tool"
    type: tool
    transition:
      next: [agent_reply, ask_user]

  # === Autonomous Planning Loop (OpenClaw-style) ===
  autonomous_plan:
    description: "Create multi-step execution plan for complex tasks"
    type: planner
    transition:
      next: [execute_step, ask_user]  # ask_user if permission needed

  execute_step:
    description: "Execute current step in the plan"
    type: executor
    transition:
      next: [execute_step, verify_progress]  # loop or verify

  verify_progress:
    description: "Check if goal achieved or continue execution"
    type: verifier
    transition:
      next: [execute_step, autonomous_plan, agent_reply, ask_user]  # continue, replan, or finish

  # === Self-Improving Skills ===
  learn_skill:
    description: "Analyze and potentially create new capabilities"
    type: skill_builder
    transition:
      next: [agent_reply, autonomous_plan]  # use skill or plan with it


```

## Next Steps

<CardGroup cols={2}>
  <Card title="Implementation Details" icon="code" href="/modules/state/implementation">
    Deep dive into state implementation
  </Card>

  <Card title="Agent Module" icon="robot" href="/modules/agent/overview">
    Learn about agent orchestration
  </Card>

  <Card title="Memory Module" icon="database" href="/modules/memory/overview">
    Explore memory integration
  </Card>

  <Card title="Model Module" icon="sparkles" href="/modules/model/overview">
    Understand LLM integration
  </Card>
</CardGroup>