From 591e02f0cf3a51611eb01ad178232408e5db51d7 Mon Sep 17 00:00:00 2001
From: Julius Degesys <julius@moonlite.ai>
Date: Tue, 3 Dec 2024 14:27:59 -0500
Subject: [PATCH] [BUG] fix reraise from validation_context (#3233)

When the type e has kwarg-only arguments, then the previous code fails.
---
 chromadb/api/models/CollectionCommon.py |  7 ++--
 chromadb/test/test_api.py               | 45 +++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/chromadb/api/models/CollectionCommon.py b/chromadb/api/models/CollectionCommon.py
index d2b3cd3789ed..675edb2f6370 100644
--- a/chromadb/api/models/CollectionCommon.py
+++ b/chromadb/api/models/CollectionCommon.py
@@ -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
 
diff --git a/chromadb/test/test_api.py b/chromadb/test/test_api.py
index ab91408c992a..4a0dc604d4b7 100644
--- a/chromadb/test/test_api.py
+++ b/chromadb/test/test_api.py
@@ -1,6 +1,7 @@
 # type: ignore
 import traceback
 import httpx
+import json
 
 import chromadb
 from chromadb.errors import ChromaError
@@ -734,6 +735,50 @@ def test_where_validation_query(client):
     with pytest.raises(ValueError, match="where"):
         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()
+    collection = client.create_collection("test")
+
+    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, response and request)
+    original_error = httpx.HTTPStatusError(
+        message="Some HTTP error",
+        request=mock_request,
+        response=mock_response
+    )
+
+    # Patch the _embed method to raise our test error
+    original_embed = collection._embed
+    def mock_embed(*args, **kwargs):
+        raise original_error
+    collection._embed = mock_embed
+
+    try:
+        # This should trigger the validation_context wrapper
+        collection.add(
+            ids=["test1"],
+            documents=["test document"]
+        )
+        pytest.fail("Expected exception was not raised")
+    # Ensures that we can still catch the original error type
+    except httpx.HTTPStatusError as e:
+        # Verify the error message was properly augmented
+        assert "Some HTTP error in add" in str(e)
+        # Verify the original keyword arguments are preserved
+        assert e.response == mock_response
+        assert e.request == mock_request
+    finally:
+        # Restore the original _embed method
+        collection._embed = original_embed
+
 
 operator_records = {
     "embeddings": [[1.1, 2.3, 3.2], [1.2, 2.24, 3.2]],