import { GoogleGenerativeAI, TaskType } from "@google/generative-ai";
let genAI: GoogleGenerativeAI | null = null;
export function initEmbeddings(apiKey: string): void {
genAI = new GoogleGenerativeAI(apiKey);
}
export async function generateEmbedding(text: string): Promise<number[]> {
if (!genAI) {
throw new Error("Embeddings not initialized. Call initEmbeddings first.");
}
const model = genAI.getGenerativeModel({ model: "text-embedding-004" });
const result = await model.embedContent({
content: { parts: [{ text }], role: "user" },
taskType: TaskType.RETRIEVAL_DOCUMENT,
});
return result.embedding.values;
}
export async function generateEmbeddings(texts: string[]): Promise<number[][]> {
if (!genAI) {
throw new Error("Embeddings not initialized. Call initEmbeddings first.");
}
const model = genAI.getGenerativeModel({ model: "text-embedding-004" });
const embeddings: number[][] = [];
for (const text of texts) {
const result = await model.embedContent({
content: { parts: [{ text }], role: "user" },
taskType: TaskType.RETRIEVAL_DOCUMENT,
});
embeddings.push(result.embedding.values);
}
return embeddings;
}
