-
Notifications
You must be signed in to change notification settings - Fork 1.2k
/
__init__.py
343 lines (281 loc) · 11.1 KB
/
__init__.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
from typing import Dict, Optional
import logging
from chromadb.api.client import Client as ClientCreator
from chromadb.api.client import AdminClient as AdminClientCreator
from chromadb.api.async_client import AsyncClient as AsyncClientCreator
from chromadb.auth.token_authn import TokenTransportHeader
import chromadb.config
from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT, Settings
from chromadb.api import AdminAPI, AsyncClientAPI, ClientAPI
from chromadb.api.models.Collection import Collection
from chromadb.api.types import (
CollectionMetadata,
Documents,
EmbeddingFunction,
Embeddings,
IDs,
Include,
Metadata,
Where,
QueryResult,
GetResult,
WhereDocument,
UpdateCollectionMetadata,
)
# Re-export types from chromadb.types
__all__ = [
"Collection",
"Metadata",
"Where",
"WhereDocument",
"Documents",
"IDs",
"Embeddings",
"EmbeddingFunction",
"Include",
"CollectionMetadata",
"UpdateCollectionMetadata",
"QueryResult",
"GetResult",
"TokenTransportHeader",
]
logger = logging.getLogger(__name__)
__settings = Settings()
__version__ = "0.5.5"
# Workaround to deal with Colab's old sqlite3 version
def is_in_colab() -> bool:
try:
import google.colab # noqa: F401
return True
except ImportError:
return False
IN_COLAB = is_in_colab()
is_client = False
try:
from chromadb.is_thin_client import is_thin_client
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
import subprocess
import sys
subprocess.check_call(
[sys.executable, "-m", "pip", "install", "pysqlite3-binary"]
)
__import__("pysqlite3")
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"
)
def configure(**kwargs) -> None: # type: ignore
"""Override Chroma's default settings, environment variables or .env files"""
global __settings
__settings = chromadb.config.Settings(**kwargs)
def get_settings() -> Settings:
return __settings
def EphemeralClient(
settings: Optional[Settings] = None,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> ClientAPI:
"""
Creates an in-memory instance of Chroma. This is useful for testing and
development, but not recommended for production use.
Args:
tenant: The tenant to use for this client. Defaults to the default tenant.
database: The database to use for this client. Defaults to the default database.
"""
if settings is None:
settings = Settings()
settings.is_persistent = False
# Make sure paramaters are the correct types -- users can pass anything.
tenant = str(tenant)
database = str(database)
return ClientCreator(settings=settings, tenant=tenant, database=database)
def PersistentClient(
path: str = "./chroma",
settings: Optional[Settings] = None,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> ClientAPI:
"""
Creates a persistent instance of Chroma that saves to disk. This is useful for
testing and development, but not recommended for production use.
Args:
path: The directory to save Chroma's data to. Defaults to "./chroma".
tenant: The tenant to use for this client. Defaults to the default tenant.
database: The database to use for this client. Defaults to the default database.
"""
if settings is None:
settings = Settings()
settings.persist_directory = path
settings.is_persistent = True
# Make sure paramaters are the correct types -- users can pass anything.
tenant = str(tenant)
database = str(database)
return ClientCreator(tenant=tenant, database=database, settings=settings)
def HttpClient(
host: str = "localhost",
port: int = 8000,
ssl: bool = False,
headers: Optional[Dict[str, str]] = None,
settings: Optional[Settings] = None,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> ClientAPI:
"""
Creates a client that connects to a remote Chroma server. This supports
many clients connecting to the same server, and is the recommended way to
use Chroma in production.
Args:
host: The hostname of the Chroma server. Defaults to "localhost".
port: The port of the Chroma server. Defaults to "8000".
ssl: Whether to use SSL to connect to the Chroma server. Defaults to False.
headers: A dictionary of headers to send to the Chroma server. Defaults to {}.
settings: A dictionary of settings to communicate with the chroma server.
tenant: The tenant to use for this client. Defaults to the default tenant.
database: The database to use for this client. Defaults to the default database.
"""
if settings is None:
settings = Settings()
# Make sure parameters are the correct types -- users can pass anything.
host = str(host)
port = int(port)
ssl = bool(ssl)
tenant = str(tenant)
database = str(database)
settings.chroma_api_impl = "chromadb.api.fastapi.FastAPI"
if settings.chroma_server_host and settings.chroma_server_host != host:
raise ValueError(
f"Chroma server host provided in settings[{settings.chroma_server_host}] is different to the one provided in HttpClient: [{host}]"
)
settings.chroma_server_host = host
if settings.chroma_server_http_port and settings.chroma_server_http_port != port:
raise ValueError(
f"Chroma server http port provided in settings[{settings.chroma_server_http_port}] is different to the one provided in HttpClient: [{port}]"
)
settings.chroma_server_http_port = port
settings.chroma_server_ssl_enabled = ssl
settings.chroma_server_headers = headers
return ClientCreator(tenant=tenant, database=database, settings=settings)
async def AsyncHttpClient(
host: str = "localhost",
port: int = 8000,
ssl: bool = False,
headers: Optional[Dict[str, str]] = None,
settings: Optional[Settings] = None,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> AsyncClientAPI:
"""
Creates an async client that connects to a remote Chroma server. This supports
many clients connecting to the same server, and is the recommended way to
use Chroma in production.
Args:
host: The hostname of the Chroma server. Defaults to "localhost".
port: The port of the Chroma server. Defaults to "8000".
ssl: Whether to use SSL to connect to the Chroma server. Defaults to False.
headers: A dictionary of headers to send to the Chroma server. Defaults to {}.
settings: A dictionary of settings to communicate with the chroma server.
tenant: The tenant to use for this client. Defaults to the default tenant.
database: The database to use for this client. Defaults to the default database.
"""
if settings is None:
settings = Settings()
# Make sure parameters are the correct types -- users can pass anything.
host = str(host)
port = int(port)
ssl = bool(ssl)
tenant = str(tenant)
database = str(database)
settings.chroma_api_impl = "chromadb.api.async_fastapi.AsyncFastAPI"
if settings.chroma_server_host and settings.chroma_server_host != host:
raise ValueError(
f"Chroma server host provided in settings[{settings.chroma_server_host}] is different to the one provided in HttpClient: [{host}]"
)
settings.chroma_server_host = host
if settings.chroma_server_http_port and settings.chroma_server_http_port != port:
raise ValueError(
f"Chroma server http port provided in settings[{settings.chroma_server_http_port}] is different to the one provided in HttpClient: [{port}]"
)
settings.chroma_server_http_port = port
settings.chroma_server_ssl_enabled = ssl
settings.chroma_server_headers = headers
return await AsyncClientCreator.create(
tenant=tenant, database=database, settings=settings
)
def CloudClient(
tenant: str,
database: str,
api_key: Optional[str] = None,
settings: Optional[Settings] = None,
*, # Following arguments are keyword-only, intended for testing only.
cloud_host: str = "api.trychroma.com",
cloud_port: int = 8000,
enable_ssl: bool = True,
) -> ClientAPI:
"""
Creates a client to connect to a tennant and database on the Chroma cloud.
Args:
tenant: The tenant to use for this client.
database: The database to use for this client.
api_key: The api key to use for this client.
"""
# If no API key is provided, try to load it from the environment variable
if api_key is None:
import os
api_key = os.environ.get("CHROMA_API_KEY")
# If the API key is still not provided, prompt the user
if api_key is None:
print(
"\033[93mDon't have an API key?\033[0m Get one at https://app.trychroma.com"
)
api_key = input("Please enter your Chroma API key: ")
if settings is None:
settings = Settings()
# Make sure paramaters are the correct types -- users can pass anything.
tenant = str(tenant)
database = str(database)
api_key = str(api_key)
cloud_host = str(cloud_host)
cloud_port = int(cloud_port)
enable_ssl = bool(enable_ssl)
settings.chroma_api_impl = "chromadb.api.fastapi.FastAPI"
settings.chroma_server_host = cloud_host
settings.chroma_server_http_port = cloud_port
# Always use SSL for cloud
settings.chroma_server_ssl_enabled = enable_ssl
settings.chroma_client_auth_provider = (
"chromadb.auth.token_authn.TokenAuthClientProvider"
)
settings.chroma_client_auth_credentials = api_key
settings.chroma_auth_token_transport_header = TokenTransportHeader.X_CHROMA_TOKEN
return ClientCreator(tenant=tenant, database=database, settings=settings)
def Client(
settings: Settings = __settings,
tenant: str = DEFAULT_TENANT,
database: str = DEFAULT_DATABASE,
) -> ClientAPI:
"""
Return a running chroma.API instance
tenant: The tenant to use for this client. Defaults to the default tenant.
database: The database to use for this client. Defaults to the default database.
"""
# Make sure paramaters are the correct types -- users can pass anything.
tenant = str(tenant)
database = str(database)
return ClientCreator(tenant=tenant, database=database, settings=settings)
def AdminClient(settings: Settings = Settings()) -> AdminAPI:
"""
Creates an admin client that can be used to create tenants and databases.
"""
return AdminClientCreator(settings=settings)