Skip to content

[ENH] Embed query strings in search api #5599

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: jai/schema-js-impl
Choose a base branch
from
Open

[ENH] Embed query strings in search api #5599

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions chromadb/api/models/AsyncCollection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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])),
Expand All @@ -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
Comment on lines +361 to +362
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[BestPractice]

Potential error handling gap: If _embed_search_string_queries() raises an exception for any search in the list comprehension, the entire operation will fail and no searches will be processed. Consider adding error handling to gracefully handle individual search embedding failures:

embedded_searches = []
for search in searches_list:
    try:
        embedded_searches.append(self._embed_search_string_queries(search))
    except Exception as e:
        logger.warning(f"Failed to embed search: {e}")
        embedded_searches.append(search)  # Use original search as fallback
Context for Agents 
[**BestPractice**]

Potential error handling gap: If `_embed_search_string_queries()` raises an exception for any search in the list comprehension, the entire operation will fail and no searches will be processed. Consider adding error handling to gracefully handle individual search embedding failures:

```python
embedded_searches = []
for search in searches_list:
    try:
        embedded_searches.append(self._embed_search_string_queries(search))
    except Exception as e:
        logger.warning(f"Failed to embed search: {e}")
        embedded_searches.append(search)  # Use original search as fallback
```

File: chromadb/api/models/AsyncCollection.py
Line: 362

]

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,
)
Expand Down
8 changes: 6 additions & 2 deletions chromadb/api/models/Collection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Comment on lines +365 to +366
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[BestPractice]

Same error handling concern applies here. If _embed_search_string_queries() fails for any search in the list, all searches will fail to process. Consider adding individual error handling as suggested for the AsyncCollection version.

Context for Agents 
[**BestPractice**]

Same error handling concern applies here. If `_embed_search_string_queries()` fails for any search in the list, all searches will fail to process. Consider adding individual error handling as suggested for the AsyncCollection version.

File: chromadb/api/models/Collection.py
Line: 366

]

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,
)
Expand Down
230 changes: 230 additions & 0 deletions chromadb/api/models/CollectionCommon.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
from uuid import UUID

from chromadb.api.types import (
EMBEDDING_KEY,
URI,
Schema,
SparseVectorIndexConfig,
Expand Down Expand Up @@ -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)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[BestPractice]

Error handling gap: Multiple embedding function calls (self._embed, self._sparse_embed, embedding_func.embed_query) can raise exceptions but aren't wrapped in try-catch blocks. If any embedding function fails (network issues, model errors, invalid input), the error will propagate uncaught and could crash the application.

Consider wrapping embedding calls:

try:
    embedding = self._embed(input=[query_text], is_query=True)
except Exception as e:
    raise ValueError(f"Failed to embed query '{query_text}': {e}") from e
Context for Agents 
[**BestPractice**]

Error handling gap: Multiple embedding function calls (`self._embed`, `self._sparse_embed`, `embedding_func.embed_query`) can raise exceptions but aren't wrapped in try-catch blocks. If any embedding function fails (network issues, model errors, invalid input), the error will propagate uncaught and could crash the application.

Consider wrapping embedding calls:
```python
try:
    embedding = self._embed(input=[query_text], is_query=True)
except Exception as e:
    raise ValueError(f"Failed to embed query '{query_text}': {e}") from e
```

File: chromadb/api/models/CollectionCommon.py
Line: 773

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(
Comment on lines +811 to +812
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[BestPractice]

Error handling gap: Similar to the main embedding case, if self._sparse_embed() fails, the error lacks context about which query and key failed. With multiple searches containing different sparse embedding keys, this makes debugging difficult.

Consider adding context:

Suggested Change
Suggested change
# Embed the query
sparse_embedding = self._sparse_embed(
# Embed the query
try:
sparse_embedding = self._sparse_embed(
input=[query_text],
sparse_embedding_function=embedding_func,
is_query=True,
)
except Exception as e:
raise ValueError(
f"Failed to embed string query '{query_text}' for sparse key '{key}': {e}"
) from e

⚡ Committable suggestion

Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.

Context for Agents 
[**BestPractice**]

Error handling gap: Similar to the main embedding case, if `self._sparse_embed()` fails, the error lacks context about which query and key failed. With multiple searches containing different sparse embedding keys, this makes debugging difficult.

Consider adding context:

<details>
<summary>Suggested Change</summary>

```suggestion
                    # Embed the query
                    try:
                        sparse_embedding = self._sparse_embed(
                            input=[query_text],
                            sparse_embedding_function=embedding_func,
                            is_query=True,
                        )
                    except Exception as e:
                        raise ValueError(
                            f"Failed to embed string query '{query_text}' for sparse key '{key}': {e}"
                        ) from e
```

⚡ **Committable suggestion**

Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.

</details>

File: chromadb/api/models/CollectionCommon.py
Line: 812

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])
Comment on lines +842 to +843
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[BestPractice]

Error handling issue: The try-except AttributeError block only catches AttributeError but the embedding function could fail with other exceptions (network errors, validation errors, etc.). Consider catching broader exceptions:

try:
    embeddings = embedding_func.embed_query(input=[query_text])
except AttributeError:
    # Fallback if embed_query doesn't exist
    embeddings = embedding_func([query_text])
except Exception as e:
    raise ValueError(
        f"Failed to embed string query '{query_text}' using embedding function: {e}"
    ) from e
Context for Agents 
[**BestPractice**]

Error handling issue: The `try-except AttributeError` block only catches `AttributeError` but the embedding function could fail with other exceptions (network errors, validation errors, etc.). Consider catching broader exceptions:

```python
try:
    embeddings = embedding_func.embed_query(input=[query_text])
except AttributeError:
    # Fallback if embed_query doesn't exist
    embeddings = embedding_func([query_text])
except Exception as e:
    raise ValueError(
        f"Failed to embed string query '{query_text}' using embedding function: {e}"
    ) from e
```

File: chromadb/api/models/CollectionCommon.py
Line: 843

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:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[BestPractice]

Potential infinite recursion: The method _embed_rank_string_queries recursively processes rank expressions but doesn't have a maximum depth check. For deeply nested or circular rank structures, this could cause a stack overflow.

Consider adding a depth limit:

def _embed_rank_string_queries(self, rank: Any, depth: int = 0, max_depth: int = 100) -> Any:
    if depth > max_depth:
        raise ValueError(f"Maximum recursion depth ({max_depth}) exceeded in rank expression")
    # ... existing logic with depth + 1 passed to recursive calls
Context for Agents 
[**BestPractice**]

Potential infinite recursion: The method `_embed_rank_string_queries` recursively processes rank expressions but doesn't have a maximum depth check. For deeply nested or circular rank structures, this could cause a stack overflow.

Consider adding a depth limit:
```python
def _embed_rank_string_queries(self, rank: Any, depth: int = 0, max_depth: int = 100) -> Any:
    if depth > max_depth:
        raise ValueError(f"Maximum recursion depth ({max_depth}) exceeded in rank expression")
    # ... existing logic with depth + 1 passed to recursive calls
```

File: chromadb/api/models/CollectionCommon.py
Line: 868

"""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,
)
21 changes: 17 additions & 4 deletions chromadb/execution/expression/operator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -897,6 +896,10 @@ def __abs__(self) -> "Abs":
"""Absolute value: abs(rank)"""
return Abs(self)

def abs(self) -> "Abs":
"""Absolute value builder: rank.abs()"""
return Abs(self)

# Builder methods for functions
def exp(self) -> "Exp":
"""Exponential: e^rank"""
Expand Down Expand Up @@ -1009,7 +1012,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
Expand All @@ -1018,16 +1024,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]",
Expand Down
Loading
Loading