From 1f8ed37f0d65e2d53bcdbc79b04740ba41b8bef4 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 | 180 ++++++++------------- packages/core/src/tests/embeddings.test.ts | 102 ++++++++++++ packages/core/src/types.ts | 14 ++ packages/core/src/voyageai.ts | 156 ++++++++++++++++++ 5 files changed, 344 insertions(+), 114 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 b1f05998b7..5005c57fce 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 49c1a4163c..c4d26d044e 100644 --- a/packages/core/src/embedding.ts +++ b/packages/core/src/embedding.ts @@ -1,51 +1,61 @@ import path from "node:path"; import { models } from "./models.ts"; -import { IAgentRuntime, ModelProviderName } from "./types.ts"; +import { EmbeddingOptions, IAgentRuntime } from "./types.ts"; import settings from "./settings.ts"; import elizaLogger from "./logger.ts"; +import { getVoyageAIEmbeddingConfig } from "./voyageai.ts"; -interface EmbeddingOptions { - model: string; - endpoint: string; - apiKey?: string; - length?: number; - isOllama?: boolean; - dimensions?: number; - provider?: string; -} +// Get embedding config based on settings +export function getEmbeddingConfig(): EmbeddingOptions { + 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/v1", + 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(); + } -// Add the embedding configuration -export const getEmbeddingConfig = () => ({ - 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", -}); + // Fallback to local BGE + return { + dimensions: 384, + model: "BGE-small-en-v1.5", + provider: "BGE", + maxInputTokens: 1000000, + }; +}; async function getRemoteEmbedding( input: string, options: EmbeddingOptions ): Promise { + elizaLogger.debug("Getting remote embedding using provider:", options.provider); + // Ensure endpoint ends with /v1 for OpenAI const baseEndpoint = options.endpoint.endsWith("/v1") ? options.endpoint @@ -54,6 +64,17 @@ async function getRemoteEmbedding( // Construct full URL const fullUrl = `${baseEndpoint}/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 || options.length, + }; + const requestOptions = { method: "POST", headers: { @@ -64,14 +85,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 { @@ -96,52 +110,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, @@ -170,39 +151,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 === "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 === "Ollama") { - return await getRemoteEmbedding(input, { - model: config.model, - endpoint: - runtime.character.modelEndpointOverride || - models[ModelProviderName.OLLAMA].endpoint, - isOllama: true, - dimensions: config.dimensions, - }); - } - - if (config.provider=="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 !== "BGE") { + return await getRemoteEmbedding(input, config); } // BGE - try local first if in Node @@ -225,6 +176,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/types.ts b/packages/core/src/types.ts index 8bb331e897..a31d5f6f76 100644 --- a/packages/core/src/types.ts +++ b/packages/core/src/types.ts @@ -240,6 +240,20 @@ export enum ModelProviderName { AKASH_CHAT_API = "akash_chat_api", } +/** + * Configuration for remote embedding providers. + */ +export interface EmbeddingOptions { + model: string; + endpoint?: string; + apiKey?: string; + length?: number; + isOllama?: boolean; + dimensions: number; + provider: string; + maxInputTokens?: number; +} + /** * Represents the current state/context of a conversation */ diff --git a/packages/core/src/voyageai.ts b/packages/core/src/voyageai.ts new file mode 100644 index 0000000000..5a1887abd4 --- /dev/null +++ b/packages/core/src/voyageai.ts @@ -0,0 +1,156 @@ +import settings from "./settings.ts"; +import { EmbeddingOptions } from "./types.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(): EmbeddingOptions { + return { + dimensions: getVoyageAIEmbeddingDimensions(), + model: getVoyageAIEmbeddingModel(), + provider: "VoyageAI", + maxInputTokens: getVoyageAIMaxInputTokens(), + endpoint: "https://api.voyageai.com/v1", + apiKey: settings.VOYAGEAI_API_KEY, + }; +}