From 9295bde28354620d69b47a5350ade8969d0cc4d4 Mon Sep 17 00:00:00 2001 From: Kyle Oswald Date: Tue, 24 Dec 2024 16:54:18 -0500 Subject: [PATCH] Add support for VoyageAI embeddings API Added support for environment variables: - USE_VOYAGEAI_EMBEDDING - VOYAGEAI_API_KEY - VOYAGEAI_EMBEDDING_DIMENSIONS - VOYAGEAI_EMBEDDING_MODEL Configuration follows existing patterns. Values for dimensions and model can be found in the VoyageAI API documentation. Some minor cleanup of the embedding.ts file. Added unit tests around embedding configuration. --- .env.example | 6 + packages/core/src/embedding.ts | 193 ++++++++------------- packages/core/src/tests/embeddings.test.ts | 102 +++++++++++ packages/core/src/voyageai.ts | 156 +++++++++++++++++ 4 files changed, 337 insertions(+), 120 deletions(-) create mode 100644 packages/core/src/tests/embeddings.test.ts create mode 100644 packages/core/src/voyageai.ts diff --git a/.env.example b/.env.example index 414f963858..3c4240da94 100644 --- a/.env.example +++ b/.env.example @@ -136,6 +136,12 @@ SMALL_ANTHROPIC_MODEL= # Default: claude-3-haiku-20240307 MEDIUM_ANTHROPIC_MODEL= # Default: claude-3-5-sonnet-20241022 LARGE_ANTHROPIC_MODEL= # Default: claude-3-5-sonnet-20241022 +# VoyageAI Configuration +VOYAGEAI_API_KEY= +USE_VOYAGEAI_EMBEDDING= # Set to TRUE for VoyageAI, leave blank for local +VOYAGEAI_EMBEDDING_MODEL= # Default: voyage-3-lite +VOYAGEAI_EMBEDDING_DIMENSIONS= # Default: 512 + # Heurist Configuration HEURIST_API_KEY= # Get from https://heurist.ai/dev-access SMALL_HEURIST_MODEL= # Default: meta-llama/llama-3-70b-instruct diff --git a/packages/core/src/embedding.ts b/packages/core/src/embedding.ts index 767b6b5673..0f1c5a7d99 100644 --- a/packages/core/src/embedding.ts +++ b/packages/core/src/embedding.ts @@ -1,24 +1,16 @@ import path from "node:path"; import { models } from "./models.ts"; -import { IAgentRuntime, ModelProviderName } from "./types.ts"; +import { IAgentRuntime } from "./types.ts"; import settings from "./settings.ts"; import elizaLogger from "./logger.ts"; - -interface EmbeddingOptions { - model: string; - endpoint: string; - apiKey?: string; - length?: number; - isOllama?: boolean; - dimensions?: number; - provider?: string; -} +import { getVoyageAIEmbeddingConfig } from "./voyageai.ts"; export const EmbeddingProvider = { OpenAI: "OpenAI", Ollama: "Ollama", GaiaNet: "GaiaNet", BGE: "BGE", + VoyageAI: "VoyageAI", } as const; export type EmbeddingProvider = @@ -29,52 +21,82 @@ export namespace EmbeddingProvider { export type Ollama = typeof EmbeddingProvider.Ollama; export type GaiaNet = typeof EmbeddingProvider.GaiaNet; export type BGE = typeof EmbeddingProvider.BGE; + export type VoyageAI = typeof EmbeddingProvider.VoyageAI; } export type EmbeddingConfig = { readonly dimensions: number; readonly model: string; readonly provider: EmbeddingProvider; + readonly endpoint?: string; + readonly apiKey?: string; + readonly maxInputTokens?: number; }; -export const getEmbeddingConfig = (): EmbeddingConfig => ({ - dimensions: - settings.USE_OPENAI_EMBEDDING?.toLowerCase() === "true" - ? 1536 // OpenAI - : settings.USE_OLLAMA_EMBEDDING?.toLowerCase() === "true" - ? 1024 // Ollama mxbai-embed-large - : settings.USE_GAIANET_EMBEDDING?.toLowerCase() === "true" - ? 768 // GaiaNet - : 384, // BGE - model: - settings.USE_OPENAI_EMBEDDING?.toLowerCase() === "true" - ? "text-embedding-3-small" - : settings.USE_OLLAMA_EMBEDDING?.toLowerCase() === "true" - ? settings.OLLAMA_EMBEDDING_MODEL || "mxbai-embed-large" - : settings.USE_GAIANET_EMBEDDING?.toLowerCase() === "true" - ? settings.GAIANET_EMBEDDING_MODEL || "nomic-embed" - : "BGE-small-en-v1.5", - provider: - settings.USE_OPENAI_EMBEDDING?.toLowerCase() === "true" - ? "OpenAI" - : settings.USE_OLLAMA_EMBEDDING?.toLowerCase() === "true" - ? "Ollama" - : settings.USE_GAIANET_EMBEDDING?.toLowerCase() === "true" - ? "GaiaNet" - : "BGE", -}); +// Get embedding config based on settings +export function getEmbeddingConfig(): EmbeddingConfig { + if (settings.USE_OPENAI_EMBEDDING?.toLowerCase() === "true") { + return { + dimensions: 1536, + model: "text-embedding-3-small", + provider: "OpenAI", + endpoint: "https://api.openai.com/v1", + apiKey: settings.OPENAI_API_KEY, + maxInputTokens: 1000000, + }; + } + if (settings.USE_OLLAMA_EMBEDDING?.toLowerCase() === "true") { + return { + dimensions: 1024, + model: settings.OLLAMA_EMBEDDING_MODEL || "mxbai-embed-large", + provider: "Ollama", + endpoint: "https://ollama.eliza.ai/", + apiKey: settings.OLLAMA_API_KEY, + maxInputTokens: 1000000, + }; + } + if (settings.USE_GAIANET_EMBEDDING?.toLowerCase() === "true") { + return { + dimensions: 768, + model: settings.GAIANET_EMBEDDING_MODEL || "nomic-embed", + provider: "GaiaNet", + endpoint: settings.SMALL_GAIANET_SERVER_URL || settings.MEDIUM_GAIANET_SERVER_URL || settings.LARGE_GAIANET_SERVER_URL, + apiKey: settings.GAIANET_API_KEY, + maxInputTokens: 1000000, + }; + } + if (settings.USE_VOYAGEAI_EMBEDDING?.toLowerCase() === "true") { + return getVoyageAIEmbeddingConfig(); + } + + // Fallback to local BGE + return { + dimensions: 384, + model: "BGE-small-en-v1.5", + provider: "BGE", + maxInputTokens: 1000000, + }; +}; async function getRemoteEmbedding( input: string, - options: EmbeddingOptions + options: EmbeddingConfig ): Promise { - // Ensure endpoint ends with /v1 for OpenAI - const baseEndpoint = options.endpoint.endsWith("/v1") - ? options.endpoint - : `${options.endpoint}${options.isOllama ? "/v1" : ""}`; + elizaLogger.debug("Getting remote embedding using provider:", options.provider); // Construct full URL - const fullUrl = `${baseEndpoint}/embeddings`; + const fullUrl = `${options.endpoint}/embeddings`; + + // jank. voyageai is the only one that doesn't use "dimensions". + const body = options.provider === "VoyageAI" ? { + input, + model: options.model, + output_dimension: options.dimensions, + } : { + input, + model: options.model, + dimensions: options.dimensions, + }; const requestOptions = { method: "POST", @@ -86,14 +108,7 @@ async function getRemoteEmbedding( } : {}), }, - body: JSON.stringify({ - input, - model: options.model, - dimensions: - options.dimensions || - options.length || - getEmbeddingConfig().dimensions, // Prefer dimensions, fallback to length - }), + body: JSON.stringify(body), }; try { @@ -118,52 +133,19 @@ async function getRemoteEmbedding( } } -export function getEmbeddingType(runtime: IAgentRuntime): "local" | "remote" { - const isNode = - typeof process !== "undefined" && - process.versions != null && - process.versions.node != null; - - // Use local embedding if: - // - Running in Node.js - // - Not using OpenAI provider - // - Not forcing OpenAI embeddings - const isLocal = - isNode && - runtime.character.modelProvider !== ModelProviderName.OPENAI && - runtime.character.modelProvider !== ModelProviderName.GAIANET && - !settings.USE_OPENAI_EMBEDDING; - - return isLocal ? "local" : "remote"; -} - export function getEmbeddingZeroVector(): number[] { - let embeddingDimension = 384; // Default BGE dimension - - if (settings.USE_OPENAI_EMBEDDING?.toLowerCase() === "true") { - embeddingDimension = 1536; // OpenAI dimension - } else if (settings.USE_OLLAMA_EMBEDDING?.toLowerCase() === "true") { - embeddingDimension = 1024; // Ollama mxbai-embed-large dimension - } - - return Array(embeddingDimension).fill(0); + // Default BGE dimension is 384 + return Array(getEmbeddingConfig().dimensions).fill(0); } /** * Gets embeddings from a remote API endpoint. Falls back to local BGE/384 * * @param {string} input - The text to generate embeddings for - * @param {EmbeddingOptions} options - Configuration options including: - * - model: The model name to use - * - endpoint: Base API endpoint URL - * - apiKey: Optional API key for authentication - * - isOllama: Whether this is an Ollama endpoint - * - dimensions: Desired embedding dimensions * @param {IAgentRuntime} runtime - The agent runtime context * @returns {Promise} Array of embedding values - * @throws {Error} If the API request fails + * @throws {Error} If the API request fails or configuration is invalid */ - export async function embed(runtime: IAgentRuntime, input: string) { elizaLogger.debug("Embedding request:", { modelProvider: runtime.character.modelProvider, @@ -192,39 +174,9 @@ export async function embed(runtime: IAgentRuntime, input: string) { const config = getEmbeddingConfig(); const isNode = typeof process !== "undefined" && process.versions?.node; - // Determine which embedding path to use - if (config.provider === EmbeddingProvider.OpenAI) { - return await getRemoteEmbedding(input, { - model: config.model, - endpoint: "https://api.openai.com/v1", - apiKey: settings.OPENAI_API_KEY, - dimensions: config.dimensions, - }); - } - - if (config.provider === EmbeddingProvider.Ollama) { - return await getRemoteEmbedding(input, { - model: config.model, - endpoint: - runtime.character.modelEndpointOverride || - models[ModelProviderName.OLLAMA].endpoint, - isOllama: true, - dimensions: config.dimensions, - }); - } - - if (config.provider == EmbeddingProvider.GaiaNet) { - return await getRemoteEmbedding(input, { - model: config.model, - endpoint: - runtime.character.modelEndpointOverride || - models[ModelProviderName.GAIANET].endpoint || - settings.SMALL_GAIANET_SERVER_URL || - settings.MEDIUM_GAIANET_SERVER_URL || - settings.LARGE_GAIANET_SERVER_URL, - apiKey: settings.GAIANET_API_KEY || runtime.token, - dimensions: config.dimensions, - }); + // Attempt remote embedding if it is configured. + if (config.provider !== EmbeddingProvider.BGE) { + return await getRemoteEmbedding(input, config); } // BGE - try local first if in Node @@ -247,6 +199,7 @@ export async function embed(runtime: IAgentRuntime, input: string) { models[runtime.character.modelProvider].endpoint, apiKey: runtime.token, dimensions: config.dimensions, + provider: config.provider, }); async function getLocalEmbedding(input: string): Promise { diff --git a/packages/core/src/tests/embeddings.test.ts b/packages/core/src/tests/embeddings.test.ts new file mode 100644 index 0000000000..c3f37f961e --- /dev/null +++ b/packages/core/src/tests/embeddings.test.ts @@ -0,0 +1,102 @@ + +import { describe, expect, vi } from "vitest"; +import { getEmbeddingConfig } from '../embedding'; +import settings from '../settings'; + +vi.mock("../settings"); +const mockedSettings = vi.mocked(settings); + +describe('getEmbeddingConfig', () => { + beforeEach(() => { + // Clear the specific mock + Object.keys(mockedSettings).forEach(key => { + delete mockedSettings[key]; + }); + + vi.clearAllMocks(); + }); + + afterEach(() => { + vi.clearAllMocks(); + }); + + it('should return BGE config by default', () => { + + mockedSettings.USE_OPENAI_EMBEDDING = 'false'; + mockedSettings.USE_OLLAMA_EMBEDDING = 'false'; + mockedSettings.USE_GAIANET_EMBEDDING = 'false'; + mockedSettings.USE_VOYAGEAI_EMBEDDING = 'false'; + + const config = getEmbeddingConfig(); + expect(config).toEqual({ + dimensions: 384, + model: 'BGE-small-en-v1.5', + provider: 'BGE', + maxInputTokens: 1000000, + }); + }); + + it('should return GaiaNet config when USE_GAIANET_EMBEDDING is true', () => { + mockedSettings.USE_GAIANET_EMBEDDING = 'true'; + mockedSettings.GAIANET_EMBEDDING_MODEL = 'test-model'; + mockedSettings.GAIANET_API_KEY = 'test-key'; + mockedSettings.SMALL_GAIANET_SERVER_URL = 'https://test.gaianet.ai'; + + const config = getEmbeddingConfig(); + expect(config).toEqual({ + dimensions: 768, + model: 'test-model', + provider: 'GaiaNet', + endpoint: 'https://test.gaianet.ai', + apiKey: 'test-key', + maxInputTokens: 1000000, + }); + }); + + + it('should return VoyageAI config when USE_VOYAGEAI_EMBEDDING is true', () => { + mockedSettings.USE_VOYAGEAI_EMBEDDING = 'true'; + mockedSettings.VOYAGEAI_API_KEY = 'test-key'; + + const config = getEmbeddingConfig(); + expect(config).toEqual({ + dimensions: 512, + model: 'voyage-3-lite', + provider: 'VoyageAI', + endpoint: 'https://api.voyageai.com/v1', + apiKey: 'test-key', + maxInputTokens: 1000000, + }); + }); + + it('should return OpenAI config when USE_OPENAI_EMBEDDING is true', () => { + mockedSettings.USE_OPENAI_EMBEDDING = 'true'; + mockedSettings.OPENAI_API_KEY = 'test-key'; + + const config = getEmbeddingConfig(); + expect(config).toEqual({ + dimensions: 1536, + model: 'text-embedding-3-small', + provider: 'OpenAI', + endpoint: 'https://api.openai.com/v1', + apiKey: 'test-key', + maxInputTokens: 1000000, + }); + }); + + it('should return Ollama config when USE_OLLAMA_EMBEDDING is true', () => { + mockedSettings.USE_OLLAMA_EMBEDDING = 'true'; + mockedSettings.OLLAMA_EMBEDDING_MODEL = 'test-model'; + mockedSettings.OLLAMA_API_KEY = 'test-key'; + + const config = getEmbeddingConfig(); + expect(config).toEqual({ + dimensions: 1024, + model: 'test-model', + provider: 'Ollama', + endpoint: 'https://ollama.eliza.ai/v1', + apiKey: 'test-key', + maxInputTokens: 1000000, + }); + }); +}); \ No newline at end of file diff --git a/packages/core/src/voyageai.ts b/packages/core/src/voyageai.ts new file mode 100644 index 0000000000..a603afa954 --- /dev/null +++ b/packages/core/src/voyageai.ts @@ -0,0 +1,156 @@ +import settings from "./settings.ts"; +import { EmbeddingConfig } from "./embedding.ts"; + +enum VoyageAIModel { + // Current models + VOYAGE_3_LARGE = 'voyage-3-large', + VOYAGE_3 = 'voyage-3', + VOYAGE_3_LITE = 'voyage-3-lite', + VOYAGE_CODE_3 = 'voyage-code-3', + VOYAGE_FINANCE_2 = 'voyage-finance-2', + VOYAGE_LAW_2 = 'voyage-law-2', + VOYAGE_CODE_2 = 'voyage-code-2', + // Legacy models + VOYAGE_MULTILINGUAL_2 = 'voyage-multilingual-2', + VOYAGE_LARGE_2_INSTRUCT = 'voyage-large-2-instruct', + VOYAGE_LARGE_2 = 'voyage-large-2', + VOYAGE_2 = 'voyage-2', + VOYAGE_LITE_02_INSTRUCT = 'voyage-lite-02-instruct', + VOYAGE_02 = 'voyage-02', + VOYAGE_01 = 'voyage-01', + VOYAGE_LITE_01 = 'voyage-lite-01', + VOYAGE_LITE_01_INSTRUCT = 'voyage-lite-01-instruct', +} + +/** + * Gets the VoyageAI embedding model to use based on settings. + * + * Checks if VOYAGEAI_EMBEDDING_MODEL is set in settings and validates that it's + * a valid model name from the VoyageraiModel enum. If no model is configured, + * defaults to VOYAGE_3_LITE. + * + * @returns {VoyageAIModel} The VoyageAI model to use for embeddings + * @throws {Error} If an invalid model name is configured in settings + */ +function getVoyageAIEmbeddingModel(): VoyageAIModel { + const modelName = settings.VOYAGEAI_EMBEDDING_MODEL ?? VoyageAIModel.VOYAGE_3_LITE; + + try { + return modelName as VoyageAIModel; + } catch { + throw new Error(`Invalid voyageai embedding model: ${modelName}`); + } +} + +/** + * Gets the embedding dimensions for the configured VoyageAI model. + * + * Each model supports specific dimension options. If VOYAGEAI_EMBEDDING_DIMENSIONS + * is set in settings, validates that it's a valid option for the model. + * Otherwise returns the default dimensions for that model. + * + * Dimensions by model: + * - VOYAGE_3_LARGE: 256, 512, 1024 (default), 2048 + * - VOYAGE_3: 1024 only + * - VOYAGE_3_LITE: 512 only + * - VOYAGE_CODE_3: 256, 512, 1024 (default), 2048 + * - VOYAGE_FINANCE_2/LAW_2: 1024 only + * - VOYAGE_CODE_2/LARGE_2: 1536 only + * - All legacy models: 1024 only + * + * @returns {number} The embedding dimensions to use + * @throws {Error} If an invalid dimension is configured for the model + * @see {@link getVoyageAIEmbeddingModel} + */ +function getVoyageAIEmbeddingDimensions(): number { + const model = getVoyageAIEmbeddingModel(); + + function validateDimensions(model: VoyageAIModel, defaultDimensions: number, validOptions: number[]) { + if (settings.VOYAGEAI_EMBEDDING_DIMENSIONS != null) { + const dim = Number(settings.VOYAGEAI_EMBEDDING_DIMENSIONS); + if (!validOptions.includes(dim)) { + throw new Error(`Invalid dimensions for ${model}: ${dim}. Valid options are: ${validOptions.join(', ')}`); + } + return dim; + } + return defaultDimensions; + } + + switch (model) { + // Current models + case VoyageAIModel.VOYAGE_3_LARGE: + return validateDimensions(model, 1024, [256, 512, 1024, 2048]); + + case VoyageAIModel.VOYAGE_3: + return validateDimensions(model, 1024, [1024]); + + case VoyageAIModel.VOYAGE_3_LITE: + return validateDimensions(model, 512, [512]); + + case VoyageAIModel.VOYAGE_CODE_3: + return validateDimensions(model, 1024, [256, 512, 1024, 2048]); + + case VoyageAIModel.VOYAGE_FINANCE_2: + case VoyageAIModel.VOYAGE_LAW_2: + return validateDimensions(model, 1024, [1024]); + + case VoyageAIModel.VOYAGE_CODE_2: + case VoyageAIModel.VOYAGE_LARGE_2: + return validateDimensions(model, 1536, [1536]); + + // Legacy models + case VoyageAIModel.VOYAGE_MULTILINGUAL_2: + case VoyageAIModel.VOYAGE_LARGE_2_INSTRUCT: + case VoyageAIModel.VOYAGE_2: + case VoyageAIModel.VOYAGE_LITE_02_INSTRUCT: + case VoyageAIModel.VOYAGE_02: + case VoyageAIModel.VOYAGE_01: + case VoyageAIModel.VOYAGE_LITE_01: + case VoyageAIModel.VOYAGE_LITE_01_INSTRUCT: + return validateDimensions(model, 1024, [1024]); + } + + // Should be unreachable. + throw new Error(`Invalid voyageai embedding model: ${model}`); +} + +/** + * Gets the maximum number of input tokens allowed for the current VoyageAI embedding model + * + * Different VoyageAI models have different token limits: + * - VOYAGE_3_LITE: 1,000,000 tokens + * - VOYAGE_3/VOYAGE_2: 320,000 tokens + * - Other models: 120,000 tokens + * + * @returns {number} The maximum number of input tokens allowed for the current model + */ +function getVoyageAIMaxInputTokens() { + switch (getVoyageAIEmbeddingModel()) { + case VoyageAIModel.VOYAGE_3_LITE: + return 1000000; + case VoyageAIModel.VOYAGE_3: + case VoyageAIModel.VOYAGE_2: + return 320000; + case VoyageAIModel.VOYAGE_3_LARGE: + case VoyageAIModel.VOYAGE_CODE_3: + case VoyageAIModel.VOYAGE_LARGE_2_INSTRUCT: + case VoyageAIModel.VOYAGE_FINANCE_2: + case VoyageAIModel.VOYAGE_MULTILINGUAL_2: + case VoyageAIModel.VOYAGE_LAW_2: + case VoyageAIModel.VOYAGE_LARGE_2: + return 120000; + default: + return 120000; // Default to most conservative limit + } +} + +export function getVoyageAIEmbeddingConfig(): EmbeddingConfig { + return { + dimensions: getVoyageAIEmbeddingDimensions(), + model: getVoyageAIEmbeddingModel(), + provider: "VoyageAI", + maxInputTokens: getVoyageAIMaxInputTokens(), + endpoint: "https://api.voyageai.com/v1", + apiKey: settings.VOYAGEAI_API_KEY, + }; +}