get the info from model settings instead of hardcoced number

elizaOS · Jan 5, 2025 · bbb1f94 · bbb1f94
1 parent e140845
commit bbb1f94
Showing 1 changed file with 21 additions and 10 deletions.
diff --git 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);