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:
Captain
2026-10-23 03:23:36 +08:00
committed by GitHub
parent 2af7d88ec7
commit 0cb5bd0a5d
3 changed files with 62 additions and 1 deletions
+3 -1
View File
@@ -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()