Skip to content

Commit

Permalink
Revert "feat: Add support for VoyageAI embeddings API"
Browse files Browse the repository at this point in the history
  • Loading branch information
@shakkernerd
shakkernerd authored Jan 14, 2025
1 parent ecfd42c commit 58f9287
Show file tree
Hide file tree
Showing 4 changed files with 80 additions and 349 deletions.
6 changes: 0 additions & 6 deletions .env.example
View file
Original file line number Diff line number Diff line change
Expand Up @@ -211,12 +211,6 @@ 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
Expand Down
165 changes: 80 additions & 85 deletions packages/core/src/embedding.ts
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,7 @@
import path from "node:path";
import { getEmbeddingModelSettings, getEndpoint } from "./models.ts";
import { IAgentRuntime, ModelProviderName } from "./types.ts";
import settings from "./settings.ts";
import elizaLogger from "./logger.ts";
import { getVoyageAIEmbeddingConfig } from "./voyageai.ts";
import { models, getEmbeddingModelSettings, getEndpoint } from "./models.ts";
import { IAgentRuntime, ModelProviderName } from "./types.ts";
import LocalEmbeddingModelManager from "./localembeddingManager.ts";

interface EmbeddingOptions {
Expand All @@ -22,93 +20,63 @@ export const EmbeddingProvider = {
GaiaNet: "GaiaNet",
Heurist: "Heurist",
BGE: "BGE",
VoyageAI: "VoyageAI",
} as const;

export type EmbeddingProviderType =
(typeof EmbeddingProvider)[keyof typeof EmbeddingProvider];

export namespace EmbeddingProvider {
export type OpenAI = typeof EmbeddingProvider.OpenAI;
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;
readonly provider: EmbeddingProviderType;
};

// 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,
};
};
export const getEmbeddingConfig = (): EmbeddingConfig => ({
dimensions:
settings.USE_OPENAI_EMBEDDING?.toLowerCase() === "true"
? getEmbeddingModelSettings(ModelProviderName.OPENAI).dimensions
: settings.USE_OLLAMA_EMBEDDING?.toLowerCase() === "true"
? getEmbeddingModelSettings(ModelProviderName.OLLAMA).dimensions
: settings.USE_GAIANET_EMBEDDING?.toLowerCase() === "true"
? getEmbeddingModelSettings(ModelProviderName.GAIANET)
.dimensions
: settings.USE_HEURIST_EMBEDDING?.toLowerCase() === "true"
? getEmbeddingModelSettings(ModelProviderName.HEURIST)
.dimensions
: 384, // BGE
model:
settings.USE_OPENAI_EMBEDDING?.toLowerCase() === "true"
? getEmbeddingModelSettings(ModelProviderName.OPENAI).name
: settings.USE_OLLAMA_EMBEDDING?.toLowerCase() === "true"
? getEmbeddingModelSettings(ModelProviderName.OLLAMA).name
: settings.USE_GAIANET_EMBEDDING?.toLowerCase() === "true"
? getEmbeddingModelSettings(ModelProviderName.GAIANET).name
: settings.USE_HEURIST_EMBEDDING?.toLowerCase() === "true"
? getEmbeddingModelSettings(ModelProviderName.HEURIST).name
: "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"
: settings.USE_HEURIST_EMBEDDING?.toLowerCase() === "true"
? "Heurist"
: "BGE",
});

async function getRemoteEmbedding(
input: string,
options: EmbeddingConfig
options: EmbeddingOptions
): Promise<number[]> {
elizaLogger.debug("Getting remote embedding using provider:", options.provider);
// Ensure endpoint ends with /v1 for OpenAI
const baseEndpoint = options.endpoint.endsWith("/v1")
? options.endpoint
: `${options.endpoint}${options.isOllama ? "/v1" : ""}`;

// Construct full URL
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 fullUrl = `${baseEndpoint}/embeddings`;

const requestOptions = {
method: "POST",
Expand All @@ -120,7 +88,14 @@ async function getRemoteEmbedding(
}
: {}),
},
body: JSON.stringify(body),
body: JSON.stringify({
input,
model: options.model,
dimensions:
options.dimensions ||
options.length ||
getEmbeddingConfig().dimensions, // Prefer dimensions, fallback to length
}),
};

try {
Expand Down Expand Up @@ -166,18 +141,44 @@ export function getEmbeddingType(runtime: IAgentRuntime): "local" | "remote" {
}

export function getEmbeddingZeroVector(): number[] {
// Default BGE dimension is 384
return Array(getEmbeddingConfig().dimensions).fill(0);
let embeddingDimension = 384; // Default BGE dimension

if (settings.USE_OPENAI_EMBEDDING?.toLowerCase() === "true") {
embeddingDimension = getEmbeddingModelSettings(
ModelProviderName.OPENAI
).dimensions; // OpenAI dimension
} else if (settings.USE_OLLAMA_EMBEDDING?.toLowerCase() === "true") {
embeddingDimension = getEmbeddingModelSettings(
ModelProviderName.OLLAMA
).dimensions; // Ollama mxbai-embed-large dimension
} else if (settings.USE_GAIANET_EMBEDDING?.toLowerCase() === "true") {
embeddingDimension = getEmbeddingModelSettings(
ModelProviderName.GAIANET
).dimensions; // GaiaNet dimension
} else if (settings.USE_HEURIST_EMBEDDING?.toLowerCase() === "true") {
embeddingDimension = getEmbeddingModelSettings(
ModelProviderName.HEURIST
).dimensions; // Heurist dimension
}

return Array(embeddingDimension).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<number[]>} Array of embedding values
* @throws {Error} If the API request fails or configuration is invalid
* @throws {Error} If the API request fails
*/

export async function embed(runtime: IAgentRuntime, input: string) {
elizaLogger.debug("Embedding request:", {
modelProvider: runtime.character.modelProvider,
Expand Down Expand Up @@ -206,11 +207,6 @@ export async function embed(runtime: IAgentRuntime, input: string) {
const config = getEmbeddingConfig();
const isNode = typeof process !== "undefined" && process.versions?.node;

// Attempt remote embedding if it is configured.
if (config.provider !== EmbeddingProvider.BGE) {
return await getRemoteEmbedding(input, config);
}

// Determine which embedding path to use
if (config.provider === EmbeddingProvider.OpenAI) {
return await getRemoteEmbedding(input, {
Expand Down Expand Up @@ -275,7 +271,6 @@ export async function embed(runtime: IAgentRuntime, input: string) {
getEndpoint(runtime.character.modelProvider),
apiKey: runtime.token,
dimensions: config.dimensions,
provider: config.provider,
});

async function getLocalEmbedding(input: string): Promise<number[]> {
Expand Down
102 changes: 0 additions & 102 deletions packages/core/src/tests/embeddings.test.ts
View file

This file was deleted.

Loading

0 comments on commit 58f9287

Please sign in to comment.