Set VectorStore.search embedding argument type (#2538)

Set VectorStore.search embedding argument type

deeplake/core/vectorstore/deeplake_vectorstore.py

@@ -301,7 +301,7 @@ def add(
 
     def search(
         self,
-        embedding_data=None,
+        embedding_data: Union[str, List[str]] = None,
         embedding_function: Optional[Callable] = None,
         embedding: Optional[Union[List[float], np.ndarray]] = None,
         k: int = 4,
@@ -344,7 +344,7 @@ def search(
 
         Args:
             embedding (Union[np.ndarray, List[float]], optional): Embedding representation for performing the search. Defaults to None. The ``embedding_data`` and ``embedding`` cannot both be specified.
-            embedding_data: Data against which the search will be performed by embedding it using the `embedding_function`. Defaults to None. The `embedding_data` and `embedding` cannot both be specified.
+            embedding_data (List[str]): Data against which the search will be performed by embedding it using the `embedding_function`. Defaults to None. The `embedding_data` and `embedding` cannot both be specified.
             embedding_function (Optional[Callable], optional): function for converting `embedding_data` into embedding. Only valid if `embedding_data` is specified
             k (int): Number of elements to return after running query. Defaults to 4.
             distance_metric (str): Type of distance metric to use for sorting the data. Avaliable options are: ``"L1", "L2", "COS", "MAX"``. Defaults to ``"COS"``.
@@ -424,11 +424,6 @@ def search(
                 embedding_data,
                 embedding_function=embedding_function or self.embedding_function,
             )
-            if isinstance(query_emb, np.ndarray):
-                assert (
-                    query_emb.ndim == 1 or query_emb.shape[0] == 1
-                ), "Query embedding must be 1-dimensional. Please consider using another embedding function for converting query string to embedding."
-
         return vector_search.search(
             query=query,
             logger=logger,

deeplake/core/vectorstore/test_deeplake_vectorstore.py

@@ -433,6 +433,24 @@ def filter_fn(x):
     assert isinstance(data.text[0].data()["value"], str)
     assert data.embedding[0].numpy().size > 0
 
+    data = vector_store.search(
+        embedding_function=embedding_fn3,
+        embedding_data="dummy",
+        return_view=True,
+        k=2,
+    )
+    assert len(data) == 2
+    assert isinstance(data.text[0].data()["value"], str)
+    assert data.embedding[0].numpy().size > 0
+
+    with pytest.raises(NotImplementedError):
+        data = vector_store.search(
+            embedding_function=embedding_fn3,
+            embedding_data=["dummy", "dummy2"],
+            return_view=True,
+            k=2,
+        )
+
     data = vector_store.search(
         filter={"metadata": {"abcdefh": 1}},
         embedding=None,
@@ -1738,7 +1756,7 @@ def test_query_dim(local_path):
     )
 
     vector_store.add(text=texts, embedding=embeddings)
-    with pytest.raises(AssertionError):
+    with pytest.raises(NotImplementedError):
         vector_store.search([texts[0], texts[0]], embedding_fn3, k=1)
 
     vector_store.search([texts[0]], embedding_fn4, k=1)

deeplake/core/vectorstore/vector_search/dataset/dataset.py

@@ -241,11 +241,17 @@ def fetch_embeddings(view, embedding_tensor: str = "embedding"):
 
 
 def get_embedding(embedding, embedding_data, embedding_function=None):
+    if isinstance(embedding_data, str):
+        embedding_data = [embedding_data]
+
     if (
         embedding is None
         and embedding_function is not None
         and embedding_data is not None
     ):
+        if len(embedding_data) > 1:
+            raise NotImplementedError("Searching batched queries is not supported yet.")
+
         embedding = embedding_function(embedding_data)  # type: ignore
 
     if embedding is not None and (