@nomyx/assistant/src/TaskExecutor.ts

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

import { PersistentState } from './PersistentState';
import { RequestHistory } from './RequestHistory';
import { ContextManager } from './ContextManager';
import { IAIProvider } from './types/provider';
import { ILogger } from './types/common';
import { ChatMessage, ChatOptions, ProviderResponse, StreamedResponse } from './types/chat';
import { Tool, StructuredPrompt } from './types';
import { MessageCollationMiddleware } from './middleware/MessageCollationMiddleware';
import { LoggingMiddleware } from './middleware/LoggingMiddleware';

type InputType = string | Record<string, any>;
type InputSchema = Record<string, 'string' | 'object' | 'number' | 'boolean'>;

export class TaskExecutionError extends Error {
  constructor(message: string, public readonly cause?: unknown) {
    super(message);
    this.name = 'TaskExecutionError';
  }
}

export class TaskExecutor {
  constructor(
    private provider: IAIProvider,
    private state: PersistentState,
    private history: RequestHistory,
    private context: ContextManager,
    private tools: Record<string, Tool>,
    private prompts: Record<string, StructuredPrompt>,
    private logger: ILogger
  ) {
    // Apply middlewares
    this.provider.use(new MessageCollationMiddleware());
    this.provider.use(new LoggingMiddleware(this.logger));
    
    this.logPrompts(); // Log prompts immediately after construction
  }

  async execute(input: InputType, promptName?: string): Promise<string> {
    try {
      const selectedPrompt = promptName ? this.prompts[promptName] : this.selectPrompt(input);
      const messages: ChatMessage[] = this.prepareMessages(input, selectedPrompt);
      const toolsForPrompt = this.prepareTools(selectedPrompt);
      
      this.logger.debug('Executing task with input', { input, selectedPrompt: selectedPrompt.name, toolsForPrompt: toolsForPrompt.map(t => t.name) });
      const response = await this.provider.chat(messages, this.getChatOptions(selectedPrompt), toolsForPrompt);
      this.logger.debug('Received response from provider', { content: response.content, toolCalls: response.toolCalls });

      if (response.toolCalls && response.toolCalls.length > 0) {
        const toolResults = await this.executeToolCalls(response.toolCalls);
        const finalResponse = await this.processFinalResponse(input, toolResults);
        return finalResponse;
      }

      return response.content;
    } catch (error: unknown) {
      this.logger.error('Error executing task', { error, input });
      throw new TaskExecutionError('Failed to execute task', error);
    }
  }

  private async executeToolCalls(toolCalls: ProviderResponse['toolCalls']): Promise<string[]> {
    if (!toolCalls) return [];

    const toolResponses: string[] = [];
    for (const toolCall of toolCalls) {
      const tool = this.tools[toolCall.name];
      if (tool) {
        this.logger.debug('Executing tool', { toolName: toolCall.name, arguments: toolCall.arguments });
        try {
          const toolResult = await tool.execute(toolCall.arguments, {} as any);
          toolResponses.push(`Tool ${toolCall.name} result: ${JSON.stringify(toolResult)}`);
          this.logger.debug('Tool execution result', { toolName: toolCall.name, result: toolResult });
        } catch (error) {
          this.logger.error('Error executing tool', { toolName: toolCall.name, error });
          toolResponses.push(`Error executing ${toolCall.name}: ${error}`);
        }
      } else {
        this.logger.warn('Tool not found for tool call', { toolCall });
        toolResponses.push(`Tool ${toolCall.name} not found`);
      }
    }
    return toolResponses;
  }

  private async processFinalResponse(originalInput: InputType, toolResults: string[]): Promise<string> {
    const processToolResultPrompt = this.prompts['processToolResult'];
    if (!processToolResultPrompt) {
      throw new Error('processToolResult prompt not found');
    }

    const finalInput = {
      toolResult: { result: toolResults.join('\n') },
      originalQuery: typeof originalInput === 'string' ? originalInput : JSON.stringify(originalInput)
    };

    const finalResponse = await this.execute(finalInput, 'processToolResult');
    this.logger.debug('Received final response from provider', { content: finalResponse });
    return finalResponse;
  }

  async *executeStream(input: InputType): AsyncIterableIterator<string> {
    try {
      const selectedPrompt = this.selectPrompt(input);
      const messages: ChatMessage[] = this.prepareMessages(input, selectedPrompt);
      const toolsForPrompt = this.prepareTools(selectedPrompt);
      
      const stream = this.provider.streamChat(messages, this.getChatOptions(selectedPrompt), toolsForPrompt);
      for await (const chunk of stream) {
        yield chunk.content;
      }
    } catch (error: unknown) {
      this.logger.error('Error executing streaming task', { error, input });
      throw new TaskExecutionError('Failed to execute streaming task', error);
    }
  }

  async *executeStreamWithDetails(input: InputType): AsyncIterableIterator<StreamedResponse> {
    try {
      const selectedPrompt = this.selectPrompt(input);
      const messages: ChatMessage[] = this.prepareMessages(input, selectedPrompt);
      const toolsForPrompt = this.prepareTools(selectedPrompt);
      
      const stream = this.provider.streamChat(messages, this.getChatOptions(selectedPrompt), toolsForPrompt);
      for await (const chunk of stream) {
        yield {
          content: chunk.content,
          toolCalls: chunk.toolCalls
        };
      }
    } catch (error: unknown) {
      this.logger.error('Error executing detailed streaming task', { error, input });
      throw new TaskExecutionError('Failed to execute detailed streaming task', error);
    }
  }

  private prepareMessages(input: InputType, prompt: StructuredPrompt): ChatMessage[] {
    const systemPrompt = this.createSystemPrompt(prompt);
    const history = this.history.getFormattedHistory();
    const context = this.context.getFormattedContext();
    
    const userMessage = this.replaceRequestPlaceholder(prompt.user_message, input, prompt.input_schema);
    
    return [
      { role: 'system', content: systemPrompt },
      { role: 'system', content: `Current context:\n${context}` },
      { role: 'system', content: `Recent conversation history:\n${history}` },
      { role: 'user', content: userMessage }
    ];
  }

  private replaceRequestPlaceholder(message: string, input: InputType, inputSchema: InputSchema): string {
    this.logger.debug('Replacing placeholders', { message, input, inputSchema });

    if (typeof input === 'string') {
      const schemaKeys = Object.keys(inputSchema);
      if (schemaKeys.length === 1 && inputSchema[schemaKeys[0]] === 'string') {
        return message.replace(`{${schemaKeys[0]}}`, input);
      } else {
        throw new Error('Input is a string, but input_schema does not match a single string field');
      }
    } else {
      return message.replace(/{(\w+)}/g, (match, key) => {
        if (key in input && key in inputSchema) {
          const value = (input as Record<string, any>)[key];
          if (typeof value === inputSchema[key] || (inputSchema[key] === 'object' && typeof value === 'object')) {
            return JSON.stringify(value);
          } else {
            throw new Error(`Type mismatch for property '${key}': expected ${inputSchema[key]}, got ${typeof value}`);
          }
        } else {
          throw new Error(`Property '${key}' not found in input object or input_schema`);
        }
      });
    }
  }

  private createSystemPrompt(prompt: StructuredPrompt): string {
    return `${this.getDefaultSystemPrompt()}\n\n${prompt.system_message}`;
  }

  private getDefaultSystemPrompt(): string {
    return `You are an AI assistant capable of performing complex multi-task activities. 
    You have access to persistent state and contextual information. 
    Always consider the provided context and history when formulating your responses. 
    If you need to perform a specific action or retrieve information, clearly state it in your response.`;
  }

  private getChatOptions(prompt: StructuredPrompt): ChatOptions {
    return {
      maxTokens: prompt.config?.max_tokens || 3600,
      temperature: prompt.config?.temperature || 0.3,
      topP: 1
    };
  }

  private selectPrompt(input: InputType): StructuredPrompt {
    this.logger.debug('Selecting prompt for input', { inputType: typeof input });

    const matchingPrompts = Object.values(this.prompts).filter(prompt => 
      this.inputMatchesSchema(input, prompt.input_schema)
    );

    if (matchingPrompts.length === 0) {
      this.logger.warn('No matching prompt found for input', { inputType: typeof input });
      throw new Error('No matching prompt found for the given input');
    }

    // If multiple prompts match, select the one with the highest priority (if set) or the first one
    const selectedPrompt = matchingPrompts.reduce((prev, current) => 
      (current.priority ?? 0) > (prev.priority ?? 0) ? current : prev
    );

    this.logger.debug('Selected prompt', { promptName: selectedPrompt.name });
    return selectedPrompt;
  }

  private inputMatchesSchema(input: InputType, schema: InputSchema): boolean {
    if (typeof input === 'string') {
      return Object.keys(schema).length === 1 && Object.values(schema)[0] === 'string';
    }

    const inputKeys = Object.keys(input);
    const schemaKeys = Object.keys(schema);
    return inputKeys.length === schemaKeys.length && inputKeys.every(key => {
      return key in schema && (typeof input[key] === schema[key] || (schema[key] === 'object' && typeof input[key] === 'object'));
    });
  }

  private prepareTools(prompt: StructuredPrompt): Tool[] {
    if (!prompt.tools || prompt.tools.length === 0) {
      return [];
    }
    
    return prompt.tools
      .map(toolName => this.tools[toolName])
      .filter((tool): tool is Tool => tool !== undefined);
  }

  private logPrompts(): void {
    this.logger.debug('Prompts in TaskExecutor:', { 
      promptCount: Object.keys(this.prompts).length,
      promptNames: Object.keys(this.prompts)
    });
    
    Object.entries(this.prompts).forEach(([promptName, prompt]) => {
      this.logger.debug(`Prompt details for ${promptName}:`, {
        name: prompt.name,
        description: prompt.description,
        inputSchema: prompt.input_schema,
        toolCount: prompt.tools.length,
        toolNames: prompt.tools,
        priority: prompt.priority ?? 'Not set'
      });
    });
    
    if (!('executeTask' in this.prompts)) {
      this.logger.warn('executeTask prompt not found in TaskExecutor');
    }
  }
}