chroma-core / chroma-core/chroma
[Feature Request]: Server side embeddings
- Dominant language
- Rust
- Stars
- 29.3k
- Forks
- 2.5k
- Avg merge
- 1d 4h
- Merged PRs (30d)
- 38
Description
### Describe the problem
At work, we have an LLM question-answering application - a "chat your data" kind of thing. We are deploying it in its own Docker container. The dev in charge of this has created an Arch Linux image to run this LLM app. The image is 500MB.
We want to use Chroma to query an embedding store to answer user questions. But the packages required to embed text - sentence-transformers and its pytorch dependency - are large. They crank the size of our image up to around 5GB. This causes problems with how long it takes us to deploy a change to the LLM app. Installing, into Arch Linux, everything needed to run deep neural network operations, takes forever! I don't know the details here, except that a 500MB image is a lot easier to work with than a 5GB one.
Therefore, we'd find it valuable to have a thin client for chroma that does not need to handle the embedding operations. We'd love to have a chroma server that has all the NN installs, and can handle embedding, so that we can keep our LLM app's container small.
I thought the recent release of the thin client would solve this, but it turns out to still be the case that embedding must be done on the client side. Since we embed new user questions at runtime, to enable similarity search, this means we'd still need the large sentence-transformers, etc., packages, in our LLM app, which brings us back to the problem of installing all those packages onto our Arch Linux image, which makes it too big.
The new chroma-client package's `collection.query` method requires an embedding to be passed. Since we get user questions at runtime, we don't have those embeddings client-side.
I'd like it if chroma had an option to embed server-side. This would make it so that our client (LLM app) image could be extremely small, and need know nothing about what an embedding is.
### Describe the proposed solution
Make it so the server-side can embed. In other words, make it so the client side can add, query, delete, create, edit, etc, without needing to have sentence-transformers installed, and does not need to pass pre-existing embeddings. My ideal client package has no knowledge of what an embedding even is.
### Alternatives considered
I've written a wrapper around chroma, that provides the functionality we're looking for:
```
import os
import json
from typing import List
from functools import wraps
from flask import Flask, request, Request
import chromadb
from chromadb.config import Settings
application = app = Flask(__name__)
LOCAL_MODE = True
if LOCAL_MODE:
os.environ["CHROMA_SERVER_HOST"] = "localhost"
os.environ["CHROMA_SERVER_PORT"] = "8000"
os.environ["CHROMA_SERVER_SSL_ENABLED"] = "false"
else:
os.environ["CHROMA_SERVER_HOST"] = "secret-site.com"
os.environ["CHROMA_SERVER_PORT"] = "1234" # actual port hidden
os.environ["CHROMA_SERVER_SSL_ENABLED"] = "true"
settings = Settings(chroma_api_impl="rest",
chroma_server_host=os.environ.get("CHROMA_SERVER_HOST"),
chroma_server_http_port=os.environ.get("CHROMA_SERVER_PORT"),
chroma_server_ssl_enabled=(os.environ.get("CHROMA_SERVER_SSL_ENABLED", "false") == "true"))
client = chromadb.Client(settings)
def _parse_metadatas(metadatas: str) -> List[dict]:
try:
return json.loads(metadatas)
except json.JSONDecodeError:
raise ValueError("metadatas must be a valid JSON string representing a list of dictionaries.")
def _parse_documents(documents: str) -> List[str]:
try:
return json.loads(documents)
except json.JSONDecodeError:
raise ValueError("documents must be a valid JSON string representing a list of strings.")
def get_collection(func):
"""
A decorator that gets the collection from the request form and passes it to the function.
If the collection does not exist, returns a 400 error.
"""
@wraps(func)
def wrapper(*args, **kwargs):
collection_name = request.form.get("collection_name")
if collection_name is None:
return "collection_name must be provided", 400
try:
collection = client.get_collection(collection_name)
except ValueError:
return f"collection {collection_name} does not exist", 400
return func(collection, *args, **kwargs)
return wrapper
@application.route("/query", methods=["POST"])
@get_collection
def query(collection):
"""
:param collection: The collection object to be queried. This is passed by the get_collection decorator.
:request body parameters:
question: The question to query.
n_results: The number of results to return.
collection_name: The name of the collection to query.
"""
question = request.form.get("question")
n_results = request.form.get("n_results")
if n_results is None:
n_results = 12
return collection.query(query_texts=[question], n_results=n_results)
@application.route("/add", methods=["POST"])
@get_collection
def add(collection):
"""
:param collection: The collection object to be queried. This is passed by the get_collection decorator.
:request body parameters:
documents: A JSON string representing a list of strings, each string being a document to add.
metadatas: A JSON string representing a list of dictionaries, each dictionary at index i
being the metadata for the document at index i from the documents list.
collection_name: The name of the collection to add the documents to.
"""
documents = request.form.get("documents")
metadatas = request.form.get("metadatas")
if documents is None:
return "documents and metadatas must both be provided; got None for documents.", 400
if metadatas is None:
return "documents and metadatas must both be provided; got None for metadatas.", 400
documents_list = _parse_documents(documents)
metadatas_list = _parse_metadatas(metadatas)
if len(documents_list) != len(metadatas_list):
return "documents and metadatas must be the same length; " \
f"got {len(documents_list)} documents but {len(metadatas_list)} metadatas.", \
400
ids = [str(i) for i in range(len(documents_list))]
collection.add(documents=documents_list, metadatas=metadatas_list, ids=ids)
return f"successfully added {len(documents_list)} documents & metadatas to collection {collection.name}", 200
@application.route("/count", methods=["POST"])
@get_collection
def count(collection):
"""
:param collection: The collection object to be queried. This is passed by the get_collection decorator.
:request body parameters:
collection_name: The name of the collection to count the documents in.
"""
return collection.count()
@application.route("/delete_collection", methods=["POST"])
def delete():
"""
:request body parameters:
collection_name: The name of the collection to delete.
"""
collection_name = request.form.get("collection_name")
if collection_name is None:
return "collection_name must be provided", 400
try:
client.delete_collection(collection_name)
return f"successfully deleted collection {collection_name}", 200
except IndexError as e:
return f"Failed to delete collection {collection_name} - perhaps it didn't exist? {e}", 204
except Exception as e:
return f"Failed to delete collection {collection_name}: {e}", 500
@application.route("/create_collection", methods=["POST"])
def create_collection():
"""
:request body parameters:
collection_name: The name of the collection to create.
"""
collection_name = request.form.get("collection_name")
if collection_name is None:
return "collection_name must be provided", 400
try:
client.create_collection(collection_name)
return f"successfully created collection {collection_name}", 200
except Exception as e:
return f"Failed to create collection - {e}", 500
if __name__ == "__main__":
application.run(port=5000)
```
an example client call to this server looks like this:
```
response = requests.post(f'{SERVER_URL}/query', data={'collection_name': "my_collection", 'question': "how do I do X?"})
```
### Importance
would make my life easier
### Additional Information
_No response_
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by reviewing the thin client’s collection.query behavior and the server API, using the Flask wrapper and its /query and /add endpoints as the stated behavior reference. The feature is done when clients can add and query documents by text without installing embedding dependencies or passing precomputed embeddings.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- arch-linux, docker, flask, python
- Domain
- ai, api, backend, database
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100