[ENH] For chroma cloud efs, extract api key from header if available to authenticate

Author: jairad26

chromadb/api/shared_system_client.py

@@ -1,4 +1,4 @@
-from typing import ClassVar, Dict
+from typing import ClassVar, Dict, Optional
 import uuid
 
 from chromadb.api import ServerAPI
@@ -94,3 +94,42 @@ def _system(self) -> System:
     def _submit_client_start_event(self) -> None:
         telemetry_client = self._system.instance(ProductTelemetryClient)
         telemetry_client.capture(ClientStartEvent())
+
+    @staticmethod
+    def get_chroma_cloud_api_key_from_clients() -> Optional[str]:
+        """
+        try to extract api key from existing client instaces by checking httpx session headers
+        if available.
+
+        Returns:
+            The first api key found, or None if no client instances have api keys set.
+        """
+        # Check FastAPI instances' session headers - this is where both paths converge
+        for system in SharedSystemClient._identifier_to_system.values():
+            try:
+                # Get the ServerAPI instance (which is FastAPI for HTTP clients)
+                server_api = system.instance(ServerAPI)
+
+                # Check if it's a FastAPI instance with a _session attribute
+                if hasattr(server_api, "_session") and hasattr(
+                    server_api._session, "headers"
+                ):
+                    if (
+                        hasattr(server_api, "_api_url")
+                        and "api.trychroma.com" not in server_api._api_url
+                    ):
+                        continue
+
+                    headers = server_api._session.headers
+                    api_key = headers.get("X-Chroma-Token") or headers.get(
+                        "x-chroma-token"
+                    )
+                    if api_key:
+                        # header value might be a string or bytes, convert to string
+                        return str(api_key)
+            except Exception:
+                # if we can't access the ServerAPI instance or it doesn't have _session,
+                # continue to the next system instance
+                continue
+
+        return None

chromadb/utils/embedding_functions/chroma_cloud_qwen_embedding_function.py

@@ -56,9 +56,19 @@ def __init__(
             )
 
         self.api_key_env_var = api_key_env_var
+        # First, try to get API key from environment variable
         self.api_key = os.getenv(api_key_env_var)
+        # If not found in env var, try to get it from existing client instances
         if not self.api_key:
-            raise ValueError(f"The {api_key_env_var} environment variable is not set.")
+            from chromadb.api.shared_system_client import SharedSystemClient
+
+            self.api_key = SharedSystemClient.get_chroma_cloud_api_key_from_clients()
+        # Raise error if still no API key found
+        if not self.api_key:
+            raise ValueError(
+                f"API key not found in environment variable {api_key_env_var} "
+                f"or in any existing client instances"
+            )
 
         self.model = model
         self.task = task

chromadb/utils/embedding_functions/chroma_cloud_splade_embedding_function.py

@@ -7,7 +7,6 @@
 from enum import Enum
 from chromadb.utils.embedding_functions.schemas import validate_config_schema
 from chromadb.utils.sparse_embedding_utils import normalize_sparse_vector
-from chromadb.base_types import SparseVector
 import os
 from typing import Union
 
@@ -36,10 +35,19 @@ def __init__(
                 "The httpx python package is not installed. Please install it with `pip install httpx`"
             )
         self.api_key_env_var = api_key_env_var
+        # First, try to get API key from environment variable
         self.api_key = os.getenv(self.api_key_env_var)
+        # If not found in env var, try to get it from existing client instances
+        if not self.api_key:
+            # Import here to avoid circular import
+            from chromadb.api.shared_system_client import SharedSystemClient
+
+            self.api_key = SharedSystemClient.get_chroma_cloud_api_key_from_clients()
+        # Raise error if still no API key found
         if not self.api_key:
             raise ValueError(
-                f"API key not found in environment variable {self.api_key_env_var}"
+                f"API key not found in environment variable {self.api_key_env_var} "
+                f"or in any existing client instances"
             )
         self.model = model
         self._api_url = "https://embed.trychroma.com/embed_sparse"