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)}`);
}
}
}
|