multivector support

python/starpoint/db.py

@@ -8,6 +8,7 @@
 import validators
 
 from starpoint import reader, writer, _utils
+from starpoint.embedding import Embedding
 
 LOGGER = logging.getLogger(__name__)
 
@@ -56,42 +57,6 @@ def delete(
             collection_name=collection_name,
         )
 
-    def column_delete(
-        self,
-        embeddings: List[List[float]],
-        document_metadatas: List[Dict[Any, Any]],
-        collection_id: Optional[str] = None,
-        collection_name: Optional[str] = None,
-    ) -> Dict[Any, Any]:
-        """Deletes documents from an existing collection by embedding and document metadata arrays.
-        The arrays are zipped together and updates the document in the order of the two arrays.
-        `column_delete()` method from [`Writer`](#writer-objects).
-
-        Args:
-            embeddings: A list of embeddings.
-                Order of the embeddings should match the document_metadatas.
-            document_metadatas: A list of metadata to be associated with embeddings.
-                Order of these metadatas should match the embeddings.
-            collection_id: The collection's id where the documents will be deleted.
-                This or the `collection_name` needs to be provided.
-            collection_name: The collection's name where the documents will be deleted.
-                This or the `collection_id` needs to be provided.
-
-        Returns:
-            dict: delete response json
-
-        Raises:
-            ValueError: If neither collection id and collection name are provided.
-            ValueError: If both collection id and collection name are provided.
-            requests.exceptions.SSLError: Failure likely due to network issues.
-        """
-        return self.writer.column_delete(
-            embeddings=embeddings,
-            document_metadatas=document_metadatas,
-            collection_id=collection_id,
-            collection_name=collection_name,
-        )
-
     def insert(
         self,
         documents: List[Dict[Any, Any]],
@@ -123,7 +88,7 @@ def insert(
 
     def column_insert(
         self,
-        embeddings: List[List[float]],
+        embeddings: List[Embedding],
         document_metadatas: List[Dict[Any, Any]],
         collection_id: Optional[str] = None,
         collection_name: Optional[str] = None,
@@ -162,7 +127,7 @@ def query(
         sql: Optional[str] = None,
         collection_id: Optional[str] = None,
         collection_name: Optional[str] = None,
-        query_embedding: Optional[List[float]] = None,
+        query_embedding: Optional[List[float] | Embedding] = None,
         params: Optional[List[Any]] = None,
         text_search_query: Optional[List[str]] = None,
         text_search_weight: Optional[float] = None,
@@ -188,11 +153,18 @@ def query(
             ValueError: If both collection id and collection name are provided.
             requests.exceptions.SSLError: Failure likely due to network issues.
         """
+
+        # check if query embedding is a float, if it is, convert to a embedding object
+        if isinstance(query_embedding, list):
+            query_embedding = Embedding(
+                vectors=query_embedding,
+                dim=len(query_embedding))
+
         return self.reader.query(
             sql=sql,
             collection_id=collection_id,
             collection_name=collection_name,
-            query_embedding=query_embedding,
+            query_embeddings=query_embedding,
             params=params,
             text_search_query=text_search_query,
             text_search_weight=text_search_weight,
@@ -259,7 +231,7 @@ def update(
 
     def column_update(
         self,
-        embeddings: List[List[float]],
+        embeddings: List[Embedding],
         document_metadatas: List[Dict[Any, Any]],
         collection_id: Optional[str] = None,
         collection_name: Optional[str] = None,

python/starpoint/embedding.py

@@ -30,6 +30,15 @@
 )
 
 
+class Embedding(object):
+    vectors: List[float]
+    dim: int
+
+    def __init__(self, vectors: List[float], dim: Optional[int] = None):
+        self.vectors = vectors
+        self.dim = len(vectors) if dim is None else dim
+
+
 class EmbeddingModel(Enum):
     MINILM = "MINI_LM"
 

python/starpoint/reader.py

@@ -11,6 +11,7 @@
     _validate_host,
 )
 
+from starpoint.embedding import Embedding
 
 LOGGER = logging.getLogger(__name__)
 
@@ -48,7 +49,7 @@ def query(
         sql: Optional[str] = None,
         collection_id: Optional[str] = None,
         collection_name: Optional[str] = None,
-        query_embedding: Optional[List[float]] = None,
+        query_embeddings: Optional[Embedding] = None,
         params: Optional[List[Any]] = None,
         text_search_query: Optional[List[str]] = None,
         text_search_weight: Optional[float] = None,
@@ -91,7 +92,7 @@ def query(
         request_data = dict(
             collection_id=collection_id,
             collection_name=collection_name,
-            query_embedding=query_embedding,
+            query_embeddings=query_embeddings,
             sql=sql,
             params=params,
             text_search_query=text_search_query,

python/starpoint/writer.py

@@ -10,6 +10,8 @@
     _validate_host,
 )
 
+from starpoint.embedding import Embedding
+
 LOGGER = logging.getLogger(__name__)
 
 # Host
@@ -100,51 +102,6 @@ def delete(
             return {}
         return response.json()
 
-    def column_delete(
-        self,
-        embeddings: List[List[float]],
-        document_metadatas: List[Dict[Any, Any]],
-        collection_id: Optional[str] = None,
-        collection_name: Optional[str] = None,
-    ) -> Dict[Any, Any]:
-        """Deletes documents from an existing collection by embedding and document metadata arrays.
-        The arrays are zipped together and updates the document in the order of the two arrays.
-
-        Args:
-            embeddings: A list of embeddings.
-                Order of the embeddings should match the document_metadatas.
-            document_metadatas: A list of metadata to be associated with embeddings.
-                Order of these metadatas should match the embeddings.
-            collection_id: The collection's id where the documents will be deleted.
-                This or the `collection_name` needs to be provided.
-            collection_name: The collection's name where the documents will be deleted.
-                This or the `collection_id` needs to be provided.
-
-        Returns:
-            dict: delete response json
-
-        Raises:
-            ValueError: If neither collection id and collection name are provided.
-            ValueError: If both collection id and collection name are provided.
-            requests.exceptions.SSLError: Failure likely due to network issues.
-        """
-        if len(embeddings) != len(document_metadatas):
-            LOGGER.warning(EMBEDDING_METADATA_LENGTH_MISMATCH_WARNING)
-
-        documents = [
-            {
-                "embedding": embedding,
-                "metadata": document_metadata,
-            }
-            for embedding, document_metadata in zip(embeddings, document_metadatas)
-        ]
-
-        return self.delete(
-            documents=documents,
-            collection_id=collection_id,
-            collection_name=collection_name,
-        )
-
     def insert(
         self,
         documents: List[Dict[Any, Any]],
@@ -214,7 +171,7 @@ def insert(
 
     def column_insert(
         self,
-        embeddings: List[List[float]],
+        embeddings: List[Embedding],
         document_metadatas: List[Dict[Any, Any]],
         collection_id: Optional[str] = None,
         collection_name: Optional[str] = None,
@@ -245,7 +202,7 @@ def column_insert(
 
         documents = [
             {
-                "embedding": embedding,
+                "embeddings": embedding,
                 "metadata": document_metadata,
             }
             for embedding, document_metadata in zip(embeddings, document_metadatas)
@@ -325,7 +282,7 @@ def update(
 
     def column_update(
         self,
-        embeddings: List[List[float]],
+        embeddings: List[Embedding],
         document_metadatas: List[Dict[Any, Any]],
         collection_id: Optional[str] = None,
         collection_name: Optional[str] = None,
@@ -356,7 +313,7 @@ def column_update(
 
         documents = [
             {
-                "embedding": embedding,
+                "embeddings": embedding,
                 "metadata": document_metadata,
             }
             for embedding, document_metadata in zip(embeddings, document_metadatas)

python/tests/test_db.py

@@ -1,6 +1,7 @@
 from tempfile import NamedTemporaryFile
 from uuid import uuid4
 from unittest.mock import MagicMock, patch
+from starpoint.embedding import Embedding
 
 import pytest
 from _pytest.monkeypatch import MonkeyPatch
@@ -30,17 +31,6 @@ def test_client_delete(mock_writer: MagicMock, mock_reader: MagicMock):
     mock_writer().delete.assert_called_once()
 
 
-@patch("starpoint.reader.Reader")
-@patch("starpoint.writer.Writer")
-def test_client_column_delete(mock_writer: MagicMock, mock_reader: MagicMock):
-    client = db.Client(api_key=uuid4())
-
-    client.column_delete(embeddings=[1.1], document_metadatas={"mock": "value"})
-
-    mock_reader.assert_called_once()  # Only called during init
-    mock_writer().column_delete.assert_called_once()
-
-
 @patch("starpoint.reader.Reader")
 @patch("starpoint.writer.Writer")
 def test_client_insert(mock_writer: MagicMock, mock_reader: MagicMock):
@@ -57,7 +47,7 @@ def test_client_insert(mock_writer: MagicMock, mock_reader: MagicMock):
 def test_client_column_insert(mock_writer: MagicMock, mock_reader: MagicMock):
     client = db.Client(api_key=uuid4())
 
-    client.column_insert(embeddings=[1.1], document_metadatas={"mock": "value"})
+    client.column_insert(embeddings=[Embedding([1.1])], document_metadatas=[{"mock": "value"}])
 
     mock_reader.assert_called_once()  # Only called during init
     mock_writer().column_insert.assert_called_once()
@@ -101,7 +91,7 @@ def test_client_update(mock_writer: MagicMock, mock_reader: MagicMock):
 def test_client_column_update(mock_writer: MagicMock, mock_reader: MagicMock):
     client = db.Client(api_key=uuid4())
 
-    client.column_update(embeddings=[1.1], document_metadatas={"mock": "value"})
+    client.column_update(embeddings=[Embedding([1.1])], document_metadatas=[{"mock": "value"}])
 
     mock_reader.assert_called_once()  # Only called during init
     mock_writer().column_update.assert_called_once()