> ## 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.
# Agent Module Overview
> Intelligent agent orchestration with state machine and tool integration
# Agent Module Overview
The Agent Module is the core orchestrator of ARKOS, coordinating the state machine, memory, LLM, and tool systems to provide intelligent conversational interactions.
## Core Concepts
The Agent class manages the complete conversation lifecycle, including context retrieval, state transitions, tool execution, and response generation.
## Architecture
```mermaid theme={null}
graph TB
subgraph "Agent Core"
AGENT[Agent]
AGENT --> FLOW[StateHandler]
AGENT --> MEM[Memory]
AGENT --> LLM[ArkModelLink]
AGENT --> TOOLS[MCPToolManager]
end
subgraph "Key Methods"
STEP[step]
STREAM[step_stream]
CALL[call_llm]
CTX[get_context]
end
subgraph "External"
API[FastAPI Server]
API --> AGENT
end
AGENT --> STEP
AGENT --> STREAM
AGENT --> CALL
AGENT --> CTX
```
## Agent Class
### Initialization
```python theme={null}
from agent_module.agent import Agent
from state_module.state_handler import StateHandler
from memory_module.memory import Memory
from model_module.ArkModelNew import ArkModelLink
from tool_module.tool_call import MCPToolManager
# Initialize components
flow = StateHandler(yaml_path="state_module/state_graph.yaml")
memory = Memory(
user_id="default_user",
session_id=None,
db_url=db_url,
use_long_term=False
)
llm = ArkModelLink(base_url="http://localhost:30000/v1")
tool_manager = MCPToolManager(mcp_config, token_store=token_store)
# Create agent
agent = Agent(
agent_id="my-agent",
flow=flow,
memory=memory,
llm=llm,
tool_manager=tool_manager
)
```
### Key Attributes
```python theme={null}
class Agent:
agent_id: str # Unique identifier
flow: StateHandler # State machine
memory: Memory # Memory system
llm: ArkModelLink # LLM interface
tool_manager: MCPToolManager # Tool manager (optional)
current_state: State # Current state in flow
system_prompt: str # System prompt with tools
available_tools: dict # Discovered MCP tools
current_user_id: str # For per-user tool auth
```
## Core Methods
### step() - Main Processing Loop
Process messages through the state machine:
```python theme={null}
async def step(self, messages, user_id: str = None):
"""
Runs the agent until reaching a terminal state.
Parameters
----------
messages : list
List of messages to process
user_id : str, optional
User ID for per-user tool authentication
Returns
-------
AIMessage
The last AI message produced
"""
self.current_user_id = user_id
self.add_context(messages)
last_ai_message = None
retry_count = 0
while not self.current_state.is_terminal:
if retry_count > MAX_ITER: # MAX_ITER = 10
break
retry_count += 1
context = self.get_context()
update = await self.current_state.run(context, self)
if update:
self.add_context([update])
if isinstance(update, AIMessage):
last_ai_message = update
if self.current_state.is_terminal:
break
# Handle state transitions
messages_list = self.memory.retrieve_short_memory(5)
if self.current_state.check_transition_ready(messages_list):
transition_dict = self.flow.get_transitions(
self.current_state, messages_list
)
transition_names = transition_dict["tt"]
if len(transition_names) == 1:
next_state_name = transition_names[0]
else:
next_state_name = await self.choose_transition(
transition_dict, messages_list
)
self.current_state = self.flow.get_state(next_state_name)
# Reset to initial state for next conversation
self.current_state = self.flow.get_state("agent_reply")
return last_ai_message
```
### step\_stream() - Streaming Response
Stream responses character by character:
```python theme={null}
async def step_stream(self, messages, user_id: str = None):
"""
Streaming version of step. Yields characters as they're produced.
Yields
------
str
Characters/chunks from each state's output
"""
self.current_user_id = user_id
self.add_context(messages)
retry_count = 0
while not self.current_state.is_terminal:
if retry_count > MAX_ITER:
yield "\n[Max iterations reached]"
break
retry_count += 1
context = self.get_context()
try:
update = await self.current_state.run(context, self)
except Exception as e:
update = AIMessage(content=f"Error: {str(e)[:200]}")
self.current_state = self.flow.get_state("agent_reply")
# Stream character by character
if update and hasattr(update, 'content') and update.content:
self.add_context([update])
for char in update.content:
yield char
if self.current_state.is_terminal:
break
# Handle transitions (same as step)
# ... transition logic ...
if not self.current_state.is_terminal:
yield "\n\n" # Separator between states
self.current_state = self.flow.get_state("agent_reply")
```
### call\_llm() - LLM Interface
Call the language model:
```python theme={null}
async def call_llm(self, context=None, json_schema=None):
"""
Agent's interface with chat model.
Parameters
----------
context : list
Messages to send to LLM
json_schema : dict, optional
JSON schema for structured output
Returns
-------
AIMessage
The LLM's response
"""
llm_response = await self.llm.generate_response(context, json_schema)
return AIMessage(content=llm_response)
```
### Context Management
```python theme={null}
def add_context(self, messages):
"""Store messages in memory."""
for message in messages:
self.memory.add_memory(message)
def get_context(self, turns=5, include_long_term=True):
"""
Retrieve context for LLM.
Returns list of messages including:
- Long-term memories (if enabled)
- Recent conversation turns
"""
short_term_mem = self.memory.retrieve_short_memory(turns)
if include_long_term:
long_term_mem = self.memory.retrieve_long_memory(context=short_term_mem)
if long_term_mem and long_term_mem.content.strip():
return [long_term_mem] + short_term_mem
return short_term_mem
```
## Transition Selection
When multiple state transitions are possible:
```python theme={null}
async def choose_transition(self, transitions_dict, messages):
"""Use LLM to select the best next state."""
transition_tuples = list(zip(
transitions_dict["tt"],
transitions_dict["td"]
))
prompt = f"""Given the conversation context and these options:
{transition_tuples}
Output the most reasonable next state.
Do not use tool result to determine the next state."""
# Create Pydantic model for structured output
NextStates = self.create_next_state_class(transition_tuples)
json_schema = {
"type": "json_schema",
"json_schema": {
"name": "class_options",
"schema": NextStates.model_json_schema(),
},
}
context_text = [SystemMessage(content=prompt)] + messages
output = await self.call_llm(context=context_text, json_schema=json_schema)
return json.loads(output.content)["next_state"]
```
## Tool Integration
### Tool Argument Classes
```python theme={null}
def fill_tool_args_class(self, tool_name: str, tool_args: Dict[str, Any]):
"""Create Pydantic model for tool call."""
ToolCall = create_model(
"ToolCall",
tool_name=(str, Field(description="Tool name to execute")),
tool_args=(Dict[str, Any], Field(default_factory=dict)),
)
return ToolCall(tool_name=tool_name, tool_args=tool_args)
async def create_tool_option_class(self):
"""Create enum of available tools for LLM selection."""
server_tool_map = await self.tool_manager.list_all_tools()
enum_members = {}
for server_name in server_tool_map:
for tool_name in server_tool_map[server_name]:
enum_members[tool_name] = tool_name
ToolEnum = Enum("ToolEnum", enum_members)
return create_model(
"ToolCall",
tool_name=(ToolEnum, Field(description="Tool to execute next"))
)
```
## Usage with FastAPI
The agent is used in the API server:
```python theme={null}
# base_module/app.py
from agent_module.agent import Agent
# Initialize agent at startup
agent = Agent(
agent_id=config.get("memory.user_id"),
flow=flow,
memory=memory,
llm=llm,
tool_manager=tool_manager,
)
@app.on_event("startup")
async def startup():
if tool_manager:
await tool_manager.initialize_servers()
agent.available_tools = await tool_manager.list_all_tools()
agent.system_prompt = base_prompt + "\n\n" + tool_prompt
@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
# ... parse request ...
if stream:
async def generate_stream():
async for chunk in agent.step_stream(context_msgs, user_id=user_id):
# Format as SSE
yield f"data: {json.dumps(data)}\n\n"
return StreamingResponse(generate_stream(), media_type="text/event-stream")
# Non-streaming
agent_response = await agent.step(context_msgs, user_id=user_id)
return JSONResponse(content=completion)
```
## Performance Timing
The agent includes built-in timing logs:
```python theme={null}
# Output during execution:
# [TIMING] add_context: 0.015s
# [TIMING] get_context: 0.102s
# [TIMING] state.run: 1.234s
# [TIMING] loop total: 1.351s
# [TIMING] step total: 2.847s
```
## Configuration
Agents are configured via `config_module/config.yaml`:
```yaml theme={null}
memory:
user_id: "default_user"
use_long_term: false
llm:
base_url: "http://localhost:30000/v1"
state:
graph_path: "state_module/state_graph.yaml"
app:
system_prompt: "You are a helpful AI assistant."
```
## Error Handling
```python theme={null}
async def step_stream(self, messages, user_id=None):
# ...
try:
update = await self.current_state.run(context, self)
except Exception as e:
print(f"[STREAM] State error: {e}")
update = AIMessage(content=f"Error: {str(e)[:200]}")
self.current_state = self.flow.get_state("agent_reply")
# ...
```
## Debug Logging
The agent logs state transitions:
```python theme={null}
print("agent.py received message")
print("agent.py CURR STATE:", self.current_state)
print("agent.py IS TERMINAL?:", self.current_state.is_terminal)
print(f"Inner loop #{retry_count + 1}")
print("REACHED TERMINAL")
print("LAST_AI_MSG", last_ai_message)
```
## Best Practices
1. **Always set user\_id**: For per-user tool authentication
2. **Monitor MAX\_ITER**: Prevent infinite loops
3. **Use streaming**: For better UX with long responses
4. **Reset state**: Agent resets to initial state after each conversation
5. **Handle errors**: Graceful degradation on state failures
## Troubleshooting
Check state machine for deadlocks:
```python theme={null}
print(f"Current state: {agent.current_state.name}")
print(f"Is terminal: {agent.current_state.is_terminal}")
```
Verify tool manager is initialized:
```python theme={null}
print(f"Available tools: {agent.available_tools.keys()}")
```
Increase the turns parameter (config.yaml):
```yaml theme={null}
short_term_turns: 50
long_term_limit: 50
```
## Next Steps
Learn about state management
Explore memory systems
Understand tool integration
See the API server