diff --git a/README.md b/README.md index 03578bf..621596a 100644 --- a/README.md +++ b/README.md @@ -130,6 +130,8 @@ cd SillyTavern-extras | `--keyphrase-model` | Load a custom key phrase extraction model.
Expects a HuggingFace model ID.
Default: [ml6team/keyphrase-extraction-distilbert-inspec](https://huggingface.co/ml6team/keyphrase-extraction-distilbert-inspec) | | `--prompt-model` | Load a custom prompt generation model.
Expects a HuggingFace model ID.
Default: [FredZhang7/anime-anything-promptgen-v2](https://huggingface.co/FredZhang7/anime-anything-promptgen-v2) | | `--embedding-model` | Load a custom text embedding model.
Expects a HuggingFace model ID.
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.
Default: `8000` | | `--sd-model` | Load a custom Stable Diffusion image generation model.
Expects a HuggingFace model ID.
Default: [ckpt/anything-v4.5-vae-swapped](https://huggingface.co/ckpt/anything-v4.5-vae-swapped)
*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.
**SLOW!** | | `--sd-remote` | Use a remote SD backend.
**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
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 ` +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=`) along with the `--chroma-host` argument. + ## API Endpoints ### Get active list `GET /api/modules` diff --git a/constants.py b/constants.py index 13ef5f6..154321f 100644 --- a/constants.py +++ b/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'] diff --git a/server.py b/server.py index fd4dc8f..091aca6 100644 --- a/server.py +++ b/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__)