Zalgorithm

FastAPI with Chroma

Notes related to using FastAPI with Chromadb.

Create a single AsyncHttpClient for the lifetime of the app

Related to the GitHub issue at https://github.com/chroma-core/chroma/issues/4296.

The Chroma HttpClient and AsyncHttpClient classes don’t provide an explicit close() method. Closing the client is handled in the background, via the httpx library (I think).

It seems that the expected pattern is to create a client for the entire lifetime of the app: https://github.com/chroma-core/chroma/issues/4296#issuecomment-2957457615.

FastAPI has a (deprecated) on_event method (from the FastAPI docs):

@app.on_event("startup")
async def startup():
    await database.connect()


@app.on_event("shutdown")
async def shutdown():
    await database.disconnect()

But on_event seems to be deprecated in favor of asynccontextmanager. Decorating a function with @asynccontextmanager turns it into an async context manager. Note that the lifespan function gets passed to app = FastAPI(lifespan=lifespan):

from contextlib import asynccontextmanager

from fastapi import FastAPI


def fake_answer_to_everything_ml_model(x: float):
    return x * 42


ml_models = {}


@asynccontextmanager
async def lifespan(app: FastAPI):
    # Load the ML model
    ml_models["answer_to_everything"] = fake_answer_to_everything_ml_model
    yield
    # Clean up the ML models and release the resources
    ml_models.clear()


app = FastAPI(lifespan=lifespan)

For my use case, I want to create a single AsyncHttpClient for the lifetime of the application:

collection_name = "zalgorithm"
chroma_client: AsyncClientAPI


@asynccontextmanager
async def lifespan(app: FastAPI):
    # startup: create a single Chroma client for entire app lifetime
    global chroma_client
    chroma_client = await chromadb.AsyncHttpClient(
        host=chroma_host, port=int(chroma_port)
    )

    yield


app = FastAPI(lifespan=lifespan)

# ...

async def get_db_connection() -> aiosqlite.Connection | Any:
    db = await aiosqlite.connect(db_path)
    db.row_factory = aiosqlite.Row

    try:
        yield db
    finally:
        await db.close()


async def get_chroma_collection() -> AsyncCollection:
    collection = await chroma_client.get_collection(name=collection_name)
    return collection


@app.post("/query", response_class=HTMLResponse)
async def query_collection(
    query: Annotated[str, Form()],
    db: aiosqlite.Connection = Depends(get_db_connection, scope="function"),
    collection: AsyncCollection = Depends(get_chroma_collection),
):
Tags: