From 56a4fa3de2e6df2dd1563cf187da2f11fdf19c29 Mon Sep 17 00:00:00 2001 From: fzowl Date: Mon, 15 Dec 2025 16:06:10 +0100 Subject: [PATCH] feat: voyageAI features --- docs/vectorizer/api-reference.md | 53 ++- docs/vectorizer/quick-start-voyage.md | 138 ++++++- projects/extension/ai/voyageai.py | 41 ++ .../extension/sql/idempotent/017-voyageai.sql | 80 +++- projects/extension/tests/test_voyageai.py | 349 ++++++++++++++++++ .../pgai/vectorizer/embedders/voyageai.py | 37 +- 6 files changed, 673 insertions(+), 25 deletions(-) diff --git a/docs/vectorizer/api-reference.md b/docs/vectorizer/api-reference.md index f78c73fdd..104936095 100644 --- a/docs/vectorizer/api-reference.md +++ b/docs/vectorizer/api-reference.md @@ -982,24 +982,59 @@ This function is used to create an embedding configuration object that is passed SELECT ai.create_vectorizer( 'my_table'::regclass, embedding => ai.embedding_voyageai( - 'voyage-3-lite', - 512, - api_key_name => "TEST_API_KEY" + 'voyage-3.5-lite', -- or 'voyage-3.5', 'voyage-3-large', etc. + 1024, -- default dimensions for voyage-3.5 models + api_key_name => "TEST_API_KEY", + output_dimension => 512 -- Optional: use 256, 512, 1024, or 2048 ), -- other parameters... ); ``` +**Example with flexible dimensions (Matryoshka embeddings):** +```sql +-- Use 256 dimensions for faster search and less storage +SELECT ai.create_vectorizer( + 'articles'::regclass, + embedding => ai.embedding_voyageai( + 'voyage-3-large', + 1024, -- Schema dimensions + output_dimension => 256 -- Actual embedding dimensions + ), + destination => ai.destination_table('articles_embeddings_256d') +); +``` + +#### Available Models + +**Current Generation (Recommended):** +| Model | Purpose | Default Dimensions | Max Tokens/Request | +|-------|---------|-------------------|-------------------| +| `voyage-3.5-lite` | Cost & latency optimized | 1024 | 1M | +| `voyage-3.5` | General-purpose optimized | 1024 | 320K | +| `voyage-3-large` | Best for general-purpose & multilingual | 1024 | 120K | +| `voyage-code-3` | Code retrieval specialized | 1024 | 120K | +| `voyage-finance-2` | Finance domain | 1024 | 120K | +| `voyage-law-2` | Legal documents | 1024 | 120K | + +**Older Models:** +| Model | Purpose | Default Dimensions | Max Tokens/Request | +|-------|---------|-------------------|-------------------| +| `voyage-3-lite` | General-purpose (older) | 512 | 120K | +| `voyage-2` | General-purpose (legacy) | - | 320K | + #### Parameters The function takes several parameters to customize the Voyage AI embedding configuration: -| Name | Type | Default | Required | Description | -|--------------|---------|------------------|----------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| model | text | - | ✔ | Specify the name of the [Voyage AI model](https://docs.voyageai.com/docs/embeddings#model-choices) to use. | -| dimensions | int | - | ✔ | Define the number of dimensions for the embedding vectors. This should match the output dimensions of the chosen model. | -| input_type | text | 'document' | ✖ | Type of the input text, null, 'query', or 'document'. | -| api_key_name | text | `VOYAGE_API_KEY` | ✖ | Set the name of the environment variable that contains the Voyage AI API key. This allows for flexible API key management without hardcoding keys in the database. On Timescale Cloud, you should set this to the name of the secret that contains the Voyage AI API key. | +| Name | Type | Default | Required | Description | +|------------------|---------|------------------|----------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| model | text | - | ✔ | Specify the name of the [Voyage AI model](https://docs.voyageai.com/docs/embeddings) to use. See table above for available models. | +| dimensions | int | - | ✔ | Define the number of dimensions for the embedding vectors. This should match the output dimensions of the chosen model (typically 1024 for voyage-3.x models). | +| input_type | text | 'document' | ✖ | Type of the input text: null, 'query', or 'document'. Setting this improves retrieval quality by allowing the model to optimize the embedding. | +| api_key_name | text | `VOYAGE_API_KEY` | ✖ | Set the name of the environment variable that contains the Voyage AI API key. This allows for flexible API key management without hardcoding keys in the database. On Timescale Cloud, you should set this to the name of the secret that contains the Voyage AI API key. | +| output_dimension | int | null | ✖ | Set the output dimension for embeddings. Supports 256, 512, 1024, or 2048 for voyage-3.x models. Lower dimensions reduce storage (up to 75%) and improve search speed with minimal accuracy loss. Uses Matryoshka embeddings technique. | +| output_dtype | text | 'float' | ✖ | Set the output data type for embeddings. Options: 'float' (default), 'int8', 'uint8', 'binary', 'ubinary'. Quantized types (int8, uint8) reduce network bandwidth and API costs. Binary types (binary, ubinary) provide maximum compression with 1/8 the dimensions. Embeddings are automatically converted to float for storage in PostgreSQL. | #### Returns diff --git a/docs/vectorizer/quick-start-voyage.md b/docs/vectorizer/quick-start-voyage.md index 979ec605c..c67876495 100644 --- a/docs/vectorizer/quick-start-voyage.md +++ b/docs/vectorizer/quick-start-voyage.md @@ -94,13 +94,69 @@ Now you can create and run a vectorizer. A vectorizer is a pgai concept, it proc 'blog'::regclass, loading => ai.loading_column('contents'), embedding => ai.embedding_voyageai( - 'voyage-3-lite', - 512 + 'voyage-3.5-lite', -- or 'voyage-3.5', 'voyage-3-large', 'voyage-code-3', etc. + 1024 -- default dimensions for voyage-3.5-lite ), destination => ai.destination_table('blog_contents_embeddings') ); ``` + **Available Voyage AI Models:** + - `voyage-3.5-lite`: Cost & latency optimized, 1024 dims (1M tokens/request) - **Recommended** + - `voyage-3.5`: General-purpose optimized, 1024 dims (320K tokens/request) + - `voyage-3-large`: Best for general-purpose & multilingual, 1024 dims (120K tokens/request) + - `voyage-code-3`: Specialized for code retrieval, 1024 dims (120K tokens/request) + - `voyage-finance-2`: Finance domain optimized, 1024 dims + - `voyage-law-2`: Legal document optimized, 1024 dims + - `voyage-3-lite`: Older model, 512 dims (120K tokens/request) + + **Flexible Dimensions (New!):** + For voyage-3.x models, you can specify `output_dimension` to reduce storage and improve performance: + ```sql + -- Use 256 dimensions for 75% storage reduction + SELECT ai.create_vectorizer( + 'blog'::regclass, + loading => ai.loading_column('contents'), + embedding => ai.embedding_voyageai( + 'voyage-3.5-lite', + 1024, -- Schema dimensions + output_dimension => 256 -- Actual embedding dimensions + ), + destination => ai.destination_table('blog_embeddings_compact') + ); + ``` + + **Dimension Trade-offs:** + - **256 dims**: Fastest search, 75% less storage, minimal accuracy loss + - **512 dims**: Balanced performance and accuracy + - **1024 dims**: Default, best accuracy (recommended for most use cases) + - **2048 dims**: Maximum accuracy for complex tasks + + **Quantization (New!):** + Use `output_dtype` to reduce network bandwidth and API costs: + ```sql + -- Use int8 quantization for 4x bandwidth reduction + SELECT ai.create_vectorizer( + 'blog'::regclass, + loading => ai.loading_column('contents'), + embedding => ai.embedding_voyageai( + 'voyage-3.5-lite', + 1024, + output_dtype => 'int8' -- Options: float, int8, uint8, binary, ubinary + ), + destination => ai.destination_table('blog_embeddings_quantized') + ); + ``` + + **Quantization Options:** + - **float**: Default, no compression (4 bytes per dimension) + - **int8**: Integer quantization, 4x smaller transfer (~1 byte per dim) + - **uint8**: Unsigned integer quantization, 4x smaller + - **binary**: Maximum compression, 32x smaller (1 bit per dim) + - **ubinary**: Unsigned binary, 32x smaller + + Note: Quantized embeddings are automatically converted to float for storage in PostgreSQL, so you get bandwidth savings but not storage savings. + 1. **Check the vectorizer worker logs** ```shell docker compose logs -f vectorizer-worker @@ -118,7 +174,7 @@ Now you can create and run a vectorizer. A vectorizer is a pgai concept, it proc ```sql SELECT chunk, - embedding <=> ai.voyageai_embed('voyage-3-lite', 'good food') as distance + embedding <=> ai.voyageai_embed('voyage-3.5-lite', 'good food') as distance FROM blog_contents_embeddings ORDER BY distance; ``` @@ -134,6 +190,78 @@ The results look like: | Cloud computing has revolutionized the way businesses operate... | 0.9131323552491029 | -That's it, you're done. You now have a table in Postgres that pgai vectorizer automatically creates -and syncs embeddings for. You can use this vectorizer for semantic search, RAG or any other AI +## Reranking with Voyage AI + +Voyage AI also provides reranking capabilities to improve search result relevance. Reranking takes your initial search results and reorders them based on relevance to your query. + +### Using the Reranker + +**Basic reranking:** +```sql +SELECT * +FROM ai.voyageai_rerank_simple( + 'rerank-2.5', + 'What are best practices for healthy eating?', + ARRAY[ + 'Maintaining a healthy diet can be challenging for busy professionals...', + 'Blogging can be a great way to share your thoughts and expertise...', + 'PostgreSQL is a powerful, open source object-relational database system...', + 'As we look towards the future, artificial intelligence continues to evolve...', + 'Cloud computing has revolutionized the way businesses operate...' + ], + api_key => 'your-api-key' +) +ORDER BY relevance_score DESC; +``` + +**Results:** +| index | document | relevance_score | +|-------|----------|-----------------| +| 0 | Maintaining a healthy diet can be challenging... | 0.9156 | +| 1 | Blogging can be a great way to share... | 0.2341 | +| 4 | Cloud computing has revolutionized... | 0.1023 | +| ... | ... | ... | + +**Limit results with top_k:** +```sql +SELECT * +FROM ai.voyageai_rerank_simple( + 'rerank-2.5-lite', + 'healthy eating', + ARRAY['...'], + api_key => 'your-api-key', + top_k => 3 +) +ORDER BY relevance_score DESC; +``` + +### Available Reranker Models + +**Current Generation (Recommended):** +| Model | Context Length | Best For | +|-------|---------------|----------| +| `rerank-2.5` | 32K tokens | Quality with multilingual/instruction support | +| `rerank-2.5-lite` | 32K tokens | Latency & quality balance | + +**Older Models:** +| Model | Context Length | Notes | +|-------|---------------|-------| +| `rerank-2` | 16K tokens | Legacy | +| `rerank-2-lite` | 8K tokens | Legacy | +| `rerank-1` | 8K tokens | Legacy | +| `rerank-lite-1` | 4K tokens | Legacy | + +### Reranker vs Semantic Search + +- **Semantic Search** (embeddings): Fast initial retrieval from large datasets +- **Reranking**: Precise relevance scoring for top-k results from semantic search + +**Typical workflow:** +1. Use semantic search to get top 100 candidates +2. Use reranker to get the most relevant 5-10 results + +--- + +That's it, you're done. You now have a table in Postgres that pgai vectorizer automatically creates +and syncs embeddings for. You can use this vectorizer for semantic search, RAG or any other AI app you can think of! If you have any questions, reach out to us on [Discord](https://discord.gg/KRdHVXAmkp). \ No newline at end of file diff --git a/projects/extension/ai/voyageai.py b/projects/extension/ai/voyageai.py index df3c29699..be444be91 100644 --- a/projects/extension/ai/voyageai.py +++ b/projects/extension/ai/voyageai.py @@ -11,12 +11,53 @@ def embed( api_key: str, input_type: str | None = None, truncation: bool | None = None, + output_dimension: int | None = None, + output_dtype: str | None = None, ) -> Generator[tuple[int, list[float]], None, None]: client = voyageai.Client(api_key=api_key) args = {} if truncation is not None: args["truncation"] = truncation + if output_dimension is not None: + args["output_dimension"] = output_dimension + if output_dtype is not None: + args["output_dtype"] = output_dtype response = client.embed(input, model=model, input_type=input_type, **args) if not hasattr(response, "embeddings"): return None yield from enumerate(response.embeddings) + + +def rerank( + model: str, + query: str, + documents: list[str], + api_key: str, + top_k: int | None = None, + truncation: bool | None = None, +) -> dict: + """ + Rerank documents using Voyage AI reranker API. + Returns the response as a dictionary for JSON serialization. + """ + client = voyageai.Client(api_key=api_key) + args = {} + if top_k is not None: + args["top_k"] = top_k + if truncation is not None: + args["truncation"] = truncation + + response = client.rerank(model=model, query=query, documents=documents, **args) + + # Convert response to dict for JSON serialization + return { + "results": [ + { + "index": r.index, + "document": r.document, + "relevance_score": r.relevance_score, + } + for r in response.results + ], + "total_tokens": response.total_tokens, + } diff --git a/projects/extension/sql/idempotent/017-voyageai.sql b/projects/extension/sql/idempotent/017-voyageai.sql index f7a7b0d82..09facc9df 100644 --- a/projects/extension/sql/idempotent/017-voyageai.sql +++ b/projects/extension/sql/idempotent/017-voyageai.sql @@ -53,11 +53,87 @@ as $python$ args["input_type"] = input_type with ai.utils.VerboseRequestTrace(plpy, "voyageai.embed()", verbose): - results = ai.voyageai.embed(model, input_texts, api_key=api_key_resolved, **args) - + results = ai.voyageai.embed(model, input_texts, api_key=api_key_resolved, **args) + for tup in results: yield tup $python$ language plpython3u immutable parallel safe security invoker set search_path to pg_catalog, pg_temp ; + +------------------------------------------------------------------------------- +-- voyageai_rerank +-- https://docs.voyageai.com/docs/reranker +create or replace function ai.voyageai_rerank +( model text +, query text +, documents text[] +, api_key text default null +, api_key_name text default null +, top_k integer default null +, truncation boolean default null +, verbose boolean default false +) returns jsonb +as $python$ + #ADD-PYTHON-LIB-DIR + import ai.voyageai + import ai.secrets + import ai.utils + import json + api_key_resolved = ai.secrets.get_secret(plpy, api_key, api_key_name, ai.voyageai.DEFAULT_KEY_NAME, SD) + + args = {} + if top_k is not None: + args["top_k"] = top_k + if truncation is not None: + args["truncation"] = truncation + + with ai.utils.VerboseRequestTrace(plpy, "voyageai.rerank()", verbose): + response = ai.voyageai.rerank(model, query, documents, api_key_resolved, **args) + + return json.dumps(response) +$python$ language plpython3u immutable parallel safe security invoker +set search_path to pg_catalog, pg_temp +; + +------------------------------------------------------------------------------- +-- voyageai_rerank_simple +-- https://docs.voyageai.com/docs/reranker +create or replace function ai.voyageai_rerank_simple +( model text +, query text +, documents text[] +, api_key text default null +, api_key_name text default null +, top_k integer default null +, truncation boolean default null +, verbose boolean default false +) returns table +( "index" int +, "document" text +, relevance_score float8 +) +as $func$ +select + x."index" +, d.document +, x.relevance_score +from pg_catalog.jsonb_to_recordset +( + ai.voyageai_rerank + ( model + , query + , documents + , api_key=>api_key + , api_key_name=>api_key_name + , top_k=>top_k + , truncation=>truncation + , verbose=>"verbose" + ) operator(pg_catalog.->) 'results' +) x("index" int, relevance_score float8) +inner join unnest(documents) with ordinality d (document, ord) +on (x."index" = (d.ord - 1)) +$func$ language sql immutable parallel safe security invoker +set search_path to pg_catalog, pg_temp +; diff --git a/projects/extension/tests/test_voyageai.py b/projects/extension/tests/test_voyageai.py index e3233577d..ad623b7e5 100644 --- a/projects/extension/tests/test_voyageai.py +++ b/projects/extension/tests/test_voyageai.py @@ -141,3 +141,352 @@ def test_voyageai_embed_with_multiple_inputs(cur, voyageai_api_key): ) result = cur.fetchone()[0] assert result == 2 + + +def test_voyageai_embed_voyage_3_5_lite(cur, voyageai_api_key): + """Test voyage-3.5-lite model (current naming).""" + cur.execute( + """ + select vector_dims + ( + ai.voyageai_embed + ( 'voyage-3.5-lite' + , 'hello world' + , api_key=>%s + ) + ) + """, + (voyageai_api_key,), + ) + actual = cur.fetchone()[0] + # voyage-3.5-lite has 1024 dimensions by default + assert actual == 1024 + + +def test_voyageai_embed_voyage_3_5(cur, voyageai_api_key): + """Test voyage-3.5 model.""" + cur.execute( + """ + select vector_dims + ( + ai.voyageai_embed + ( 'voyage-3.5' + , 'hello world' + , api_key=>%s + ) + ) + """, + (voyageai_api_key,), + ) + actual = cur.fetchone()[0] + # voyage-3.5 has 1024 dimensions by default + assert actual == 1024 + + +def test_voyageai_embed_voyage_3_large(cur, voyageai_api_key): + """Test voyage-3-large model.""" + cur.execute( + """ + select vector_dims + ( + ai.voyageai_embed + ( 'voyage-3-large' + , 'hello world' + , api_key=>%s + ) + ) + """, + (voyageai_api_key,), + ) + actual = cur.fetchone()[0] + # voyage-3-large has 1024 dimensions by default + assert actual == 1024 + + +def test_voyageai_embed_with_output_dimension_256(cur, voyageai_api_key): + """Test output_dimension parameter with 256 dimensions.""" + cur.execute( + """ + select vector_dims + ( + ai.voyageai_embed + ( 'voyage-3-large' + , 'hello world' + , api_key=>%s + , output_dimension=>256 + ) + ) + """, + (voyageai_api_key,), + ) + actual = cur.fetchone()[0] + assert actual == 256 + + +def test_voyageai_embed_with_output_dimension_512(cur, voyageai_api_key): + """Test output_dimension parameter with 512 dimensions.""" + cur.execute( + """ + select vector_dims + ( + ai.voyageai_embed + ( 'voyage-3.5-lite' + , 'hello world' + , api_key=>%s + , output_dimension=>512 + ) + ) + """, + (voyageai_api_key,), + ) + actual = cur.fetchone()[0] + assert actual == 512 + + +def test_voyageai_embed_with_output_dimension_2048(cur, voyageai_api_key): + """Test output_dimension parameter with 2048 dimensions.""" + cur.execute( + """ + select vector_dims + ( + ai.voyageai_embed + ( 'voyage-3.5' + , 'hello world' + , api_key=>%s + , output_dimension=>2048 + ) + ) + """, + (voyageai_api_key,), + ) + actual = cur.fetchone()[0] + assert actual == 2048 + + +def test_voyageai_embed_output_dimension_with_multiple_inputs(cur, voyageai_api_key): + """Test output_dimension works with multiple inputs.""" + cur.execute( + """ + select count(*) from ai.voyageai_embed + ( 'voyage-3-large' + , ARRAY['hello world', 'goodbye world', 'test embedding'] + , api_key=>%s + , output_dimension=>256 + ) + """, + (voyageai_api_key,), + ) + result = cur.fetchone()[0] + assert result == 3 + + +def test_voyageai_embed_with_output_dtype_float(cur, voyageai_api_key): + """Test output_dtype parameter with float (default).""" + cur.execute( + """ + select vector_dims + ( + ai.voyageai_embed + ( 'voyage-3-large' + , 'hello world' + , api_key=>%s + , output_dtype=>'float' + ) + ) + """, + (voyageai_api_key,), + ) + actual = cur.fetchone()[0] + assert actual == 1024 + + +def test_voyageai_embed_with_output_dtype_int8(cur, voyageai_api_key): + """Test output_dtype parameter with int8 quantization.""" + cur.execute( + """ + select vector_dims + ( + ai.voyageai_embed + ( 'voyage-3-large' + , 'hello world' + , api_key=>%s + , output_dtype=>'int8' + ) + ) + """, + (voyageai_api_key,), + ) + actual = cur.fetchone()[0] + # int8 still returns 1024 dimensions (quantized values converted to float for storage) + assert actual == 1024 + + +def test_voyageai_embed_with_output_dtype_uint8(cur, voyageai_api_key): + """Test output_dtype parameter with uint8 quantization.""" + cur.execute( + """ + select vector_dims + ( + ai.voyageai_embed + ( 'voyage-3.5-lite' + , 'hello world' + , api_key=>%s + , output_dtype=>'uint8' + ) + ) + """, + (voyageai_api_key,), + ) + actual = cur.fetchone()[0] + assert actual == 1024 + + +def test_voyageai_embed_with_output_dtype_binary(cur, voyageai_api_key): + """Test output_dtype parameter with binary quantization.""" + cur.execute( + """ + select vector_dims + ( + ai.voyageai_embed + ( 'voyage-3-large' + , 'hello world' + , api_key=>%s + , output_dimension=>1024 + , output_dtype=>'binary' + ) + ) + """, + (voyageai_api_key,), + ) + actual = cur.fetchone()[0] + # binary returns 1024/8 = 128 packed integers, converted to float for storage + assert actual == 128 + + +def test_voyageai_embed_with_output_dtype_ubinary(cur, voyageai_api_key): + """Test output_dtype parameter with ubinary quantization.""" + cur.execute( + """ + select vector_dims + ( + ai.voyageai_embed + ( 'voyage-3.5' + , 'hello world' + , api_key=>%s + , output_dimension=>512 + , output_dtype=>'ubinary' + ) + ) + """, + (voyageai_api_key,), + ) + actual = cur.fetchone()[0] + # ubinary returns 512/8 = 64 packed integers + assert actual == 64 + + +def test_voyageai_embed_output_dtype_with_dimension(cur, voyageai_api_key): + """Test output_dtype combined with output_dimension.""" + cur.execute( + """ + select vector_dims + ( + ai.voyageai_embed + ( 'voyage-3-large' + , 'hello world' + , api_key=>%s + , output_dimension=>256 + , output_dtype=>'int8' + ) + ) + """, + (voyageai_api_key,), + ) + actual = cur.fetchone()[0] + assert actual == 256 + + +def test_voyageai_rerank(cur, voyageai_api_key): + """Test basic reranking functionality.""" + cur.execute( + """ + with x as + ( + select ai.voyageai_rerank + ( 'rerank-2.5' + , 'How long does it take for two programmers to work on something?' + , array + [ $$Good programmers don't just write programs. They build a working vocabulary.$$ + , 'One of the best programming skills you can have is knowing when to walk away for awhile.' + , 'What one programmer can do in one month, two programmers can do in two months.' + , 'how much wood would a woodchuck chuck if a woodchuck could chuck wood?' + ] + , api_key=>%s + ) as actual + ) + select y."index" as actual + from x + cross join lateral jsonb_to_recordset(x.actual->'results') y("index" int, relevance_score float8) + order by y.relevance_score desc + limit 1 + """, + (voyageai_api_key,), + ) + actual = cur.fetchone()[0] + # The document at index 2 is most relevant to the query + assert actual == 2 + + +def test_voyageai_rerank_with_top_k(cur, voyageai_api_key): + """Test reranking with top_k parameter.""" + cur.execute( + """ + with x as + ( + select ai.voyageai_rerank + ( 'rerank-2.5-lite' + , 'programming best practices' + , array + [ $$Good programmers don't just write programs. They build a working vocabulary.$$ + , 'One of the best programming skills you can have is knowing when to walk away for awhile.' + , 'What one programmer can do in one month, two programmers can do in two months.' + , 'how much wood would a woodchuck chuck if a woodchuck could chuck wood?' + ] + , api_key=>%s + , top_k=>2 + ) as actual + ) + select jsonb_array_length(x.actual->'results') as result_count + from x + """, + (voyageai_api_key,), + ) + actual = cur.fetchone()[0] + # Should only return top 2 results + assert actual == 2 + + +def test_voyageai_rerank_simple(cur, voyageai_api_key): + """Test simplified reranking interface.""" + cur.execute( + """ + select x."index" as actual + from ai.voyageai_rerank_simple + ( 'rerank-2.5' + , 'How long does it take for two programmers to work on something?' + , array + [ $$Good programmers don't just write programs. They build a working vocabulary.$$ + , 'One of the best programming skills you can have is knowing when to walk away for awhile.' + , 'What one programmer can do in one month, two programmers can do in two months.' + , 'how much wood would a woodchuck chuck if a woodchuck could chuck wood?' + ] + , api_key=>%s + ) x + order by relevance_score desc + limit 1 + """, + (voyageai_api_key,), + ) + actual = cur.fetchone()[0] + # The document at index 2 is most relevant + assert actual == 2 diff --git a/projects/pgai/pgai/vectorizer/embedders/voyageai.py b/projects/pgai/pgai/vectorizer/embedders/voyageai.py index f851cb104..b40ffb5a3 100644 --- a/projects/pgai/pgai/vectorizer/embedders/voyageai.py +++ b/projects/pgai/pgai/vectorizer/embedders/voyageai.py @@ -1,5 +1,5 @@ from collections.abc import AsyncGenerator, Callable -from typing import Literal +from typing import Any, Literal from pydantic import BaseModel from tokenizers import Tokenizer @@ -18,16 +18,16 @@ def voyage_max_tokens_per_batch(model: str) -> int: # According to https://docs.voyageai.com/docs/embeddings: # The total number of tokens in the list is at most: - # - 1M for voyage-3-lite - # - 320K for voyage-3 and voyage-2 - # - 120K for voyage-3-large, voyage-code-3, voyage-large-2-instruct, voyage-finance-2, voyage-multilingual-2, voyage-law-2, and voyage-large-2 + # - 1M for voyage-3.5-lite and voyage-3-lite + # - 320K for voyage-3.5, voyage-3, and voyage-2 + # - 120K for voyage-3-large, voyage-code-3, voyage-large-2-instruct, voyage-finance-2, voyage-multilingual-2, voyage-law-2, voyage-large-2, and voyage-3-lite match model: - case "voyage-3-lite": + case "voyage-3.5-lite" | "voyage-3-lite": return 1_000_000 - case "voyage-2" | "voyage-3": + case "voyage-3.5" | "voyage-2" | "voyage-3": return 320_000 case _: - return 120_000 # NOTE: This is conservative, but there probably won't be new Voyage models, so... + return 120_000 # Conservative default for specialized and older models def voyage_token_counter( @@ -54,12 +54,22 @@ class VoyageAI(ApiKeyMixin, BaseModel, Embedder): model (str): The name of the Voyage AU model used for embeddings. input_type ("document" | "query" | None): Set the input type of the items to be embedded. If set, improves retrieval quality. + output_dimension (int | None): Set the output dimension for embeddings. + Supports 256, 512, 1024, or 2048 for newer models (voyage-3.x). + Lower dimensions reduce storage and improve search speed with slight + accuracy trade-off (Matryoshka embeddings). + output_dtype (str | None): Set the output data type for embeddings. + Supports "float" (default), "int8", "uint8", "binary", "ubinary". + Quantized types reduce network bandwidth and API costs. Embeddings + are automatically converted to float for storage in PostgreSQL. """ implementation: Literal["voyageai"] model: str input_type: Literal["document"] | Literal["query"] | None = None + output_dimension: int | None = None + output_dtype: str | None = None @override async def embed( @@ -100,10 +110,19 @@ async def call_embed_api(self, documents: list[str]) -> EmbeddingResponse: # Note: deferred import to avoid import overhead import voyageai + # Build API call parameters + params: dict[str, Any] = { + "model": self.model, + "input_type": self.input_type, + } + if self.output_dimension is not None: + params["output_dimension"] = self.output_dimension + if self.output_dtype is not None: + params["output_dtype"] = self.output_dtype + response = await voyageai.AsyncClient(api_key=self._api_key).embed( documents, - model=self.model, - input_type=self.input_type, + **params, ) usage = Usage( prompt_tokens=response.total_tokens,