diff --git a/examples/gemini.ts b/examples/gemini/embedding.ts similarity index 76% rename from examples/gemini.ts rename to examples/gemini/embedding.ts index 416a586f98d90646bfab8e9d04dbec4f33433bba..6ecc2692ba67f4faa4b8c3dfbe88724b7a63d85e 100644 --- a/examples/gemini.ts +++ b/examples/gemini/embedding.ts @@ -1,11 +1,11 @@ -import { GEMINI_MODEL, GeminiEmbedding } from "llamaindex"; +import { GEMINI_EMBEDDING_MODEL, GeminiEmbedding } from "llamaindex"; async function main() { if (!process.env.GOOGLE_API_KEY) { throw new Error("Please set the GOOGLE_API_KEY environment variable."); } const embedModel = new GeminiEmbedding({ - model: GEMINI_MODEL.GEMINI_PRO, + model: GEMINI_EMBEDDING_MODEL.EMBEDDING_001, }); const texts = ["hello", "world"]; const embeddings = await embedModel.getTextEmbeddingsBatch(texts); diff --git a/packages/core/src/embeddings/GeminiEmbedding.ts b/packages/core/src/embeddings/GeminiEmbedding.ts index fa5eb0cb582186f6712cbf7fac98faa80b311ce7..1b4fda6bcf37bfe6270286ec85f1edceb670209e 100644 --- a/packages/core/src/embeddings/GeminiEmbedding.ts +++ b/packages/core/src/embeddings/GeminiEmbedding.ts @@ -1,27 +1,21 @@ -import { - GEMINI_MODEL, - GeminiSessionStore, - type GeminiConfig, - type GeminiSession, -} from "../llm/gemini.js"; +import { GeminiSessionStore, type GeminiSession } from "../llm/gemini.js"; import { BaseEmbedding } from "./types.js"; +export enum GEMINI_EMBEDDING_MODEL { + EMBEDDING_001 = "embedding-001", + TEXT_EMBEDDING_004 = "text-embedding-004", +} + /** * GeminiEmbedding is an alias for Gemini that implements the BaseEmbedding interface. */ export class GeminiEmbedding extends BaseEmbedding { - model: GEMINI_MODEL; - temperature: number; - topP: number; - maxTokens?: number; + model: GEMINI_EMBEDDING_MODEL; session: GeminiSession; - constructor(init?: GeminiConfig) { + constructor(init?: Partial<GeminiEmbedding>) { super(); - this.model = init?.model ?? GEMINI_MODEL.GEMINI_PRO; - this.temperature = init?.temperature ?? 0.1; - this.topP = init?.topP ?? 1; - this.maxTokens = init?.maxTokens ?? undefined; + this.model = init?.model ?? GEMINI_EMBEDDING_MODEL.EMBEDDING_001; this.session = init?.session ?? GeminiSessionStore.get(); } diff --git a/packages/core/src/llm/gemini.ts b/packages/core/src/llm/gemini.ts index 7e498c0fc0ca7f2b681042c98a50e36ac62fe50d..413e4e2a350b5e91c4a09da2129cc42c2604f58f 100644 --- a/packages/core/src/llm/gemini.ts +++ b/packages/core/src/llm/gemini.ts @@ -32,8 +32,6 @@ type GeminiSessionOptions = { export enum GEMINI_MODEL { GEMINI_PRO = "gemini-pro", GEMINI_PRO_VISION = "gemini-pro-vision", - EMBEDDING_001 = "embedding-001", - AQA = "aqa", GEMINI_PRO_LATEST = "gemini-1.5-pro-latest", } @@ -44,16 +42,12 @@ export interface GeminiModelInfo { export const GEMINI_MODEL_INFO_MAP: Record<GEMINI_MODEL, GeminiModelInfo> = { [GEMINI_MODEL.GEMINI_PRO]: { contextWindow: 30720 }, [GEMINI_MODEL.GEMINI_PRO_VISION]: { contextWindow: 12288 }, - [GEMINI_MODEL.EMBEDDING_001]: { contextWindow: 2048 }, - [GEMINI_MODEL.AQA]: { contextWindow: 7168 }, [GEMINI_MODEL.GEMINI_PRO_LATEST]: { contextWindow: 10 ** 6 }, }; const SUPPORT_TOOL_CALL_MODELS: GEMINI_MODEL[] = [ GEMINI_MODEL.GEMINI_PRO, GEMINI_MODEL.GEMINI_PRO_VISION, - GEMINI_MODEL.EMBEDDING_001, - GEMINI_MODEL.AQA, ]; const DEFAULT_GEMINI_PARAMS = {