-
Notifications
You must be signed in to change notification settings - Fork 1.9k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
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.
- Loading branch information
Showing
5 changed files
with
344 additions
and
114 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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, | ||
}); | ||
}); | ||
}); |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.