Merge branch 'main' into feature/collection-dimensionality-exposed

.gitattributes

@@ -0,0 +1 @@
+*_pb2.py* linguist-generated

.github/workflows/chroma-cluster-test.yml

@@ -16,12 +16,8 @@ jobs:
       matrix:
         python: ['3.7']
         platform: [ubuntu-latest]
-        testfile: ["--ignore-glob 'chromadb/test/property/*' --ignore='chromadb/test/test_cli.py'",
-                   "chromadb/test/property/test_add.py",
-                   "chromadb/test/property/test_collections.py",
-                   "chromadb/test/property/test_embeddings.py",
-                   "chromadb/test/property/test_filtering.py",
-                   "chromadb/test/property/test_persist.py"]
+        testfile: ["chromadb/test/ingest/test_producer_consumer.py",
+                   "chromadb/test/segment/distributed/test_memberlist_provider.py",]
     runs-on: ${{ matrix.platform }}
     steps:
     - name: Checkout
@@ -32,6 +28,14 @@ jobs:
         python-version: ${{ matrix.python }}
     - name: Install test dependencies
       run: python -m pip install -r requirements.txt && python -m pip install -r requirements_dev.txt
+    - name: Start minikube
+      id: minikube
+      uses: medyagh/setup-minikube@latest
+      with:
+        minikube-version: latest
+        kubernetes-version: latest
+        driver: docker
+        addons: ingress, ingress-dns
+        start-args: '--profile chroma-test'
     - name: Integration Test
       run: bin/cluster-test.sh ${{ matrix.testfile }}
-      continue-on-error: true # Mark the job as successful even if the tests fail for now (Xfail)

.gitignore

@@ -14,7 +14,10 @@ index_data
 # Default configuration for persist_directory in chromadb/config.py
 # Currently it's located in "./chroma/"
 chroma/
+chroma_test_data/
+server.htpasswd
 
+.venv
 venv
 .env
 .chroma

.pre-commit-config.yaml

@@ -1,4 +1,4 @@
-exclude: 'chromadb/proto/chroma_pb2\.(py|pyi|py_grpc\.py)' # Generated files
+exclude: 'chromadb/proto/(chroma_pb2|coordinator_pb2)\.(py|pyi|py_grpc\.py)' # Generated files
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v4.4.0
@@ -32,5 +32,5 @@ repos:
     rev: "v1.2.0"
     hooks:
       - id: mypy
-        args: [--strict, --ignore-missing-imports, --follow-imports=silent, --disable-error-code=type-abstract]
-        additional_dependencies: ["types-requests", "pydantic", "overrides", "hypothesis", "pytest", "pypika", "numpy", "types-protobuf"]
+        args: [--strict, --ignore-missing-imports, --follow-imports=silent, --disable-error-code=type-abstract, --config-file=./mypy.ini]
+        additional_dependencies: ["types-requests", "pydantic", "overrides", "hypothesis", "pytest", "pypika", "numpy", "types-protobuf", "kubernetes"]

.vscode/settings.json

@@ -35,5 +35,10 @@
     "--no-pretty",
     "--strict",
     "--disable-error-code=type-abstract"
-  ]
+  ],
+  "protoc": {
+    "options": [
+      "--proto_path=idl/",
+    ]
+  }
 }

README.md

@@ -88,7 +88,7 @@ What are embeddings?
 - __Technical__: An embedding is the latent-space position of a document at a layer of a deep neural network. For models trained specifically to embed data, this is the last layer.
 - __A small example__: If you search your photos for "famous bridge in San Francisco". By embedding this query and comparing it to the embeddings of your photos and their metadata - it should return photos of the Golden Gate Bridge.
 
-Embeddings databases (also known as **vector databases**) store embeddings and allow you to search by nearest neighbors rather than by substrings like a traditional database. By default, Chroma uses [Sentence Transformers](https://docs.trychroma.com/embeddings#default-sentence-transformers) to embed for you but you can also use OpenAI embeddings, Cohere (multilingual) embeddings, or your own.
+Embeddings databases (also known as **vector databases**) store embeddings and allow you to search by nearest neighbors rather than by substrings like a traditional database. By default, Chroma uses [Sentence Transformers](https://docs.trychroma.com/embeddings#sentence-transformers) to embed for you but you can also use OpenAI embeddings, Cohere (multilingual) embeddings, or your own.
 
 ## Get involved
 

bin/cluster-test.sh

@@ -3,14 +3,50 @@
 set -e
 
 function cleanup {
-  docker compose -f docker-compose.cluster.test.yml down --rmi local --volumes
+  # Restore the previous kube context
+  kubectl config use-context $PREV_CHROMA_KUBE_CONTEXT
+  # Kill the tunnel process
+  kill $TUNNEL_PID
+  minikube delete -p chroma-test
 }
 
 trap cleanup EXIT
 
-docker compose -f docker-compose.cluster.test.yml up -d --wait
+# Save the current kube context into a variable
+export PREV_CHROMA_KUBE_CONTEXT=$(kubectl config current-context)
+
+# Create a new minikube cluster for the test
+minikube start -p chroma-test
+
+# Add the ingress addon to the cluster
+minikube addons enable ingress -p chroma-test
+minikube addons enable ingress-dns -p chroma-test
+
+# Setup docker to build inside the minikube cluster and build the image
+eval $(minikube -p chroma-test docker-env)
+docker build -t server:latest -f Dockerfile .
+
+# Apply the kubernetes manifests
+kubectl apply -f k8s/deployment
+kubectl apply -f k8s/crd
+kubectl apply -f k8s/cr
+kubectl apply -f k8s/test
+
+# Wait for the pods in the chroma namespace to be ready
+kubectl wait --namespace chroma --for=condition=Ready pods --all --timeout=300s
+
+# Run mini kube tunnel in the background to expose the service
+minikube tunnel -p chroma-test &
+TUNNEL_PID=$!
+
+# Wait for the tunnel to be ready. There isn't an easy way to check if the tunnel is ready. So we just wait for 10 seconds
+sleep 10
 
 export CHROMA_CLUSTER_TEST_ONLY=1
+export CHROMA_SERVER_HOST=$(kubectl get svc server -n chroma -o=jsonpath='{.status.loadBalancer.ingress[0].ip}')
+export PULSAR_BROKER_URL=$(kubectl get svc pulsar -n chroma -o=jsonpath='{.status.loadBalancer.ingress[0].ip}')
+echo "Chroma Server is running at port $CHROMA_SERVER_HOST"
+echo "Pulsar Broker is running at port $PULSAR_BROKER_URL"
 
 echo testing: python -m pytest "$@"
 python -m pytest "$@"

chromadb/__init__.py

@@ -35,15 +35,15 @@
     "QueryResult",
     "GetResult",
 ]
-from chromadb.telemetry.events import ClientStartEvent
-from chromadb.telemetry import Telemetry
+from chromadb.telemetry.product.events import ClientStartEvent
+from chromadb.telemetry.product import ProductTelemetryClient
 
 
 logger = logging.getLogger(__name__)
 
 __settings = Settings()
 
-__version__ = "0.4.13"
+__version__ = "0.4.14"
 
 # Workaround to deal with Colab's old sqlite3 version
 try:
@@ -56,12 +56,14 @@
 is_client = False
 try:
     from chromadb.is_thin_client import is_thin_client  # type: ignore
+
     is_client = is_thin_client
 except ImportError:
     is_client = False
 
 if not is_client:
     import sqlite3
+
     if sqlite3.sqlite_version_info < (3, 35, 0):
         if IN_COLAB:
             # In Colab, hotswap to pysqlite-binary if it's too old
@@ -75,8 +77,11 @@
             sys.modules["sqlite3"] = sys.modules.pop("pysqlite3")
         else:
             raise RuntimeError(
-                "\033[91mYour system has an unsupported version of sqlite3. Chroma requires sqlite3 >= 3.35.0.\033[0m\n"
-                "\033[94mPlease visit https://docs.trychroma.com/troubleshooting#sqlite to learn how to upgrade.\033[0m"
+                "\033[91mYour system has an unsupported version of sqlite3. Chroma \
+                    requires sqlite3 >= 3.35.0.\033[0m\n"
+                "\033[94mPlease visit \
+                    https://docs.trychroma.com/troubleshooting#sqlite to learn how \
+                    to upgrade.\033[0m"
             )
 
 
@@ -147,12 +152,11 @@ def Client(settings: Settings = __settings) -> API:
 
     system = System(settings)
 
-    telemetry_client = system.instance(Telemetry)
+    product_telemetry_client = system.instance(ProductTelemetryClient)
     api = system.instance(API)
 
     system.start()
 
-    # Submit event for client start
-    telemetry_client.capture(ClientStartEvent())
+    product_telemetry_client.capture(ClientStartEvent())
 
     return api

chromadb/api/fastapi.py

@@ -31,7 +31,12 @@
 from chromadb.auth.providers import RequestsClientAuthProtocolAdapter
 from chromadb.auth.registry import resolve_provider
 from chromadb.config import Settings, System
-from chromadb.telemetry import Telemetry
+from chromadb.telemetry.opentelemetry import (
+    OpenTelemetryClient,
+    OpenTelemetryGranularity,
+    trace_method,
+)
+from chromadb.telemetry.product import ProductTelemetryClient
 from urllib.parse import urlparse, urlunparse, quote
 
 logger = logging.getLogger(__name__)
@@ -51,7 +56,8 @@ def _validate_host(host: str) -> None:
         if "/" in host and (not host.startswith("http")):
             raise ValueError(
                 "Invalid URL. "
-                "Seems that you are trying to pass URL as a host but without specifying the protocol. "
+                "Seems that you are trying to pass URL as a host but without \
+                    specifying the protocol. "
                 "Please add http:// or https:// to the host."
             )
 
@@ -94,7 +100,8 @@ def __init__(self, system: System):
         system.settings.require("chroma_server_host")
         system.settings.require("chroma_server_http_port")
 
-        self._telemetry_client = self.require(Telemetry)
+        self._opentelemetry_client = self.require(OpenTelemetryClient)
+        self._product_telemetry_client = self.require(ProductTelemetryClient)
         self._settings = system.settings
 
         self._api_url = FastAPI.resolve_url(
@@ -130,13 +137,15 @@ def __init__(self, system: System):
         if self._header is not None:
             self._session.headers.update(self._header)
 
+    @trace_method("FastAPI.heartbeat", OpenTelemetryGranularity.OPERATION)
     @override
     def heartbeat(self) -> int:
         """Returns the current server time in nanoseconds to check if the server is alive"""
         resp = self._session.get(self._api_url)
         raise_chroma_error(resp)
         return int(resp.json()["nanosecond heartbeat"])
 
+    @trace_method("FastAPI.list_collections", OpenTelemetryGranularity.OPERATION)
     @override
     def list_collections(self) -> Sequence[Collection]:
         """Returns a list of all collections"""
@@ -149,6 +158,7 @@ def list_collections(self) -> Sequence[Collection]:
 
         return collections
 
+    @trace_method("FastAPI.create_collection", OpenTelemetryGranularity.OPERATION)
     @override
     def create_collection(
         self,
@@ -176,6 +186,7 @@ def create_collection(
             metadata=resp_json["metadata"],
         )
 
+    @trace_method("FastAPI.get_collection", OpenTelemetryGranularity.OPERATION)
     @override
     def get_collection(
         self,
@@ -195,6 +206,9 @@ def get_collection(
             metadata=resp_json["metadata"],
         )
 
+    @trace_method(
+        "FastAPI.get_or_create_collection", OpenTelemetryGranularity.OPERATION
+    )
     @override
     def get_or_create_collection(
         self,
@@ -207,6 +221,7 @@ def get_or_create_collection(
             name, metadata, embedding_function, get_or_create=True
         )
 
+    @trace_method("FastAPI._modify", OpenTelemetryGranularity.OPERATION)
     @override
     def _modify(
         self,
@@ -222,12 +237,14 @@ def _modify(
         )
         raise_chroma_error(resp)
 
+    @trace_method("FastAPI.delete_collection", OpenTelemetryGranularity.OPERATION)
     @override
     def delete_collection(self, name: str) -> None:
         """Deletes a collection"""
         resp = self._session.delete(self._api_url + "/collections/" + name)
         raise_chroma_error(resp)
 
+    @trace_method("FastAPI._count", OpenTelemetryGranularity.OPERATION)
     @override
     def _count(self, collection_id: UUID) -> int:
         """Returns the number of embeddings in the database"""
@@ -237,6 +254,7 @@ def _count(self, collection_id: UUID) -> int:
         raise_chroma_error(resp)
         return cast(int, resp.json())
 
+    @trace_method("FastAPI._peek", OpenTelemetryGranularity.OPERATION)
     @override
     def _dimensions(self, collection_id: UUID) -> int:
         """Returns the dimensionality of the embeddings in the collection"""
@@ -255,6 +273,7 @@ def _peek(self, collection_id: UUID, n: int = 10) -> GetResult:
             include=["embeddings", "documents", "metadatas"],
         )
 
+    @trace_method("FastAPI._get", OpenTelemetryGranularity.OPERATION)
     @override
     def _get(
         self,
@@ -297,6 +316,7 @@ def _get(
             documents=body.get("documents", None),
         )
 
+    @trace_method("FastAPI._delete", OpenTelemetryGranularity.OPERATION)
     @override
     def _delete(
         self,
@@ -316,6 +336,7 @@ def _delete(
         raise_chroma_error(resp)
         return cast(IDs, resp.json())
 
+    @trace_method("FastAPI._submit_batch", OpenTelemetryGranularity.ALL)
     def _submit_batch(
         self,
         batch: Tuple[
@@ -339,6 +360,7 @@ def _submit_batch(
         )
         return resp
 
+    @trace_method("FastAPI._add", OpenTelemetryGranularity.ALL)
     @override
     def _add(
         self,
@@ -359,6 +381,7 @@ def _add(
         raise_chroma_error(resp)
         return True
 
+    @trace_method("FastAPI._update", OpenTelemetryGranularity.ALL)
     @override
     def _update(
         self,
@@ -380,6 +403,7 @@ def _update(
         resp.raise_for_status()
         return True
 
+    @trace_method("FastAPI._upsert", OpenTelemetryGranularity.ALL)
     @override
     def _upsert(
         self,
@@ -401,6 +425,7 @@ def _upsert(
         resp.raise_for_status()
         return True
 
+    @trace_method("FastAPI._query", OpenTelemetryGranularity.ALL)
     @override
     def _query(
         self,
@@ -436,13 +461,15 @@ def _query(
             documents=body.get("documents", None),
         )
 
+    @trace_method("FastAPI.reset", OpenTelemetryGranularity.ALL)
     @override
     def reset(self) -> bool:
         """Resets the database"""
         resp = self._session.post(self._api_url + "/reset")
         raise_chroma_error(resp)
         return cast(bool, resp.json())
 
+    @trace_method("FastAPI.get_version", OpenTelemetryGranularity.OPERATION)
     @override
     def get_version(self) -> str:
         """Returns the version of the server"""
@@ -456,6 +483,7 @@ def get_settings(self) -> Settings:
         return self._settings
 
     @property
+    @trace_method("FastAPI.max_batch_size", OpenTelemetryGranularity.OPERATION)
     @override
     def max_batch_size(self) -> int:
         if self._max_batch_size == -1: