[ENH] return chroma-trace-id header, include trace ID in thrown errors

chromadb/api/base_http_client.py

@@ -85,6 +85,10 @@ def _raise_chroma_error(resp: httpx.Response) -> None:
  if body["error"] in errors.error_types:
  chroma_error = errors.error_types[body["error"]](body["message"])
 
+ trace_id = resp.headers.get("chroma-trace-id")
+ if trace_id:
+ chroma_error.trace_id = trace_id
+
  except BaseException:
  pass
 
@@ -94,4 +98,7 @@ def _raise_chroma_error(resp: httpx.Response) -> None:
  try:
  resp.raise_for_status()
  except httpx.HTTPStatusError:
+ trace_id = resp.headers.get("chroma-trace-id")
+ if trace_id:
+ raise Exception(f"{resp.text} (trace ID: {trace_id})")
  raise (Exception(resp.text))

chromadb/api/segment.py

@@ -248,7 +248,7 @@ def get_collection(
  if existing:
  return existing[0]
  else:
- raise ValueError(f"Collection {name} does not exist.")
+ raise InvalidCollectionException(f"Collection {name} does not exist.")
 
  @trace_method("SegmentAPI.list_collection", OpenTelemetryGranularity.OPERATION)
  @override

chromadb/errors.py

@@ -1,9 +1,11 @@
 from abc import abstractmethod
-from typing import Dict, Type
+from typing import Dict, Optional, Type
 from overrides import overrides, EnforceOverrides
 
 
 class ChromaError(Exception, EnforceOverrides):
+ trace_id: Optional[str] = None
+
  def code(self) -> int:
  """Return an appropriate HTTP response code for this error"""
  return 400 # Bad Request

chromadb/server/fastapi/__init__.py

@@ -57,6 +57,7 @@
 import logging
 from chromadb.telemetry.product.events import ServerStartEvent
 from chromadb.utils.fastapi import fastapi_json_response, string_to_uuid as _uuid
+from opentelemetry import trace
 from chromadb.telemetry.opentelemetry.fastapi import instrument_fastapi
 from chromadb.types import Database, Tenant
 from chromadb.telemetry.product import ServerContext, ProductTelemetryClient
@@ -81,6 +82,15 @@ def use_route_names_as_operation_ids(app: _FastAPI) -> None:
  route.operation_id = route.name
 
 
+async def add_trace_id_to_response_middleware(
+ request: Request, call_next: Callable[[Request], Any]
+) -> Response:
+ trace_id = trace.get_current_span().get_span_context().trace_id
+ response = await call_next(request)
+ response.headers["Chroma-Trace-Id"] = format(trace_id, "x")
+ return response
+
+
 async def catch_exceptions_middleware(
  request: Request, call_next: Callable[[Request], Any]
 ) -> Response:
@@ -105,7 +115,7 @@ async def check_http_version_middleware(
 D = TypeVar("D", bound=BaseModel, contravariant=True)
 
 
-def validate_model(model: Type[D], data: Any) -> D:
+def validate_model(model: Type[D], data: Any) -> D: # type: ignore
  """Used for backward compatibility with Pydantic 1.x"""
  try:
  return model.model_validate(data) # pydantic 2.x
@@ -156,6 +166,7 @@ def __init__(self, settings: Settings):
 
  self._app.middleware("http")(check_http_version_middleware)
  self._app.middleware("http")(catch_exceptions_middleware)
+ self._app.middleware("http")(add_trace_id_to_response_middleware)
  self._app.add_middleware(
  CORSMiddleware,
  allow_headers=["*"],

chromadb/test/test_api.py

@@ -3,6 +3,7 @@
 import httpx
 
 import chromadb
+from chromadb.errors import ChromaError
 from chromadb.api.fastapi import FastAPI
 from chromadb.api.types import QueryResult, EmbeddingFunction, Document
 from chromadb.config import Settings
@@ -525,6 +526,18 @@ def test_add_a_collection(client):
  collection = client.get_collection("testspace2")
 
 
+def test_error_includes_trace_id(client):
+ if client._system.settings.chroma_server_host is None:
+ pytest.skip("Trace IDs are only attached for HTTP requests")
+
+ client.reset()
+
+ with pytest.raises(ChromaError) as error:
+ client.get_collection("testspace2")
+
+ assert error.value.trace_id is not None
+
+
 def test_list_collections(client):
  client.reset()
  client.create_collection("testspace")