Skip to content

Commit

Permalink
[BUG] fix reraise from validation_context (chroma-core#3233)
Browse files Browse the repository at this point in the history
When the embedding function raises an error with a kwarg-only constructor
arguments, then the validation_context wrapper code fails and raises a
TypeError.
  • Loading branch information
JuliusDegesys committed Dec 3, 2024
1 parent c199996 commit 29e4064
Show file tree
Hide file tree
Showing 2 changed files with 41 additions and 2 deletions.
7 changes: 5 additions & 2 deletions chromadb/api/models/CollectionCommon.py
Original file line number Diff line number Diff line change
Expand Up @@ -89,8 +89,11 @@ def wrapper(self: Any, *args: Any, **kwargs: Any) -> T:
try:
return func(self, *args, **kwargs)
except Exception as e:
msg = f"{str(e)} in {name}."
raise type(e)(msg).with_traceback(e.__traceback__)
if e.args:
e.args = (f"{e.args[0]} in {name}.",) + e.args[1:]
else:
e.args = (f"{type(e)} in {name}.",)
raise

return wrapper

Expand Down
36 changes: 36 additions & 0 deletions chromadb/test/test_api.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
# type: ignore
import traceback
import httpx
import json

import chromadb
from chromadb.errors import ChromaError
Expand Down Expand Up @@ -735,6 +736,41 @@ def test_where_validation_query(client):
collection.query(query_embeddings=[0, 0, 0], where={"value": {"nested": "5"}})


def test_validation_context(client):
"""Test that the validation_context decorator properly re-raises exceptions
with keyword-only arguments"""
client.reset()

mock_request = httpx.Request("POST", "https://embedding.com")
mock_response = httpx.Response(
status_code=500,
content=json.dumps({"error": "test error"}).encode(),
request=mock_request,
)

# An error with keyword-only arguments (namely, request and response)
ef_error = httpx.HTTPStatusError(
message="Some HTTP error", request=mock_request, response=mock_response
)

class MockEmbeddingFunction:
def __call__(self, input):
raise ef_error

ef = MockEmbeddingFunction()
collection = client.create_collection("test", embedding_function=ef)

with pytest.raises(
httpx.HTTPStatusError, match="Some HTTP error in add."
) as exc_info:
# This should trigger the validation_context wrapper
collection.add(ids=["test1"], documents=["test document"])

# Verify the original keyword arguments are preserved
assert exc_info.value.response == mock_response
assert exc_info.value.request == mock_request


operator_records = {
"embeddings": [[1.1, 2.3, 3.2], [1.2, 2.24, 3.2]],
"ids": ["id1", "id2"],
Expand Down

0 comments on commit 29e4064

Please sign in to comment.