learning_ai_common_plat/packages/llm/src/providers/mock.ts

/**
 * Mock LLM provider — for testing.
 *
 * Returns pre-configured responses or a default echo response.
 * Supports vision content, streaming, and embedding.
 */

import type {
  ChatCompletionRequest,
  ChatCompletionResponse,
  EmbeddingRequest,
  EmbeddingResponse,
  LLMProvider,
} from '../types.js';
import { getMessageText } from '../types.js';

export class MockLLMProvider implements LLMProvider {
  private responses: ChatCompletionResponse[] = [];
  private embeddingResponses: EmbeddingResponse[] = [];
  public calls: ChatCompletionRequest[] = [];
  public embedCalls: EmbeddingRequest[] = [];

  constructor(responses?: ChatCompletionResponse[]) {
    if (responses) this.responses = [...responses];
  }

  isConfigured(): boolean {
    return true;
  }

  /** Add a chat response to the queue. */
  addResponse(response: ChatCompletionResponse): void {
    this.responses.push(response);
  }

  /** Add an embedding response to the queue. */
  addEmbeddingResponse(response: EmbeddingResponse): void {
    this.embeddingResponses.push(response);
  }

  async chatCompletion(req: ChatCompletionRequest): Promise<ChatCompletionResponse> {
    this.calls.push(req);

    if (this.responses.length > 0) {
      return this.responses.shift()!;
    }

    // Default echo response — handles both string and multipart content
    const lastMessage = req.messages[req.messages.length - 1];
    const text = lastMessage ? getMessageText(lastMessage) : '(empty)';
    return {
      content: `Mock response to: ${text}`,
      model: req.model ?? 'mock-model',
      finishReason: 'stop',
      usage: { promptTokens: 10, completionTokens: 10, totalTokens: 20 },
    };
  }

  async *chatCompletionStream(req: ChatCompletionRequest): AsyncIterable<string> {
    this.calls.push(req);

    if (this.responses.length > 0) {
      const resp = this.responses.shift()!;
      // Yield word-by-word to simulate streaming
      const words = resp.content.split(' ');
      for (const word of words) {
        yield word + ' ';
      }
      return;
    }

    const lastMessage = req.messages[req.messages.length - 1];
    const text = lastMessage ? getMessageText(lastMessage) : '(empty)';
    const words = `Mock response to: ${text}`.split(' ');
    for (const word of words) {
      yield word + ' ';
    }
  }

  async embed(req: EmbeddingRequest): Promise<EmbeddingResponse> {
    this.embedCalls.push(req);

    if (this.embeddingResponses.length > 0) {
      return this.embeddingResponses.shift()!;
    }

    // Default: return deterministic pseudo-embeddings (dimension 8 for testing)
    const inputs = Array.isArray(req.input) ? req.input : [req.input];
    const embeddings = inputs.map(text => {
      // Simple hash-based deterministic vector for testing
      const vec = new Array(8).fill(0);
      for (let i = 0; i < text.length; i++) {
        vec[i % 8] += text.charCodeAt(i) / 1000;
      }
      // Normalize
      const mag = Math.sqrt(vec.reduce((sum, v) => sum + v * v, 0)) || 1;
      return vec.map(v => v / mag);
    });

    return {
      embeddings,
      model: req.model ?? 'mock-embedding-model',
      usage: {
        promptTokens: inputs.join(' ').split(/\s+/).length,
        completionTokens: 0,
        totalTokens: inputs.join(' ').split(/\s+/).length,
      },
    };
  }

  /** Reset call history and responses. */
  reset(): void {
    this.calls = [];
    this.embedCalls = [];
    this.responses = [];
    this.embeddingResponses = [];
  }
}