[ENH] Bundle CLI in python package (#3986)

## Description of changes
* Add bindings for the CLI.
* Currently only using them for the "run" command in the python CLI app.

## Test plan
- [x] Tests pass locally with `pytest` for python, `yarn test` for js,
`cargo test` for rust

## Documentation Changes
N/A

Author: itaismith

.github/workflows/_python-tests.yml

@@ -44,6 +44,18 @@ jobs:
         uses: ./.github/actions/python
         with:
           python-version: ${{ matrix.python }}
+      - name: Setup Rust
+        uses: ./.github/actions/rust
+        with:
+          github-token: ${{ github.token }}
+      - name: Build Rust bindings
+        uses: PyO3/maturin-action@v1
+        with:
+          command: build
+          sccache: true
+      - name: Install built wheel
+        shell: bash
+        run: pip install --no-index --find-links target/wheels/ chromadb
       - name: Test
         run: python -m pytest ${{ matrix.parallelized && '-n auto' || '' }} ${{ matrix.test-globs }}
         shell: bash
@@ -125,6 +137,18 @@ jobs:
       uses: actions/checkout@v4
     - name: Set up Python (${{ matrix.python }})
       uses: ./.github/actions/python
+    - name: Setup Rust
+      uses: ./.github/actions/rust
+      with:
+        github-token: ${{ github.token }}
+    - name: Build Rust bindings
+      uses: PyO3/maturin-action@v1
+      with:
+        command: build
+        sccache: true
+    - name: Install built wheel
+      shell: bash
+      run: pip install --no-index --find-links target/wheels/ chromadb
     - name: Integration Test
       run: bin/python-integration-test ${{ matrix.test-globs }}
       shell: bash

Cargo.lock

@@ -75,6 +75,7 @@ chroma-sysdb = { path = "rust/sysdb" }
 chroma-tracing = { path = "rust/tracing" }
 chroma-types = { path = "rust/types" }
 chroma-sqlite = { path = "rust/sqlite" }
+chroma-cli = { path = "rust/cli" }
 mdac = { path = "rust/mdac" }
 wal3 = { path = "rust/wal3" }
 worker = { path = "rust/worker" }

Cargo.toml

@@ -1,5 +1,6 @@
 from typing import Optional
 
+import chromadb_rust_bindings
 from rich.console import Console
 from rich.progress import Progress, SpinnerColumn, TextColumn
 import typer.rich_utils
@@ -20,20 +21,16 @@
 utils_app = typer.Typer(short_help="Use maintenance utilities")
 app.add_typer(utils_app, name="utils")
 
-_logo = """
-                \033[38;5;069m(((((((((    \033[38;5;203m(((((\033[38;5;220m####
-             \033[38;5;069m(((((((((((((\033[38;5;203m(((((((((\033[38;5;220m#########
-           \033[38;5;069m(((((((((((((\033[38;5;203m(((((((((((\033[38;5;220m###########
-         \033[38;5;069m((((((((((((((\033[38;5;203m((((((((((((\033[38;5;220m############
-        \033[38;5;069m(((((((((((((\033[38;5;203m((((((((((((((\033[38;5;220m#############
-        \033[38;5;069m(((((((((((((\033[38;5;203m((((((((((((((\033[38;5;220m#############
-         \033[38;5;069m((((((((((((\033[38;5;203m(((((((((((((\033[38;5;220m##############
-         \033[38;5;069m((((((((((((\033[38;5;203m((((((((((((\033[38;5;220m##############
-           \033[38;5;069m((((((((((\033[38;5;203m(((((((((((\033[38;5;220m#############
-             \033[38;5;069m((((((((\033[38;5;203m((((((((\033[38;5;220m##############
-                \033[38;5;069m(((((\033[38;5;203m((((    \033[38;5;220m#########\033[0m
 
-    """
+def build_cli_args(**kwargs):
+    args = []
+    for key, value in kwargs.items():
+        if isinstance(value, bool):
+            if value:
+                args.append(f"--{key}")
+        elif value is not None:
+            args.extend([f"--{key}", str(value)])
+    return args
 
 
 @app.command()  # type: ignore
@@ -44,54 +41,16 @@ def run(
     host: Annotated[
         Optional[str], typer.Option(help="The host to listen to. Default: localhost")
     ] = "localhost",
-    log_path: Annotated[
-        Optional[str], typer.Option(help="The path to the log file.")
-    ] = "chroma.log",
     port: int = typer.Option(8000, help="The port to run the server on."),
-    test: bool = typer.Option(False, help="Test mode.", show_envvar=False, hidden=True),
 ) -> None:
     """Run a chroma server"""
-    console = Console()
-
-    print("\033[1m")  # Bold logo
-    print(_logo)
-    print("\033[1m")  # Bold
-    print("Running Chroma")
-    print("\033[0m")  # Reset
-
-    console.print(f"[bold]Saving data to:[/bold] [green]{path}[/green]")
-    console.print(
-        f"[bold]Connect to chroma at:[/bold] [green]http://{host}:{port}[/green]"
-    )
-    console.print(
-        "[bold]Getting started guide[/bold]: [blue]https://docs.trychroma.com/getting-started[/blue]\n\n"
-    )
-
-    # set ENV variable for PERSIST_DIRECTORY to path
-    os.environ["IS_PERSISTENT"] = "True"
-    os.environ["PERSIST_DIRECTORY"] = path
-    os.environ["CHROMA_SERVER_NOFILE"] = "65535"
-    os.environ["CHROMA_CLI"] = "True"
-
-    # get the path where chromadb is installed
-    chromadb_path = os.path.dirname(os.path.realpath(__file__))
-
-    # this is the path of the CLI, we want to move up one directory
-    chromadb_path = os.path.dirname(chromadb_path)
-    log_config = set_log_file_path(f"{chromadb_path}/log_config.yml", f"{log_path}")
-    config = {
-        "app": "chromadb.app:app",
-        "host": host,
-        "port": port,
-        "workers": 1,
-        "log_config": log_config,  # Pass the modified log_config dictionary
-        "timeout_keep_alive": 30,
-    }
-
-    if test:
-        return
-
-    uvicorn.run(**config)
+    cli_args = ["chroma", "run"]
+    cli_args.extend(build_cli_args(
+        path=path,
+        host=host,
+        port=port,
+    ))
+    chromadb_rust_bindings.cli(cli_args)
 
 
 @utils_app.command()  # type: ignore

chromadb/cli/cli.py

@@ -1,12 +1,15 @@
 import multiprocessing
 import multiprocessing.context
+import os
+import time
 from multiprocessing.synchronize import Event
 
 from typer.testing import CliRunner
 
+import chromadb
 from chromadb.api.client import Client
 from chromadb.api.models.Collection import Collection
-from chromadb.cli.cli import app
+from chromadb.cli.cli import app, build_cli_args
 from chromadb.cli.utils import set_log_file_path
 from chromadb.config import Settings, System
 from chromadb.db.base import get_sql
@@ -18,21 +21,26 @@
 
 runner = CliRunner()
 
+def start_app(args: list[str]) -> None:
+    runner.invoke(app, args)
 
 def test_app() -> None:
-    result = runner.invoke(
-        app,
-        [
-            "run",
-            "--path",
-            "chroma_test_data",
-            "--port",
-            "8001",
-            "--test",
-        ],
-    )
-    assert "chroma_test_data" in result.stdout
-    assert "8001" in result.stdout
+    kwargs = {"path": "chroma_test_data", "port": 8001}
+    args = ["run"]
+    args.extend(build_cli_args(**kwargs))
+    print(args)
+    server_process = multiprocessing.Process(target=start_app, args=(args,))
+    server_process.start()
+    time.sleep(5)
+
+    host = os.getenv("CHROMA_SERVER_HOST", kwargs.get("host", "localhost"))
+    port = os.getenv("CHROMA_SERVER_HTTP_PORT", kwargs.get("port", 8000))
+
+    client = chromadb.HttpClient(host=host, port=port)
+    heartbeat = client.heartbeat()
+    server_process.terminate()
+    server_process.join()
+    assert heartbeat > 0
 
 
 def test_utils_set_log_file_path() -> None:

chromadb/test/test_cli.py

@@ -0,0 +1,43 @@
+mod commands;
+mod utils;
+
+use crate::commands::run::{run, RunArgs};
+use clap::{Parser, Subcommand};
+
+#[derive(Subcommand, Debug)]
+enum Command {
+    Docs,
+    Run(RunArgs),
+    Support,
+}
+
+#[derive(Parser, Debug)]
+#[command(name = "chroma")]
+#[command(version = "1.0.0")]
+#[command(about = "A CLI for Chroma", long_about = None)]
+struct Cli {
+    #[command(subcommand)]
+    command: Command,
+}
+
+pub fn chroma_cli(args: Vec<String>) {
+    let cli = Cli::parse_from(args);
+
+    match cli.command {
+        Command::Docs => {
+            let url = "https://docs.trychroma.com";
+            if webbrowser::open(url).is_err() {
+                eprintln!("Error: Failed to open the browser. Visit {}.", url);
+            }
+        }
+        Command::Run(args) => {
+            run(args);
+        }
+        Command::Support => {
+            let url = "https://discord.gg/MMeYNTmh3x";
+            if webbrowser::open(url).is_err() {
+                eprintln!("Error: Failed to open the browser. Visit {}.", url);
+            }
+        }
+    }
+}

rust/cli/src/lib.rs

@@ -1,42 +1,9 @@
-mod commands;
-mod utils;
-use crate::commands::run::{run, RunArgs};
-use clap::{Parser, Subcommand};
+#![windows_subsystem = "console"]
 
-#[derive(Subcommand, Debug)]
-enum Command {
-    Docs,
-    Run(RunArgs),
-    Support,
-}
-
-#[derive(Parser, Debug)]
-#[command(name = "chroma")]
-#[command(version = "0.0.1")]
-#[command(about = "A CLI for Chroma", long_about = None)]
-struct Cli {
-    #[command(subcommand)]
-    command: Command,
-}
+use chroma_cli::chroma_cli;
+use std::env;
 
 fn main() {
-    let cli = Cli::parse();
-
-    match cli.command {
-        Command::Docs => {
-            let url = "https://docs.trychroma.com";
-            if webbrowser::open(url).is_err() {
-                eprintln!("Error: Failed to open the browser. Visit {}.", url);
-            }
-        }
-        Command::Run(args) => {
-            run(args);
-        }
-        Command::Support => {
-            let url = "https://discord.gg/MMeYNTmh3x";
-            if webbrowser::open(url).is_err() {
-                eprintln!("Error: Failed to open the browser. Visit {}.", url);
-            }
-        }
-    }
+    let args: Vec<String> = env::args().collect();
+    chroma_cli(args)
 }

rust/cli/src/main.rs

@@ -28,5 +28,6 @@ chroma-cache = { workspace = true }
 chroma-system = { workspace = true }
 chroma-types = { workspace = true, features = ["pyo3"] }
 chroma-error = { workspace = true, features = ["validator"] }
+chroma-cli = { workspace = true }
 mdac = { workspace = true }
 

rust/python_bindings/Cargo.toml

@@ -1,5 +1,6 @@
 use crate::errors::{ChromaPyResult, WrappedPyErr, WrappedSerdeJsonError, WrappedUuidError};
 use chroma_cache::FoyerCacheConfig;
+use chroma_cli::chroma_cli;
 use chroma_config::{registry::Registry, Configurable};
 use chroma_frontend::{
     executor::config::{ExecutorConfig, LocalExecutorConfig},
@@ -22,7 +23,9 @@ use chroma_types::{
     ListCollectionsRequest, ListDatabasesRequest, Metadata, QueryResponse, UpdateCollectionRequest,
     UpdateMetadata,
 };
-use pyo3::{exceptions::PyValueError, pyclass, pymethods, types::PyAnyMethods, PyObject, Python};
+use pyo3::{
+    exceptions::PyValueError, pyclass, pyfunction, pymethods, types::PyAnyMethods, PyObject, Python,
+};
 use std::time::SystemTime;
 
 const DEFAULT_DATABASE: &str = "default_database";
@@ -48,6 +51,20 @@ impl PythonBindingsConfig {
     }
 }
 
+#[pyfunction]
+#[pyo3(signature = (py_args=None))]
+#[allow(dead_code)]
+pub fn cli(py_args: Option<Vec<String>>) -> ChromaPyResult<()> {
+    let args = py_args.unwrap_or_else(|| std::env::args().collect());
+    let args = if args.is_empty() {
+        vec!["chroma".to_string()]
+    } else {
+        args
+    };
+    chroma_cli(args);
+    Ok(())
+}
+
 //////////////////////// PyMethods Implementation ////////////////////////
 #[pymethods]
 impl Bindings {

rust/python_bindings/src/bindings.rs

@@ -1,6 +1,7 @@
 mod bindings;
 mod errors;
 
+use crate::bindings::cli;
 use bindings::{Bindings, PythonBindingsConfig};
 use pyo3::prelude::*;
 
@@ -21,6 +22,8 @@ fn chromadb_rust_bindings(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<MigrationMode>()?;
     m.add_class::<MigrationHash>()?;
 
+    m.add_function(wrap_pyfunction!(cli, m)?)?;
+
     // Log config classes
     // TODO
     Ok(())