playbook

/**
 * EmbeddingClient — fetches embeddings from LM Studio /v1/embeddings.
 * Adds E5-Instruct query:/passage: prefixes for cross-lingual retrieval.
 */

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

interface EmbeddingResponse {
  data: Array<{ embedding: number[]; index: number }>;
  model: string;
}

export class EmbeddingClient {
  private config: Required<EmbeddingConfig>;

  constructor(config: EmbeddingConfig) {
    this.config = { timeout: 30_000, ...config };
  }

  /**
   * Embed passages (documents being indexed). Adds "passage: " prefix for E5 models.
   */
  async embedPassages(texts: string[]): Promise<Float32Array[]> {
    return this.embed(texts.map((t) => `passage: ${t}`));
  }

  /**
   * Embed a search query. Adds "query: " prefix for E5 models.
   */
  async embedQuery(text: string): Promise<Float32Array> {
    const [vec] = await this.embed([`query: ${text}`]);
    return vec;
  }

  private async embed(texts: string[]): Promise<Float32Array[]> {
    if (texts.length === 0) return [];

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

    try {
      const response = await fetch(`${this.config.baseUrl}/v1/embeddings`, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ model: this.config.model, input: texts }),
        signal: controller.signal,
      });

      if (!response.ok) {
        const body = await response.text();
        throw new Error(`LM Studio /v1/embeddings error ${response.status}: ${body}`);
      }

      const data = (await response.json()) as EmbeddingResponse;
      return data.data
        .sort((a, b) => a.index - b.index)
        .map((item) => new Float32Array(item.embedding));
    } catch (err) {
      if (err instanceof Error && err.name === "AbortError") {
        throw new Error(`Embedding request timed out (${this.config.timeout}ms)`);
      }
      throw err;
    } finally {
      clearTimeout(tid);
    }
  }

  getModelName(): string {
    return this.config.model;
  }

  async getDimension(): Promise<number> {
    const [vec] = await this.embed(["test"]);
    return vec.length;
  }
}