File size: 1,215 Bytes
c09f67c | 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 47 48 49 50 51 52 53 54 | import { createGoogleGenerativeAI } from "@ai-sdk/google";
import { embed, embedMany } from "ai";
const GOOGLE_API_KEY = process.env.GOOGLE_GENERATIVE_AI_API_KEY!;
const google = createGoogleGenerativeAI({
apiKey: GOOGLE_API_KEY,
});
const EMBEDDING_CONFIG = {
model: google.textEmbedding("gemini-embedding-001"),
providerOptions: {
google: {
outputDimensionality: 768,
taskType: "SEMANTIC_SIMILARITY",
},
},
modelName: "gemini-embedding-001",
};
export class Embed {
public async embedMany(content: string[]): Promise<{
embeddings: number[][];
model: string;
}> {
const { embeddings } = await embedMany({
model: EMBEDDING_CONFIG.model,
values: content,
providerOptions: EMBEDDING_CONFIG.providerOptions,
});
return {
embeddings,
model: EMBEDDING_CONFIG.modelName,
};
}
public async embed(content: string): Promise<{
embedding: number[];
model: string;
}> {
const { embedding } = await embed({
model: EMBEDDING_CONFIG.model,
value: content,
providerOptions: EMBEDDING_CONFIG.providerOptions,
});
return {
embedding,
model: EMBEDDING_CONFIG.modelName,
};
}
}
|