@nomyx/assistant/src/AIAssistant.ts

A powerful assistant library and cli for your AI projects. works with Vertex AI (Claude and Gemini)

import { EventEmitter } from 'events';
import { v4 as uuidv4 } from 'uuid';
import * as dotenv from 'dotenv';
import * as path from 'path';
import {
  ILogger,
  IAIProvider,
  ExecutionOptions,
  Tool,
  StructuredPrompt,
  StreamedResponse
} from './types';
import { PersistentState } from './PersistentState';
import { RequestHistory } from './RequestHistory';
import { ContextManager } from './ContextManager';
import { TaskExecutor } from './TaskExecutor';
import { AIProviderFactory } from './ai/AIProviderFactory';
import { ProviderInitializer } from './providers/ProviderInitializer';
import { ToolAndPromptLoader } from './utils/ToolAndPromptLoader';
import { Middleware } from './types/middleware';

dotenv.config();

export class AIAssistant extends EventEmitter {
  private state: PersistentState;
  private history: RequestHistory;
  private context: ContextManager;
  private taskExecutor!: TaskExecutor;
  private tools: Record<string, Tool> = {};
  private prompts: Record<string, StructuredPrompt> = {};
  private provider!: IAIProvider;
  private providerInitializer: ProviderInitializer;
  private toolAndPromptLoader: ToolAndPromptLoader;
  private middlewares: Middleware[] = [];

  constructor(
    private logger: ILogger,
    private options: ExecutionOptions & {
      maxHistoryLength?: number;
      initialWorkingDirectory?: string;
      providerName?: string;
      toolsPath?: string;
      promptsPath?: string;
    } = {}
  ) {
    super();
    this.state = new PersistentState(options.statePersistenceDir, logger);
    this.history = new RequestHistory(options.maxHistoryLength || 5);
    this.context = new ContextManager(logger);
    this.providerInitializer = new ProviderInitializer(logger);
    this.toolAndPromptLoader = new ToolAndPromptLoader(logger);
  }

  async initialize(): Promise<void> {
    this.initializeProvider(this.options.providerName);
    await this.initializeToolsAndPrompts(this.options.toolsPath, this.options.promptsPath);

    this.logger.debug('Prompts before creating TaskExecutor:', { promptNames: Object.keys(this.prompts) });

    this.taskExecutor = new TaskExecutor(
      this.provider,
      this.state,
      this.history,
      this.context,
      this.tools,
      this.prompts,
      this.logger
    );

    this.updateWorkingDirectory(this.options.initialWorkingDirectory || process.cwd());
    this.logger.info('AIAssistant initialized');
  }

  private initializeProvider(providerName?: string): void {
    const config = this.providerInitializer.initializeProvider(providerName);
    const factory = new AIProviderFactory(this.logger);
    this.provider = factory.createProvider(config);
    this.logger.info(`Initialized AI provider: ${this.provider.constructor.name} with config: ${JSON.stringify(config)}`);
  }

  private async initializeToolsAndPrompts(toolsPath?: string, promptsPath?: string): Promise<void> {
    const { tools, prompts } = await this.toolAndPromptLoader.initializeToolsAndPrompts(toolsPath, promptsPath);
    this.tools = tools;
    this.prompts = prompts;
  }

  async processRequest(input: string): Promise<string> {
    const requestId = uuidv4();
    this.emit('requestStart', { requestId, input });
    this.logger.info(`Processing request: ${requestId}`, { input });

    try {
      let result = await this.taskExecutor.execute(input);
      
      // Apply middlewares
      for (const middleware of this.middlewares) {
        result = await middleware.process(result);
      }

      this.history.addEntry(input, result);
      this.emit('requestComplete', { requestId, result });
      this.logger.info(`Request completed: ${requestId}`);
      return result;
    } catch (error: unknown) {
      this.emit('requestError', { requestId, error });
      this.logger.error(`Error processing request: ${requestId}`, { error });
      throw error;
    }
  }

  async *processRequestStream(input: string): AsyncIterableIterator<string> {
    const requestId = uuidv4();
    this.emit('requestStart', { requestId, input });
    this.logger.info(`Processing streaming request: ${requestId}`, { input });

    try {
      const stream = this.taskExecutor.executeStream(input);
      let fullResponse = '';

      for await (const chunk of stream) {
        fullResponse += chunk;
        yield chunk;
      }

      // Apply middlewares to the full response
      for (const middleware of this.middlewares) {
        fullResponse = await middleware.process(fullResponse);
      }

      this.history.addEntry(input, fullResponse);
      this.emit('requestComplete', { requestId, result: fullResponse });
      this.logger.info(`Streaming request completed: ${requestId}`);
    } catch (error: unknown) {
      this.emit('requestError', { requestId, error });
      this.logger.error(`Error processing streaming request: ${requestId}`, { error });
      throw error;
    }
  }

  async *processRequestStreamWithDetails(input: string): AsyncIterableIterator<StreamedResponse> {
    const requestId = uuidv4();
    this.emit('requestStart', { requestId, input });
    this.logger.info(`Processing detailed streaming request: ${requestId}`, { input });

    try {
      const stream = this.taskExecutor.executeStreamWithDetails(input);
      let fullResponse = '';

      for await (const chunk of stream) {
        if (chunk.content) {
          fullResponse += chunk.content;
        }
        yield chunk;
      }

      // Apply middlewares to the full response
      for (const middleware of this.middlewares) {
        fullResponse = await middleware.process(fullResponse);
      }

      this.history.addEntry(input, fullResponse);
      this.emit('requestComplete', { requestId, result: fullResponse });
      this.logger.info(`Detailed streaming request completed: ${requestId}`);
    } catch (error: unknown) {
      this.emit('requestError', { requestId, error });
      this.logger.error(`Error processing detailed streaming request: ${requestId}`, { error });
      throw error;
    }
  }

  updateWorkingDirectory(newWorkingDirectory: string): void {
    const normalizedPath = path.normalize(newWorkingDirectory);
    this.context.updateContext({
      currentWorkingDirectory: normalizedPath,
      parentDirectory: path.dirname(normalizedPath),
      directoryName: path.basename(normalizedPath)
    });
    this.logger.info(`Updated working directory: ${normalizedPath}`);
  }

  getProvider(): IAIProvider {
    return this.provider;
  }

  addMiddleware(middleware: Middleware): void {
    this.middlewares.push(middleware);
    this.logger.info(`Added middleware: ${middleware.constructor.name}`);
  }
}