File size: 1,504 Bytes
d4abe4b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3e12ae4
 
d4abe4b
 
 
 
 
 
 
 
 
3e12ae4
 
d4abe4b
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
import { Embeddings } from '@langchain/core/embeddings';
import { GoogleGenerativeAI } from '@google/generative-ai';
import { config } from '../utils/config.js';
import { logger } from '../utils/logger.js';

export class GeminiEmbedding extends Embeddings {
  private genAI: GoogleGenerativeAI;
  private modelName: string;

  constructor() {
    super({});
    this.genAI = new GoogleGenerativeAI(config.gemini.apiKey);
    this.modelName = config.gemini.embeddingModel;
  }

  async embedDocuments(texts: string[]): Promise<number[][]> {
    try {
      const model = this.genAI.getGenerativeModel({ model: this.modelName });
      
      const embeddings = await Promise.all(
        texts.map(async (text) => {
          const result = await model.embedContent(text);
          return result.embedding.values;
        })
      );
      
      return embeddings;
    } catch (error) {
      logger.error({ error }, 'Error generating embeddings');
      throw new Error(`Embedding error: ${error instanceof Error ? error.message : String(error)}`);
    }
  }

  async embedQuery(text: string): Promise<number[]> {
    try {
      const model = this.genAI.getGenerativeModel({ model: this.modelName });
      const result = await model.embedContent(text);
      return result.embedding.values;
    } catch (error) {
      logger.error({ error }, 'Error generating query embedding');
      throw new Error(`Query embedding error: ${error instanceof Error ? error.message : String(error)}`);
    }
  }
}