user-docs

/**
 * Embeddings API Client
 * Connects to an OpenAI-compatible embeddings endpoint.
 */

import { ragDebug } from '../utils/ragLogger.js';
import { createHash } from 'crypto';

export interface EmbeddingConfig {
  baseUrl: string;
  model: string;
  apiKey?: string;
  timeout?: number;
}

export interface EmbeddingResponse {
  object: 'list';
  data: Array<{
    object: 'embedding';
    embedding: number[];
    index: number;
  }>;
  model: string;
  usage: {
    prompt_tokens: number;
    total_tokens: number;
  };
}

export class EmbeddingsClient {
  private config: EmbeddingConfig;
  private cache = new Map<string, Float32Array>();

  constructor(config: EmbeddingConfig) {
    this.config = {
      timeout: 120000,
      ...config,
    };
    ragDebug('RAG Embeddings', `Client configured - Model: "${this.config.model}"`);
    ragDebug('RAG Embeddings', `Embedding API URL: ${this.config.baseUrl}`);
  }

  /**
   * Generate embeddings for a batch of texts
   */
  async embed(texts: string[]): Promise<Float32Array[]> {
    if (texts.length === 0) {
      return [];
    }

    const results = new Array<Float32Array>(texts.length);
    const missing: Array<{ text: string; index: number; key: string }> = [];
    for (let index = 0; index < texts.length; index++) {
      const text = texts[index];
      const key = this.cacheKey(text);
      const cached = this.cache.get(key);
      if (cached) {
        results[index] = cached;
      } else {
        missing.push({ text, index, key });
      }
    }

    if (missing.length === 0) return results;

    ragDebug('RAG Embeddings', `Sending request to embedding API - Model: "${this.config.model}", Texts: ${missing.length}`);
    if (missing.length === 1) {
      ragDebug('RAG Embeddings', `Text preview: "${missing[0].text.substring(0, 50)}..."`);
    }

    const requestTexts = missing.map(item => item.text);

    const controller = new AbortController();
    const timeoutId = setTimeout(() => controller.abort(), this.config.timeout);

    try {
      const response = await fetch(`${this.config.baseUrl}/v1/embeddings`, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          ...this.authHeaders(),
        },
        body: JSON.stringify({
          model: this.config.model,
          input: requestTexts.length === 1 ? requestTexts[0] : requestTexts,
        }),
        signal: controller.signal,
      });

      if (!response.ok) {
        const error = await response.text();
        throw new Error(`API error: ${response.status} - ${error}`);
      }

      const data: EmbeddingResponse = await response.json();
      
      ragDebug('RAG Embeddings', `Embeddings generated successfully - Model: "${data.model}", Dimension: ${data.data[0]?.embedding.length || 0}`);
      ragDebug('RAG Embeddings', `Token usage: ${data.usage.total_tokens} tokens`);
      
      // Convert to Float32Array and sort by index
      const generated = data.data
        .sort((a, b) => a.index - b.index)
        .map(item => new Float32Array(item.embedding));
      for (let generatedIndex = 0; generatedIndex < generated.length; generatedIndex++) {
        const missingItem = missing[generatedIndex];
        const embedding = generated[generatedIndex];
        this.cache.set(missingItem.key, embedding);
        results[missingItem.index] = embedding;
      }

      return results;
    } catch (error) {
      if (error instanceof Error && error.name === 'AbortError') {
        throw new Error(`Embedding request timed out after ${this.config.timeout}ms`);
      }
      throw error;
    } finally {
      clearTimeout(timeoutId);
    }
  }

  private cacheKey(text: string): string {
    return `${this.config.model}:${createHash('sha256').update(text).digest('hex')}`;
  }

  /**
   * Generate embedding for a single text
   */
  async embedSingle(text: string): Promise<Float32Array> {
    const [embedding] = await this.embed([text]);
    return embedding;
  }

  /**
   * Get the dimension of embeddings from this model
   */
  async getDimension(): Promise<number> {
    const testEmbedding = await this.embedSingle('Testing...');
    return testEmbedding.length;
  }

  /**
   * Health check for LM Studio API
   */
  async healthCheck(): Promise<boolean> {
    try {
      const response = await fetch(`${this.config.baseUrl}/v1/models`, {
        method: 'GET',
        headers: this.authHeaders(),
        signal: AbortSignal.timeout(5000),
      });
      return response.ok;
    } catch {
      return false;
    }
  }

  private authHeaders(): Record<string, string> {
    const apiKey = this.config.apiKey?.trim();
    return apiKey ? { Authorization: `Bearer ${apiKey}` } : {};
  }
}