Limit transparent vector encoding to write upserts.

src/turbopuffer/_utils/_transform.py

@@ -32,6 +32,7 @@
     strip_annotated_type,
 )
 from ..lib.vector import b64encode_vector
+from ..types.namespace_write_params import NamespaceWriteParams
 
 _T = TypeVar("_T")
 
@@ -115,7 +116,7 @@ def _vector_keys_from_write_body(data: object) -> frozenset[str]:
 
 
 # turbopuffer: Simple transform without expensive type introspection.
-def _turbopuffer_transform(obj: object, *, vector_keys: frozenset[str]) -> object:
+def _turbopuffer_transform(obj: object) -> object:
     if obj is None or isinstance(obj, (int, float, bool, str)):
         return obj
     if isinstance(obj, dict):
@@ -124,22 +125,41 @@ def _turbopuffer_transform(obj: object, *, vector_keys: frozenset[str]) -> objec
             # Strip Omit and NotGiven values
             if isinstance(v, (Omit, NotGiven)):
                 continue
-            if k in vector_keys:
-                result[k] = _encode_vector(v)
-            else:
-                result[k] = _turbopuffer_transform(v, vector_keys=vector_keys)
+            result[k] = _turbopuffer_transform(v)
         return result
     if isinstance(obj, list):
-        return [_turbopuffer_transform(i, vector_keys=vector_keys) for i in cast(List[object], obj)]
+        return [_turbopuffer_transform(i) for i in cast(List[object], obj)]
     if isinstance(obj, tuple):
-        return tuple(_turbopuffer_transform(i, vector_keys=vector_keys) for i in cast(tuple[object, ...], obj))
+        return tuple(_turbopuffer_transform(i) for i in cast(tuple[object, ...], obj))
     return obj
 
 
+def _encode_write_vectors(data: object) -> object:
+    """Base64-encode vector attrs on write bodies (top-level document fields only)."""
+    if not isinstance(data, dict):
+        return data
+    body = cast(dict[str, object], data)
+    vector_keys = _vector_keys_from_write_body(body)
+    result = dict(body)
+
+    def encode_attrs(attrs: dict[str, object]) -> dict[str, object]:
+        return {k: _encode_vector(v) if k in vector_keys else v for k, v in attrs.items()}
+
+    rows = result.get("upsert_rows")
+    if is_iterable(rows) and not isinstance(rows, (str, bytes, dict)):
+        result["upsert_rows"] = [
+            encode_attrs(cast(dict[str, object], row)) if isinstance(row, dict) else row for row in cast(Any, rows)
+        ]
+    columns = result.get("upsert_columns")
+    if isinstance(columns, dict):
+        result["upsert_columns"] = encode_attrs(cast(dict[str, object], columns))
+    return result
+
+
 # Wrapper over _transform_recursive providing fake types
 def transform(
     data: _T,
-    expected_type: object,  # noqa: ARG001 - kept for API compatibility
+    expected_type: object,
 ) -> _T:
     """Transform dictionaries based off of type information from the given type, for example:
 
@@ -157,7 +177,10 @@ class Params(TypedDict, total=False):
     It should be noted that the transformations that this function does are not represented in the type system.
     """
     # turbopuffer: Use simple vector encoding instead of generic type-based transform.
-    return cast(_T, _turbopuffer_transform(data, vector_keys=_vector_keys_from_write_body(data)))
+    data = cast(_T, _turbopuffer_transform(data))
+    if expected_type is NamespaceWriteParams:
+        return cast(_T, _encode_write_vectors(data))
+    return data
 
 
 @lru_cache(maxsize=8096)
@@ -358,7 +381,7 @@ async def async_maybe_transform(
 
 async def async_transform(
     data: _T,
-    expected_type: object,  # noqa: ARG001 - kept for API compatibility
+    expected_type: object,
 ) -> _T:
     """Transform dictionaries based off of type information from the given type, for example:
 
@@ -376,7 +399,10 @@ class Params(TypedDict, total=False):
     It should be noted that the transformations that this function does are not represented in the type system.
     """
     # turbopuffer: Use simple vector encoding instead of generic type-based transform.
-    return cast(_T, _turbopuffer_transform(data, vector_keys=_vector_keys_from_write_body(data)))
+    data = cast(_T, _turbopuffer_transform(data))
+    if expected_type is NamespaceWriteParams:
+        return cast(_T, _encode_write_vectors(data))
+    return data
 
 
 async def _async_transform_recursive(

tests/custom/test_vectors.py

@@ -10,10 +10,9 @@
 from turbopuffer.types import (
     Row,
     Vector,
-    RowParam,
-    ColumnsParam,
     QueryBilling,
     VectorEncoding,
+    NamespaceWriteParams,
     NamespaceQueryResponse,
     namespace_query_params,
 )
@@ -566,11 +565,13 @@ def test_transparent_vector_encoding():
     # Due to the nature of the hack, there's a high risk of a future refactoring
     # to the Stainless specification breaking the fast path.
 
-    transformed = transform({"id": 1, "vector": [0.1, 0.2, 0.3]}, RowParam)
-    assert transformed == {"id": 1, "vector": "zczMPc3MTD6amZk+"}
+    transformed = transform({"upsert_rows": [{"id": 1, "vector": [0.1, 0.2, 0.3]}]}, NamespaceWriteParams)
+    assert transformed == {"upsert_rows": [{"id": 1, "vector": "zczMPc3MTD6amZk+"}]}
 
-    transformed = transform({"id": [1, 2], "vector": [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]}, ColumnsParam)
-    assert transformed == {"id": [1, 2], "vector": ["zczMPc3MTD6amZk+", "zczMPgAAAD+amRk/"]}
+    transformed = transform(
+        {"upsert_columns": {"id": [1, 2], "vector": [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]}}, NamespaceWriteParams
+    )
+    assert transformed == {"upsert_columns": {"id": [1, 2], "vector": ["zczMPc3MTD6amZk+", "zczMPgAAAD+amRk/"]}}
 
     # Named vector columns are encoded when schema sets ann; ann: False does not encode.
     body = {
@@ -589,7 +590,7 @@ def test_transparent_vector_encoding():
             "other_embedding": {"type": "[3]f32", "ann": False},
         },
     }
-    transformed = transform(body, object)
+    transformed = transform(body, NamespaceWriteParams)
     assert transformed == {
         "upsert_rows": [
             {
@@ -612,7 +613,7 @@ def test_transparent_vector_encoding():
         "upsert_rows": [{"id": 1, "title_embedding": [0.1, 0.2, 0.3]}],
         "schema": {"title_embedding": {"type": "[3]f32", "ann": dict[str, object]()}},
     }
-    transformed = transform(body_empty_ann, object)
+    transformed = transform(body_empty_ann, NamespaceWriteParams)
     assert transformed == {
         "upsert_rows": [{"id": 1, "title_embedding": "zczMPc3MTD6amZk+"}],
         "schema": {"title_embedding": {"type": "[3]f32", "ann": {}}},
@@ -625,11 +626,13 @@ async def test_transparent_vector_encoding_async():
     # Due to the nature of the hack, there's a high risk of a future refactoring
     # to the Stainless specification breaking the fast path.
 
-    transformed = await async_transform({"id": 1, "vector": [0.1, 0.2, 0.3]}, RowParam)
-    assert transformed == {"id": 1, "vector": "zczMPc3MTD6amZk+"}
+    transformed = await async_transform({"upsert_rows": [{"id": 1, "vector": [0.1, 0.2, 0.3]}]}, NamespaceWriteParams)
+    assert transformed == {"upsert_rows": [{"id": 1, "vector": "zczMPc3MTD6amZk+"}]}
 
-    transformed = await async_transform({"id": [1, 2], "vector": [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]}, ColumnsParam)
-    assert transformed == {"id": [1, 2], "vector": ["zczMPc3MTD6amZk+", "zczMPgAAAD+amRk/"]}
+    transformed = await async_transform(
+        {"upsert_columns": {"id": [1, 2], "vector": [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]}}, NamespaceWriteParams
+    )
+    assert transformed == {"upsert_columns": {"id": [1, 2], "vector": ["zczMPc3MTD6amZk+", "zczMPgAAAD+amRk/"]}}
 
     body = {
         "upsert_rows": [
@@ -647,7 +650,7 @@ async def test_transparent_vector_encoding_async():
             "other_embedding": {"type": "[3]f32", "ann": False},
         },
     }
-    transformed = await async_transform(body, object)
+    transformed = await async_transform(body, NamespaceWriteParams)
     assert transformed == {
         "upsert_rows": [
             {