Skip to content

Commit bbb1f94

Browse files
committed
get the info from model settings instead of hardcoced number
1 parent e140845 commit bbb1f94

File tree

1 file changed

+21
-10
lines changed

1 file changed

+21
-10
lines changed

packages/core/src/embedding.ts

+21-10
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
import path from "node:path";
2-
import { getEndpoint } from "./models.ts";
2+
import {
3+
getEmbeddingModelSettings,
4+
getEndpoint,
5+
getModelSettings,
6+
} from "./models.ts";
37
import { IAgentRuntime, ModelProviderName } from "./types.ts";
48
import settings from "./settings.ts";
59
import elizaLogger from "./logger.ts";
@@ -33,19 +37,20 @@ export type EmbeddingConfig = {
3337
export const getEmbeddingConfig = (): EmbeddingConfig => ({
3438
dimensions:
3539
settings.USE_OPENAI_EMBEDDING?.toLowerCase() === "true"
36-
? 1536 // OpenAI
40+
? getEmbeddingModelSettings(ModelProviderName.OPENAI).dimensions
3741
: settings.USE_OLLAMA_EMBEDDING?.toLowerCase() === "true"
38-
? 1024 // Ollama mxbai-embed-large
42+
? getEmbeddingModelSettings(ModelProviderName.OLLAMA).dimensions
3943
: settings.USE_GAIANET_EMBEDDING?.toLowerCase() === "true"
40-
? 768 // GaiaNet
44+
? getEmbeddingModelSettings(ModelProviderName.GAIANET)
45+
.dimensions
4146
: 384, // BGE
4247
model:
4348
settings.USE_OPENAI_EMBEDDING?.toLowerCase() === "true"
44-
? "text-embedding-3-small"
49+
? getEmbeddingModelSettings(ModelProviderName.OPENAI).name
4550
: settings.USE_OLLAMA_EMBEDDING?.toLowerCase() === "true"
46-
? settings.OLLAMA_EMBEDDING_MODEL || "mxbai-embed-large"
51+
? getEmbeddingModelSettings(ModelProviderName.OLLAMA).name
4752
: settings.USE_GAIANET_EMBEDDING?.toLowerCase() === "true"
48-
? settings.GAIANET_EMBEDDING_MODEL || "nomic-embed"
53+
? getEmbeddingModelSettings(ModelProviderName.GAIANET).name
4954
: "BGE-small-en-v1.5",
5055
provider:
5156
settings.USE_OPENAI_EMBEDDING?.toLowerCase() === "true"
@@ -134,11 +139,17 @@ export function getEmbeddingZeroVector(): number[] {
134139
let embeddingDimension = 384; // Default BGE dimension
135140

136141
if (settings.USE_OPENAI_EMBEDDING?.toLowerCase() === "true") {
137-
embeddingDimension = 1536; // OpenAI dimension
142+
embeddingDimension = getEmbeddingModelSettings(
143+
ModelProviderName.OPENAI
144+
).dimensions; // OpenAI dimension
138145
} else if (settings.USE_OLLAMA_EMBEDDING?.toLowerCase() === "true") {
139-
embeddingDimension = 1024; // Ollama mxbai-embed-large dimension
146+
embeddingDimension = getEmbeddingModelSettings(
147+
ModelProviderName.OLLAMA
148+
).dimensions; // Ollama mxbai-embed-large dimension
140149
} else if (settings.USE_GAIANET_EMBEDDING?.toLowerCase() === "true") {
141-
embeddingDimension = 768; // GaiaNet dimension
150+
embeddingDimension = getEmbeddingModelSettings(
151+
ModelProviderName.GAIANET
152+
).dimensions; // GaiaNet dimension
142153
}
143154

144155
return Array(embeddingDimension).fill(0);

0 commit comments

Comments
 (0)