AI API

The AI API provides access to Trae IDE's artificial intelligence capabilities, enabling intelligent code generation, analysis, and assistance features in your extensions.

Overview

The AI API allows you to:

• Generate code using AI models
• Analyze code for patterns and issues
• Provide intelligent suggestions and completions
• Integrate with chat-based AI assistance
• Customize AI behavior for specific use cases

Core Interfaces

AIProvider

Main interface for accessing AI capabilities.

interface AIProvider {
  readonly models: readonly AIModel[];
  readonly isAvailable: boolean;
  
  generateCode(request: CodeGenerationRequest): Promise<CodeGenerationResponse>;
  analyzeCode(request: CodeAnalysisRequest): Promise<CodeAnalysisResponse>;
  chat(request: ChatRequest): Promise<ChatResponse>;
  complete(request: CompletionRequest): Promise<CompletionResponse>;
  explain(request: ExplanationRequest): Promise<ExplanationResponse>;
}

AIModel

Represents an available AI model.

interface AIModel {
  readonly id: string;
  readonly name: string;
  readonly description: string;
  readonly capabilities: AICapability[];
  readonly maxTokens: number;
  readonly supportedLanguages: string[];
}

enum AICapability {
  CodeGeneration = 'code-generation',
  CodeAnalysis = 'code-analysis',
  Chat = 'chat',
  Completion = 'completion',
  Explanation = 'explanation',
  Translation = 'translation',
  Refactoring = 'refactoring'
}

Getting Started

Accessing the AI Provider

import { TraeAPI } from '@trae/api';

// Get the AI provider
const aiProvider = TraeAPI.ai.getProvider();

if (aiProvider.isAvailable) {
  console.log('AI is available');
  console.log('Available models:', aiProvider.models.map(m => m.name));
} else {
  console.log('AI is not available');
}

Checking Model Capabilities

function findModelWithCapability(capability: AICapability): AIModel | undefined {
  return aiProvider.models.find(model => 
    model.capabilities.includes(capability)
  );
}

// Find a model that can generate code
const codeGenModel = findModelWithCapability(AICapability.CodeGeneration);
if (codeGenModel) {
  console.log(`Using ${codeGenModel.name} for code generation`);
}

Code Generation

Basic Code Generation

interface CodeGenerationRequest {
  prompt: string;
  language: string;
  context?: CodeContext;
  model?: string;
  maxTokens?: number;
  temperature?: number;
  stopSequences?: string[];
}

interface CodeGenerationResponse {
  code: string;
  explanation?: string;
  confidence: number;
  model: string;
  usage: TokenUsage;
}

// Generate a function
const request: CodeGenerationRequest = {
  prompt: 'Create a function that calculates the factorial of a number',
  language: 'typescript',
  context: {
    currentFile: editor.document.uri,
    cursorPosition: editor.selection.active,
    selectedText: editor.document.getText(editor.selection)
  }
};

const response = await aiProvider.generateCode(request);
console.log('Generated code:', response.code);
console.log('Explanation:', response.explanation);

Advanced Code Generation

// Generate code with specific requirements
const advancedRequest: CodeGenerationRequest = {
  prompt: 'Create a React component for a user profile card',
  language: 'typescript',
  context: {
    currentFile: editor.document.uri,
    imports: [
      'import React from "react";',
      'import { User } from "./types";'
    ],
    existingCode: editor.document.getText(),
    projectStructure: await getProjectStructure()
  },
  maxTokens: 1000,
  temperature: 0.7,
  stopSequences: ['\n\n\n']
};

const response = await aiProvider.generateCode(advancedRequest);

// Insert the generated code
if (response.confidence > 0.8) {
  await editor.edit(editBuilder => {
    editBuilder.insert(editor.selection.active, response.code);
  });
} else {
  // Show the code for review
  TraeAPI.window.showInformationMessage(
    `Generated code (confidence: ${response.confidence})`,
    'Insert', 'Review'
  ).then(choice => {
    if (choice === 'Insert') {
      editor.edit(editBuilder => {
        editBuilder.insert(editor.selection.active, response.code);
      });
    } else if (choice === 'Review') {
      showCodeReviewPanel(response);
    }
  });
}

Code Context

interface CodeContext {
  currentFile?: TraeAPI.Uri;
  cursorPosition?: TraeAPI.Position;
  selectedText?: string;
  imports?: string[];
  existingCode?: string;
  projectStructure?: ProjectStructure;
  dependencies?: string[];
  framework?: string;
  codeStyle?: CodeStyle;
}

interface ProjectStructure {
  files: string[];
  directories: string[];
  packageJson?: any;
  tsConfig?: any;
}

interface CodeStyle {
  indentation: 'spaces' | 'tabs';
  indentSize: number;
  quotes: 'single' | 'double';
  semicolons: boolean;
  trailingCommas: boolean;
}

Code Analysis

Analyzing Code Quality

interface CodeAnalysisRequest {
  code: string;
  language: string;
  analysisType: AnalysisType[];
  context?: CodeContext;
}

enum AnalysisType {
  Quality = 'quality',
  Security = 'security',
  Performance = 'performance',
  Maintainability = 'maintainability',
  Bugs = 'bugs',
  CodeSmells = 'code-smells',
  Complexity = 'complexity'
}

interface CodeAnalysisResponse {
  issues: CodeIssue[];
  metrics: CodeMetrics;
  suggestions: CodeSuggestion[];
  overallScore: number;
}

// Analyze current file
const editor = TraeAPI.window.activeTextEditor;
if (editor) {
  const analysisRequest: CodeAnalysisRequest = {
    code: editor.document.getText(),
    language: editor.document.languageId,
    analysisType: [AnalysisType.Quality, AnalysisType.Security, AnalysisType.Performance]
  };
  
  const analysis = await aiProvider.analyzeCode(analysisRequest);
  
  // Show issues in problems panel
  const diagnostics: TraeAPI.Diagnostic[] = analysis.issues.map(issue => ({
    range: new TraeAPI.Range(
      new TraeAPI.Position(issue.line - 1, issue.column),
      new TraeAPI.Position(issue.line - 1, issue.column + issue.length)
    ),
    message: issue.message,
    severity: mapSeverity(issue.severity),
    source: 'AI Analysis',
    code: issue.code
  }));
  
  TraeAPI.languages.createDiagnosticCollection('ai-analysis')
    .set(editor.document.uri, diagnostics);
}

Code Issue Types

interface CodeIssue {
  type: AnalysisType;
  severity: IssueSeverity;
  message: string;
  line: number;
  column: number;
  length: number;
  code?: string;
  suggestion?: string;
  fixable: boolean;
}

enum IssueSeverity {
  Error = 'error',
  Warning = 'warning',
  Info = 'info',
  Hint = 'hint'
}

interface CodeMetrics {
  complexity: number;
  maintainabilityIndex: number;
  linesOfCode: number;
  duplicatedLines: number;
  testCoverage?: number;
}

interface CodeSuggestion {
  type: 'refactor' | 'optimize' | 'simplify' | 'modernize';
  description: string;
  before: string;
  after: string;
  impact: 'low' | 'medium' | 'high';
}

Chat Integration

Basic Chat

interface ChatRequest {
  message: string;
  context?: ChatContext;
  model?: string;
  conversationId?: string;
}

interface ChatContext {
  currentFile?: TraeAPI.Uri;
  selectedCode?: string;
  projectContext?: ProjectContext;
  previousMessages?: ChatMessage[];
}

interface ChatResponse {
  message: string;
  actions?: ChatAction[];
  conversationId: string;
  model: string;
}

// Start a chat conversation
const chatRequest: ChatRequest = {
  message: 'How can I optimize this React component?',
  context: {
    currentFile: editor.document.uri,
    selectedCode: editor.document.getText(editor.selection),
    projectContext: {
      framework: 'react',
      typescript: true,
      dependencies: ['react', '@types/react']
    }
  }
};

const chatResponse = await aiProvider.chat(chatRequest);
console.log('AI Response:', chatResponse.message);

// Handle suggested actions
chatResponse.actions?.forEach(action => {
  if (action.type === 'code-edit') {
    // Show code edit suggestion
    showCodeEditSuggestion(action);
  } else if (action.type === 'file-create') {
    // Show file creation suggestion
    showFileCreationSuggestion(action);
  }
});

Chat Actions

interface ChatAction {
  type: ChatActionType;
  title: string;
  description?: string;
  data: any;
}

enum ChatActionType {
  CodeEdit = 'code-edit',
  FileCreate = 'file-create',
  FileDelete = 'file-delete',
  RunCommand = 'run-command',
  InstallPackage = 'install-package',
  OpenFile = 'open-file',
  SearchCode = 'search-code'
}

// Handle chat actions
function handleChatAction(action: ChatAction): void {
  switch (action.type) {
    case ChatActionType.CodeEdit:
      applyCodeEdit(action.data);
      break;
    case ChatActionType.FileCreate:
      createFile(action.data.path, action.data.content);
      break;
    case ChatActionType.RunCommand:
      TraeAPI.commands.executeCommand(action.data.command, ...action.data.args);
      break;
    case ChatActionType.InstallPackage:
      installPackage(action.data.package);
      break;
  }
}

Code Completion

Intelligent Completions

interface CompletionRequest {
  document: TraeAPI.TextDocument;
  position: TraeAPI.Position;
  context: CompletionContext;
  maxSuggestions?: number;
}

interface CompletionContext {
  triggerKind: CompletionTriggerKind;
  triggerCharacter?: string;
  linePrefix: string;
  lineSuffix: string;
}

enum CompletionTriggerKind {
  Invoked = 1,
  TriggerCharacter = 2,
  TriggerForIncompleteCompletions = 3
}

interface CompletionResponse {
  suggestions: CompletionSuggestion[];
  isIncomplete: boolean;
}

interface CompletionSuggestion {
  label: string;
  kind: CompletionItemKind;
  detail?: string;
  documentation?: string;
  insertText: string;
  range?: TraeAPI.Range;
  confidence: number;
  sortText?: string;
}

// Register AI completion provider
TraeAPI.languages.registerCompletionItemProvider('typescript', {
  async provideCompletionItems(
    document: TraeAPI.TextDocument,
    position: TraeAPI.Position,
    token: TraeAPI.CancellationToken,
    context: TraeAPI.CompletionContext
  ): Promise<TraeAPI.CompletionList> {
    const line = document.lineAt(position);
    const linePrefix = line.text.substring(0, position.character);
    const lineSuffix = line.text.substring(position.character);
    
    const completionRequest: CompletionRequest = {
      document,
      position,
      context: {
        triggerKind: context.triggerKind,
        triggerCharacter: context.triggerCharacter,
        linePrefix,
        lineSuffix
      },
      maxSuggestions: 10
    };
    
    const response = await aiProvider.complete(completionRequest);
    
    const items = response.suggestions.map(suggestion => {
      const item = new TraeAPI.CompletionItem(suggestion.label, suggestion.kind);
      item.detail = suggestion.detail;
      item.documentation = suggestion.documentation;
      item.insertText = suggestion.insertText;
      item.range = suggestion.range;
      item.sortText = suggestion.sortText || suggestion.confidence.toString();
      return item;
    });
    
    return new TraeAPI.CompletionList(items, response.isIncomplete);
  }
}, '.', '(', '[');

Code Explanation

Explaining Code

interface ExplanationRequest {
  code: string;
  language: string;
  explanationType: ExplanationType;
  context?: CodeContext;
}

enum ExplanationType {
  Overview = 'overview',
  LineByLine = 'line-by-line',
  Concepts = 'concepts',
  Performance = 'performance',
  Security = 'security'
}

interface ExplanationResponse {
  explanation: string;
  sections: ExplanationSection[];
  relatedConcepts: string[];
  suggestions: string[];
}

interface ExplanationSection {
  title: string;
  content: string;
  codeRange?: TraeAPI.Range;
  importance: 'low' | 'medium' | 'high';
}

// Explain selected code
const editor = TraeAPI.window.activeTextEditor;
if (editor && !editor.selection.isEmpty) {
  const selectedCode = editor.document.getText(editor.selection);
  
  const explanationRequest: ExplanationRequest = {
    code: selectedCode,
    language: editor.document.languageId,
    explanationType: ExplanationType.Overview,
    context: {
      currentFile: editor.document.uri,
      cursorPosition: editor.selection.active
    }
  };
  
  const explanation = await aiProvider.explain(explanationRequest);
  
  // Show explanation in a webview
  showExplanationPanel(explanation);
}

Advanced Features

Custom AI Models

// Register a custom AI model
class CustomAIModel implements AIModel {
  readonly id = 'custom-model-v1';
  readonly name = 'Custom Code Model';
  readonly description = 'A specialized model for our codebase';
  readonly capabilities = [AICapability.CodeGeneration, AICapability.CodeAnalysis];
  readonly maxTokens = 4096;
  readonly supportedLanguages = ['typescript', 'javascript', 'python'];
  
  async generateCode(request: CodeGenerationRequest): Promise<CodeGenerationResponse> {
    // Custom implementation
    const response = await this.callCustomAPI(request);
    return {
      code: response.generatedCode,
      explanation: response.explanation,
      confidence: response.confidence,
      model: this.id,
      usage: response.tokenUsage
    };
  }
  
  private async callCustomAPI(request: CodeGenerationRequest): Promise<any> {
    // Implementation details
  }
}

// Register the custom model
TraeAPI.ai.registerModel(new CustomAIModel());

AI-Powered Refactoring

class AIRefactoring {
  static async extractFunction(
    editor: TraeAPI.TextEditor,
    selection: TraeAPI.Selection
  ): Promise<void> {
    const selectedCode = editor.document.getText(selection);
    
    const request: CodeGenerationRequest = {
      prompt: `Extract the following code into a well-named function:\n\n${selectedCode}`,
      language: editor.document.languageId,
      context: {
        currentFile: editor.document.uri,
        selectedText: selectedCode,
        existingCode: editor.document.getText()
      }
    };
    
    const response = await TraeAPI.ai.getProvider().generateCode(request);
    
    // Apply the refactoring
    await editor.edit(editBuilder => {
      // Replace selected code with function call
      editBuilder.replace(selection, response.code.functionCall);
      
      // Insert function definition
      const insertPosition = this.findBestInsertPosition(editor.document);
      editBuilder.insert(insertPosition, response.code.functionDefinition);
    });
  }
  
  static async renameVariable(
    editor: TraeAPI.TextEditor,
    position: TraeAPI.Position
  ): Promise<void> {
    const wordRange = editor.document.getWordRangeAtPosition(position);
    if (!wordRange) return;
    
    const currentName = editor.document.getText(wordRange);
    const context = this.getVariableContext(editor.document, position);
    
    const request: CodeGenerationRequest = {
      prompt: `Suggest a better name for the variable '${currentName}' based on its usage context`,
      language: editor.document.languageId,
      context: {
        currentFile: editor.document.uri,
        selectedText: currentName,
        existingCode: context
      }
    };
    
    const response = await TraeAPI.ai.getProvider().generateCode(request);
    const suggestedName = response.code.trim();
    
    // Show rename suggestion
    const choice = await TraeAPI.window.showInformationMessage(
      `Rename '${currentName}' to '${suggestedName}'?`,
      'Yes', 'No'
    );
    
    if (choice === 'Yes') {
      await TraeAPI.commands.executeCommand('editor.action.rename', [
        editor.document.uri,
        position,
        suggestedName
      ]);
    }
  }
}

AI Code Review

class AICodeReview {
  static async reviewChanges(changes: TraeAPI.TextDocumentChangeEvent[]): Promise<void> {
    for (const change of changes) {
      const changedCode = this.extractChangedCode(change);
      
      const analysisRequest: CodeAnalysisRequest = {
        code: changedCode,
        language: change.document.languageId,
        analysisType: [
          AnalysisType.Quality,
          AnalysisType.Security,
          AnalysisType.Performance,
          AnalysisType.Bugs
        ]
      };
      
      const analysis = await TraeAPI.ai.getProvider().analyzeCode(analysisRequest);
      
      // Show review comments
      this.showReviewComments(change.document, analysis);
    }
  }
  
  static async reviewPullRequest(files: string[]): Promise<ReviewResult> {
    const reviews: FileReview[] = [];
    
    for (const file of files) {
      const document = await TraeAPI.workspace.openTextDocument(file);
      const content = document.getText();
      
      const request: CodeAnalysisRequest = {
        code: content,
        language: document.languageId,
        analysisType: [AnalysisType.Quality, AnalysisType.Security]
      };
      
      const analysis = await TraeAPI.ai.getProvider().analyzeCode(request);
      
      reviews.push({
        file,
        issues: analysis.issues,
        suggestions: analysis.suggestions,
        score: analysis.overallScore
      });
    }
    
    return {
      files: reviews,
      overallScore: this.calculateOverallScore(reviews),
      summary: this.generateSummary(reviews)
    };
  }
}

interface FileReview {
  file: string;
  issues: CodeIssue[];
  suggestions: CodeSuggestion[];
  score: number;
}

interface ReviewResult {
  files: FileReview[];
  overallScore: number;
  summary: string;
}

Configuration

AI Settings

interface AIConfiguration {
  defaultModel: string;
  maxTokens: number;
  temperature: number;
  enableCodeGeneration: boolean;
  enableCodeAnalysis: boolean;
  enableChat: boolean;
  enableCompletion: boolean;
  customModels: CustomModelConfig[];
  apiKeys: Record<string, string>;
}

interface CustomModelConfig {
  id: string;
  name: string;
  endpoint: string;
  apiKey: string;
  capabilities: AICapability[];
}

// Get AI configuration
const config = TraeAPI.workspace.getConfiguration('trae.ai');
const aiConfig: AIConfiguration = {
  defaultModel: config.get('defaultModel', 'gpt-4'),
  maxTokens: config.get('maxTokens', 2048),
  temperature: config.get('temperature', 0.7),
  enableCodeGeneration: config.get('enableCodeGeneration', true),
  enableCodeAnalysis: config.get('enableCodeAnalysis', true),
  enableChat: config.get('enableChat', true),
  enableCompletion: config.get('enableCompletion', true),
  customModels: config.get('customModels', []),
  apiKeys: config.get('apiKeys', {})
};

Error Handling

AI API Errors

enum AIErrorCode {
  ModelNotAvailable = 'MODEL_NOT_AVAILABLE',
  RateLimitExceeded = 'RATE_LIMIT_EXCEEDED',
  InvalidRequest = 'INVALID_REQUEST',
  NetworkError = 'NETWORK_ERROR',
  AuthenticationError = 'AUTHENTICATION_ERROR',
  TokenLimitExceeded = 'TOKEN_LIMIT_EXCEEDED'
}

class AIError extends Error {
  constructor(
    public code: AIErrorCode,
    message: string,
    public details?: any
  ) {
    super(message);
    this.name = 'AIError';
  }
}

// Handle AI errors
try {
  const response = await aiProvider.generateCode(request);
  // Use response
} catch (error) {
  if (error instanceof AIError) {
    switch (error.code) {
      case AIErrorCode.RateLimitExceeded:
        TraeAPI.window.showWarningMessage(
          'AI rate limit exceeded. Please try again later.'
        );
        break;
      case AIErrorCode.ModelNotAvailable:
        TraeAPI.window.showErrorMessage(
          `AI model not available: ${error.message}`
        );
        break;
      default:
        TraeAPI.window.showErrorMessage(
          `AI error: ${error.message}`
        );
    }
  } else {
    TraeAPI.window.showErrorMessage(
      `Unexpected error: ${error.message}`
    );
  }
}

Best Practices

Performance

• Cache AI responses when appropriate
• Use streaming for long responses
• Implement request debouncing for real-time features
• Monitor token usage and costs

User Experience

• Provide loading indicators for AI operations
• Allow users to cancel long-running requests
• Show confidence scores for AI suggestions
• Provide fallback options when AI is unavailable

Security

• Never send sensitive data to external AI services
• Validate AI responses before applying changes
• Implement proper authentication for AI services
• Respect user privacy preferences

Examples

AI-Powered Code Assistant

export class AICodeAssistant {
  private aiProvider: AIProvider;
  
  constructor() {
    this.aiProvider = TraeAPI.ai.getProvider();
  }
  
  async suggestFix(diagnostic: TraeAPI.Diagnostic, document: TraeAPI.TextDocument): Promise<void> {
    const problemCode = document.getText(diagnostic.range);
    const context = this.getCodeContext(document, diagnostic.range);
    
    const request: CodeGenerationRequest = {
      prompt: `Fix this ${diagnostic.severity} issue: ${diagnostic.message}\n\nCode: ${problemCode}`,
      language: document.languageId,
      context
    };
    
    const response = await this.aiProvider.generateCode(request);
    
    // Show quick fix
    const action = new TraeAPI.CodeAction(
      `AI Fix: ${diagnostic.message}`,
      TraeAPI.CodeActionKind.QuickFix
    );
    
    action.edit = new TraeAPI.WorkspaceEdit();
    action.edit.replace(document.uri, diagnostic.range, response.code);
    
    return action;
  }
  
  async generateTests(functionRange: TraeAPI.Range, document: TraeAPI.TextDocument): Promise<string> {
    const functionCode = document.getText(functionRange);
    
    const request: CodeGenerationRequest = {
      prompt: `Generate comprehensive unit tests for this function:\n\n${functionCode}`,
      language: document.languageId,
      context: {
        currentFile: document.uri,
        existingCode: document.getText()
      }
    };
    
    const response = await this.aiProvider.generateCode(request);
    return response.code;
  }
}