From 9b8cbe6a30798fb974f0ab3912f63de209911b01 Mon Sep 17 00:00:00 2001 From: Trayan Azarov Date: Fri, 2 Feb 2024 19:58:51 +0200 Subject: [PATCH 1/9] [BUG]: FTS delete (#1664) Without FTS clean-up, the SQLite file grows over time as collections are added and removed: No clean-up: ![image](https://github.com/chroma-core/chroma/assets/1157440/fdea6f36-8a7f-4d5d-8810-40feea9c7c91) With clean-up: ![image](https://github.com/chroma-core/chroma/assets/1157440/c1ea757f-3cb7-4eca-a090-ebfe0c45c067) ## Description of changes *Summarize the changes made by this PR.* - Improvements & Bug fixes - Enabled secure-delete option in FTS - Removing entries for deleted segments from FTS ## Test plan *How are these changes tested?* - [ ] Tests pass locally with `pytest` for python, `yarn test` for js ## Documentation Changes *Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the [docs repository](https://github.com/chroma-core/docs)?* --- chromadb/segment/impl/metadata/sqlite.py | 17 +++++++++++++++++ chromadb/test/segment/test_metadata.py | 20 ++++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/chromadb/segment/impl/metadata/sqlite.py b/chromadb/segment/impl/metadata/sqlite.py index 7f6bae488131..2e5af88b0d05 100644 --- a/chromadb/segment/impl/metadata/sqlite.py +++ b/chromadb/segment/impl/metadata/sqlite.py @@ -573,6 +573,7 @@ def _where_doc_criterion( def delete(self) -> None: t = Table("embeddings") t1 = Table("embedding_metadata") + t2 = Table("embedding_fulltext_search") q0 = ( self._db.querybuilder() .from_(t1) @@ -603,7 +604,23 @@ def delete(self) -> None: ) ) ) + q_fts = ( + self._db.querybuilder() + .from_(t2) + .delete() + .where( + t2.rowid.isin( + self._db.querybuilder() + .from_(t) + .select(t.id) + .where( + t.segment_id == ParameterValue(self._db.uuid_to_db(self._id)) + ) + ) + ) + ) with self._db.tx() as cur: + cur.execute(*get_sql(q_fts)) cur.execute(*get_sql(q0)) cur.execute(*get_sql(q)) diff --git a/chromadb/test/segment/test_metadata.py b/chromadb/test/segment/test_metadata.py index ef6400b210ef..1f03d6350f48 100644 --- a/chromadb/test/segment/test_metadata.py +++ b/chromadb/test/segment/test_metadata.py @@ -657,3 +657,23 @@ def test_delete_segment( res = cur.execute(sql, params) # assert that the segment is gone assert len(res.fetchall()) == 0 + + fts_t = Table("embedding_fulltext_search") + q_fts = ( + _db.querybuilder() + .from_(fts_t) + .select() + .where( + fts_t.rowid.isin( + _db.querybuilder() + .from_(t) + .select(t.id) + .where(t.segment_id == ParameterValue(_db.uuid_to_db(_id))) + ) + ) + ) + sql, params = get_sql(q_fts) + with _db.tx() as cur: + res = cur.execute(sql, params) + # assert that all FTS rows are gone + assert len(res.fetchall()) == 0 From 64b63bb897cec7131d52e9912688c952db2e7c1c Mon Sep 17 00:00:00 2001 From: tonis Date: Sat, 3 Feb 2024 01:43:50 +0700 Subject: [PATCH 2/9] add auth transport header to docker compose (#1653) ## Description of changes This just adds a missing env variable to docker-compose --- docker-compose.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/docker-compose.yml b/docker-compose.yml index 3e023109458a..4cd653b01c7f 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -22,6 +22,7 @@ services: - CHROMA_SERVER_AUTH_CREDENTIALS_FILE=${CHROMA_SERVER_AUTH_CREDENTIALS_FILE} - CHROMA_SERVER_AUTH_CREDENTIALS=${CHROMA_SERVER_AUTH_CREDENTIALS} - CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER=${CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER} + - CHROMA_SERVER_AUTH_TOKEN_TRANSPORT_HEADER=${CHROMA_SERVER_AUTH_TOKEN_TRANSPORT_HEADER} - PERSIST_DIRECTORY=${PERSIST_DIRECTORY:-/chroma/chroma} - CHROMA_OTEL_EXPORTER_ENDPOINT=${CHROMA_OTEL_EXPORTER_ENDPOINT} - CHROMA_OTEL_EXPORTER_HEADERS=${CHROMA_OTEL_EXPORTER_HEADERS} From efc16a20fa4870d05e1ae751ee3ef84d8e35cefd Mon Sep 17 00:00:00 2001 From: nicolasgere Date: Tue, 6 Feb 2024 12:33:11 -0800 Subject: [PATCH 3/9] Tilt setup for local dev (#1688) Use tilt for local end to end dev. It is a v1 (Without debugging instruction, optimized docker image etc) --------- Co-authored-by: nicolas --- DEVELOP.md | 9 ++++ Tiltfile | 30 ++++++++++++ k8s/dev/coordinator.yaml | 42 ++++++++++++++++ k8s/dev/pulsar.yaml | 45 ++++++++++++++++++ k8s/dev/server.yaml | 52 ++++++++++++++++++++ k8s/dev/setup.yaml | 100 +++++++++++++++++++++++++++++++++++++++ k8s/dev/worker.yaml | 40 ++++++++++++++++ 7 files changed, 318 insertions(+) create mode 100644 Tiltfile create mode 100644 k8s/dev/coordinator.yaml create mode 100644 k8s/dev/pulsar.yaml create mode 100644 k8s/dev/server.yaml create mode 100644 k8s/dev/setup.yaml create mode 100644 k8s/dev/worker.yaml diff --git a/DEVELOP.md b/DEVELOP.md index f034e07bed38..05357f29e60a 100644 --- a/DEVELOP.md +++ b/DEVELOP.md @@ -50,6 +50,15 @@ api = chromadb.HttpClient(host="localhost", port="8000") print(api.heartbeat()) ``` +## Local dev setup for distributed chroma +We use tilt for providing local dev setup. Tilt is an open source project +##### Requirement +- Docker +- Local Kubernetes cluster (Recommended: [OrbStack](https://orbstack.dev/) for mac, [Kind](https://kind.sigs.k8s.io/) for linux) +- [Tilt](https://docs.tilt.dev/) + +For starting the distributed Chroma in the workspace, use `tilt up`. It will create all the required resources and build the necessary Docker image in the current kubectl context. +Once done, it will expose Chroma on port 8000. You can also visit the Tilt dashboard UI at http://localhost:10350/. To clean and remove all the resources created by Tilt, use `tilt down`. ## Testing diff --git a/Tiltfile b/Tiltfile new file mode 100644 index 000000000000..7be3d4ca594f --- /dev/null +++ b/Tiltfile @@ -0,0 +1,30 @@ +docker_build('coordinator', + context='.', + dockerfile='./go/coordinator/Dockerfile' +) + +docker_build('server', + context='.', + dockerfile='./Dockerfile', +) + +docker_build('worker', + context='.', + dockerfile='./rust/worker/Dockerfile' +) + + +k8s_yaml(['k8s/dev/setup.yaml']) +k8s_resource( + objects=['chroma:Namespace', 'memberlist-reader:ClusterRole', 'memberlist-reader:ClusterRoleBinding', 'pod-list-role:Role', 'pod-list-role-binding:RoleBinding', 'memberlists.chroma.cluster:CustomResourceDefinition','worker-memberlist:MemberList'], + new_name='k8s_setup', + labels=["infrastructure"] +) +k8s_yaml(['k8s/dev/pulsar.yaml']) +k8s_resource('pulsar', resource_deps=['k8s_setup'], labels=["infrastructure"]) +k8s_yaml(['k8s/dev/server.yaml']) +k8s_resource('server', resource_deps=['k8s_setup'],labels=["chroma"], port_forwards=8000 ) +k8s_yaml(['k8s/dev/coordinator.yaml']) +k8s_resource('coordinator', resource_deps=['pulsar', 'server'], labels=["chroma"]) +k8s_yaml(['k8s/dev/worker.yaml']) +k8s_resource('worker', resource_deps=['coordinator'],labels=["chroma"]) diff --git a/k8s/dev/coordinator.yaml b/k8s/dev/coordinator.yaml new file mode 100644 index 000000000000..ce897d44c82b --- /dev/null +++ b/k8s/dev/coordinator.yaml @@ -0,0 +1,42 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: coordinator + namespace: chroma +spec: + replicas: 1 + selector: + matchLabels: + app: coordinator + template: + metadata: + labels: + app: coordinator + spec: + containers: + - command: + - "chroma" + - "coordinator" + - "--pulsar-admin-url=http://pulsar.chroma:8080" + - "--pulsar-url=pulsar://pulsar.chroma:6650" + - "--notifier-provider=pulsar" + image: coordinator + imagePullPolicy: IfNotPresent + name: coordinator + ports: + - containerPort: 50051 + name: grpc +--- +apiVersion: v1 +kind: Service +metadata: + name: coordinator + namespace: chroma +spec: + ports: + - name: grpc + port: 50051 + targetPort: grpc + selector: + app: coordinator + type: ClusterIP \ No newline at end of file diff --git a/k8s/dev/pulsar.yaml b/k8s/dev/pulsar.yaml new file mode 100644 index 000000000000..4038ecda2093 --- /dev/null +++ b/k8s/dev/pulsar.yaml @@ -0,0 +1,45 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: pulsar + namespace: chroma +spec: + replicas: 1 + selector: + matchLabels: + app: pulsar + template: + metadata: + labels: + app: pulsar + spec: + containers: + - name: pulsar + image: apachepulsar/pulsar + command: [ "/pulsar/bin/pulsar", "standalone" ] + ports: + - containerPort: 6650 + - containerPort: 8080 + volumeMounts: + - name: pulsardata + mountPath: /pulsar/data + volumes: + - name: pulsardata + emptyDir: {} +--- +apiVersion: v1 +kind: Service +metadata: + name: pulsar + namespace: chroma +spec: + ports: + - name: pulsar-port + port: 6650 + targetPort: 6650 + - name: admin-port + port: 8080 + targetPort: 8080 + selector: + app: pulsar + type: ClusterIP \ No newline at end of file diff --git a/k8s/dev/server.yaml b/k8s/dev/server.yaml new file mode 100644 index 000000000000..9d76314e693e --- /dev/null +++ b/k8s/dev/server.yaml @@ -0,0 +1,52 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: server + namespace: chroma +spec: + replicas: 2 + selector: + matchLabels: + app: server + template: + metadata: + labels: + app: server + spec: + containers: + - name: server + image: server + imagePullPolicy: IfNotPresent + ports: + - containerPort: 8000 + volumeMounts: + - name: chroma + mountPath: /test + env: + - name: IS_PERSISTENT + value: "TRUE" + - name: CHROMA_PRODUCER_IMPL + value: "chromadb.ingest.impl.pulsar.PulsarProducer" + - name: CHROMA_CONSUMER_IMPL + value: "chromadb.ingest.impl.pulsar.PulsarConsumer" + - name: CHROMA_SEGMENT_MANAGER_IMPL + value: "chromadb.segment.impl.manager.distributed.DistributedSegmentManager" + - name: PULSAR_BROKER_URL + value: "pulsar.chroma" + - name: PULSAR_BROKER_PORT + value: "6650" + - name: PULSAR_ADMIN_PORT + value: "8080" + - name: ALLOW_RESET + value: "TRUE" + - name: CHROMA_SYSDB_IMPL + value: "chromadb.db.impl.grpc.client.GrpcSysDB" + - name: CHROMA_SERVER_GRPC_PORT + value: "50051" + - name: CHROMA_COORDINATOR_HOST + value: "coordinator.chroma" + volumes: + - name: chroma + emptyDir: {} + + diff --git a/k8s/dev/setup.yaml b/k8s/dev/setup.yaml new file mode 100644 index 000000000000..d9e1d95cc151 --- /dev/null +++ b/k8s/dev/setup.yaml @@ -0,0 +1,100 @@ +kind: Namespace +apiVersion: v1 +metadata: + name: chroma +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: memberlist-reader +rules: +- apiGroups: + - chroma.cluster + resources: + - memberlists + verbs: + - get + - list + - watch + - create + - update + - patch + - delete +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + name: memberlist-reader +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: ClusterRole + name: memberlist-reader +subjects: +- kind: ServiceAccount + name: default + namespace: chroma +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: Role +metadata: + namespace: chroma + name: pod-list-role +rules: +- apiGroups: [""] + resources: ["pods"] + verbs: ["get", "list", "watch"] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: RoleBinding +metadata: + name: pod-list-role-binding + namespace: chroma +subjects: +- kind: ServiceAccount + name: default + namespace: chroma +roleRef: + kind: Role + name: pod-list-role + apiGroup: rbac.authorization.k8s.io +--- +apiVersion: apiextensions.k8s.io/v1 +kind: CustomResourceDefinition +metadata: + name: memberlists.chroma.cluster +spec: + group: chroma.cluster + versions: + - name: v1 + served: true + storage: true + schema: + openAPIV3Schema: + type: object + properties: + spec: + type: object + properties: + members: + type: array + items: + type: object + properties: + url: # Rename to ip + type: string + pattern: '^((25[0-5]|(2[0-4]|1\d|[1-9]|)\d)\.?\b){4}$' + scope: Namespaced + names: + plural: memberlists + singular: memberlist + kind: MemberList + shortNames: + - ml +--- +apiVersion: chroma.cluster/v1 +kind: MemberList +metadata: + name: worker-memberlist + namespace: chroma +spec: + members: \ No newline at end of file diff --git a/k8s/dev/worker.yaml b/k8s/dev/worker.yaml new file mode 100644 index 000000000000..82b4c9d905ba --- /dev/null +++ b/k8s/dev/worker.yaml @@ -0,0 +1,40 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: worker + namespace: chroma +spec: + replicas: 1 + selector: + matchLabels: + app: worker + template: + metadata: + labels: + app: worker + member-type: worker + spec: + containers: + - name: worker + image: worker + imagePullPolicy: IfNotPresent + command: ["cargo", "run"] + ports: + - containerPort: 50051 + volumeMounts: + - name: chroma + mountPath: /index_data + env: + - name: CHROMA_WORKER__PULSAR_URL + value: pulsar://pulsar.chroma:6650 + - name: CHROMA_WORKER__PULSAR_NAMESPACE + value: default + - name: CHROMA_WORKER__PULSAR_TENANT + value: default + - name: CHROMA_WORKER__MY_IP + valueFrom: + fieldRef: + fieldPath: status.podIP + volumes: + - name: chroma + emptyDir: {} \ No newline at end of file From a62cfb07f8882f9d2a585ff265abd96ea411fadc Mon Sep 17 00:00:00 2001 From: Trayan Azarov Date: Wed, 7 Feb 2024 18:47:47 +0200 Subject: [PATCH 4/9] [ENH][SEC]: CIP-01022024 SSL Verify Client Config (#1604) ## Description of changes *Summarize the changes made by this PR.* - New functionality - New CIP to introduce SSL verify flag to support custom PKIs or to accept self-signed certs for testing and experimentation purposes ## Test plan *How are these changes tested?* - [x] Tests pass locally with `pytest` for python ## Documentation Changes CIP document in the PR. --- chromadb/api/fastapi.py | 2 + chromadb/config.py | 4 +- chromadb/test/conftest.py | 72 ++++++++++- chromadb/test/openssl.cnf | 12 ++ chromadb/test/test_api.py | 43 +++++++ .../CIP-01022024_SSL_Verify_Client_Config.md | 68 ++++++++++ .../CIP-01022024-test_self_signed.ipynb | 119 ++++++++++++++++++ 7 files changed, 318 insertions(+), 2 deletions(-) create mode 100644 chromadb/test/openssl.cnf create mode 100644 docs/cip/CIP-01022024_SSL_Verify_Client_Config.md create mode 100644 docs/cip/assets/CIP-01022024-test_self_signed.ipynb diff --git a/chromadb/api/fastapi.py b/chromadb/api/fastapi.py index d3d1a8a4e7ea..1ee7a45af541 100644 --- a/chromadb/api/fastapi.py +++ b/chromadb/api/fastapi.py @@ -138,6 +138,8 @@ def __init__(self, system: System): self._session = requests.Session() if self._header is not None: self._session.headers.update(self._header) + if self._settings.chroma_server_ssl_verify is not None: + self._session.verify = self._settings.chroma_server_ssl_verify @trace_method("FastAPI.heartbeat", OpenTelemetryGranularity.OPERATION) @override diff --git a/chromadb/config.py b/chromadb/config.py index 61b789d0eee6..e9ceffc5dd02 100644 --- a/chromadb/config.py +++ b/chromadb/config.py @@ -4,7 +4,7 @@ import os from abc import ABC from graphlib import TopologicalSorter -from typing import Optional, List, Any, Dict, Set, Iterable +from typing import Optional, List, Any, Dict, Set, Iterable, Union from typing import Type, TypeVar, cast from overrides import EnforceOverrides @@ -122,6 +122,8 @@ class Settings(BaseSettings): # type: ignore chroma_server_headers: Optional[Dict[str, str]] = None chroma_server_http_port: Optional[str] = None chroma_server_ssl_enabled: Optional[bool] = False + # the below config value is only applicable to Chroma HTTP clients + chroma_server_ssl_verify: Optional[Union[bool, str]] = None chroma_server_api_default_path: Optional[str] = "/api/v1" chroma_server_grpc_port: Optional[str] = None # eg ["http://localhost:3000"] diff --git a/chromadb/test/conftest.py b/chromadb/test/conftest.py index 087cb2271bd8..34a1b040dd19 100644 --- a/chromadb/test/conftest.py +++ b/chromadb/test/conftest.py @@ -3,6 +3,7 @@ import os import shutil import socket +import subprocess import tempfile import time from typing import ( @@ -47,7 +48,6 @@ ) hypothesis.settings.load_profile(os.getenv("HYPOTHESIS_PROFILE", "dev")) - NOT_CLUSTER_ONLY = os.getenv("CHROMA_CLUSTER_TEST_ONLY") != "1" @@ -58,6 +58,35 @@ def skip_if_not_cluster() -> pytest.MarkDecorator: ) +def generate_self_signed_certificate() -> None: + config_path = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "openssl.cnf" + ) + print(f"Config path: {config_path}") # Debug print to verify path + if not os.path.exists(config_path): + raise FileNotFoundError(f"Config file not found at {config_path}") + subprocess.run( + [ + "openssl", + "req", + "-x509", + "-newkey", + "rsa:4096", + "-keyout", + "serverkey.pem", + "-out", + "servercert.pem", + "-days", + "365", + "-nodes", + "-subj", + "/CN=localhost", + "-config", + config_path, + ] + ) + + def find_free_port() -> int: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.bind(("", 0)) @@ -77,6 +106,8 @@ def _run_server( chroma_server_authz_provider: Optional[str] = None, chroma_server_authz_config_file: Optional[str] = None, chroma_server_authz_config: Optional[Dict[str, Any]] = None, + chroma_server_ssl_certfile: Optional[str] = None, + chroma_server_ssl_keyfile: Optional[str] = None, ) -> None: """Run a Chroma server locally""" if is_persistent and persist_directory: @@ -123,6 +154,8 @@ def _run_server( port=port, log_level="error", timeout_keep_alive=30, + ssl_keyfile=chroma_server_ssl_keyfile, + ssl_certfile=chroma_server_ssl_certfile, ) @@ -152,6 +185,8 @@ def _fastapi_fixture( chroma_server_authz_provider: Optional[str] = None, chroma_server_authz_config_file: Optional[str] = None, chroma_server_authz_config: Optional[Dict[str, Any]] = None, + chroma_server_ssl_certfile: Optional[str] = None, + chroma_server_ssl_keyfile: Optional[str] = None, ) -> Generator[System, None, None]: """Fixture generator that launches a server in a separate process, and yields a fastapi client connect to it""" @@ -171,6 +206,8 @@ def _fastapi_fixture( Optional[str], Optional[str], Optional[Dict[str, Any]], + Optional[str], + Optional[str], ] = ( port, False, @@ -183,6 +220,8 @@ def _fastapi_fixture( chroma_server_authz_provider, chroma_server_authz_config_file, chroma_server_authz_config, + chroma_server_ssl_certfile, + chroma_server_ssl_keyfile, ) persist_directory = None if is_persistent: @@ -199,6 +238,8 @@ def _fastapi_fixture( chroma_server_authz_provider, chroma_server_authz_config_file, chroma_server_authz_config, + chroma_server_ssl_certfile, + chroma_server_ssl_keyfile, ) proc = ctx.Process(target=_run_server, args=args, daemon=True) proc.start() @@ -210,6 +251,8 @@ def _fastapi_fixture( chroma_client_auth_provider=chroma_client_auth_provider, chroma_client_auth_credentials=chroma_client_auth_credentials, chroma_client_auth_token_transport_header=chroma_client_auth_token_transport_header, + chroma_server_ssl_verify=chroma_server_ssl_certfile, + chroma_server_ssl_enabled=True if chroma_server_ssl_certfile else False, ) system = System(settings) api = system.instance(ServerAPI) @@ -231,6 +274,15 @@ def fastapi_persistent() -> Generator[System, None, None]: return _fastapi_fixture(is_persistent=True) +def fastapi_ssl() -> Generator[System, None, None]: + generate_self_signed_certificate() + return _fastapi_fixture( + is_persistent=False, + chroma_server_ssl_certfile="./servercert.pem", + chroma_server_ssl_keyfile="./serverkey.pem", + ) + + def basic_http_client() -> Generator[System, None, None]: settings = Settings( chroma_api_impl="chromadb.api.fastapi.FastAPI", @@ -400,6 +452,11 @@ def system_fixtures_wrong_auth() -> List[Callable[[], Generator[System, None, No return fixtures +def system_fixtures_ssl() -> List[Callable[[], Generator[System, None, None]]]: + fixtures = [fastapi_ssl] + return fixtures + + @pytest.fixture(scope="module", params=system_fixtures_wrong_auth()) def system_wrong_auth( request: pytest.FixtureRequest, @@ -412,6 +469,11 @@ def system(request: pytest.FixtureRequest) -> Generator[ServerAPI, None, None]: yield next(request.param()) +@pytest.fixture(scope="module", params=system_fixtures_ssl()) +def system_ssl(request: pytest.FixtureRequest) -> Generator[ServerAPI, None, None]: + yield next(request.param()) + + @pytest.fixture(scope="module", params=system_fixtures_auth()) def system_auth(request: pytest.FixtureRequest) -> Generator[ServerAPI, None, None]: yield next(request.param()) @@ -432,6 +494,14 @@ def client(system: System) -> Generator[ClientAPI, None, None]: client.clear_system_cache() +@pytest.fixture(scope="function") +def client_ssl(system_ssl: System) -> Generator[ClientAPI, None, None]: + system_ssl.reset_state() + client = ClientCreator.from_system(system_ssl) + yield client + client.clear_system_cache() + + @pytest.fixture(scope="function") def api_wrong_cred( system_wrong_auth: System, diff --git a/chromadb/test/openssl.cnf b/chromadb/test/openssl.cnf new file mode 100644 index 000000000000..11704076bd47 --- /dev/null +++ b/chromadb/test/openssl.cnf @@ -0,0 +1,12 @@ +[req] +distinguished_name = req_distinguished_name +x509_extensions = usr_cert + +[req_distinguished_name] +CN = localhost + +[usr_cert] +subjectAltName = @alt_names + +[alt_names] +DNS.1 = localhost \ No newline at end of file diff --git a/chromadb/test/test_api.py b/chromadb/test/test_api.py index 36a82205e45c..cb88ed2bb77e 100644 --- a/chromadb/test/test_api.py +++ b/chromadb/test/test_api.py @@ -1,5 +1,7 @@ # type: ignore +import traceback import requests +from urllib3.connectionpool import InsecureRequestWarning import chromadb from chromadb.api.fastapi import FastAPI @@ -360,6 +362,7 @@ def test_modify_error_on_existing_name(api): with pytest.raises(Exception): c2.modify(name="testspace") + def test_modify_warn_on_DF_change(api, caplog): api.reset() @@ -368,6 +371,7 @@ def test_modify_warn_on_DF_change(api, caplog): with pytest.raises(Exception, match="not supported") as e: collection.modify(metadata={"hnsw:space": "cosine"}) + def test_metadata_cru(api): api.reset() metadata_a = {"a": 1, "b": 2} @@ -1437,6 +1441,7 @@ def test_invalid_embeddings(api): # test to make sure update shows exception for bad dimensionality + def test_dimensionality_exception_update(api): api.reset() collection = api.create_collection("test_dimensionality_update_exception") @@ -1446,8 +1451,10 @@ def test_dimensionality_exception_update(api): collection.update(**bad_dimensionality_records) assert "dimensionality" in str(e.value) + # test to make sure upsert shows exception for bad dimensionality + def test_dimensionality_exception_upsert(api): api.reset() collection = api.create_collection("test_dimensionality_upsert_exception") @@ -1456,3 +1463,39 @@ def test_dimensionality_exception_upsert(api): with pytest.raises(Exception) as e: collection.upsert(**bad_dimensionality_records) assert "dimensionality" in str(e.value) + + +def test_ssl_self_signed(client_ssl): + if os.environ.get("CHROMA_INTEGRATION_TEST_ONLY"): + pytest.skip("Skipping test for integration test") + client_ssl.heartbeat() + + +def test_ssl_self_signed_without_ssl_verify(client_ssl): + if os.environ.get("CHROMA_INTEGRATION_TEST_ONLY"): + pytest.skip("Skipping test for integration test") + client_ssl.heartbeat() + _port = client_ssl._server._settings.chroma_server_http_port + with pytest.raises(ValueError) as e: + chromadb.HttpClient(ssl=True, port=_port) + stack_trace = traceback.format_exception( + type(e.value), e.value, e.value.__traceback__ + ) + client_ssl.clear_system_cache() + assert "CERTIFICATE_VERIFY_FAILED" in "".join(stack_trace) + + +def test_ssl_self_signed_with_verify_false(client_ssl): + if os.environ.get("CHROMA_INTEGRATION_TEST_ONLY"): + pytest.skip("Skipping test for integration test") + client_ssl.heartbeat() + _port = client_ssl._server._settings.chroma_server_http_port + with pytest.warns(InsecureRequestWarning) as record: + client = chromadb.HttpClient( + ssl=True, + port=_port, + settings=chromadb.Settings(chroma_server_ssl_verify=False), + ) + client.heartbeat() + client_ssl.clear_system_cache() + assert "Unverified HTTPS request" in str(record[0].message) diff --git a/docs/cip/CIP-01022024_SSL_Verify_Client_Config.md b/docs/cip/CIP-01022024_SSL_Verify_Client_Config.md new file mode 100644 index 000000000000..2448af11c88e --- /dev/null +++ b/docs/cip/CIP-01022024_SSL_Verify_Client_Config.md @@ -0,0 +1,68 @@ +# CIP-01022024 SSL Verify Client Config + +## Status + +Current Status: `Under Discussion` + +## Motivation + +The motivation for this change is to enhance security and flexibility in Chroma's client API. Users need the ability to +configure SSL contexts to trust custom CA certificates or self-signed certificates, which is not straightforward with +the current setup. This capability is crucial for organizations that operate their own CA or for developers who need to +test their applications in environments where certificates from a recognized CA are not available or practical. + +The suggested change entails a server-side certificate be available, but this CIP does not prescribe how such +certificate should be configured or obtained. In our testing, we used a self-signed certificate generated with +`openssl` and configured the client to trust the certificate. We also experiment with a SSL-terminated proxy server. +Both of approaches yielded the same results. + +> **IMPORTANT:** It should be noted that we do not recommend or encourage the use of self-signed certificates in +> production environments. + +We also provide a sample notebook that to help the reader run a local Chroma server with a self-signed certificate and +configure the client to trust the certificate. The notebook can be found +in [assets/CIP-01022024-test_self_signed.ipynb](./assets/CIP-01022024-test_self_signed.ipynb). + +## Public Interfaces + +> **Note:** The following changes are only applicable to Chroma HttpClient. + +New settings variable `chroma_server_ssl_verify` accepting either a boolean or a path to a certificate file. If the +value is a path to a certificate file, the file will be used to verify the server's certificate. If the value is a +boolean, the SSL certificate verification can be bypassed (`false`) or enforced (`true`). + +The value is passed as `verify` parameter to `requests.Session` of the `FastAPI` client. See +requests [documentation](https://requests.readthedocs.io/en/latest/user/advanced/#ssl-cert-verification) for +more details. + +Example Usage: + +```python +import chromadb +from chromadb import Settings +client = chromadb.HttpClient(host="localhost",port="8443",ssl=True, settings=Settings(chroma_server_ssl_verify='./servercert.pem')) +# or with boolean +client = chromadb.HttpClient(host="localhost",port="8443",ssl=True, settings=Settings(chroma_server_ssl_verify=False)) +``` + +### Resources + +- https://requests.readthedocs.io/en/latest/api/#requests.request +- https://www.geeksforgeeks.org/ssl-certificate-verification-python-requests/ + +## Proposed Changes + +The proposed changes are mentioned in the public interfaces. + +## Compatibility, Deprecation, and Migration Plan + +The change is not backward compatible from client's perspective as the lack of the feature in prior clients will cause +an error when passing the new settings parameter. Server-side is not affected by this change. + +## Test Plan + +API tests with SSL verification enabled and a self-signed certificate. + +## Rejected Alternatives + +N/A diff --git a/docs/cip/assets/CIP-01022024-test_self_signed.ipynb b/docs/cip/assets/CIP-01022024-test_self_signed.ipynb new file mode 100644 index 000000000000..d607b51824b7 --- /dev/null +++ b/docs/cip/assets/CIP-01022024-test_self_signed.ipynb @@ -0,0 +1,119 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "source": [ + "# Generate a Certificate\n", + "\n", + "```bash\n", + "openssl req -new -newkey rsa:2048 -sha256 -days 365 -nodes -x509 \\\n", + " -keyout ./serverkey.pem \\\n", + " -out ./servercert.pem \\\n", + " -subj \"/O=Chroma/C=US\" \\\n", + " -config chromadb/test/openssl.cnf\n", + "```\n", + "\n", + "> Note: The above command should be executed at the root of the repo (openssl.cnf uses relative path)\n" + ], + "metadata": { + "collapsed": false + }, + "id": "faa8cefb6825fe83" + }, + { + "cell_type": "markdown", + "source": [ + "# Start the server\n", + "\n", + "```bash\n", + "uvicorn chromadb.app:app --workers 1 --host 0.0.0.0 --port 8443 \\\n", + " --proxy-headers --log-config chromadb/log_config.yml --ssl-keyfile ./serverkey.pem --ssl-certfile ./servercert.pem\n", + "```" + ], + "metadata": { + "collapsed": false + }, + "id": "e084285e11c3747d" + }, + { + "cell_type": "markdown", + "source": [ + "# Test with cert as SSL verify string" + ], + "metadata": { + "collapsed": false + }, + "id": "130df9c0a6d67b52" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "initial_id", + "metadata": { + "collapsed": true + }, + "outputs": [], + "source": [ + "from chromadb import Settings\n", + "import chromadb\n", + "client = chromadb.HttpClient(host=\"localhost\",port=\"8443\",ssl=True, settings=Settings(chroma_server_ssl_verify='./servercert.pem'))\n", + "print(client.heartbeat())" + ] + }, + { + "cell_type": "markdown", + "source": [ + "# Test with cert as SSL verify boolean" + ], + "metadata": { + "collapsed": false + }, + "id": "8223d0100df06ec4" + }, + { + "cell_type": "code", + "outputs": [], + "source": [ + "from chromadb import Settings\n", + "import chromadb\n", + "client = chromadb.HttpClient(host=\"localhost\",port=\"8443\",ssl=True, settings=Settings(chroma_server_ssl_verify=False))\n", + "print(client.heartbeat())" + ], + "metadata": { + "collapsed": false + }, + "id": "f7cf299721741c1", + "execution_count": null + }, + { + "cell_type": "code", + "outputs": [], + "source": [], + "metadata": { + "collapsed": false + }, + "id": "6231ac2ac38383c2" + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 2 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython2", + "version": "2.7.6" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From 691ac3fa59be1cbe59553685b2d2cb793bab2e49 Mon Sep 17 00:00:00 2001 From: Mikhail Merkulov Date: Wed, 7 Feb 2024 18:48:28 +0200 Subject: [PATCH 5/9] Dockerized chroma arguments customization (#1658) ## Description of changes *Summarize the changes made by this PR.* - New functionality - Added an ability to customize the default arguments that are passed from `docker run` or `docker compose` `command` field to `uvicorn chromadb.app:app`. I needed it to be able to customize the port because in certain scenarios it cannot be change (i.e. ECS where internal port is proxies as is). The default arguments are not changed: `--workers 1 --host 0.0.0.0 --port 8000 --proxy-headers --log-config chromadb/log_config.yml --timeout-keep-alive 30` - Added ENV variables for basic customization with default values: ``` CHROMA_HOST_ADDR="0.0.0.0" CHROMA_HOST_PORT=8000 CHROMA_WORKERS=1 CHROMA_LOG_CONFIG="chromadb/log_config.yml" CHROMA_TIMEOUT_KEEP_ALIVE=30 ``` ## Test plan *How are these changes tested?* - Tested locally using `docker build` and `docker run` commands - Tested customization in `docker-compose` - now it works as expected. ## Documentation Changes *Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the [docs repository](https://github.com/chroma-core/docs)?* TODO: Deployment docs needs to be updated to cover container arguments customization. --- Dockerfile | 11 ++++++++++- bin/docker_entrypoint.sh | 12 +++++++++++- docker-compose.test-auth.yml | 2 +- docker-compose.test.yml | 2 +- docker-compose.yml | 2 +- 5 files changed, 24 insertions(+), 5 deletions(-) diff --git a/Dockerfile b/Dockerfile index 1f90733edbb7..c871a4cd8c70 100644 --- a/Dockerfile +++ b/Dockerfile @@ -25,6 +25,15 @@ COPY --from=builder /install /usr/local COPY ./bin/docker_entrypoint.sh /docker_entrypoint.sh COPY ./ /chroma +RUN chmod +x /docker_entrypoint.sh + +ENV CHROMA_HOST_ADDR "0.0.0.0" +ENV CHROMA_HOST_PORT 8000 +ENV CHROMA_WORKERS 1 +ENV CHROMA_LOG_CONFIG "chromadb/log_config.yml" +ENV CHROMA_TIMEOUT_KEEP_ALIVE 30 + EXPOSE 8000 -CMD ["/docker_entrypoint.sh"] +ENTRYPOINT ["/docker_entrypoint.sh"] +CMD [ "--workers ${CHROMA_WORKERS} --host ${CHROMA_HOST_ADDR} --port ${CHROMA_HOST_PORT} --proxy-headers --log-config ${CHROMA_LOG_CONFIG} --timeout-keep-alive ${CHROMA_TIMEOUT_KEEP_ALIVE}"] \ No newline at end of file diff --git a/bin/docker_entrypoint.sh b/bin/docker_entrypoint.sh index e6f2df70be87..e9498b4fd7ca 100755 --- a/bin/docker_entrypoint.sh +++ b/bin/docker_entrypoint.sh @@ -1,5 +1,15 @@ #!/bin/bash +set -e export IS_PERSISTENT=1 export CHROMA_SERVER_NOFILE=65535 -exec uvicorn chromadb.app:app --workers 1 --host 0.0.0.0 --port 8000 --proxy-headers --log-config chromadb/log_config.yml --timeout-keep-alive 30 +args="$@" + +if [[ $args =~ ^uvicorn.* ]]; then + echo "Starting server with args: $(eval echo "$args")" + echo -e "\033[31mWARNING: Please remove 'uvicorn chromadb.app:app' from your command line arguments. This is now handled by the entrypoint script." + exec $(eval echo "$args") +else + echo "Starting 'uvicorn chromadb.app:app' with args: $(eval echo "$args")" + exec uvicorn chromadb.app:app $(eval echo "$args") +fi diff --git a/docker-compose.test-auth.yml b/docker-compose.test-auth.yml index 259d4c54e79a..d3297b5a04fc 100644 --- a/docker-compose.test-auth.yml +++ b/docker-compose.test-auth.yml @@ -11,7 +11,7 @@ services: dockerfile: Dockerfile volumes: - chroma-data:/chroma/chroma - command: uvicorn chromadb.app:app --workers 1 --host 0.0.0.0 --port 8000 --log-config chromadb/log_config.yml --timeout-keep-alive 30 + command: "--workers 1 --host 0.0.0.0 --port 8000 --proxy-headers --log-config chromadb/log_config.yml --timeout-keep-alive 30" environment: - ANONYMIZED_TELEMETRY=False - ALLOW_RESET=True diff --git a/docker-compose.test.yml b/docker-compose.test.yml index c8cae63b3eba..4384bad1982a 100644 --- a/docker-compose.test.yml +++ b/docker-compose.test.yml @@ -11,7 +11,7 @@ services: dockerfile: Dockerfile volumes: - chroma-data:/chroma/chroma - command: uvicorn chromadb.app:app --workers 1 --host 0.0.0.0 --port 8000 --log-config chromadb/log_config.yml --timeout-keep-alive 30 + command: "--workers 1 --host 0.0.0.0 --port 8000 --proxy-headers --log-config chromadb/log_config.yml --timeout-keep-alive 30" environment: - ANONYMIZED_TELEMETRY=False - ALLOW_RESET=True diff --git a/docker-compose.yml b/docker-compose.yml index 4cd653b01c7f..20d096569070 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -15,7 +15,7 @@ services: # Default configuration for persist_directory in chromadb/config.py # Read more about deployments: https://docs.trychroma.com/deployment - chroma-data:/chroma/chroma - command: uvicorn chromadb.app:app --reload --workers 1 --host 0.0.0.0 --port 8000 --log-config chromadb/log_config.yml --timeout-keep-alive 30 + command: "--workers 1 --host 0.0.0.0 --port 8000 --proxy-headers --log-config chromadb/log_config.yml --timeout-keep-alive 30" environment: - IS_PERSISTENT=TRUE - CHROMA_SERVER_AUTH_PROVIDER=${CHROMA_SERVER_AUTH_PROVIDER} From c665838b0d143e2c2ceb82c4ade7404dc98124ff Mon Sep 17 00:00:00 2001 From: Viktor Due <66885944+DueViktor@users.noreply.github.com> Date: Wed, 7 Feb 2024 19:05:55 +0100 Subject: [PATCH 6/9] FIPS Compliance (#1673) ## Description of changes Close #1672 ## Test plan *How are these changes tested?* - [x] Tests pass locally with `pytest` for python, `yarn test` for js ## Documentation Changes *Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the [docs repository](https://github.com/chroma-core/docs)?* --- chromadb/db/migrations.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/chromadb/db/migrations.py b/chromadb/db/migrations.py index 951cb762c36b..97ef029092ab 100644 --- a/chromadb/db/migrations.py +++ b/chromadb/db/migrations.py @@ -1,3 +1,4 @@ +import sys from typing import Sequence from typing_extensions import TypedDict, NotRequired from importlib_resources.abc import Traversable @@ -253,7 +254,7 @@ def _read_migration_file(file: MigrationFile, hash_alg: str) -> Migration: sql = file["path"].read_text() if hash_alg == "md5": - hash = hashlib.md5(sql.encode("utf-8")).hexdigest() + hash = hashlib.md5(sql.encode("utf-8"), usedforsecurity=False).hexdigest() if sys.version_info >= (3, 9) else hashlib.md5(sql.encode("utf-8")).hexdigest() elif hash_alg == "sha256": hash = hashlib.sha256(sql.encode("utf-8")).hexdigest() else: From 01369afe1d00d844829281c89175e279e1591f44 Mon Sep 17 00:00:00 2001 From: Trayan Azarov Date: Fri, 9 Feb 2024 21:53:08 +0200 Subject: [PATCH 7/9] [BUG]: Disallowing 0-dimensional embeddings (#1702) Refs: #1698 ## Description of changes *Summarize the changes made by this PR.* - Improvements & Bug fixes - Added validation for 0-dimensional embeddings ## Test plan *How are these changes tested?* - [x] Tests pass locally with `pytest` for python, `yarn test` for js ## Documentation Changes N/A --- chromadb/api/types.py | 6 +++++- chromadb/test/property/test_embeddings.py | 7 +++++++ clients/js/test/add.collections.test.ts | 14 ++++++++++++++ 3 files changed, 26 insertions(+), 1 deletion(-) diff --git a/chromadb/api/types.py b/chromadb/api/types.py index 7781c4225725..0054f283e8d0 100644 --- a/chromadb/api/types.py +++ b/chromadb/api/types.py @@ -476,7 +476,11 @@ def validate_embeddings(embeddings: Embeddings) -> Embeddings: raise ValueError( f"Expected each embedding in the embeddings to be a list, got {embeddings}" ) - for embedding in embeddings: + for i,embedding in enumerate(embeddings): + if len(embedding) == 0: + raise ValueError( + f"Expected each embedding in the embeddings to be a non-empty list, got empty embedding at pos {i}" + ) if not all( [ isinstance(value, (int, float)) and not isinstance(value, bool) diff --git a/chromadb/test/property/test_embeddings.py b/chromadb/test/property/test_embeddings.py index cfb2c93fa524..bf3e882184ff 100644 --- a/chromadb/test/property/test_embeddings.py +++ b/chromadb/test/property/test_embeddings.py @@ -455,3 +455,10 @@ def test_autocasting_validate_embeddings_incompatible_types( validate_embeddings(Collection._normalize_embeddings(embds)) assert "Expected each value in the embedding to be a int or float" in str(e) + + +def test_0dim_embedding_validation() -> None: + embds = [[]] + with pytest.raises(ValueError) as e: + validate_embeddings(embds) + assert "Expected each embedding in the embeddings to be a non-empty list" in str(e) \ No newline at end of file diff --git a/clients/js/test/add.collections.test.ts b/clients/js/test/add.collections.test.ts index cb89fa8dbe06..7ac271ff98e9 100644 --- a/clients/js/test/add.collections.test.ts +++ b/clients/js/test/add.collections.test.ts @@ -100,3 +100,17 @@ test('It should return an error when inserting duplicate IDs in the same batch', expect(e.message).toMatch('duplicates') } }) + + +test('should error on empty embedding', async () => { + await chroma.reset() + const collection = await chroma.createCollection({ name: "test" }); + const ids = ["id1"] + const embeddings = [[]] + const metadatas = [{ test: 'test1', 'float_value': 0.1 }] + try { + await collection.add({ ids, embeddings, metadatas }); + } catch (e: any) { + expect(e.message).toMatch('got empty embedding at pos') + } +}) \ No newline at end of file From 4b656d9b48373ef272bcb7282c76682e595d6374 Mon Sep 17 00:00:00 2001 From: Trayan Azarov Date: Wed, 14 Feb 2024 20:32:07 +0200 Subject: [PATCH 8/9] [ENH]: Chroma python client orjson serialization (#1705) ## Description of changes *Summarize the changes made by this PR.* - Improvements & Bug fixes - Faster serialization of requests for the HttpClient ## Test plan *How are these changes tested?* - [x] Tests pass locally with `pytest` for python ## Documentation Changes N/A # Perf Benchmark ![image](https://github.com/chroma-core/chroma/assets/1157440/20307cd4-042a-46f3-86df-c11311bc1a7c) - The test was conducted with HttpClient from `main` and from this PR. - Batches of 10, 100, 1000, and 10000 were used to test (the times for generating the batches are discounted) - A mock server was used, which ignored parsing JSON at the server side to reduce latency # Refs - https://showmax.engineering/articles/json-python-libraries-overview - https://github.com/ijl/orjson - https://catnotfoundnear.github.io/finding-the-fastest-python-json-library-on-all-python-versions-8-compared.html --- chromadb/api/fastapi.py | 32 ++++++++++++++++---------------- clients/python/pyproject.toml | 1 + clients/python/requirements.txt | 1 + pyproject.toml | 1 + requirements.txt | 1 + server.htpasswd | 1 - 6 files changed, 20 insertions(+), 17 deletions(-) delete mode 100644 server.htpasswd diff --git a/chromadb/api/fastapi.py b/chromadb/api/fastapi.py index 1ee7a45af541..a10fdfaf02d9 100644 --- a/chromadb/api/fastapi.py +++ b/chromadb/api/fastapi.py @@ -1,4 +1,4 @@ -import json +import orjson as json import logging from typing import Optional, cast, Tuple from typing import Sequence @@ -147,7 +147,7 @@ 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"]) + return int(json.loads(resp.text)["nanosecond heartbeat"]) @trace_method("FastAPI.create_database", OpenTelemetryGranularity.OPERATION) @override @@ -177,7 +177,7 @@ def get_database( params={"tenant": tenant}, ) raise_chroma_error(resp) - resp_json = resp.json() + resp_json = json.loads(resp.text) return Database( id=resp_json["id"], name=resp_json["name"], tenant=resp_json["tenant"] ) @@ -198,7 +198,7 @@ def get_tenant(self, name: str) -> Tenant: self._api_url + "/tenants/" + name, ) raise_chroma_error(resp) - resp_json = resp.json() + resp_json = json.loads(resp.text) return Tenant(name=resp_json["name"]) @trace_method("FastAPI.list_collections", OpenTelemetryGranularity.OPERATION) @@ -221,7 +221,7 @@ def list_collections( }, ) raise_chroma_error(resp) - json_collections = resp.json() + json_collections = json.loads(resp.text) collections = [] for json_collection in json_collections: collections.append(Collection(self, **json_collection)) @@ -239,7 +239,7 @@ def count_collections( params={"tenant": tenant, "database": database}, ) raise_chroma_error(resp) - return cast(int, resp.json()) + return cast(int, json.loads(resp.text)) @trace_method("FastAPI.create_collection", OpenTelemetryGranularity.OPERATION) @override @@ -268,7 +268,7 @@ def create_collection( params={"tenant": tenant, "database": database}, ) raise_chroma_error(resp) - resp_json = resp.json() + resp_json = json.loads(resp.text) return Collection( client=self, id=resp_json["id"], @@ -302,7 +302,7 @@ def get_collection( self._api_url + "/collections/" + name if name else str(id), params=_params ) raise_chroma_error(resp) - resp_json = resp.json() + resp_json = json.loads(resp.text) return Collection( client=self, name=resp_json["name"], @@ -381,7 +381,7 @@ def _count( self._api_url + "/collections/" + str(collection_id) + "/count" ) raise_chroma_error(resp) - return cast(int, resp.json()) + return cast(int, json.loads(resp.text)) @trace_method("FastAPI._peek", OpenTelemetryGranularity.OPERATION) @override @@ -434,7 +434,7 @@ def _get( ) raise_chroma_error(resp) - body = resp.json() + body = json.loads(resp.text) return GetResult( ids=body["ids"], embeddings=body.get("embeddings", None), @@ -462,7 +462,7 @@ def _delete( ) raise_chroma_error(resp) - return cast(IDs, resp.json()) + return cast(IDs, json.loads(resp.text)) @trace_method("FastAPI._submit_batch", OpenTelemetryGranularity.ALL) def _submit_batch( @@ -586,7 +586,7 @@ def _query( ) raise_chroma_error(resp) - body = resp.json() + body = json.loads(resp.text) return QueryResult( ids=body["ids"], @@ -604,7 +604,7 @@ def reset(self) -> bool: """Resets the database""" resp = self._session.post(self._api_url + "/reset") raise_chroma_error(resp) - return cast(bool, resp.json()) + return cast(bool, json.loads(resp.text)) @trace_method("FastAPI.get_version", OpenTelemetryGranularity.OPERATION) @override @@ -612,7 +612,7 @@ def get_version(self) -> str: """Returns the version of the server""" resp = self._session.get(self._api_url + "/version") raise_chroma_error(resp) - return cast(str, resp.json()) + return cast(str, json.loads(resp.text)) @override def get_settings(self) -> Settings: @@ -626,7 +626,7 @@ def max_batch_size(self) -> int: if self._max_batch_size == -1: resp = self._session.get(self._api_url + "/pre-flight-checks") raise_chroma_error(resp) - self._max_batch_size = cast(int, resp.json()["max_batch_size"]) + self._max_batch_size = cast(int, json.loads(resp.text)["max_batch_size"]) return self._max_batch_size @@ -637,7 +637,7 @@ def raise_chroma_error(resp: requests.Response) -> None: chroma_error = None try: - body = resp.json() + body = json.loads(resp.text) if "error" in body: if body["error"] in errors.error_types: chroma_error = errors.error_types[body["error"]](body["message"]) diff --git a/clients/python/pyproject.toml b/clients/python/pyproject.toml index b62c002d095c..edd0c00d7cf5 100644 --- a/clients/python/pyproject.toml +++ b/clients/python/pyproject.toml @@ -26,6 +26,7 @@ dependencies = [ 'typing_extensions >= 4.5.0', 'tenacity>=8.2.3', 'PyYAML>=6.0.0', + 'orjson>=3.9.12', ] [tool.black] diff --git a/clients/python/requirements.txt b/clients/python/requirements.txt index 1242bf7d7e0f..b977b03f064a 100644 --- a/clients/python/requirements.txt +++ b/clients/python/requirements.txt @@ -9,3 +9,4 @@ PyYAML>=6.0.0 requests >= 2.28 tenacity>=8.2.3 typing_extensions >= 4.5.0 +orjson>=3.9.12 diff --git a/pyproject.toml b/pyproject.toml index 01aa0d8663bc..d425e77952d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,7 @@ dependencies = [ 'tenacity>=8.2.3', 'PyYAML>=6.0.0', 'mmh3>=4.0.1', + 'orjson>=3.9.12', ] [tool.black] diff --git a/requirements.txt b/requirements.txt index 3e99734fd3a8..6a1b1fb966f2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -25,3 +25,4 @@ tqdm>=4.65.0 typer>=0.9.0 typing_extensions>=4.5.0 uvicorn[standard]==0.18.3 +orjson>=3.9.12 \ No newline at end of file diff --git a/server.htpasswd b/server.htpasswd deleted file mode 100644 index 77f277a399ba..000000000000 --- a/server.htpasswd +++ /dev/null @@ -1 +0,0 @@ -admin:$2y$05$e5sRb6NCcSH3YfbIxe1AGu2h5K7OOd982OXKmd8WyQ3DRQ4MvpnZS From da68516f72bc8d23e0811a23c1d6085143f4e2df Mon Sep 17 00:00:00 2001 From: Trayan Azarov Date: Wed, 14 Feb 2024 21:17:28 +0200 Subject: [PATCH 9/9] [BUG]: Adding validation check for "chroma:document" in metadata. (#1718) Closes: #1717 ## Description of changes *Summarize the changes made by this PR.* - Improvements & Bug fixes - Fixed an issue where if the metadata key is set to `chroma:document` it is either ignore when inserting or overrides the actual document when updating records by `id` ## Test plan *How are these changes tested?* - [x] Tests pass locally with `pytest` for python, `yarn test` for js ## Documentation Changes N/A --- chromadb/api/types.py | 8 ++++++-- chromadb/test/segment/test_metadata.py | 9 +++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/chromadb/api/types.py b/chromadb/api/types.py index 0054f283e8d0..347461718fde 100644 --- a/chromadb/api/types.py +++ b/chromadb/api/types.py @@ -20,7 +20,7 @@ # Re-export types from chromadb.types __all__ = ["Metadata", "Where", "WhereDocument", "UpdateCollectionMetadata"] - +META_KEY_CHROMA_DOCUMENT = "chroma:document" T = TypeVar("T") OneOrMany = Union[T, List[T]] @@ -265,6 +265,10 @@ def validate_metadata(metadata: Metadata) -> Metadata: if len(metadata) == 0: raise ValueError(f"Expected metadata to be a non-empty dict, got {metadata}") for key, value in metadata.items(): + if key == META_KEY_CHROMA_DOCUMENT: + raise ValueError( + f"Expected metadata to not contain the reserved key {META_KEY_CHROMA_DOCUMENT}" + ) if not isinstance(key, str): raise TypeError( f"Expected metadata key to be a str, got {key} which is a {type(key)}" @@ -476,7 +480,7 @@ def validate_embeddings(embeddings: Embeddings) -> Embeddings: raise ValueError( f"Expected each embedding in the embeddings to be a list, got {embeddings}" ) - for i,embedding in enumerate(embeddings): + for i, embedding in enumerate(embeddings): if len(embedding) == 0: raise ValueError( f"Expected each embedding in the embeddings to be a non-empty list, got empty embedding at pos {i}" diff --git a/chromadb/test/segment/test_metadata.py b/chromadb/test/segment/test_metadata.py index 1f03d6350f48..2126c6d1febc 100644 --- a/chromadb/test/segment/test_metadata.py +++ b/chromadb/test/segment/test_metadata.py @@ -3,6 +3,8 @@ import tempfile import pytest from typing import Generator, List, Callable, Iterator, Dict, Optional, Union, Sequence + +from chromadb.api.types import validate_metadata from chromadb.config import System, Settings from chromadb.db.base import ParameterValue, get_sql from chromadb.db.impl.sqlite import SqliteDB @@ -677,3 +679,10 @@ def test_delete_segment( res = cur.execute(sql, params) # assert that all FTS rows are gone assert len(res.fetchall()) == 0 + + +def test_metadata_validation_forbidden_key() -> None: + with pytest.raises(ValueError, match="chroma:document"): + validate_metadata( + {"chroma:document": "this is not the document you are looking for"} + )