Add tests around include_attributes=True

tests/test_vectors.py

@@ -201,6 +201,24 @@ def check_result(row, expected):
     for i in range(len(vector_set)):  # Use VectorResult in index mode
         check_result(vector_set[i], expected[i])
 
+    # Test query with all attributes
+    vector_set = ns.query(
+        top_k=5,
+        vector=[0.8, 0.7],
+        distance_metric='euclidean_squared',
+        include_vectors=True,
+        include_attributes=True
+    )
+    expected = [
+        tpuf.VectorRow(id=7, vector=[0.7, 0.7], dist=0.01, attributes={'hello': 'world'}),
+        tpuf.VectorRow(id=10, vector=[1.0, 1.0], dist=0.13, attributes={'test': 'cols'}),
+        tpuf.VectorRow(id=11, vector=[1.1, 1.1], dist=0.25, attributes={'test': 'cols'}),
+        tpuf.VectorRow(id=3, vector=[0.3, 0.3], dist=0.41, attributes={'test': 'cols', 'key1': 'three', 'key2': 'c'}),
+        tpuf.VectorRow(id=6, vector=[0.3, 0.3], dist=0.41, attributes={'test': 'cols', 'key1': 'three', 'key2': 'c'}),
+    ]
+    for i in range(len(vector_set)):  # Use VectorResult in index mode
+        check_result(vector_set[i], expected[i])
+
     # Test query with typed query
     vector_set = ns.query(tpuf.VectorQuery(
         top_k=5,

turbopuffer/backend.py

@@ -52,6 +52,8 @@ def make_api_request(self,
                 # before = time.monotonic()
                 json_payload = tpuf.dump_json_bytes(payload)
                 # print('Json time:', time.monotonic() - before)
+            elif isinstance(payload, bytes):
+                json_payload = payload
             else:
                 raise ValueError(f'Unsupported POST payload type: {type(payload)}')
 

turbopuffer/error.py

@@ -10,4 +10,4 @@ class APIError(TurbopufferError):
     def __init__(self, status_code: int, status_name: str, message: str):
         self.status_code = status_code
         self.status_name = status_name
-        super().__init__(f'{status_name}: {message}')
+        super().__init__(f'{status_name} (HTTP {status_code}): {message}')

turbopuffer/namespace.py

@@ -1,3 +1,4 @@
+import json
 import sys
 from turbopuffer.vectors import Cursor, VectorResult, VectorColumns, VectorRow, batch_iter
 from turbopuffer.backend import Backend

turbopuffer/query.py

@@ -21,7 +21,7 @@ class VectorQuery:
     distance_metric: Optional[str] = None
     top_k: int = 10
     include_vectors: bool = False
-    include_attributes: Optional[List[str]] = None
+    include_attributes: Optional[Union[List[str], bool]] = None
     filters: Optional[Dict[str, List[FilterTuple]]] = None
 
     def from_dict(source: dict) -> 'VectorQuery':
@@ -41,8 +41,9 @@ def __post_init__(self):
                     raise ValueError(f'VectorQuery.vector must a 1d-array, got {self.vector.ndim} dimensions')
             elif not isinstance(self.vector, list):
                 raise ValueError('VectorQuery.vector must be a list, got:', type(self.vector))
-        if self.include_attributes is not None and not isinstance(self.include_attributes, list) and not isinstance(self.include_attributes, bool):
-            raise ValueError('VectorQuery.include_attributes must be a list or bool, got:', type(self.include_attributes))
+        if self.include_attributes is not None:
+            if not isinstance(self.include_attributes, list) and not isinstance(self.include_attributes, bool):
+                raise ValueError('VectorQuery.include_attributes must be a list or bool, got:', type(self.include_attributes))
         if self.filters is not None:
             if not isinstance(self.filters, dict):
                 raise ValueError('VectorQuery.filters must be a dict, got:', type(self.filters))