Skip to content

Feat: VoyageAI integration and embedding models #8443

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 8 commits into
base: master
Choose a base branch
from
Open

Feat: VoyageAI integration and embedding models #8443

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
3fe48a4
VoyageAI integration and embedding models
fzowl Mar 5, 2025
4a1c435
VoyageAI integration and embedding models
fzowl Mar 5, 2025
5658352
VoyageAI integration and embedding models
fzowl Mar 20, 2025
d9b46e9
VoyageAI integration and embedding models
fzowl Mar 23, 2025
d5a81b2
Adding Voyage 3.5 models
fzowl May 21, 2025
215265f
Merge branch 'master' into master
fzowl May 21, 2025
fdc6866
Adding Voyage 3.5 models
fzowl May 21, 2025
9567158
Merge branch 'geldata:master' into master
fzowl Jul 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
133 changes: 132 additions & 1 deletion edb/lib/ext/ai.edgeql
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ CREATE EXTENSION PACKAGE ai VERSION '1.0' {
create module ext::ai;

create scalar type ext::ai::ProviderAPIStyle
extending enum<OpenAI, Anthropic>;
extending enum<OpenAI, Anthropic, VoyageAI>;

create abstract type ext::ai::ProviderConfig extending cfg::ConfigObject {
create required property name: std::str {
Expand Down Expand Up @@ -137,6 +137,27 @@ CREATE EXTENSION PACKAGE ai VERSION '1.0' {
};
};

create type ext::ai::VoyageAIProviderConfig extending ext::ai::ProviderConfig {
alter property name {
set protected := true;
set default := 'builtin::voyageai';
};

alter property display_name {
set protected := true;
set default := 'VoyageAI';
};

alter property api_url {
set default := 'https://api.voyageai.com/v1'
};

alter property api_style {
set protected := true;
set default := ext::ai::ProviderAPIStyle.VoyageAI;
};
};

create type ext::ai::Config extending cfg::ExtensionConfig {
create required property indexer_naptime: std::duration {
set default := <std::duration>'10s';
Expand Down Expand Up @@ -527,6 +548,116 @@ CREATE EXTENSION PACKAGE ai VERSION '1.0' {
ext::ai::text_gen_model_context_window := "200000";
};

# VoyageAI models
create abstract type ext::ai::Voyage3LargeEmbedModel
extending ext::ai::EmbeddingModel
{
alter annotation
ext::ai::model_name := "voyage-3-large";
alter annotation
ext::ai::model_provider := "builtin::voyageai";
alter annotation
ext::ai::embedding_model_max_input_tokens := "32000";
alter annotation
ext::ai::embedding_model_max_batch_tokens := "32000";
alter annotation
ext::ai::embedding_model_max_output_dimensions := "2048";
alter annotation
ext::ai::embedding_model_supports_shortening := "true";
};

create abstract type ext::ai::VoyageCode3EmbedModel
extending ext::ai::EmbeddingModel
{
alter annotation
ext::ai::model_name := "voyage-code-3";
alter annotation
ext::ai::model_provider := "builtin::voyageai";
alter annotation
ext::ai::embedding_model_max_input_tokens := "32000";
alter annotation
ext::ai::embedding_model_max_batch_tokens := "32000";
alter annotation
ext::ai::embedding_model_max_output_dimensions := "2048";
alter annotation
ext::ai::embedding_model_supports_shortening := "true";
};

create abstract type ext::ai::Voyage3EmbedModel
extending ext::ai::EmbeddingModel
{
alter annotation
ext::ai::model_name := "voyage-3";
alter annotation
ext::ai::model_provider := "builtin::voyageai";
alter annotation
ext::ai::embedding_model_max_input_tokens := "32000";
alter annotation
ext::ai::embedding_model_max_batch_tokens := "32000";
alter annotation
ext::ai::embedding_model_max_output_dimensions := "1024";
};

create abstract type ext::ai::Voyage3LiteEmbedModel
extending ext::ai::EmbeddingModel
{
alter annotation
ext::ai::model_name := "voyage-3-lite";
alter annotation
ext::ai::model_provider := "builtin::voyageai";
alter annotation
ext::ai::embedding_model_max_input_tokens := "32000";
alter annotation
ext::ai::embedding_model_max_batch_tokens := "32000";
alter annotation
ext::ai::embedding_model_max_output_dimensions := "512";
};

create abstract type ext::ai::VoyageFinance2EmbedModel
extending ext::ai::EmbeddingModel
{
alter annotation
ext::ai::model_name := "voyage-finance-2";
alter annotation
ext::ai::model_provider := "builtin::voyageai";
alter annotation
ext::ai::embedding_model_max_input_tokens := "32000";
alter annotation
ext::ai::embedding_model_max_batch_tokens := "32000";
alter annotation
ext::ai::embedding_model_max_output_dimensions := "1024";
};

create abstract type ext::ai::VoyageLaw2EmbedModel
extending ext::ai::EmbeddingModel
{
alter annotation
ext::ai::model_name := "voyage-law-2";
alter annotation
ext::ai::model_provider := "builtin::voyageai";
alter annotation
ext::ai::embedding_model_max_input_tokens := "16000";
alter annotation
ext::ai::embedding_model_max_batch_tokens := "16000";
alter annotation
ext::ai::embedding_model_max_output_dimensions := "1024";
};

create abstract type ext::ai::VoyageCode2EmbedModel
extending ext::ai::EmbeddingModel
{
alter annotation
ext::ai::model_name := "voyage-code-2";
alter annotation
ext::ai::model_provider := "builtin::voyageai";
alter annotation
ext::ai::embedding_model_max_input_tokens := "16000";
alter annotation
ext::ai::embedding_model_max_batch_tokens := "16000";
alter annotation
ext::ai::embedding_model_max_output_dimensions := "1536";
};

create scalar type ext::ai::DistanceFunction
extending enum<Cosine, InnerProduct, L2>;

Expand Down
46 changes: 46 additions & 0 deletions edb/server/protocol/ai_ext.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,7 @@ class BadRequestError(AIExtError):
class ApiStyle(s_enum.StrEnum):
OpenAI = 'OpenAI'
Anthropic = 'Anthropic'
VoyageAI = 'VoyageAI'


class Tokenizer(abc.ABC):
Expand Down Expand Up @@ -1062,6 +1063,10 @@ async def _generate_embeddings(
return await _generate_openai_embeddings(
provider, model_name, inputs, shortening, user, http_client
)
elif provider.api_style == ApiStyle.VoyageAI:
return await _generate_voyageai_embeddings(
provider, model_name, inputs, http_client
)
else:
raise RuntimeError(
f"unsupported model provider API style: {provider.api_style}, "
Expand Down Expand Up @@ -1124,6 +1129,47 @@ async def _generate_openai_embeddings(
)


async def _generate_voyageai_embeddings(
provider: ProviderConfig,
model_name: str,
inputs: list[str],
http_client: http.HttpClient,
) -> EmbeddingsResult:

headers = {
"Authorization": f"Bearer {provider.secret}",
}
client = http_client.with_context(
headers=headers,
base_url=provider.api_url,
)

params: dict[str, Any] = {"input": inputs, "model": model_name, "output_format": "float", "truncation": False}

result = await client.post(
"/embeddings",
json=params,
)

error = None
if result.status_code >= 400:
error = rs.Error(
message=(
f"API call to generate embeddings failed with status "
f"{result.status_code}: {result.text}"
),
retry=(
# If the request fails with 429 - too many requests, it can be
# retried
result.status_code == 429
),
)

return EmbeddingsResult(
data=(error if error else EmbeddingsData(result.bytes())),
)


def _read_openai_header_field(
result: Any,
field_names: list[str],
Expand Down