[ENH] Embed query strings in search api

jairad26 · jairad26 · commit e7655bcd6ff3 · 2025-10-16T13:18:29.000-07:00
diff --git a/chromadb/api/models/AsyncCollection.py b/chromadb/api/models/AsyncCollection.py
@@ -344,7 +344,7 @@ async def search(
 
             # Single search
             result = await collection.search(search)
-            
+
             # Multiple searches at once
             searches = [
                 Search().where(K("type") == "article").rank(Knn(query=[0.1, 0.2])),
@@ -357,9 +357,14 @@ async def search(
         if searches_list is None:
             searches_list = []
 
+        # Embed any string queries in Knn objects
+        embedded_searches = [
+            self._embed_search_string_queries(search) for search in searches_list
+        ]
+
         return await self._client._search(
             collection_id=self.id,
-            searches=cast(List[Search], searches_list),
+            searches=cast(List[Search], embedded_searches),
             tenant=self.tenant,
             database=self.database,
         )
diff --git a/chromadb/api/models/Collection.py b/chromadb/api/models/Collection.py
@@ -22,7 +22,6 @@
 )
 from chromadb.api.collection_configuration import UpdateCollectionConfiguration
 from chromadb.execution.expression.plan import Search
-from typing import cast, List
 
 import logging
 
@@ -362,9 +361,14 @@ def search(
         if searches_list is None:
             searches_list = []
 
+        # Embed any string queries in Knn objects
+        embedded_searches = [
+            self._embed_search_string_queries(search) for search in searches_list
+        ]
+
         return self._client._search(
             collection_id=self.id,
-            searches=cast(List[Search], searches_list),
+            searches=cast(List[Search], embedded_searches),
             tenant=self.tenant,
             database=self.database,
         )
diff --git a/chromadb/api/models/CollectionCommon.py b/chromadb/api/models/CollectionCommon.py
@@ -17,6 +17,7 @@
 from uuid import UUID
 
 from chromadb.api.types import (
+    EMBEDDING_KEY,
     URI,
     Schema,
     SparseVectorIndexConfig,
@@ -741,3 +742,232 @@ def _sparse_embed(
         if is_query:
             return sparse_embedding_function.embed_query(input=input)
         return sparse_embedding_function(input=input)
+
+    def _embed_knn_string_queries(self, knn: Any) -> Any:
+        """Embed string queries in Knn objects using the appropriate embedding function.
+
+        Args:
+            knn: A Knn object that may have a string query
+
+        Returns:
+            A Knn object with the string query replaced by an embedding
+
+        Raises:
+            ValueError: If the query is a string but no embedding function is available
+        """
+        from chromadb.execution.expression.operator import Knn
+
+        if not isinstance(knn, Knn):
+            return knn
+
+        # If query is not a string, nothing to do
+        if not isinstance(knn.query, str):
+            return knn
+
+        query_text = knn.query
+        key = knn.key
+
+        # Handle main embedding field
+        if key == EMBEDDING_KEY:
+            # Use the collection's main embedding function
+            embedding = self._embed(input=[query_text], is_query=True)
+            if not embedding or len(embedding) != 1:
+                raise ValueError(
+                    "Embedding function returned unexpected number of embeddings"
+                )
+            # Return a new Knn with the embedded query
+            return Knn(
+                query=embedding[0],
+                key=knn.key,
+                limit=knn.limit,
+                default=knn.default,
+                return_rank=knn.return_rank,
+            )
+
+        # Handle metadata field with potential sparse embedding
+        schema = self.schema
+        if schema is None or key not in schema.key_overrides:
+            raise ValueError(
+                f"Cannot embed string query for key '{key}': "
+                f"key not found in schema. Please provide an embedded vector or "
+                f"configure an embedding function for this key in the schema."
+            )
+
+        value_type = schema.key_overrides[key]
+
+        # Check for sparse vector with embedding function
+        if value_type.sparse_vector is not None:
+            sparse_index = value_type.sparse_vector.sparse_vector_index
+            if sparse_index is not None and sparse_index.enabled:
+                config = sparse_index.config
+                if config.embedding_function is not None:
+                    embedding_func = config.embedding_function
+                    if not isinstance(embedding_func, SparseEmbeddingFunction):
+                        embedding_func = cast(
+                            SparseEmbeddingFunction[Any], embedding_func
+                        )
+                    validate_sparse_embedding_function(embedding_func)
+
+                    # Embed the query
+                    sparse_embedding = self._sparse_embed(
+                        input=[query_text],
+                        sparse_embedding_function=embedding_func,
+                        is_query=True,
+                    )
+
+                    if not sparse_embedding or len(sparse_embedding) != 1:
+                        raise ValueError(
+                            "Sparse embedding function returned unexpected number of embeddings"
+                        )
+
+                    # Return a new Knn with the sparse embedding
+                    return Knn(
+                        query=sparse_embedding[0],
+                        key=knn.key,
+                        limit=knn.limit,
+                        default=knn.default,
+                        return_rank=knn.return_rank,
+                    )
+
+        # Check for dense vector with embedding function (float_list)
+        if value_type.float_list is not None:
+            vector_index = value_type.float_list.vector_index
+            if vector_index is not None and vector_index.enabled:
+                config = vector_index.config
+                if config.embedding_function is not None:
+                    embedding_func = config.embedding_function
+                    validate_embedding_function(embedding_func)
+
+                    # Embed the query using the schema's embedding function
+                    try:
+                        embeddings = embedding_func.embed_query(input=[query_text])
+                    except AttributeError:
+                        # Fallback if embed_query doesn't exist
+                        embeddings = embedding_func([query_text])
+
+                    if not embeddings or len(embeddings) != 1:
+                        raise ValueError(
+                            "Embedding function returned unexpected number of embeddings"
+                        )
+
+                    # Return a new Knn with the dense embedding
+                    return Knn(
+                        query=embeddings[0],
+                        key=knn.key,
+                        limit=knn.limit,
+                        default=knn.default,
+                        return_rank=knn.return_rank,
+                    )
+
+        raise ValueError(
+            f"Cannot embed string query for key '{key}': "
+            f"no embedding function configured for this key in the schema. "
+            f"Please provide an embedded vector or configure an embedding function."
+        )
+
+    def _embed_rank_string_queries(self, rank: Any) -> Any:
+        """Recursively embed string queries in Rank expressions.
+
+        Args:
+            rank: A Rank expression that may contain Knn objects with string queries
+
+        Returns:
+            A Rank expression with all string queries embedded
+        """
+        # Import here to avoid circular dependency
+        from chromadb.execution.expression.operator import (
+            Knn,
+            Abs,
+            Div,
+            Exp,
+            Log,
+            Max,
+            Min,
+            Mul,
+            Sub,
+            Sum,
+            Val,
+            Rrf,
+        )
+
+        if rank is None:
+            return None
+
+        # Base case: Knn - embed if it has a string query
+        if isinstance(rank, Knn):
+            return self._embed_knn_string_queries(rank)
+
+        # Base case: Val - no embedding needed
+        if isinstance(rank, Val):
+            return rank
+
+        # Recursive cases: walk through child ranks
+        if isinstance(rank, Abs):
+            return Abs(self._embed_rank_string_queries(rank.rank))
+
+        if isinstance(rank, Div):
+            return Div(
+                self._embed_rank_string_queries(rank.left),
+                self._embed_rank_string_queries(rank.right),
+            )
+
+        if isinstance(rank, Exp):
+            return Exp(self._embed_rank_string_queries(rank.rank))
+
+        if isinstance(rank, Log):
+            return Log(self._embed_rank_string_queries(rank.rank))
+
+        if isinstance(rank, Max):
+            return Max([self._embed_rank_string_queries(r) for r in rank.ranks])
+
+        if isinstance(rank, Min):
+            return Min([self._embed_rank_string_queries(r) for r in rank.ranks])
+
+        if isinstance(rank, Mul):
+            return Mul([self._embed_rank_string_queries(r) for r in rank.ranks])
+
+        if isinstance(rank, Sub):
+            return Sub(
+                self._embed_rank_string_queries(rank.left),
+                self._embed_rank_string_queries(rank.right),
+            )
+
+        if isinstance(rank, Sum):
+            return Sum([self._embed_rank_string_queries(r) for r in rank.ranks])
+
+        if isinstance(rank, Rrf):
+            return Rrf(
+                ranks=[self._embed_rank_string_queries(r) for r in rank.ranks],
+                k=rank.k,
+                weights=rank.weights,
+                normalize=rank.normalize,
+            )
+
+        # Unknown rank type - return as is
+        return rank
+
+    def _embed_search_string_queries(self, search: Any) -> Any:
+        """Embed string queries in a Search object.
+
+        Args:
+            search: A Search object that may contain Knn objects with string queries
+
+        Returns:
+            A Search object with all string queries embedded
+        """
+        # Import here to avoid circular dependency
+        from chromadb.execution.expression.plan import Search
+
+        if not isinstance(search, Search):
+            return search
+
+        # Embed the rank expression if it exists
+        embedded_rank = self._embed_rank_string_queries(search._rank)
+
+        # Create a new Search with the embedded rank
+        return Search(
+            where=search._where,
+            rank=embedded_rank,
+            limit=search._limit,
+            select=search._select,
+        )
diff --git a/chromadb/execution/expression/operator.py b/chromadb/execution/expression/operator.py
@@ -1,5 +1,4 @@
 from dataclasses import dataclass, field
-from enum import Enum
 from typing import Optional, List, Dict, Set, Any, Union
 
 import numpy as np
@@ -1009,7 +1008,10 @@ class Knn(Rank):
     """KNN-based ranking
 
     Args:
-        query: The query vector for KNN search (dense, sparse, or numpy array)
+        query: The query for KNN search. Can be:
+               - A string (will be automatically embedded using the collection's embedding function)
+               - A dense vector (list or numpy array)
+               - A sparse vector (SparseVector dict)
         key: The embedding key to search against. Can be:
              - "#embedding" (default) - searches the main embedding field
              - A metadata field name (e.g., "my_custom_field") - searches that metadata field
@@ -1018,16 +1020,23 @@ class Knn(Rank):
         return_rank: If True, return the rank position (0, 1, 2, ...) instead of distance (default: False)
 
     Examples:
-        # Search main embeddings (equivalent forms)
+        # Search with string query (automatically embedded)
+        Knn(query="hello world")  # Will use collection's embedding function
+
+        # Search main embeddings with vectors (equivalent forms)
         Knn(query=[0.1, 0.2])  # Uses default key="#embedding"
         Knn(query=[0.1, 0.2], key=K.EMBEDDING)
         Knn(query=[0.1, 0.2], key="#embedding")
 
-        # Search sparse embeddings stored in metadata
+        # Search sparse embeddings stored in metadata with string
+        Knn(query="hello world", key="custom_embedding")  # Will use schema's embedding function
+
+        # Search sparse embeddings stored in metadata with vector
         Knn(query=my_vector, key="custom_embedding")  # Example: searches a metadata field
     """
 
     query: Union[
+        str,
         List[float],
         SparseVector,
         "NDArray[np.float32]",
diff --git a/chromadb/test/api/test_schema_e2e.py b/chromadb/test/api/test_schema_e2e.py
diff --git a/chromadb/test/test_api.py b/chromadb/test/test_api.py
