prompt-enchancer

import {
  EnhancementVariant,
  GenerationOptions,
  Enhancer,
  EnhancementResult,
  PluginConfig,
  LLMClient,
  EnhancerError,
} from "./types";
import { ConfigService } from "./config";

export class Generator {
  private configService: ConfigService;
  private enhancers: Map<string, Enhancer> = new Map();
  private cache: Map<string, EnhancementVariant[]> = new Map();

  constructor(
    private llmClient: LLMClient,
    config?: Partial<PluginConfig>,
  ) {
    this.configService = new ConfigService(config);
  }

  registerEnhancer(enhancer: Enhancer): void {
    this.enhancers.set(enhancer.id, enhancer);
  }

  async generate(
    originalPrompt: string,
    options?: Partial<GenerationOptions>,
  ): Promise<EnhancementVariant[]> {
    const generationOptions: GenerationOptions = {
      input: originalPrompt,
      temperature: this.configService.getConfig().temperature,
      maxRetries: 2,
      timeout: 30000,
      ...options,
    };

    const cacheKey = this.getCacheKey(originalPrompt);
    const cached = this.cache.get(cacheKey);
    if (cached) return cached;

    const enabledEnhancerIds = this.configService.getEnabledEnhancerIds();
    const results: EnhancementVariant[] = [];

    const batchSize = 3;
    for (let i = 0; i < enabledEnhancerIds.length; i += batchSize) {
      const batch = enabledEnhancerIds.slice(i, i + batchSize);
      const batchResults = await Promise.allSettled(
        batch.map(async (id) => {
          const enhancer = this.enhancers.get(id);
          return enhancer
            ? this.generateWithRetry(enhancer, generationOptions)
            : null;
        }),
      );

      for (const result of batchResults) {
        if (result.status === "fulfilled" && result.value) {
          results.push(result.value);
        }
      }
    }

    this.cache.set(cacheKey, results);
    return results.slice(0, this.configService.getConfig().maxVariants);
  }

  private async generateWithRetry(
    enhancer: Enhancer,
    options: GenerationOptions,
    retryCount: number = 0,
  ): Promise<EnhancementVariant | null> {
    try {
      return await enhancer.generate(options.input, options);
    } catch (error) {
      if (retryCount < (options.maxRetries || 2)) {
        await new Promise((r) => setTimeout(r, Math.pow(2, retryCount) * 1000));
        return this.generateWithRetry(enhancer, options, retryCount + 1);
      }
      throw new EnhancerError(
        enhancer.id,
        options.input,
        retryCount,
        String(error),
      );
    }
  }

  private getCacheKey(prompt: string): string {
    let hash = 0;
    for (let i = 0; i < prompt.length; i++)
      hash = (hash << 5) - hash + prompt.charCodeAt(i);
    return `prompt_${hash}`;
  }
}

export async function generate(prompt: string, context: any): Promise<string> {
  const llmClient: LLMClient = {
    call: async (opts) =>
      await context.model.generate(opts.userPrompt, {
        systemPrompt: opts.systemPrompt,
        temperature: opts.temperature,
      }),
    isAvailable: async () => true,
    getModelInfo: async () => ({
      name: "n/a",
      version: "1.0",
      maxContextLength: 8192,
      features: [],
    }),
  };

  const generator = new Generator(llmClient);
  const variants = await generator.generate(prompt);
  return variants.map((v) => `=== ${v.label} ===\n${v.prompt}`).join("\n\n");
}