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

# Agent System

> LangGraph agent workflow and tool execution

## Overview

Aurora's AI agent is built on **LangGraph**, a framework for building stateful, multi-step workflows with large language models. The agent can execute cloud operations, search knowledge bases, run commands, and more through a tool-based architecture.

## LangGraph Architecture

### State Graph

The agent workflow is implemented as a directed graph with nodes representing operations and edges defining the flow.

```python theme={null}
from langgraph.graph import StateGraph, START, END

workflow = StateGraph(State)
workflow.add_node("agentic_tool_flow", agent.agentic_tool_flow)
workflow.add_edge(START, "agentic_tool_flow")
workflow.add_edge("agentic_tool_flow", END)
app = workflow.compile(checkpointer=memory)
```

Reference: `server/chat/backend/agent/workflow.py:103-113`

### State Management

The workflow state stores conversation context and execution metadata:

```python theme={null}
class State(TypedDict):
    messages: List[BaseMessage]      # Chat history
    user_id: str                     # Authenticated user
    session_id: str                  # Chat session ID
    question: str                    # User query
    provider_preference: List[str]   # Connected cloud providers
    selected_project_id: str         # Active project/subscription
    mode: str                        # "agent" or "ask" mode
    attachments: List[Dict]          # File uploads
    session_files: List[Dict]        # S3 storage files
```

Reference: `server/chat/backend/agent/utils/state.py`

### Memory Persistence

LangGraph uses a memory saver to persist conversation context across requests:

```python theme={null}
from chat.backend.agent.utils.safe_memory_saver import SafeMemorySaver

memory = SafeMemorySaver()
config = RunnableConfig({"configurable": {"thread_id": session_id}})
app = workflow.compile(checkpointer=memory)
```

Reference: `server/chat/backend/agent/workflow.py:60-64`

***

## Agent Execution Flow

### 1. Initialize Connection

User connects via WebSocket with authentication:

```json theme={null}
{
  "type": "init",
  "user_id": "user-123"
}
```

Reference: `server/main_chatbot.py:94-133`

### 2. Receive Query

User sends a question with context:

```json theme={null}
{
  "type": "query",
  "query": "List all GKE clusters in my project",
  "user_id": "user-123",
  "session_id": "session-456",
  "mode": "agent",
  "model": "openai/gpt-4",
  "provider_preference": ["gcp"]
}
```

Reference: `server/main_chatbot.py:806-839`

### 3. Set User Context

The agent sets thread-local context for tools:

```python theme={null}
from chat.backend.agent.tools.cloud_tools import set_user_context

set_user_context(
    user_id=state.user_id,
    session_id=state.session_id,
    provider_preference=provider_preference,
    selected_project_id=selected_project_id,
    state=state,
    mode=mode
)
```

Reference: `server/chat/backend/agent/agent.py:200-208`

### 4. Build System Prompt

Dynamic system prompt based on connected providers and mode:

```python theme={null}
from chat.backend.agent.prompt.prompt_builder import build_prompt_segments

segments = build_prompt_segments(
    provider_preference=provider_preference,
    mode=mode,
    has_zip_reference=has_zip,
    state=state
)

system_prompt = assemble_system_prompt(segments)
```

Reference: `server/chat/backend/agent/agent.py:244-250`

### 5. Load Available Tools

Tools are filtered based on mode and connected providers:

```python theme={null}
from chat.backend.agent.tools.cloud_tools import get_cloud_tools

tools = get_cloud_tools(
    provider_preference=provider_preference,
    mode=mode,
    state=state
)
```

Reference: `server/chat/backend/agent/agent.py:168`

### 6. Execute LangGraph Workflow

The workflow streams events as the agent thinks and acts:

```python theme={null}
async for event_type, event_data in workflow.stream(state):
    if event_type == "token":
        # Stream LLM tokens to frontend
        await websocket.send(json.dumps({
            "type": "message",
            "data": {"text": token_text, "is_chunk": True}
        }))
    elif event_type == "messages":
        # Handle tool calls from LLM
        if msg_chunk.tool_calls:
            for tool_call in msg_chunk.tool_calls:
                # Execute tool and return result
```

Reference: `server/main_chatbot.py:331-470`

### 7. Tool Execution

When the LLM decides to use a tool:

```python theme={null}
@tool
def list_gke_clusters(project_id: str, location: str = "-") -> str:
    """List all GKE clusters in a GCP project."""
    # Tool execution logic
    return json.dumps({"status": "success", "clusters": [...]})
```

The tool result is added to the conversation and the LLM continues reasoning.

### 8. Stream Response

Final answer is streamed token-by-token to the frontend:

```json theme={null}
{
  "type": "message",
  "data": {
    "text": "Found 3 GKE clusters:\n1. prod-cluster-1\n2. staging-cluster\n3. dev-cluster",
    "is_chunk": false,
    "is_complete": true
  }
}
```

Reference: `server/main_chatbot.py:439-470`

***

## Agent Tools

Aurora provides 30+ tools across multiple categories:

### Cloud Operations

| Tool                         | Purpose               | File             |
| ---------------------------- | --------------------- | ---------------- |
| `list_gke_clusters`          | List GKE clusters     | `cloud_tools.py` |
| `list_gcp_compute_instances` | List GCP VMs          | `cloud_tools.py` |
| `list_aws_ec2_instances`     | List AWS EC2          | `cloud_tools.py` |
| `list_azure_vms`             | List Azure VMs        | `cloud_tools.py` |
| `get_gcp_logs`               | Query GCP logs        | `cloud_tools.py` |
| `iac_write`                  | Write Terraform       | `iac_tool.py`    |
| `iac_deploy`                 | Deploy infrastructure | `iac_tool.py`    |

### Knowledge & Search

| Tool                    | Purpose                       | File                            |
| ----------------------- | ----------------------------- | ------------------------------- |
| `knowledge_base_search` | Semantic search uploaded docs | `knowledge_base_search_tool.py` |
| `confluence_search`     | Search Confluence wikis       | `confluence_search_tool.py`     |
| `web_search`            | Search the internet           | `web_search_tool.py`            |

### Command Execution

| Tool                  | Purpose                    | File                    |
| --------------------- | -------------------------- | ----------------------- |
| `run_kubectl_command` | Execute kubectl in pods    | `cloud_tools.py`        |
| `terminal_exec`       | Run bash commands          | `terminal_exec_tool.py` |
| `cloud_exec`          | Execute cloud CLI commands | `cloud_exec_tool.py`    |

### Source Control

| Tool                 | Purpose                 | File                       |
| -------------------- | ----------------------- | -------------------------- |
| `github_search_code` | Search GitHub repos     | `github_rca_tool.py`       |
| `github_get_file`    | Read file content       | `github_rca_tool.py`       |
| `github_commit`      | Commit and push changes | `github_commit_tool.py`    |
| `github_apply_fix`   | Apply code fixes        | `github_apply_fix_tool.py` |

### Monitoring & Observability

| Tool                     | Purpose                 | File                  |
| ------------------------ | ----------------------- | --------------------- |
| `splunk_search`          | Search Splunk logs      | `splunk_tool.py`      |
| `dynatrace_query`        | Query Dynatrace metrics | `dynatrace_tool.py`   |
| `coroot_analyze`         | Analyze Coroot data     | `coroot_tool.py`      |
| `jenkins_get_build_logs` | Fetch Jenkins logs      | `jenkins_rca_tool.py` |

### File Operations

| Tool                  | Purpose              | File               |
| --------------------- | -------------------- | ------------------ |
| `extract_zip_file`    | Extract ZIP archives | `zip_file_tool.py` |
| `read_terraform_file` | Read Terraform files | `iac_tool.py`      |

Reference: `server/chat/backend/agent/tools/`

***

## Tool Execution Pattern

### Tool Definition

Tools are defined using LangChain's `@tool` decorator:

```python theme={null}
from langchain.agents import tool
from chat.backend.agent.tools.cloud_tools import get_user_context

@tool
def list_gke_clusters(project_id: str, location: str = "-") -> str:
    """List all GKE clusters in a GCP project.
    
    Args:
        project_id: GCP project ID
        location: GCP region or zone (default: all)
    
    Returns:
        JSON string with cluster details
    """
    # Get user context
    ctx = get_user_context()
    user_id = ctx["user_id"]
    
    # Execute operation
    from connectors.gcp_connector.gcp.projects import list_gke_clusters
    clusters = list_gke_clusters(user_id, project_id, location)
    
    return json.dumps({"status": "success", "clusters": clusters})
```

### Tool Context Access

Tools access user context via thread-local storage:

```python theme={null}
from chat.backend.agent.tools.cloud_tools import get_user_context

ctx = get_user_context()
user_id = ctx["user_id"]
session_id = ctx["session_id"]
provider = ctx["provider_preference"]
mode = ctx["mode"]  # "agent" or "ask"
```

Reference: `server/chat/backend/agent/tools/cloud_tools.py:58-98`

### Tool Output Streaming

Tools can send real-time updates via WebSocket:

```python theme={null}
from chat.backend.agent.tools.cloud_tools import get_websocket_sender

sender = get_websocket_sender()
if sender:
    await sender({
        "type": "tool_progress",
        "data": {"message": "Deploying cluster..."}
    })
```

Reference: `server/main_chatbot.py:176-215`

### Error Handling

Tools return structured JSON with status:

```python theme={null}
# Success
return json.dumps({
    "status": "success",
    "data": {...}
})

# Failure
return json.dumps({
    "status": "failed",
    "error": "Permission denied",
    "message": "User lacks IAM permissions"
})
```

***

## Access Control

### Agent vs Ask Mode

Aurora supports two operational modes:

**Agent Mode** (Full Access):

* Can execute infrastructure changes
* Can commit code to GitHub
* Can deploy Terraform resources
* Can run destructive operations

**Ask Mode** (Read-Only):

* Can read cloud resources
* Can search knowledge bases
* Cannot modify infrastructure
* Cannot commit code

Reference: `server/chat/backend/agent/access.py`

### Tool Filtering

```python theme={null}
from chat.backend.agent.access import ModeAccessController

if not ModeAccessController.is_tool_allowed(mode, tool_name):
    return {"error": "Tool not available in Ask mode"}
```

Reference: `server/main_chatbot.py:867-876`

***

## LLM Integration

### Multi-Provider Support

Aurora supports multiple LLM providers:

```python theme={null}
from chat.backend.agent.llm import LLMManager

llm_manager = LLMManager()
llm = llm_manager.get_llm(
    model="openai/gpt-4",
    streaming=True,
    temperature=0.0
)
```

Supported providers:

* **OpenRouter**: Access 100+ models (default)
* **OpenAI**: GPT-3.5, GPT-4, GPT-4o
* **Anthropic**: Claude 3.5 Sonnet, Claude 3 Opus
* **Google**: Gemini Pro, Gemini 2.0 (with thinking)

Reference: `server/chat/backend/agent/llm.py`

### Prompt Caching

Aurora uses prefix caching to reduce latency and costs:

```python theme={null}
from chat.backend.agent.utils.prefix_cache import PrefixCacheManager

cache_manager = PrefixCacheManager()
cache_manager.register_breakpoint(
    "system_prompt",
    content=system_prompt,
    priority=100
)
```

Reference: `server/chat/backend/agent/utils/prefix_cache.py`

***

## MCP Integration

Aurora supports **Model Context Protocol (MCP)** for extending agent capabilities:

### MCP Preloader

MCP servers are preloaded on startup for faster response times:

```python theme={null}
from chat.backend.agent.tools.mcp_preloader import start_mcp_preloader

mcp_preloader = start_mcp_preloader()
```

Reference: `server/main_compute.py:102-106`

### Dynamic Tool Loading

MCP tools are dynamically loaded based on user connections:

```python theme={null}
from chat.backend.agent.tools.mcp_tools import get_mcp_tools_for_user

mcp_tools = get_mcp_tools_for_user(user_id)
all_tools = base_tools + mcp_tools
```

Reference: `server/chat/backend/agent/tools/mcp_tools.py`

***

## Streaming Architecture

### Token Streaming

LLM responses stream token-by-token for perceived speed:

```python theme={null}
async for event_type, event_data in workflow.stream(state):
    if event_type == "token":
        await websocket.send(json.dumps({
            "type": "message",
            "data": {"text": token, "is_chunk": True}
        }))
```

Reference: `server/main_chatbot.py:341-360`

### Tool Call Streaming

Tool execution status is streamed in real-time:

```python theme={null}
# Tool call started
await websocket.send(json.dumps({
    "type": "tool_call",
    "data": {
        "tool_name": "list_gke_clusters",
        "input": {"project_id": "my-project"},
        "status": "running"
    }
}))

# Tool call completed
await websocket.send(json.dumps({
    "type": "tool_result",
    "data": {
        "tool_name": "list_gke_clusters",
        "result": {...},
        "status": "success"
    }
}))
```

Reference: `server/main_chatbot.py:386-414`

***

## Cancellation & Cleanup

### Workflow Cancellation

Users can cancel in-progress workflows:

```json theme={null}
{
  "type": "control",
  "action": "cancel",
  "session_id": "session-456"
}
```

The system:

1. Cancels the asyncio task
2. Waits for ongoing tool calls to complete
3. Consolidates message chunks
4. Saves context for resumption
5. Sends END status to frontend

Reference: `server/main_chatbot.py:709-800`

### Terraform Cleanup

Terraform state is cleaned up after deployment:

```python theme={null}
from utils.terraform.terraform_cleanup import cleanup_terraform_directory

cleanup_terraform_directory(user_id, session_id)
```

Reference: `server/chat/backend/agent/agent.py:40-78`

***

## Related Documentation

* [Architecture Overview](/architecture/overview)
* [Service Components](/architecture/services)
* [Data Flow](/architecture/data-flow)