@nomyx/assistant/src/prompts/PromptExecution.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 fs from 'fs/promises';
import path from 'path';
import Ajv from 'ajv';
import { ILogger, PromptDefinition } from '../types/common';
import { ExecutionOptions, ExecutionState } from '../types/execution';
import { StandardizedToolCall, StandardTool, IToolManager } from '../types/tool';
import { IAIProvider } from '../types/provider';
import { ChatMessage, ChatOptions } from '../types/chat';
import { Parser } from 'expr-eval';

interface PromptRegistry {
  get(name: string): StructuredPrompt;
  getPromptsSortedByPriority(): StructuredPrompt[];
  getPromptsByConfig(config: Partial<NonNullable<PromptDefinition['config']>>): StructuredPrompt[];
}

interface StructuredPrompt {
  definition: PromptDefinition;
  render(variables: Record<string, string>, availablePrompts: PromptDefinition[]): string;
  getPriority(): number | undefined;
  getConfig(): PromptDefinition['config'] | undefined;
}

export class PromptExecution extends EventEmitter {
  private state: ExecutionState;
  private ajv: any;
  private toolManager: IToolManager;

  constructor(
    private registry: PromptRegistry,
    private initialPrompt: StructuredPrompt,
    private provider: IAIProvider,
    private availablePrompts: PromptDefinition[],
    private options: ExecutionOptions = {},
    private logger: ILogger
  ) {
    super();
    this.state = {
      id: uuidv4(),
      variables: {},
      history: [],
      currentPrompt: initialPrompt.definition.name
    };
    this.ajv = new Ajv();
    this.toolManager = {
      registerTool: (tool: StandardTool) => {},
      getTool: (name: string) => undefined,
      getAllTools: () => []
    };
  }

  async execute(initialReplacements: Record<string, string> = {}): Promise<void> {
    this.state.variables = { ...this.state.variables, ...initialReplacements };
    
    while (true) {
      try {
        await this.loadState();
        const currentPrompt = this.registry.get(this.state.currentPrompt);
        const { shouldContinue, nextPrompt } = await this.executePromptWithRetry(currentPrompt);
        if (!shouldContinue) break;
        if (nextPrompt) this.state.currentPrompt = nextPrompt.definition.name;
        await this.saveState();
      } catch (error) {
        this.emit('error', error);
        break;
      }
    }

    this.emit('executionComplete', this.state);
  }

  registerTool(tool: StandardTool): void {
    this.toolManager.registerTool(tool);
  }

  private async executePromptWithRetry(prompt: StructuredPrompt): Promise<{ shouldContinue: boolean; nextPrompt?: StructuredPrompt }> {
    let retries = 0;
    while (retries <= (this.options.maxRetries || 3)) {
      try {
        return await this.executePromptWithTimeout(prompt);
      } catch (error) {
        if (error instanceof Error && error.name === 'TimeoutError') {
          this.emit('promptTimeout', { prompt: prompt.definition.name, retries });
        } else {
          this.emit('promptError', { prompt: prompt.definition.name, error, retries });
        }
        retries++;
        if (retries <= (this.options.maxRetries || 3)) {
          await new Promise(resolve => setTimeout(resolve, this.options.retryDelayMs || 1000));
        } else {
          throw error;
        }
      }
    }
    throw new Error('Max retries exceeded');
  }

  private async executePromptWithTimeout(prompt: StructuredPrompt): Promise<{ shouldContinue: boolean; nextPrompt?: StructuredPrompt }> {
    return new Promise((resolve, reject) => {
      const timeout = setTimeout(() => {
        reject(new Error('Prompt execution timed out'));
      }, this.options.timeoutMs || 30000);

      this.executePrompt(prompt)
        .then(resolve)
        .catch(reject)
        .finally(() => clearTimeout(timeout));
    });
  }

  private async executePrompt(prompt: StructuredPrompt): Promise<{ shouldContinue: boolean; nextPrompt?: StructuredPrompt }> {
    let renderedPrompt = prompt.render(this.state.variables, this.availablePrompts);
    
    this.emit('promptStart', { prompt: prompt.definition.name, renderedPrompt });

    const messages: ChatMessage[] = [{ role: 'user', content: renderedPrompt }];
    const chatOptions: ChatOptions = {
      ...this.options,
      ...prompt.getConfig()
    };
    const response = await this.provider.chat(messages, chatOptions);
    
    let parsedResponse: any;
    try {
      parsedResponse = JSON.parse(response.content);
    } catch (error) {
      throw new Error(`Invalid JSON response for prompt: ${prompt.definition.name}`);
    }

    if (!this.validateResponse(parsedResponse, prompt.definition.responseFormat)) {
      throw new Error(`Invalid response format for prompt: ${prompt.definition.name}`);
    }

    this.state.history.push({ prompt: renderedPrompt, response: parsedResponse });
    this.emit('promptComplete', { prompt: prompt.definition.name, response: parsedResponse });

    // Process tool calls
    if (response.toolCalls) {
      for (const toolCall of response.toolCalls) {
        await this.processToolCall(toolCall);
      }
    }

    // Update state variables
    this.state.variables = { ...this.state.variables, ...parsedResponse.response };

    // Handle next steps
    if (parsedResponse.next) {
      if (parsedResponse.next.if) {
        const nextPromptName = this.evaluateCondition(parsedResponse.next.if) ? parsedResponse.next.then : parsedResponse.next.else;
        return { shouldContinue: true, nextPrompt: this.registry.get(nextPromptName) };
      } else if (parsedResponse.next.each) {
        for (const item of parsedResponse.next.each) {
          const nextPromptName = this.evaluateCondition(item.if) ? item.then : item.else;
          await this.executePrompt(this.registry.get(nextPromptName));
        }
        return { shouldContinue: true };
      }
    }

    return { shouldContinue: false };
  }

  private async processToolCall(toolCall: StandardizedToolCall): Promise<void> {
    this.emit('toolCallStart', toolCall);

    try {
      const tool = this.toolManager.getTool(toolCall.name);
      if (!tool) {
        throw new Error(`Tool not found: ${toolCall.name}`);
      }

      const result = await tool.execute(toolCall.arguments, { user: { id: this.state.id, name: 'User' }, session: { id: this.state.id, timestamp: new Date().toISOString() } });
      this.state.variables[toolCall.name] = result;

      this.emit('toolCallComplete', { toolCall, result });
    } catch (error) {
      this.emit('toolCallError', { toolCall, error });
      throw error;
    }
  }

  private validateResponse(response: any, schema: string): boolean {
    const validate = this.ajv.compile(JSON.parse(schema));
    return validate(response);
  }
  
  private evaluateCondition(condition: any): boolean {
    const parser = new Parser();
    const expr = parser.parse(condition);
    return expr.evaluate(this.state.variables);
  }

  private async saveState(): Promise<void> {
    if (this.options.statePersistenceDir) {
      const filePath = path.join(this.options.statePersistenceDir, `${this.state.id}.json`);
      await fs.writeFile(filePath, JSON.stringify(this.state), 'utf8');
    }
  }

  private async loadState(): Promise<void> {
    if (this.options.statePersistenceDir) {
      const filePath = path.join(this.options.statePersistenceDir, `${this.state.id}.json`);
      try {
        const data = await fs.readFile(filePath, 'utf8');
        this.state = JSON.parse(data);
      } catch (error) {
        if (error instanceof Error && (error as NodeJS.ErrnoException).code !== 'ENOENT') {
          throw error;
        }
      }
    }
  }
}