frontend/prompt-service-client.ts

/**
 * Prompt Service Client
 * Handles communication with the Python prompt service for real model generation
 */

export interface PromptComparisonRequest {
  type: 'prompt_comparison';
  prompt_a: string;
  prompt_b: string;
}

export interface SinglePromptRequest {
  type: 'single_prompt';
  prompt: string;
  prompt_id?: string;
  max_tokens?: number;
  temperature?: number;
}

export interface AttentionTrace {
  type: 'attention';
  layer: string;
  weights: number[][];
  max_weight: number;
  entropy?: number;
  prompt_id: string;
  comparison_group?: 'prompt_a' | 'prompt_b';
  timestamp: number;
}

export interface ComparisonSummary {
  type: 'prompt_comparison';
  prompt_a: {
    id: string;
    text: string;
    generated: string;
    num_traces: number;
  };
  prompt_b: {
    id: string;
    text: string;
    generated: string;
    num_traces: number;
  };
  timestamp: number;
}

export class PromptServiceClient {
  private ws: WebSocket | null = null;
  private readonly url: string;
  private messageHandlers: Map<string, (data: Record<string, unknown>) => void> = new Map();
  private connectionPromise: Promise<void> | null = null;

  constructor(url: string = 'ws://localhost:8767') {
    this.url = url;
  }

  /**
   * Connect to the prompt service
   */
  async connect(): Promise<void> {
    if (this.ws?.readyState === WebSocket.OPEN) {
      return;
    }

    if (this.connectionPromise) {
      return this.connectionPromise;
    }

    this.connectionPromise = new Promise((resolve, reject) => {
      const timeout = setTimeout(() => {
        if (this.ws?.readyState !== WebSocket.OPEN) {
          console.log('⏱️ Connection timeout - using demo mode');
          this.ws?.close();
          reject(new Error('Connection timeout'));
        }
      }, 3000); // 3 second timeout

      try {
        this.ws = new WebSocket(this.url);

        this.ws.onopen = () => {
          clearTimeout(timeout);
          console.log('✅ Connected to prompt service');
          resolve();
        };

        this.ws.onerror = (error) => {
          clearTimeout(timeout);
          console.log('⚠️ Prompt service not available - will use demo mode');
          reject(new Error('Prompt service not available'));
        };

        this.ws.onclose = () => {
          console.log('🔌 Disconnected from prompt service');
          this.ws = null;
          this.connectionPromise = null;
        };

        this.ws.onmessage = (event) => {
          try {
            const data = JSON.parse(event.data);
            this.handleMessage(data);
          } catch (error) {
            console.error('Error parsing message:', error);
          }
        };
      } catch (error) {
        reject(error);
      }
    });

    return this.connectionPromise;
  }

  /**
   * Handle incoming messages from the prompt service
   */
  private handleMessage(data: Record<string, unknown>) {
    // Notify all registered handlers
    this.messageHandlers.forEach((handler) => {
      handler(data);
    });
  }

  /**
   * Register a message handler
   */
  onMessage(id: string, handler: (data: Record<string, unknown>) => void) {
    this.messageHandlers.set(id, handler);
  }

  /**
   * Unregister a message handler
   */
  offMessage(id: string) {
    this.messageHandlers.delete(id);
  }

  /**
   * Request prompt comparison from the service
   */
  async comparePrompts(promptA: string, promptB: string): Promise<void> {
    await this.connect();

    if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
      throw new Error('Not connected to prompt service');
    }

    const request: PromptComparisonRequest = {
      type: 'prompt_comparison',
      prompt_a: promptA,
      prompt_b: promptB
    };

    this.ws.send(JSON.stringify(request));
  }

  /**
   * Generate for a single prompt
   */
  async generateSingle(
    prompt: string,
    options?: {
      prompt_id?: string;
      max_tokens?: number;
      temperature?: number;
    }
  ): Promise<void> {
    await this.connect();

    if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
      throw new Error('Not connected to prompt service');
    }

    const request: SinglePromptRequest = {
      type: 'single_prompt',
      prompt,
      ...options
    };

    this.ws.send(JSON.stringify(request));
  }

  /**
   * Disconnect from the service
   */
  disconnect() {
    if (this.ws) {
      this.ws.close();
      this.ws = null;
    }
  }

  /**
   * Check if connected
   */
  isConnected(): boolean {
    return this.ws?.readyState === WebSocket.OPEN;
  }
}