Chroma memory

.env.template

@@ -118,6 +118,14 @@ OPENAI_API_KEY=your-openai-api-key
 # MILVUS_SECURE=
 # MILVUS_COLLECTION=autogpt
 
+### CHROMADB
+# CHROMA_DB_DIRECTORY - ChromaDB directory if using local instance
+# CHROMA_SERVER_HOST - ChromaDB server host if using remote instance
+# CHROMA_SERVER_PORT - ChromaDB server port if using remote instance
+CHROMA_SERVER_HOST=
+CHROMA_SERVER_PORT=
+CHROMA_DB_DIRECTORY=your-chroma-db-directory
+
 ################################################################################
 ### IMAGE GENERATION PROVIDER
 ################################################################################

autogpt/config/config.py

@@ -80,7 +80,13 @@ def __init__(self) -> None:
         self.pinecone_api_key = os.getenv("PINECONE_API_KEY")
         self.pinecone_region = os.getenv("PINECONE_ENV")
 
-        self.weaviate_host = os.getenv("WEAVIATE_HOST")
+        #chroma configuration
+        self.chroma_db_directory = os.getenv("CHROMA_DB_DIRECTORY")
+        self.chroma_server_host = os.getenv("CHROMA_SERVER_HOST")
+        self.chroma_server_port = os.getenv("CHROMA_SERVER_PORT")
+
+        #weaviate configuration
+        self.weaviate_host  = os.getenv("WEAVIATE_HOST")
         self.weaviate_port = os.getenv("WEAVIATE_PORT")
         self.weaviate_protocol = os.getenv("WEAVIATE_PROTOCOL", "http")
         self.weaviate_username = os.getenv("WEAVIATE_USERNAME", None)

autogpt/memory/__init__.py

@@ -1,6 +1,7 @@
 from autogpt.logs import logger
 from autogpt.memory.local import LocalCache
 from autogpt.memory.no_memory import NoMemory
+from autogpt.memory.chroma import ChromaMemory
 
 # List of supported memory backends
 # Add a backend to this list if the import attempt is successful
@@ -34,6 +35,14 @@
 except ImportError:
     MilvusMemory = None
 
+try:
+    from autogpt.memory.chroma import ChromaMemory
+
+    supported_memory.append("chroma")
+except ImportError:
+    print("Chroma not installed. Skipping import.")
+    ChromaMemory = None
+
 
 def get_memory(cfg, init=False):
     memory = None
@@ -71,6 +80,16 @@ def get_memory(cfg, init=False):
             )
         else:
             memory = MilvusMemory(cfg)
+    elif cfg.memory_backend == "chroma":
+        if not ChromaMemory:
+            print(
+                "Error: Chroma is not installed. Please install chroma-sdk to"
+                " use Chroma as a memory backend."
+            )
+        else:
+            memory = ChromaMemory(cfg)
+            if init:
+                memory.clear()
     elif cfg.memory_backend == "no_memory":
         memory = NoMemory(cfg)
 

autogpt/memory/chroma.py

@@ -0,0 +1,49 @@
+import chromadb
+from chromadb.config import Settings
+from chromadb.utils import embedding_functions
+import uuid
+import os
+from datetime import datetime
+import random
+from memory.base import MemoryProviderSingleton
+
+
+class ChromaMemory(MemoryProviderSingleton):
+    def __init__(self, cfg):
+        if cfg.chroma_db_directory is not None:
+            self.db = chromadb.Client(Settings(chroma_db_impl="duckdb+parquet", persist_directory=cfg.chroma_db_directory))
+        elif cfg.chroma_server_host is not None:
+            self.db = chromadb.Client(Settings(chroma_db_impl="rest", chroma_server_host=cfg.chroma_server_host, chroma_server_http_port=cfg.chroma_server_port))
+        else:
+            self.db = chromadb.Client(Settings(chroma_db_impl="duckdb+parquet", persist_directory="./chromadb"))
+        self.embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction("paraphrase-MiniLM-L6-v2")
+        self.collection = self.db.get_or_create_collection(name="autogpt", embedding_function=self.embedding_function)
+
+    def add(self, data):
+        _timestamp = (datetime.now() - datetime(1970, 1, 1)).total_seconds()
+        if self.collection.add(documents=[data], metadatas=[{"timestamp": _timestamp}], id=[str(uuid.uuid4())]):
+            return True
+        return False
+
+    def remove(self, data):
+        _deleted_id = self.collection.remove(data=[str(uuid.uuid4())])
+        _result = f"Deleted {_deleted_id}"
+        return _result
+
+    def get(self, data):
+        return self.collection.query(query_texts=[data], n_results=1)
+
+    def get_relevant(self, data, num_relevant=5):
+        try:
+            results = self.collection.query(query_texts=[data], n_results=num_relevant)
+        except Exception as e:
+            print(f"Error querying ChromaDB: {e}")
+            return None
+        return results
+
+    def clear(self):
+        self.db.reset()
+        return "Cleared"
+
+    def get_stats(self):
+        return self.collection.count()

requirements.txt

@@ -4,6 +4,7 @@ distro==1.8.0
 openai==0.27.2
 playsound==1.2.2
 python-dotenv==1.0.0
+chromadb==0.3.21
 pyyaml==6.0
 readability-lxml==0.8.1
 requests