stack/apps/gateway/src/memory/embedding.service.ts

import { Injectable, Logger } from '@nestjs/common';
import type { EmbeddingProvider } from '@mosaic/memory';

const DEFAULT_MODEL = 'text-embedding-3-small';
const DEFAULT_DIMENSIONS = 1536;

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

/**
 * Generates embeddings via the OpenAI-compatible embeddings API.
 * Supports OpenAI, Azure OpenAI, and any provider with a compatible endpoint.
 */
@Injectable()
export class EmbeddingService implements EmbeddingProvider {
  private readonly logger = new Logger(EmbeddingService.name);
  private readonly apiKey: string | undefined;
  private readonly baseUrl: string;
  private readonly model: string;

  readonly dimensions = DEFAULT_DIMENSIONS;

  constructor() {
    this.apiKey = process.env['OPENAI_API_KEY'];
    this.baseUrl = process.env['EMBEDDING_API_URL'] ?? 'https://api.openai.com/v1';
    this.model = process.env['EMBEDDING_MODEL'] ?? DEFAULT_MODEL;
  }

  get available(): boolean {
    return !!this.apiKey;
  }

  async embed(text: string): Promise<number[]> {
    const results = await this.embedBatch([text]);
    return results[0]!;
  }

  async embedBatch(texts: string[]): Promise<number[][]> {
    if (!this.apiKey) {
      this.logger.warn('No OPENAI_API_KEY configured — returning zero vectors');
      return texts.map(() => new Array<number>(this.dimensions).fill(0));
    }

    const response = await fetch(`${this.baseUrl}/embeddings`, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        Authorization: `Bearer ${this.apiKey}`,
      },
      body: JSON.stringify({
        model: this.model,
        input: texts,
        dimensions: this.dimensions,
      }),
    });

    if (!response.ok) {
      const body = await response.text();
      this.logger.error(`Embedding API error: ${response.status} ${body}`);
      throw new Error(`Embedding API returned ${response.status}`);
    }

    const json = (await response.json()) as EmbeddingResponse;
    return json.data.sort((a, b) => a.index - b.index).map((d) => d.embedding);
  }
}