pyrox.dev
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 fetchFromGitHub,
6 fetchurl,
7
8 # build inputs
9 cargo,
10 openssl,
11 pkg-config,
12 protobuf,
13 rustc,
14 rustPlatform,
15 zstd-c,
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 pybase64,
37 pydantic,
38 pypika,
39 pyyaml,
40 requests,
41 tenacity,
42 tokenizers,
43 tqdm,
44 typer,
45 typing-extensions,
46 uvicorn,
47
48 # optional dependencies
49 chroma-hnswlib,
50
51 # tests
52 hnswlib,
53 hypothesis,
54 pandas,
55 psutil,
56 pytest-asyncio,
57 pytest-xdist,
58 pytestCheckHook,
59 sqlite,
60 starlette,
61 writableTmpDirAsHomeHook,
62
63 # passthru
64 nixosTests,
65 nix-update-script,
66}:
67
68buildPythonPackage rec {
69 pname = "chromadb";
70 version = "1.1.0";
71 pyproject = true;
72
73 src = fetchFromGitHub {
74 owner = "chroma-core";
75 repo = "chroma";
76 tag = version;
77 hash = "sha256-RVXMjniqZ0zUVhdgcYHFgYV1WrNZzBLW9jdrvV8AnRU=";
78 };
79
80 cargoDeps = rustPlatform.fetchCargoVendor {
81 inherit src;
82 name = "${pname}-${version}-vendor";
83 hash = "sha256-owy+6RttjVDCfsnn7MLuMn9/esHPwb7Z7jXqJ4IHfaE=";
84 };
85
86 # Can't use fetchFromGitHub as the build expects a zipfile
87 swagger-ui = fetchurl {
88 url = "https://github.com/swagger-api/swagger-ui/archive/refs/tags/v5.22.0.zip";
89 hash = "sha256-H+kXxA/6rKzYA19v7Zlx2HbIg/DGicD5FDIs0noVGSk=";
90 };
91
92 postPatch = ''
93 # Nixpkgs is taking the version from `chromadb_rust_bindings` which is versioned independently
94 substituteInPlace pyproject.toml \
95 --replace-fail "dynamic = [\"version\"]" "version = \"${version}\""
96 '';
97
98 pythonRelaxDeps = [
99 "fastapi"
100 "posthog"
101 ];
102
103 build-system = [
104 rustPlatform.maturinBuildHook
105 ];
106
107 nativeBuildInputs = [
108 cargo
109 pkg-config
110 protobuf
111 rustc
112 rustPlatform.cargoSetupHook
113 ];
114
115 buildInputs = [
116 openssl
117 zstd-c
118 ];
119
120 dependencies = [
121 bcrypt
122 build
123 fastapi
124 grpcio
125 httpx
126 importlib-resources
127 jsonschema
128 kubernetes
129 mmh3
130 numpy
131 onnxruntime
132 opentelemetry-api
133 opentelemetry-exporter-otlp-proto-grpc
134 opentelemetry-instrumentation-fastapi
135 opentelemetry-sdk
136 orjson
137 overrides
138 posthog
139 pybase64
140 pydantic
141 pypika
142 pyyaml
143 requests
144 tenacity
145 tokenizers
146 tqdm
147 typer
148 typing-extensions
149 uvicorn
150 ];
151
152 optional-dependencies = {
153 dev = [ chroma-hnswlib ];
154 };
155
156 nativeCheckInputs = [
157 chroma-hnswlib
158 hnswlib
159 hypothesis
160 pandas
161 psutil
162 pytest-asyncio
163 pytest-xdist
164 pytestCheckHook
165 sqlite
166 starlette
167 writableTmpDirAsHomeHook
168 ];
169
170 # Disable on aarch64-linux due to broken onnxruntime
171 # https://github.com/microsoft/onnxruntime/issues/10038
172 pythonImportsCheck = lib.optionals doCheck [ "chromadb" ];
173
174 # Test collection breaks on aarch64-linux
175 doCheck = with stdenv.buildPlatform; !(isAarch && isLinux);
176
177 env = {
178 ZSTD_SYS_USE_PKG_CONFIG = true;
179 SWAGGER_UI_DOWNLOAD_URL = "file://${swagger-ui}";
180 };
181
182 pytestFlags = [
183 "-v"
184 "-Wignore:DeprecationWarning"
185 "-Wignore:PytestCollectionWarning"
186 ];
187
188 # Skip the distributed and integration tests
189 # See https://github.com/chroma-core/chroma/issues/5315
190 preCheck = ''
191 (($(ulimit -n) < 1024)) && ulimit -n 1024
192 export CHROMA_RUST_BINDINGS_TEST_ONLY=1
193 '';
194
195 enabledTestPaths = [
196 "chromadb/test"
197 ];
198
199 disabledTests = [
200 # Failure in name resolution
201 "test_collection_query_with_invalid_collection_throws"
202 "test_collection_update_with_invalid_collection_throws"
203 "test_default_embedding"
204 "test_persist_index_loading"
205
206 # Deadlocks intermittently
207 "test_app"
208
209 # Depends on specific floating-point precision
210 "test_base64_conversion_is_identity_f16"
211
212 # No such file or directory: 'openssl'
213 "test_ssl_self_signed_without_ssl_verify"
214 "test_ssl_self_signed"
215 ];
216
217 disabledTestPaths = [
218 # Tests require network access
219 "chromadb/test/distributed"
220 "chromadb/test/ef"
221 "chromadb/test/property/test_cross_version_persist.py"
222 "chromadb/test/stress"
223
224 # Excessively slow
225 "chromadb/test/property/test_add.py"
226 "chromadb/test/property/test_persist.py"
227
228 # ValueError: An instance of Chroma already exists for ephemeral with different settings
229 "chromadb/test/test_chroma.py"
230 ];
231
232 __darwinAllowLocalNetworking = true;
233
234 passthru = {
235 tests = {
236 inherit (nixosTests) chromadb;
237 };
238
239 updateScript = nix-update-script {
240 # we have to update both the python hash and the cargo one,
241 # so use nix-update-script
242 extraArgs = [
243 "--version-regex"
244 "([0-9].+)"
245 ];
246 };
247 };
248
249 meta = {
250 description = "AI-native open-source embedding database";
251 homepage = "https://github.com/chroma-core/chroma";
252 changelog = "https://github.com/chroma-core/chroma/releases/tag/${version}";
253 license = lib.licenses.asl20;
254 maintainers = with lib.maintainers; [
255 fab
256 sarahec
257 ];
258 mainProgram = "chroma";
259 };
260}