Azad Client Documentation

Overview

The Azad Client is a high-level TypeScript client that facilitates communication between your application and an AI assistant through WebSocket connections. It manages task execution, tool calling, message streaming, and provides a structured interface for handling various AI events.

Architecture

The Azad Client architecture consists of several key components:

1. WebSocketClient: Low-level client that handles the WebSocket connection and message passing
2. AzadClient: High-level client that orchestrates task execution and event handling
3. LocalEnvironment: Manages tool registration, execution, and provides a bridge between tools and the client
4. Tools: Individual capabilities that can be registered with the environment and called by the AI assistant

┌─────────────┐      ┌────────────────┐      ┌───────────────┐
│  AzadClient │◄────►│ WebSocketClient│◄────►│ Remote Server │
└──────┬──────┘      └────────────────┘      └───────────────┘
       │
       ▼
┌──────────────┐     ┌───────────┐
│ Environment  │◄───►│   Tools   │
└──────────────┘     └───────────┘

Key Interfaces

AzadClient

The AzadClient class is the main entry point for working with the Azad system.

class AzadClient {
  constructor(workingDirectory: string);

  // Connection management
  async connect(): Promise<void>;
  async disconnect(): Promise<void>;

  // Task execution
  async executeTask(options: ExecuteTaskOptions): Promise<void>;

  // Task management
  public loadTask(task: Task): void;
  public answerToolCall(toolCallId: string, result: ApprovalResponse): void;
}

ExecuteTaskOptions

Options for executing a task through the Azad Client:

interface ExecuteTaskOptions {
  // Model configuration
  model: {
    id: string;
    apiKey: string;
  };

  // Tool and task options
  tools: AnyToolSchema[];
  toolDialect?: 'xml' | 'json' | 'native';
  thinkingBudgetTokens?: number;
  maxSteps?: number;
  system?: string;
  maxRetry?: number;
  task?: (ImagePart | TextPart)[];
  initialTask?: (ImagePart | TextPart)[];

  // Event callbacks
  onText?: TextStreamCallback;
  onReasoning?: ReasoningStreamCallback;
  onStepStart?: StepStartCallback;
  onStepEnd?: StepEndCallback;
  onToolExecutionStart?: ToolExecutionStartCallback;
  onToolCallParamsUpdate?: ToolCallParamsUpdateCallback;
  onToolExecutionEnd?: ToolExecutionEndCallback;
  onToolApproval?: ToolApprovalCallback;
  onNetworkError?: (error: Error) => void;
}

WebSocketClient

The WebSocketClient handles low-level WebSocket communication:

class WebSocketClient {
  constructor(url: string, environment: Environment);

  // Connection management
  async connect(): Promise<void>;
  async disconnect(): Promise<void>;

  // API methods
  async abortStep(
    args: Omit<AbortStepMessage['data'], 'operation'>
  ): Promise<any>;
  async getLoadedTask(
    args: Omit<GetLoadedTaskMessage['data'], 'operation'>
  ): Promise<DataResponse<Task | null>>;
  async loadTask(
    args: Omit<LoadTaskMessage['data'], 'operation'>
  ): Promise<MessageResponse | ErrorResponse>;

  // Streaming API methods
  async *runServerTool(
    args: Omit<RunServerToolMessage['data'], 'operation'>
  ): AsyncGenerator<AINetworkEvent, void>;
  async *stepTask(
    args: Omit<StepTaskMessage['data'], 'operation'>
  ): AsyncGenerator<AINetworkEvent, void>;
}

Usage Examples

Basic Task Execution

import { AzadClient } from './server/azad-client';
import { executeCommandTool } from './environment/tools/definitions/execute-command';

// Initialize the client with the current working directory
const client = new AzadClient(process.cwd());

// Execute a task
await client.executeTask({
  model: {
    id: 'gpt-4',
    apiKey: 'your-api-key',
  },
  tools: [executeCommandTool],
  toolDialect: 'xml',
  task: 'List all the files in the current directory and create a summary of their contents.',

  // Text streaming callback
  onText: (chunk, accumulated) => {
    console.log('New text:', chunk);
  },

  // Tool approval callback
  onToolApproval: async (toolName, toolCallId, args) => {
    console.log(`Tool approval requested: ${toolName}`);
    console.log('Arguments:', args);

    // Automatically approve all tools
    return { approved: true, images: [], feedback: null };
  },
});

// Disconnect when done
await client.disconnect();

Advanced Task Execution with Callbacks

const client = new AzadClient(process.cwd());

await client.executeTask({
  model: {
    id: 'gpt-4',
    apiKey: 'your-api-key',
  },
  tools: [executeCommandTool, readFileTool, writeFileTool],
  toolDialect: 'xml',
  maxSteps: 10,
  system: 'You are a helpful coding assistant.',
  task: 'Analyze the package.json file and upgrade all dependencies to their latest versions.',

  // Full event callbacks
  onText: (chunk, accumulated) => {
    process.stdout.write(chunk);
  },

  onReasoning: (chunk, accumulated) => {
    process.stderr.write(chalk.gray(chunk));
  },

  onStepStart: (assistantMessage) => {
    console.log('\n--- Step started ---');
  },

  onStepEnd: (task, newMessages) => {
    console.log('\n--- Step completed ---');
  },

  onToolExecutionStart: (toolCall) => {
    console.log(`\n[Tool] Executing: ${toolCall.tool_name}`);
  },

  onToolCallParamsUpdate: (toolCall) => {
    console.log(`[Tool] Params updated for: ${toolCall.tool_name}`);
  },

  onToolExecutionEnd: (toolCall, result) => {
    console.log(`[Tool] Execution completed: ${toolCall.tool_name}`);
    console.log(`[Tool] Success: ${result.success}`);
  },

  onToolApproval: async (toolName, toolCallId, args) => {
    // Prompt the user for approval
    const answer = await promptUser(`Approve tool ${toolName}? (y/n)`);
    return {
      approved: answer.toLowerCase() === 'y',
      images: [],
      feedback: answer.toLowerCase() !== 'y' ? 'Tool rejected by user' : null,
    };
  },

  onNetworkError: (error) => {
    console.error(`Error: `, error);
  },
});

Task and Message Structure

Tasks in the Azad system have a specific structure:

interface Task {
  id: string;
  content: Array<TaskContent>;
  state: 'running' | 'completed' | 'failed';
  messages: Array<Message>;
}

interface Message {
  role: 'user' | 'assistant' | 'system' | 'taskconfig';
  id: string;
  task_id: string;
  started_ts: number;
  content: Array<MessageContent>;
  // Additional properties depending on the role
}

Event Types

The Azad Client processes various event types during task execution:

1. Connection Events: network_attempt_connection, ai_first_token, network_connection_established, etc.
2. Text Events: text_start, text_chunk, text_end
3. Reasoning Events: reasoning_start, reasoning_chunk, reasoning_end
4. Tool Events: tool_name, parameter_start, parameter_chunk, parameter_end, parameters_complete, tool_ready
5. Task Events: task_message, step_end

Error Handling

The Azad Client includes built-in error handling with retry capabilities:

// Configure retry behavior
await client.executeTask({
  // Other options...
  maxRetry: 3, // Retry up to 3 times on failure
});

Implementation Details

Connection Management

The WebSocketClient handles the low-level WebSocket connection:

1. Connection: Establishes a WebSocket connection to the specified URL
2. Message Queuing: Queues messages when not connected and sends them when the connection is established
3. Reconnection: Handles reconnection with retry logic
4. Event Handling: Processes incoming messages and routes them to the appropriate handlers

Task Execution Flow

1. Task Creation: A task is created with the provided options
2. Tool Registration: Tools are registered with the LocalEnvironment
3. Task Configuration: The task is configured with model settings, tool metadata, etc.
4. Execution: The task is executed by sending a stepTask request to the server
5. Event Processing: Events are processed as they arrive and routed to the appropriate callbacks
6. Task Completion: The task is marked as completed when the execution is finished

Tool Approval and Execution

1. Tool Call: The AI assistant generates a tool call
2. Parameter Collection: Parameters are collected and accumulated
3. Tool Approval: The tool call is sent for approval via the onToolApproval callback
4. Tool Execution: If approved, the tool is executed by the LocalEnvironment
5. Result Processing: The tool result is processed and sent back to the assistant

Advanced Topics

Custom Tools

Custom tools can be registered with the Azad Client by creating a ToolSchema and passing it to the executeTask method. See the Tool Creation Guide for details.

Multi-step Tasks

Complex tasks may require multiple steps to complete. The Azad Client tracks the task state and messages across steps, allowing for continuous interaction.

Parallel Tool Execution

The Azad system supports parallel tool execution, allowing multiple tools to be executed simultaneously in one assistant response for improved performance.

• Note execution is sequental and awaited (Tool > Tool Result, ...)

Best Practices

1. Resource Management: Always disconnect the client when done to free up resources
2. Error Handling: Use the maxRetry option to handle transient errors
3. Tool Approval: Implement proper tool approval to ensure security
4. Environment Details: Provide accurate working directory information for proper tool execution
5. Callback Management: Keep callback functions lightweight to avoid blocking the event processing loop