Merge remote-tracking branch 'origin/main' into gatekeeper/fix-1982-1

python/python/lancedb/embeddings/base.py

@@ -183,7 +183,11 @@ class EmbeddingFunction(BaseModel, ABC):
     def VectorField(self, **kwargs):
         """
         Creates a pydantic Field that can automatically annotate
-        the target vector column for this embedding function
+        the target vector column for this embedding function.
+
+        The field can be annotated as ``list[float]`` for compatibility with
+        static type checkers. LanceDB will infer the fixed vector dimension from
+        this embedding function.
         """
         return Field(json_schema_extra={"vector_column_for": self}, **kwargs)
 

python/python/lancedb/pydantic.py

@@ -99,6 +99,17 @@ def Vector(
     ...     pa.field("url", pa.utf8(), False),
     ...     pa.field("embeddings", pa.list_(pa.float32(), 768))
     ... ])
+
+    Notes
+    -----
+    ``Vector`` creates a type dynamically, so calls such as ``Vector(768)`` are
+    not valid static type annotations. For an embedding field, use the standard
+    ``list[float]`` annotation when running mypy; ``VectorField`` supplies the
+    fixed dimension to LanceDB::
+
+        class MyModel(LanceModel):
+            text: str = embeddings.SourceField()
+            vector: list[float] = embeddings.VectorField()
     """
 
     # TODO: make a public parameterized type.
@@ -379,6 +390,10 @@ def _unwrap_optional_annotation(annotation: Any) -> Any | None:
 
 def _pydantic_to_arrow_type(field: FieldInfo) -> pa.DataType:
     """Convert a Pydantic FieldInfo to Arrow DataType"""
+    embedding_vector_type = _embedding_vector_to_arrow_type(field)
+    if embedding_vector_type is not None:
+        return embedding_vector_type
+
     unwrapped = _unwrap_optional_annotation(field.annotation)
     if unwrapped is not None:
         return _pydantic_type_to_arrow_type(unwrapped, field)
@@ -392,8 +407,32 @@ def _pydantic_to_arrow_type(field: FieldInfo) -> pa.DataType:
     return _pydantic_type_to_arrow_type(field.annotation, field)
 
 
+def _embedding_vector_to_arrow_type(field: FieldInfo) -> pa.DataType | None:
+    """Infer a fixed-size vector type from ``VectorField`` metadata."""
+    if not _is_embedding_vector_annotation(field):
+        return None
+
+    function = get_extras(field, "vector_column_for")
+    return pa.list_(pa.float32(), function.ndims())
+
+
+def _is_embedding_vector_annotation(field: FieldInfo) -> bool:
+    if get_extras(field, "vector_column_for") is None:
+        return False
+
+    annotation = _unwrap_optional_annotation(field.annotation)
+    if annotation is None:
+        annotation = field.annotation
+
+    origin = getattr(annotation, "__origin__", None)
+    args = getattr(annotation, "__args__", ())
+    return origin is list and args == (float,)
+
+
 def is_nullable(field: FieldInfo) -> bool:
     """Check if a Pydantic FieldInfo is nullable."""
+    if _is_embedding_vector_annotation(field):
+        return True
     if _unwrap_optional_annotation(field.annotation) is not None:
         return True
     if isinstance(field.annotation, (_GenericAlias, GenericAlias)):

python/python/tests/test_pydantic.py

@@ -9,6 +9,7 @@ from typing import List, Optional, Tuple
 import pyarrow as pa
 import pydantic
 import pytest
+from lancedb.conftest import MockTextEmbeddingFunction
 from lancedb.pydantic import (
     PYDANTIC_VERSION,
     LanceModel,
@@ -426,6 +427,25 @@ def test_bare_vector_raises_clear_error():
         exec("class TestModel(LanceModel):\n    vector: Vector", namespace)
 
 
+def test_embedding_vector_list_annotation():
+    embedding = MockTextEmbeddingFunction.create()
+
+    class StaticTypingModel(LanceModel):
+        text: str = embedding.SourceField()
+        vector: list[float] = embedding.VectorField()
+
+    schema = pydantic_to_schema(StaticTypingModel)
+    assert schema == pa.schema(
+        [
+            pa.field("text", pa.utf8(), False),
+            pa.field("vector", pa.list_(pa.float32(), embedding.ndims()), True),
+        ]
+    )
+
+    model = StaticTypingModel(text="hello", vector=[0.0] * embedding.ndims())
+    assert model.vector == [0.0] * embedding.ndims()
+
+
 def test_fixed_size_list_field():
     class TestModel(pydantic.BaseModel):
         vec: Vector(16)