wip on different input vectors

Author: fcostaoliveira

engine/clients/pgvector/search.py

@@ -37,6 +37,26 @@ def search_one(cls, vector, meta_conditions, top) -> List[Tuple[int, float]]:
         if "hnsw_ef" in cls.search_params:
             cls.cur.execute(f"SET hnsw.ef_search = {cls.search_params['hnsw_ef']}")
 
+        # Ensure vector is in the correct format for pgvector
+        try:
+            if isinstance(vector, bytes):
+                # If vector is bytes, it might be serialized - try to convert
+                # First try to interpret as float32 bytes
+                try:
+                    import struct
+                    num_floats = len(vector) // 4  # 4 bytes per float32
+                    vector_array = np.array(struct.unpack(f'{num_floats}f', vector), dtype=np.float32)
+                except struct.error:
+                    # If that fails, try to decode as numpy array
+                    vector_array = np.frombuffer(vector, dtype=np.float32)
+            elif isinstance(vector, np.ndarray):
+                vector_array = vector.astype(np.float32)
+            else:
+                # Convert list to numpy array
+                vector_array = np.array(vector, dtype=np.float32)
+        except Exception as e:
+            raise ValueError(f"Failed to convert vector to proper format. Vector type: {type(vector)}, Error: {e}")
+
         if cls.distance == Distance.COSINE:
             query = f"SELECT id, embedding <=> %s AS _score FROM items ORDER BY _score LIMIT {top};"
         elif cls.distance == Distance.L2:
@@ -46,7 +66,7 @@ def search_one(cls, vector, meta_conditions, top) -> List[Tuple[int, float]]:
 
         cls.cur.execute(
             query,
-            (np.array(vector),),
+            (vector_array,),
         )
         return cls.cur.fetchall()
 

engine/clients/pgvector/upload.py

@@ -50,11 +50,14 @@ def init_client(cls, host, distance, connection_params, upload_params):
     def upload_batch(
         cls, ids: List[int], vectors: List[list], metadata: Optional[List[dict]]
     ):
-        vectors = np.array(vectors)
-
         # Copy is faster than insert
         with cls.cur.copy("COPY items (id, embedding) FROM STDIN") as copy:
-            for i, embedding in zip(ids, vectors):
+            for i, vector in zip(ids, vectors):
+                # Convert vector to numpy array with float32 dtype for pgvector
+                if isinstance(vector, np.ndarray):
+                    embedding = vector.astype(np.float32)
+                else:
+                    embedding = np.array(vector, dtype=np.float32)
                 copy.write_row((i, embedding))
 
     @classmethod