tjh.dev
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 fetchFromGitHub,
6 fetchurl,
7
8 # build inputs
9 cargo,
10 pkg-config,
11 protobuf,
12 rustc,
13 rustPlatform,
14 pkgs, # zstd hidden by python3Packages.zstd
15 openssl,
16
17 # dependencies
18 bcrypt,
19 build,
20 fastapi,
21 grpcio,
22 httpx,
23 importlib-resources,
24 jsonschema,
25 kubernetes,
26 mmh3,
27 numpy,
28 onnxruntime,
29 opentelemetry-api,
30 opentelemetry-exporter-otlp-proto-grpc,
31 opentelemetry-instrumentation-fastapi,
32 opentelemetry-sdk,
33 orjson,
34 overrides,
35 posthog,
36 pydantic,
37 pypika,
38 pyyaml,
39 requests,
40 tenacity,
41 tokenizers,
42 tqdm,
43 typer,
44 typing-extensions,
45 uvicorn,
46
47 # optional dependencies
48 chroma-hnswlib,
49
50 # tests
51 hnswlib,
52 hypothesis,
53 pandas,
54 psutil,
55 pytest-asyncio,
56 pytest-xdist,
57 pytestCheckHook,
58 sqlite,
59 starlette,
60 writableTmpDirAsHomeHook,
61
62 # passthru
63 nixosTests,
64 nix-update-script,
65}:
66
67buildPythonPackage rec {
68 pname = "chromadb";
69 version = "1.0.12";
70 pyproject = true;
71
72 src = fetchFromGitHub {
73 owner = "chroma-core";
74 repo = "chroma";
75 tag = version;
76 hash = "sha256-Q4PhJTRNzJeVx6DIPWirnI9KksNb8vfOtqb/q9tSK3c=";
77 };
78
79 cargoDeps = rustPlatform.fetchCargoVendor {
80 inherit src;
81 name = "${pname}-${version}-vendor";
82 hash = "sha256-+Ea2aRrsBGfVCLdOF41jeMehJhMurc8d0UKrpR6ndag=";
83 };
84
85 # Can't use fetchFromGitHub as the build expects a zipfile
86 swagger-ui = fetchurl {
87 url = "https://github.com/swagger-api/swagger-ui/archive/refs/tags/v5.22.0.zip";
88 hash = "sha256-H+kXxA/6rKzYA19v7Zlx2HbIg/DGicD5FDIs0noVGSk=";
89 };
90
91 patches = [
92 # The fastapi servers can't set up their networking in the test environment, so disable for testing
93 ./disable-fastapi-fixtures.patch
94 ];
95
96 postPatch = ''
97 # Nixpkgs is taking the version from `chromadb_rust_bindings` which is versioned independently
98 substituteInPlace pyproject.toml \
99 --replace-fail "dynamic = [\"version\"]" "version = \"${version}\""
100 '';
101
102 pythonRelaxDeps = [
103 "fastapi"
104 ];
105
106 build-system = [
107 rustPlatform.maturinBuildHook
108 ];
109
110 nativeBuildInputs = [
111 cargo
112 pkg-config
113 protobuf
114 rustc
115 rustPlatform.cargoSetupHook
116 ];
117
118 buildInputs = [
119 openssl
120 pkgs.zstd
121 ];
122
123 dependencies = [
124 bcrypt
125 build
126 fastapi
127 grpcio
128 httpx
129 importlib-resources
130 jsonschema
131 kubernetes
132 mmh3
133 numpy
134 onnxruntime
135 opentelemetry-api
136 opentelemetry-exporter-otlp-proto-grpc
137 opentelemetry-instrumentation-fastapi
138 opentelemetry-sdk
139 orjson
140 overrides
141 posthog
142 pydantic
143 pypika
144 pyyaml
145 requests
146 tenacity
147 tokenizers
148 tqdm
149 typer
150 typing-extensions
151 uvicorn
152 ];
153
154 optional-dependencies = {
155 dev = [ chroma-hnswlib ];
156 };
157
158 nativeCheckInputs = [
159 chroma-hnswlib
160 hnswlib
161 hypothesis
162 pandas
163 psutil
164 pytest-asyncio
165 pytest-xdist
166 pytestCheckHook
167 sqlite
168 starlette
169 writableTmpDirAsHomeHook
170 ];
171
172 # Disable on aarch64-linux due to broken onnxruntime
173 # https://github.com/microsoft/onnxruntime/issues/10038
174 pythonImportsCheck = lib.optionals doCheck [ "chromadb" ];
175
176 # Test collection breaks on aarch64-linux
177 doCheck = with stdenv.buildPlatform; !(isAarch && isLinux);
178
179 env = {
180 ZSTD_SYS_USE_PKG_CONFIG = true;
181 SWAGGER_UI_DOWNLOAD_URL = "file://${swagger-ui}";
182 };
183
184 pytestFlags = [
185 "-x" # these are slow tests, so stop on the first failure
186 "-v"
187 "-Wignore:DeprecationWarning"
188 "-Wignore:PytestCollectionWarning"
189 ];
190
191 preCheck = ''
192 (($(ulimit -n) < 1024)) && ulimit -n 1024
193 export HOME=$(mktemp -d)
194 '';
195
196 disabledTests = [
197 # Tests are flaky / timing sensitive
198 "test_fastapi_server_token_authn_allows_when_it_should_allow"
199 "test_fastapi_server_token_authn_rejects_when_it_should_reject"
200
201 # Issue with event loop
202 "test_http_client_bw_compatibility"
203
204 # httpx ReadError
205 "test_not_existing_collection_delete"
206
207 # Tests launch a server and try to connect to it
208 # These either have https connection errors or name resolution errors
209 "test_collection_query_with_invalid_collection_throws"
210 "test_collection_update_with_invalid_collection_throws"
211 "test_default_embedding"
212 "test_invalid_index_params"
213 "test_peek"
214 "test_persist_index_loading"
215 "test_query_id_filtering_e2e"
216 "test_query_id_filtering_medium_dataset"
217 "test_query_id_filtering_small_dataset"
218 "test_ssl_self_signed_without_ssl_verify"
219 "test_ssl_self_signed"
220
221 # Apparent race condition with sqlite
222 # See https://github.com/chroma-core/chroma/issues/4661
223 "test_multithreaded_get_or_create"
224 ];
225
226 disabledTestPaths = [
227 # Tests require network access
228 "bin/rust_python_compat_test.py"
229 "chromadb/test/configurations/test_collection_configuration.py"
230 "chromadb/test/ef/test_default_ef.py"
231 "chromadb/test/ef/test_onnx_mini_lm_l6_v2.py"
232 "chromadb/test/ef/test_voyageai_ef.py"
233 "chromadb/test/property/"
234 "chromadb/test/property/test_cross_version_persist.py"
235 "chromadb/test/stress/"
236 "chromadb/test/test_api.py"
237
238 # Tests time out (waiting for server)
239 "chromadb/test/test_cli.py"
240
241 # Cannot find protobuf file while loading test
242 "chromadb/test/distributed/test_log_failover.py"
243
244 # ValueError: An instance of Chroma already exists for ephemeral with different settings
245 "chromadb/test/test_chroma.py"
246 "chromadb/test/ef/test_multimodal_ef.py"
247 ];
248
249 __darwinAllowLocalNetworking = true;
250
251 passthru = {
252 tests = {
253 inherit (nixosTests) chromadb;
254 };
255
256 updateScript = nix-update-script {
257 # we have to update both the python hash and the cargo one,
258 # so use nix-update-script
259 extraArgs = [
260 "--version-regex"
261 "([0-9].+)"
262 ];
263 };
264 };
265
266 meta = {
267 description = "AI-native open-source embedding database";
268 homepage = "https://github.com/chroma-core/chroma";
269 changelog = "https://github.com/chroma-core/chroma/releases/tag/${version}";
270 license = lib.licenses.asl20;
271 maintainers = with lib.maintainers; [
272 fab
273 sarahec
274 ];
275 mainProgram = "chroma";
276 };
277}