Skip to content

Commit 7d412ae

Browse files
jeffchuberjeffchuber
authored
[ENH] initial CLI (chroma-core#1032)
This proposes an initial CLI The CLI is installed when you installed `pip install chromadb`. You then call the CLI with `chroma run --path <persist_dir> --port <port>` where path and port are optional. This also adds `chroma help` and `chroma docs` as convenience links - but I'm open to removing those. To make this easy - I added `typer` (by the author of FastAPI). I'm not sure this is the tool that we want to commit to for a fuller featured CLI, but given the extremely minimal footprint of this - I don't think it's a one way door. <img width="1477" alt="Screenshot 2023-08-23 at 4 59 54 PM" src="https://github.com/chroma-core/chroma/assets/891664/30374228-d303-41e1-8e9e-188b7f8532d4"> *** #### TODO - [x] test in fresh env - i think i need to add `typer` as a req - [ ] consider expanding the test to make sure the service is actually running - [x] hide the test option from the typer UI - [x] linking to a getting started guide could be interesting at the top of the logs
1 parent 8e96730 commit 7d412ae

15 files changed

+175
-41
lines changed

‎.github/workflows/chroma-integration-test.yml

+2-1
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,9 @@ jobs:
1616
matrix:
1717
python: ['3.7']
1818
platform: [ubuntu-latest, windows-latest]
19-
testfile: ["--ignore-glob 'chromadb/test/property/*'",
19+
testfile: ["--ignore-glob 'chromadb/test/property/*' --ignore='chromadb/test/test_cli.py'",
2020
"chromadb/test/property/test_add.py",
21+
"chromadb/test/test_cli.py",
2122
"chromadb/test/property/test_collections.py",
2223
"chromadb/test/property/test_cross_version_persist.py",
2324
"chromadb/test/property/test_embeddings.py",

‎DEVELOP.md

+1-1
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ print(api.heartbeat())
4343

4444
3. With a persistent backend and a small frontend client
4545

46-
Run `docker-compose up -d --build`
46+
Run `chroma run --path /chroma_db_path`
4747
```python
4848
import chromadb
4949
api = chromadb.HttpClient(host="localhost", port="8000")

‎README.md

+1-1
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@
3535
```bash
3636
pip install chromadb # python client
3737
# for javascript, npm install chromadb!
38-
# for client-server mode, docker-compose up -d --build
38+
# for client-server mode, chroma run --path /chroma_db_path
3939
```
4040

4141
The core API is only 4 functions (run our [💡 Google Colab](https://colab.research.google.com/drive/1QEzFyqnoFxq7LUGyP1vzR4iLt9PpCDXv?usp=sharing) or [Replit template](https://replit.com/@swyx/BasicChromaStarter?v=1)):

‎chromadb/cli/__init__.py

Whitespace-only changes.

‎chromadb/cli/cli.py

+87
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
import typer
2+
import uvicorn
3+
import os
4+
import webbrowser
5+
6+
app = typer.Typer()
7+
8+
_logo = """
9+
\033[38;5;069m((((((((( \033[38;5;203m(((((\033[38;5;220m####
10+
\033[38;5;069m(((((((((((((\033[38;5;203m(((((((((\033[38;5;220m#########
11+
\033[38;5;069m(((((((((((((\033[38;5;203m(((((((((((\033[38;5;220m###########
12+
\033[38;5;069m((((((((((((((\033[38;5;203m((((((((((((\033[38;5;220m############
13+
\033[38;5;069m(((((((((((((\033[38;5;203m((((((((((((((\033[38;5;220m#############
14+
\033[38;5;069m(((((((((((((\033[38;5;203m((((((((((((((\033[38;5;220m#############
15+
\033[38;5;069m((((((((((((\033[38;5;203m(((((((((((((\033[38;5;220m##############
16+
\033[38;5;069m((((((((((((\033[38;5;203m((((((((((((\033[38;5;220m##############
17+
\033[38;5;069m((((((((((\033[38;5;203m(((((((((((\033[38;5;220m#############
18+
\033[38;5;069m((((((((\033[38;5;203m((((((((\033[38;5;220m##############
19+
\033[38;5;069m(((((\033[38;5;203m(((( \033[38;5;220m#########\033[0m
20+
21+
"""
22+
23+
24+
@app.command() # type: ignore
25+
def run(
26+
path: str = typer.Option(
27+
"./chroma_data", help="The path to the file or directory."
28+
),
29+
port: int = typer.Option(8000, help="The port to run the server on."),
30+
test: bool = typer.Option(False, help="Test mode.", show_envvar=False, hidden=True),
31+
) -> None:
32+
"""Run a chroma server"""
33+
34+
print("\033[1m") # Bold logo
35+
print(_logo)
36+
print("\033[1m") # Bold
37+
print("Running Chroma")
38+
print("\033[0m") # Reset
39+
40+
typer.echo(f"\033[1mSaving data to\033[0m: \033[32m{path}\033[0m")
41+
typer.echo(
42+
f"\033[1mConnect to chroma at\033[0m: \033[32mhttp://localhost:{port}\033[0m"
43+
)
44+
typer.echo(
45+
"\033[1mGetting started guide\033[0m: https://docs.trychroma.com/getting-started\n\n"
46+
)
47+
48+
# set ENV variable for PERSIST_DIRECTORY to path
49+
os.environ["IS_PERSISTENT"] = "True"
50+
os.environ["PERSIST_DIRECTORY"] = path
51+
52+
# get the path where chromadb is installed
53+
chromadb_path = os.path.dirname(os.path.realpath(__file__))
54+
55+
# this is the path of the CLI, we want to move up one directory
56+
chromadb_path = os.path.dirname(chromadb_path)
57+
58+
config = {
59+
"app": "chromadb.app:app",
60+
"host": "0.0.0.0",
61+
"port": port,
62+
"workers": 1,
63+
"log_config": f"{chromadb_path}/log_config.yml",
64+
}
65+
66+
if test:
67+
return
68+
69+
uvicorn.run(**config)
70+
71+
72+
@app.command() # type: ignore
73+
def help() -> None:
74+
"""Opens help url in your browser"""
75+
76+
webbrowser.open("https://discord.gg/MMeYNTmh3x")
77+
78+
79+
@app.command() # type: ignore
80+
def docs() -> None:
81+
"""Opens docs url in your browser"""
82+
83+
webbrowser.open("https://docs.trychroma.com")
84+
85+
86+
if __name__ == "__main__":
87+
app()

‎chromadb/log_config.yml

+37
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
version: 1
2+
disable_existing_loggers: False
3+
formatters:
4+
default:
5+
"()": uvicorn.logging.DefaultFormatter
6+
format: '%(levelprefix)s [%(asctime)s] %(message)s'
7+
use_colors: null
8+
datefmt: '%d-%m-%Y %H:%M:%S'
9+
access:
10+
"()": uvicorn.logging.AccessFormatter
11+
format: '%(levelprefix)s [%(asctime)s] %(client_addr)s - "%(request_line)s" %(status_code)s'
12+
datefmt: '%d-%m-%Y %H:%M:%S'
13+
handlers:
14+
default:
15+
formatter: default
16+
class: logging.StreamHandler
17+
stream: ext://sys.stderr
18+
access:
19+
formatter: access
20+
class: logging.StreamHandler
21+
stream: ext://sys.stdout
22+
console:
23+
class: logging.StreamHandler
24+
stream: ext://sys.stdout
25+
formatter: default
26+
file:
27+
class : logging.handlers.RotatingFileHandler
28+
filename: chroma.log
29+
formatter: default
30+
loggers:
31+
root:
32+
level: WARN
33+
handlers: [console, file]
34+
chromadb:
35+
level: DEBUG
36+
uvicorn:
37+
level: INFO

‎chromadb/test/test_cli.py

+21
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
from typer.testing import CliRunner
2+
3+
from chromadb.cli.cli import app
4+
5+
runner = CliRunner()
6+
7+
8+
def test_app() -> None:
9+
result = runner.invoke(
10+
app,
11+
[
12+
"run",
13+
"--path",
14+
"chroma_test_data",
15+
"--port",
16+
"8001",
17+
"--test",
18+
],
19+
)
20+
assert "chroma_test_data" in result.stdout
21+
assert "8001" in result.stdout

‎clients/js/package.json

+4-1
Original file line numberDiff line numberDiff line change
@@ -29,14 +29,17 @@
2929
"dist"
3030
],
3131
"scripts": {
32-
"test": "run-s db:clean db:run test:runfull db:clean db:run-auth test:runfull-authonly db:clean",
32+
"test": "run-s db:clean db:cleanauth db:run test:runfull db:clean db:run-auth test:runfull-authonly db:cleanauth",
33+
"testnoauth": "run-s db:clean db:run test:runfull db:clean",
34+
"testauth": "run-s db:cleanauth db:run-auth test:runfull-authonly db:cleanauth",
3335
"test:set-port": "cross-env URL=localhost:8001",
3436
"test:run": "jest --runInBand --testPathIgnorePatterns=test/auth.basic.test.ts",
3537
"test:run-auth": "jest --runInBand --testPathPattern=test/auth.basic.test.ts",
3638
"test:runfull": "PORT=8001 jest --runInBand --testPathIgnorePatterns=test/auth.basic.test.ts",
3739
"test:runfull-authonly": "PORT=8001 jest --runInBand --testPathPattern=test/auth.basic.test.ts",
3840
"test:update": "run-s db:clean db:run && jest --runInBand --updateSnapshot && run-s db:clean",
3941
"db:clean": "cd ../.. && CHROMA_PORT=8001 docker-compose -f docker-compose.test.yml down --volumes",
42+
"db:cleanauth": "cd ../.. && CHROMA_PORT=8001 docker-compose -f docker-compose.test-auth.yml down --volumes",
4043
"db:run": "cd ../.. && CHROMA_PORT=8001 docker-compose -f docker-compose.test.yml up --detach && sleep 5",
4144
"db:run-auth": "cd ../.. && CHROMA_PORT=8001 docker-compose -f docker-compose.test-auth.yml up --detach && sleep 5",
4245
"clean": "rimraf dist",

‎docker-compose.test-auth.yml

+1-1
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ services:
1212
volumes:
1313
- ./:/chroma
1414
- test_index_data:/index_data
15-
command: uvicorn chromadb.app:app --workers 1 --host 0.0.0.0 --port 8000 --log-config log_config.yml
15+
command: uvicorn chromadb.app:app --workers 1 --host 0.0.0.0 --port 8000 --log-config chromadb/log_config.yml
1616
environment:
1717
- ANONYMIZED_TELEMETRY=False
1818
- ALLOW_RESET=True

‎docker-compose.test.yml

+1-1
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ services:
1212
volumes:
1313
- ./:/chroma
1414
- test_index_data:/index_data
15-
command: uvicorn chromadb.app:app --workers 1 --host 0.0.0.0 --port 8000 --log-config log_config.yml
15+
command: uvicorn chromadb.app:app --workers 1 --host 0.0.0.0 --port 8000 --log-config chromadb/log_config.yml
1616
environment:
1717
- ANONYMIZED_TELEMETRY=False
1818
- ALLOW_RESET=True

‎docker-compose.yml

+1-1
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ services:
1414
- ./:/chroma
1515
# Be aware that indexed data are located in "/chroma/chroma/"
1616
# Default configuration for persist_directory in chromadb/config.py
17-
command: uvicorn chromadb.app:app --reload --workers 1 --host 0.0.0.0 --port 8000 --log-config log_config.yml
17+
command: uvicorn chromadb.app:app --reload --workers 1 --host 0.0.0.0 --port 8000 --log-config chromadb/log_config.yml
1818
environment:
1919
- IS_PERSISTENT=TRUE
2020
- CHROMA_SERVER_AUTH_PROVIDER=${CHROMA_SERVER_AUTH_PROVIDER}

‎examples/use_with/cohere/cohere_js.js

+10-10
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ First run Chroma
77
```
88
git clone [email protected]:chroma-core/chroma.git
99
cd chroma
10-
docker-compose up -d --build
10+
chroma run --path /chroma_db_path
1111
```
1212
1313
Then install chroma and cohere
@@ -61,20 +61,20 @@ const main = async () => {
6161
});
6262

6363
// # 나는 오렌지를 좋아한다 is "I like oranges" in Korean
64-
multilingual_texts = [ 'Hello from Cohere!', 'مرحبًا من كوهير!',
65-
'Hallo von Cohere!', 'Bonjour de Cohere!',
66-
'¡Hola desde Cohere!', 'Olá do Cohere!',
67-
'Ciao da Cohere!', '您好，来自 Cohere！',
68-
'कोहेरे से नमस्ते!', '나는 오렌지를 좋아한다' ]
64+
multilingual_texts = ['Hello from Cohere!', 'مرحبًا من كوهير!',
65+
'Hallo von Cohere!', 'Bonjour de Cohere!',
66+
'¡Hola desde Cohere!', 'Olá do Cohere!',
67+
'Ciao da Cohere!', '您好，来自 Cohere！',
68+
'कोहेरे से नमस्ते!', '나는 오렌지를 좋아한다']
6969

7070
let ids = Array.from({ length: multilingual_texts.length }, (_, i) => String(i));
7171

7272
await collection.add({
73-
ids:ids,
74-
documents:multilingual_texts
75-
})
73+
ids: ids,
74+
documents: multilingual_texts
75+
})
7676

77-
console.log(await collection.query({queryTexts:["citrus"], nResults:1}))
77+
console.log(await collection.query({ queryTexts: ["citrus"], nResults: 1 }))
7878

7979
}
8080

‎log_config.yml

-23
This file was deleted.

‎pyproject.toml

+8-1
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,8 @@ dependencies = [
3232
'overrides >= 7.3.1',
3333
'importlib-resources',
3434
'graphlib_backport >= 1.0.3; python_version < "3.9"',
35-
'bcrypt >= 4.0.1'
35+
'bcrypt >= 4.0.1',
36+
'typer >= 0.9.0',
3637
]
3738

3839
[tool.black]
@@ -43,6 +44,9 @@ target-version = ['py36', 'py37', 'py38', 'py39', 'py310']
4344
[tool.pytest.ini_options]
4445
pythonpath = ["."]
4546

47+
[project.scripts]
48+
chroma = "chromadb.cli.cli:app"
49+
4650
[project.urls]
4751
"Homepage" = "https://github.com/chroma-core/chroma"
4852
"Bug Tracker" = "https://github.com/chroma-core/chroma/issues"
@@ -56,3 +60,6 @@ local_scheme="no-local-version"
5660

5761
[tool.setuptools]
5862
packages = ["chromadb"]
63+
64+
[tool.setuptools.package-data]
65+
chromadb = ["*.yml"]

‎requirements.txt

+1
Original file line numberDiff line numberDiff line change
@@ -14,5 +14,6 @@ pypika==0.48.9
1414
requests==2.28.1
1515
tokenizers==0.13.2
1616
tqdm==4.65.0
17+
typer>=0.9.0
1718
typing_extensions==4.5.0
1819
uvicorn[standard]==0.18.3

0 commit comments

Comments
 (0)