mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-23 12:12:27 +00:00
fix(python): accept VoyageAI text batch containers (#4558)
Direct VoyageAI `compute_source_embeddings()` calls reject `list[str]` and NumPy text batches before reaching the provider, although both belong to the declared `TEXT` input type. For example, `["first", "second"]` raises `ValueError: Input type <class 'list'> not allowed with text model.` Normalize these containers to a list before the existing string-entry validation. The regression cases exercise the public embedding API with a stub provider: input order, list/NumPy batches, Arrow controls, empty batches, and rejection of non-text entries and ordinary two-dimensional arrays. The VoyageAI integration page now documents the accepted text containers. Verified on current main `7b3ba4c` with a locally compiled native extension (0.41.0b7, Python 3.12.12): - VoyageAI suite: 7 failed / 20 passed before the fix; 27 passed afterward. - Adjacent embedding suite: 13 passed, 4 skipped, 8 slow cases deselected. - Python CI's pinned Ruff 0.9.9 lint/format checks pass; configured Pyright reports zero errors/warnings; all applicable pre-commit hooks pass. The provider is stubbed; no live Voyage API request or model-quality evaluation was performed. AI disclosure: I used Codex for investigation, implementation, tests, review, and drafting this description.
This commit is contained in:
@@ -45,6 +45,8 @@ Supported parameters (to be passed in `create` method) are:
|
||||
| `output_dimension` | `int` | `None` | Output embedding dimension. Only supported by `voyage-multimodal-3.5`. Valid options: 256, 512, 1024 (default), 2048. |
|
||||
|
||||
|
||||
For text models, `compute_source_embeddings()` accepts a single string, a list of strings, a one-dimensional NumPy array, or a PyArrow `Array` or `ChunkedArray`. Batch inputs retain their order.
|
||||
|
||||
Usage Example:
|
||||
|
||||
<CodeGroup>
|
||||
@@ -61,4 +63,4 @@ The `voyage-multimodal-3.5` model can embed text alongside images. You can use i
|
||||
<CodeBlock filename="Python" language="Python" icon="python">
|
||||
{PyEmbeddingVoyageaiMultimodal}
|
||||
</CodeBlock>
|
||||
</CodeGroup>
|
||||
</CodeGroup>
|
||||
|
||||
@@ -145,6 +145,8 @@ def sanitize_text_input(inputs: TEXT) -> List[str]:
|
||||
"""
|
||||
if isinstance(inputs, str):
|
||||
inputs = [inputs]
|
||||
elif isinstance(inputs, (list, np.ndarray)):
|
||||
inputs = list(inputs)
|
||||
elif isinstance(inputs, pa.Array):
|
||||
inputs = inputs.to_pylist()
|
||||
elif isinstance(inputs, pa.ChunkedArray):
|
||||
|
||||
@@ -6,6 +6,8 @@
|
||||
These tests verify model registration and configuration without requiring API calls.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
@@ -122,3 +124,58 @@ class TestVoyageAIModelRegistration:
|
||||
assert "voyage-4" not in func.multimodal_embedding_models
|
||||
assert "voyage-4-lite" not in func.multimodal_embedding_models
|
||||
assert "voyage-4-large" not in func.multimodal_embedding_models
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"inputs",
|
||||
[
|
||||
["first", "second"],
|
||||
np.array(["first", "second"]),
|
||||
pa.array(["first", "second"]),
|
||||
pa.chunked_array([["first"], ["second"]]),
|
||||
],
|
||||
ids=["list", "numpy", "arrow", "chunked-arrow"],
|
||||
)
|
||||
def test_text_source_embeddings_accept_batch_containers(inputs):
|
||||
from lancedb.embeddings.voyageai import VoyageAIEmbeddingFunction
|
||||
|
||||
client = MagicMock()
|
||||
client.tokenize.return_value = [[1], [2]]
|
||||
client.embed.return_value.embeddings = [[0.1, 0.2], [0.3, 0.4]]
|
||||
VoyageAIEmbeddingFunction.client = client
|
||||
func = get_registry().get("voyageai").create(name="voyage-3")
|
||||
|
||||
assert func.compute_source_embeddings(inputs) == [[0.1, 0.2], [0.3, 0.4]]
|
||||
client.embed.assert_called_once_with(
|
||||
texts=["first", "second"], model="voyage-3", input_type="document"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("inputs", [[], np.array([], dtype=str)])
|
||||
def test_text_source_embeddings_accept_empty_batch(inputs):
|
||||
from lancedb.embeddings.voyageai import VoyageAIEmbeddingFunction
|
||||
|
||||
client = MagicMock()
|
||||
VoyageAIEmbeddingFunction.client = client
|
||||
func = get_registry().get("voyageai").create(name="voyage-3")
|
||||
|
||||
assert func.compute_source_embeddings(inputs) == []
|
||||
client.tokenize.assert_not_called()
|
||||
client.embed.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"inputs",
|
||||
[["text", 42], np.array(["text", None]), np.array([["first", "second"]])],
|
||||
)
|
||||
def test_text_source_embeddings_reject_non_text_batch_entries(inputs):
|
||||
from lancedb.embeddings.voyageai import VoyageAIEmbeddingFunction
|
||||
|
||||
client = MagicMock()
|
||||
VoyageAIEmbeddingFunction.client = client
|
||||
func = get_registry().get("voyageai").create(name="voyage-3")
|
||||
|
||||
with pytest.raises(ValueError, match="Each input should be str"):
|
||||
func.compute_source_embeddings(inputs)
|
||||
client.tokenize.assert_not_called()
|
||||
client.embed.assert_not_called()
|
||||
|
||||
Reference in New Issue
Block a user