Merge pull request #39 from bdashore3/main

Add support for ChromaDB instances on the local network

README.md

@@ -130,6 +130,8 @@ cd SillyTavern-extras
 | `--keyphrase-model`      | Load a custom key phrase extraction model.<br>Expects a HuggingFace model ID.<br>Default: [ml6team/keyphrase-extraction-distilbert-inspec](https://huggingface.co/ml6team/keyphrase-extraction-distilbert-inspec) |
 | `--prompt-model`         | Load a custom prompt generation model.<br>Expects a HuggingFace model ID.<br>Default: [FredZhang7/anime-anything-promptgen-v2](https://huggingface.co/FredZhang7/anime-anything-promptgen-v2) |
 | `--embedding-model`      | Load a custom text embedding model.<br>Expects a HuggingFace model ID.<br>Default: [sentence-transformers/all-mpnet-base-v2](https://huggingface.co/sentence-transformers/all-mpnet-base-v2) |
+| `--chroma-host`          | Specifies a host IP for a remote ChromaDB server. |
+| `--chroma-port`          | Specifies an HTTP port for a remote ChromaDB server.<br>Default: `8000` |
 | `--sd-model`             | Load a custom Stable Diffusion image generation model.<br>Expects a HuggingFace model ID.<br>Default: [ckpt/anything-v4.5-vae-swapped](https://huggingface.co/ckpt/anything-v4.5-vae-swapped)<br>*Must have VAE pre-baked in PyTorch format or the output will look drab!* |
 | `--sd-cpu`               | Force the Stable Diffusion generation pipeline to run on the CPU.<br>**SLOW!** |
 | `--sd-remote`            | Use a remote SD backend.<br>**Supported APIs: [sd-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui)**  |
@@ -138,6 +140,35 @@ cd SillyTavern-extras
 | `--sd-remote-ssl`        | Use SSL for the remote SD backend<br>Default: **False** |
 | `--sd-remote-auth`       | Specify the `username:password` for the remote SD backend (if required) |
 
+## ChromaDB
+ChromaDB is a blazing fast and open source database that is used for long-term memory when chatting with characters. It can be run in-memory or on a local server on your LAN.
+
+NOTE: You should **NOT** run ChromaDB on a cloud server. There are no methods for authentication (yet), so unless you want to expose an unauthenticated ChromaDB to the world, run this on a local server in your LAN.
+
+### In-memory setup
+Run the extras server with the `chromadb` module enabled.
+
+### Remote setup
+Prerequisites: Docker, Docker compose (make sure you're running in rootless mode with the systemd service enabled if on Linux)
+
+Steps:
+
+1. Run `git clone https://github.com/chroma-core/chroma chromadb` and `cd chromadb`
+2. Run `docker-compose up -d --build` to build ChromaDB. This may take a long time depending on your system
+3. Once the build process is finished, ChromaDB should be running in the background. You can check with the command `docker ps`
+4. On your client machine, specify your local server ip in the `--chroma-host` argument (ex. `--chroma-host=192.168.1.10`)
+  
+
+If you are running ChromaDB on the same machine as SillyTavern, you will have to change the port of one of the services. To do this for ChromaDB:
+
+1. Run `docker ps` to get the container ID and then `docker container stop <container ID>`
+2. Enter the ChromaDB git repository `cd chromadb`
+3. Open `docker-compose.yml` and look for the line starting with `uvicorn chromadb.app:app`
+4. Change the `--port` argument to whatever port you want.
+5. Look for the `ports` category and change the occurrences of `8000` to whatever port you chose in step 4.
+6. Save and exit. Then run `docker-compose up --detach`
+7. On your client machine, make sure to specity the `--chroma-port` argument (ex. `--chroma-port=<your-port-here>`) along with the `--chroma-host` argument.
+
 ## API Endpoints
 ### Get active list
 `GET /api/modules`

constants.py

@@ -11,6 +11,7 @@ DEFAULT_SD_MODEL = "ckpt/anything-v4.5-vae-swapped"
 DEFAULT_EMBEDDING_MODEL = "sentence-transformers/all-mpnet-base-v2"
 DEFAULT_REMOTE_SD_HOST = "127.0.0.1"
 DEFAULT_REMOTE_SD_PORT = 7860
+DEFAULT_CHROMA_PORT = 8000
 SILERO_SAMPLES_PATH = "tts_samples"
 SILERO_SAMPLE_TEXT = "The quick brown fox jumps over the lazy dog"
 # ALL_MODULES = ['caption', 'summarize', 'classify', 'keywords', 'prompt', 'sd']

server.py

@@ -62,6 +62,8 @@ parser.add_argument(
 )
 parser.add_argument("--prompt-model", help="Load a custom prompt generation model")
 parser.add_argument("--embedding-model", help="Load a custom text embedding model")
+parser.add_argument("--chroma-host", help="Host IP for a remote ChromaDB instance")
+parser.add_argument("--chroma-port", help="HTTP port for a remote ChromaDB instance (defaults to 8000)")
 
 sd_group = parser.add_mutually_exclusive_group()
 
@@ -241,12 +243,35 @@ if "chromadb" in modules:
     from chromadb.config import Settings
     from sentence_transformers import SentenceTransformer
 
-    # disable chromadb telemetry
+    # Assume that the user wants in-memory unless a host is specified
+    # Also disable chromadb telemetry
     posthog.capture = lambda *args, **kwargs: None
-    chromadb_client = chromadb.Client(Settings(anonymized_telemetry=False))
+    if args.chroma_host is None:
+        chromadb_client = chromadb.Client(Settings(anonymized_telemetry=False))
+        print("ChromaDB is running in-memory. It will be cleared when the server is restarted!")
+    else:
+        chroma_port=(
+            args.chroma_port if args.chroma_port else DEFAULT_CHROMA_PORT
+        )
+        chromadb_client = chromadb.Client(
+            Settings(
+                anonymized_telemetry=False,
+                chroma_api_impl="rest",
+                chroma_server_host=args.chroma_host,
+                chroma_server_http_port=chroma_port
+            )
+        )
+        print(f"ChromaDB is remotely configured at {args.chroma_host}:{chroma_port}")
+
     chromadb_embedder = SentenceTransformer(embedding_model)
     chromadb_embed_fn = chromadb_embedder.encode
 
+    # Check if the db is connected and running, otherwise tell the user
+    try:
+        chromadb_client.heartbeat()
+        print("Successfully pinged ChromaDB! Your client is successfully connected.")
+    except:
+        print("Could not ping ChromaDB! If you are running remotely, please check your host and port!")
 
 # Flask init
 app = Flask(__name__)