import { EmbeddingModel } from "../types";
import { OpenAIEmbeddings } from "@langchain/openai";

export class OpenAIEmbeddingModel implements EmbeddingModel {
  private embeddings: OpenAIEmbeddings;

  constructor(apiKey: string, modelName: string = "text-embedding-3-small") {
    this.embeddings = new OpenAIEmbeddings({
      openAIApiKey: apiKey,
      modelName,
      batchSize: 1000, // Process embeddings in batches
    });
  }

  async embed(text: string): Promise<number[]> {
    try {
      const result = await this.embeddings.embedQuery(text);
      return result;
    } catch (error) {
      throw new Error(`Failed to generate embedding: ${error}`);
    }
  }

  async embedBatch(texts: string[]): Promise<number[][]> {
    try {
      const results = await this.embeddings.embedDocuments(texts);
      return results;
    } catch (error) {
      throw new Error(`Failed to generate batch embeddings: ${error}`);
    }
  }

  getDimensions(): number {
    // text-embedding-3-small has 1536 dimensions
    return 1536;
  }
}