diff --git a/packages/core/src/embedding.ts b/packages/core/src/embedding.ts index 117853df37b..33124e61df2 100644 --- a/packages/core/src/embedding.ts +++ b/packages/core/src/embedding.ts @@ -1,5 +1,9 @@ import path from "node:path"; -import { getEndpoint } from "./models.ts"; +import { + getEmbeddingModelSettings, + getEndpoint, + getModelSettings, +} from "./models.ts"; import { IAgentRuntime, ModelProviderName } from "./types.ts"; import settings from "./settings.ts"; import elizaLogger from "./logger.ts"; @@ -33,19 +37,20 @@ export type EmbeddingConfig = { export const getEmbeddingConfig = (): EmbeddingConfig => ({ dimensions: settings.USE_OPENAI_EMBEDDING?.toLowerCase() === "true" - ? 1536 // OpenAI + ? getEmbeddingModelSettings(ModelProviderName.OPENAI).dimensions : settings.USE_OLLAMA_EMBEDDING?.toLowerCase() === "true" - ? 1024 // Ollama mxbai-embed-large + ? getEmbeddingModelSettings(ModelProviderName.OLLAMA).dimensions : settings.USE_GAIANET_EMBEDDING?.toLowerCase() === "true" - ? 768 // GaiaNet + ? getEmbeddingModelSettings(ModelProviderName.GAIANET) + .dimensions : 384, // BGE model: settings.USE_OPENAI_EMBEDDING?.toLowerCase() === "true" - ? "text-embedding-3-small" + ? getEmbeddingModelSettings(ModelProviderName.OPENAI).name : settings.USE_OLLAMA_EMBEDDING?.toLowerCase() === "true" - ? settings.OLLAMA_EMBEDDING_MODEL || "mxbai-embed-large" + ? getEmbeddingModelSettings(ModelProviderName.OLLAMA).name : settings.USE_GAIANET_EMBEDDING?.toLowerCase() === "true" - ? settings.GAIANET_EMBEDDING_MODEL || "nomic-embed" + ? getEmbeddingModelSettings(ModelProviderName.GAIANET).name : "BGE-small-en-v1.5", provider: settings.USE_OPENAI_EMBEDDING?.toLowerCase() === "true" @@ -134,11 +139,17 @@ export function getEmbeddingZeroVector(): number[] { let embeddingDimension = 384; // Default BGE dimension if (settings.USE_OPENAI_EMBEDDING?.toLowerCase() === "true") { - embeddingDimension = 1536; // OpenAI dimension + embeddingDimension = getEmbeddingModelSettings( + ModelProviderName.OPENAI + ).dimensions; // OpenAI dimension } else if (settings.USE_OLLAMA_EMBEDDING?.toLowerCase() === "true") { - embeddingDimension = 1024; // Ollama mxbai-embed-large dimension + embeddingDimension = getEmbeddingModelSettings( + ModelProviderName.OLLAMA + ).dimensions; // Ollama mxbai-embed-large dimension } else if (settings.USE_GAIANET_EMBEDDING?.toLowerCase() === "true") { - embeddingDimension = 768; // GaiaNet dimension + embeddingDimension = getEmbeddingModelSettings( + ModelProviderName.GAIANET + ).dimensions; // GaiaNet dimension } return Array(embeddingDimension).fill(0);