@megabrain/agent-sdk

Official SDK for developing custom agents for the MegaBrain platform. Create powerful AI agents with document processing, language detection, and extensible service integration.

Features

• 🚀 Easy Agent Development: Simple base class with common functionality
• 🌍 Multi-language Support: Automatic language detection and response adaptation
• 📄 Document Processing: Built-in file reference extraction and document search
• 🔧 Service Integration: Abstract services for LLM, document search, web search, and data analysis
• ✅ Type Safety: Full TypeScript support with strict typing
• 🧪 Testing Utilities: Comprehensive testing framework with mock services
• 📦 Validation: Built-in validation for manifests and contexts
• 🛡️ Error Handling: Structured error types and handling

Installation

npm install @megabrain/agent-sdk
# or
yarn add @megabrain/agent-sdk

Quick Start

Creating Your First Agent

import { BaseAgent, SDKAgentContext, AgentResponse, AgentServices } from '@megabrain/agent-sdk';

export class MyCustomAgent extends BaseAgent {
  constructor(services: AgentServices) {
    super(services);
  }

  async process(context: SDKAgentContext): Promise<AgentResponse> {
    // Validate the context
    this.validateContext(context);

    // Extract file references and clean query
    const { filenames, cleanQuery } = this.extractFileReferences(context.query);
    const query = cleanQuery || context.query;

    // Use services to process the query
    if (this.hasDocumentService() && filenames.length > 0) {
      const documents = await this.services.document.search(query);
      const context = documents.map(doc => doc.content).join('\n');
      
      if (this.hasLLMService()) {
        const response = await this.services.llm.complete(
          `Answer based on this context: ${context}\n\nQuestion: ${query}`
        );
        
        return this.createSuccessResponse(response, `Found ${documents.length} relevant documents`);
      }
    }

    // Simple response without additional services
    return this.createSuccessResponse(
      `You asked: "${query}"`,
      'Processed successfully'
    );
  }
}

Creating an Agent Plugin

import { createAgentPlugin, AgentServices } from '@megabrain/agent-sdk';

export const myAgentPlugin = createAgentPlugin({
  name: "my-custom-agent",
  description: "A custom agent that processes queries with document context",
  version: "1.0.0",
  author: "Your Name",
  createAgent: (services: AgentServices) => new MyCustomAgent(services),
  manifest: {
    name: "my-custom-agent",
    version: "1.0.0",
    description: "A custom agent that processes queries with document context",
    author: "Your Name",
    capabilities: [
      {
        type: "document-processing",
        description: "Processes documents and answers questions based on content"
      }
    ],
    requirements: {
      minSDKVersion: "1.0.0",
      services: ["document", "llm"],
      permissions: ["read-documents"]
    }
  }
});

Core Concepts

Base Agent

All agents extend the BaseAgent class which provides:

• Language Detection: Automatic detection of user language with detectLanguage()
• File Reference Extraction: Parse file references like file:document.pdf with extractFileReferences()
• Message Formatting: Format conversation history with formatMessagesForCompletion()
• Validation: Context validation with validateContext()
• Error Handling: Standardized error responses with createErrorResponse()

Services

Agents receive services through dependency injection:

interface AgentServices {
  document: DocumentService;    // Document search and processing
  llm: LLMService;             // Large Language Model interactions
  web?: WebSearchService;      // Web search (optional)
  dataAnalysis?: DataAnalysisService; // Data analysis (optional)
}

Context

Every agent receives an SDKAgentContext with:

interface SDKAgentContext {
  query: string;              // User's query
  sdkVersion: string;         // SDK version
  messages: ChatMessage[];    // Conversation history
  fileRefs?: string[];        // Referenced files
  userId?: number;            // User identifier
  agentId?: string;           // Agent identifier
  // ... and more
}

Advanced Usage

Document Processing Agent

export class DocumentAgent extends BaseAgent {
  async process(context: SDKAgentContext): Promise<AgentResponse> {
    const { filenames, cleanQuery } = this.extractFileReferences(context.query);
    
    if (!this.hasDocumentService()) {
      throw new ServiceUnavailableError("document");
    }

    // Search for relevant documents
    const documents = await this.services.document.search(cleanQuery, {
      limit: 5,
      fileTypes: ['pdf', 'txt', 'docx']
    });

    if (documents.length === 0) {
      return this.createSuccessResponse(
        "I couldn't find any relevant documents for your query.",
        "No documents found"
      );
    }

    // Generate response using LLM with document context
    if (this.hasLLMService()) {
      const documentContext = documents
        .map(doc => `Document: ${doc.metadata.filename}\n${doc.content}`)
        .join('\n\n');

      const messages = this.formatMessagesForCompletion(
        "You are a helpful assistant. Answer questions based on the provided documents.",
        cleanQuery,
        context.messages
      );

      const response = await this.services.llm.chat(messages, {
        maxTokens: context.maxTokens || 500
      });

      return this.createSuccessResponse(
        response,
        `Analyzed ${documents.length} documents`,
        { documentsCount: documents.length, sources: documents.map(d => d.metadata.filename) }
      );
    }

    // Fallback without LLM
    return this.createSuccessResponse(
      `Found ${documents.length} relevant documents: ${documents.map(d => d.metadata.filename).join(', ')}`,
      `Document search completed`
    );
  }
}

Streaming Agent

export class StreamingAgent extends BaseAgent {
  async process(context: SDKAgentContext): Promise<AgentResponse> {
    if (!context.onProgress) {
      return this.createErrorResponse(new Error("Streaming not supported in this context"));
    }

    const response = await this.services.llm.complete(context.query, {
      stream: true,
      maxTokens: context.maxTokens
    });

    // Simulate streaming (in real implementation, this would stream from LLM)
    const words = response.split(' ');
    let currentResponse = '';

    for (const word of words) {
      currentResponse += word + ' ';
      context.onProgress(currentResponse.trim());
      await new Promise(resolve => setTimeout(resolve, 100)); // Simulate delay
    }

    return this.createSuccessResponse(response, "Streaming completed");
  }
}

Testing

The SDK includes comprehensive testing utilities:

import { 
  AgentTestRunner, 
  createMockServices, 
  createMockContext,
  MockDocumentService 
} from '@megabrain/agent-sdk/testing';

// Create test runner with mock services
const testRunner = new AgentTestRunner(createMockServices({
  includeDocument: true,
  includeLLM: true
}));

// Run individual test
const result = await testRunner.testAgent(MyCustomAgent, createMockContext({
  query: "What is artificial intelligence?",
  maxTokens: 100
}));

console.log('Test result:', result);

// Run test suite
await testRunner.runTestSuite(MyCustomAgent, [
  {
    name: "Basic query test",
    context: createMockContext({ query: "Hello world" }),
    expectedSuccess: true
  },
  {
    name: "Empty query test", 
    context: createMockContext({ query: "" }),
    expectedSuccess: false
  }
]);

Validation

The SDK provides comprehensive validation:

import { 
  validateAgentManifest, 
  validateSDKAgentContext,
  validateRequiredServices 
} from '@megabrain/agent-sdk';

// Validate manifest
try {
  const manifest = validateAgentManifest({
    name: "my-agent",
    version: "1.0.0",
    description: "My custom agent",
    author: "Developer",
    capabilities: [],
    requirements: {
      minSDKVersion: "1.0.0",
      services: ["document"],
      permissions: []
    }
  });
  console.log('Valid manifest:', manifest);
} catch (error) {
  console.error('Invalid manifest:', error.message);
}

// Validate context
try {
  const context = validateSDKAgentContext({
    query: "Test query",
    sdkVersion: "1.0.0",
    messages: []
  });
  console.log('Valid context:', context);
} catch (error) {
  console.error('Invalid context:', error.message);
}

Error Handling

The SDK provides structured error types:

import { 
  SDKError, 
  AgentValidationError, 
  ServiceUnavailableError,
  isSDKError,
  wrapError 
} from '@megabrain/agent-sdk';

try {
  // Your agent code
} catch (error) {
  if (isSDKError(error)) {
    console.error(`SDK Error [${error.code}]:`, error.message);
    console.error('Context:', error.context);
  } else {
    // Wrap unknown errors
    const sdkError = wrapError(error);
    throw sdkError;
  }
}

Utilities

The SDK includes many utility functions:

import { 
  delay, 
  retryWithBackoff, 
  withTimeout,
  generateUUID,
  formatFileSize,
  sanitizeInput 
} from '@megabrain/agent-sdk';

// Retry with exponential backoff
const result = await retryWithBackoff(async () => {
  return await someUnreliableOperation();
}, 3, 1000);

// Add timeout to promises
const response = await withTimeout(
  longRunningOperation(),
  5000, // 5 second timeout
  "Operation timed out"
);

// Generate UUID
const id = generateUUID();

// Format file sizes
const sizeStr = formatFileSize(1024 * 1024); // "1 MB"

// Sanitize user input
const clean = sanitizeInput(userInput);

API Reference

BaseAgent

Methods

• process(context: SDKAgentContext): Promise<AgentResponse> - Abstract Main processing method
• validateContext(context: SDKAgentContext): void - Validates agent context
• detectLanguage(text: string): string | null - Detects text language
• extractFileReferences(query: string): FileReferenceResult - Extracts file references
• formatMessagesForCompletion(systemPrompt: string, query: string, messages?: Array<{role: string; content: string}>): ChatMessage[] - Formats messages for LLM
• createSuccessResponse(answer: string, context?: string, metadata?: Record<string, unknown>): AgentResponse - Creates success response
• createErrorResponse(error: Error, context?: string): AgentResponse - Creates error response

Service Checks

• hasDocumentService(): boolean - Check if document service is available
• hasLLMService(): boolean - Check if LLM service is available
• hasWebSearchService(): boolean - Check if web search service is available
• hasDataAnalysisService(): boolean - Check if data analysis service is available

License

MIT © MegaBrain Team

Support

• Documentation
• GitHub Issues
• Discord Community

Contributing

We welcome contributions! Please see our Contributing Guide for details.