Chroma memory

.env.template

@@ -17,4 +17,7 @@ OPENAI_AZURE_EMBEDDINGS_DEPLOYMENT_ID=deployment-id-for-azure-embeddigs
 IMAGE_PROVIDER=dalle
 HUGGINGFACE_API_TOKEN=
 USE_MAC_OS_TTS=False
+CHROMA_SERVER_HOST=
+CHROMA_SERVER_PORT=
+CHROMA_DB_DIRECTORY=your-chroma-db-directory
 MEMORY_BACKEND=local

requirements.txt

@@ -3,6 +3,7 @@ colorama==0.4.6
 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

scripts/config.py

@@ -42,6 +42,10 @@ def __init__(self):
         self.fast_token_limit = int(os.getenv("FAST_TOKEN_LIMIT", 4000))
         self.smart_token_limit = int(os.getenv("SMART_TOKEN_LIMIT", 8000))
 
+        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", 8000)
+
         self.openai_api_key = os.getenv("OPENAI_API_KEY")
         self.use_azure = False
         self.use_azure = os.getenv("USE_AZURE") == 'True'
@@ -144,3 +148,15 @@ def set_pinecone_region(self, value: str):
     def set_debug_mode(self, value: bool):
         """Set the debug mode value."""
         self.debug_mode = value
+
+    def set_chroma_db_directory(self, value: str):
+        """Set the Chroma DB directory value."""
+        self.chroma_db_directory = value
+
+    def set_chroma_server_host(self, value: str):
+        """Set the Chroma server host value."""
+        self.chroma_server_host = value
+
+    def set_chroma_server_port(self, value: str):
+        """Set the Chroma server port value."""
+        self.chroma_server_port = value

scripts/memory/__init__.py

@@ -17,6 +17,12 @@
 except ImportError:
     print("Pinecone not installed. Skipping import.")
     PineconeMemory = None
+
+try:
+    from memory.chroma import ChromaMemory
+except ImportError:
+    print("Chroma not installed. Skipping import.")
+    ChromaMemory = None
 
 def get_memory(cfg, init=False):
     memory = None
@@ -34,6 +40,14 @@ def get_memory(cfg, init=False):
                   " use Redis as a memory backend.")
         else:
             memory = RedisMemory(cfg)
+    elif cfg.memory_backend == "chroma":
+        if not ChromaMemory:
+            print("Error: Chroma is not installed. Please install chroma to"
+                  " use Chroma as a memory backend.")
+        else:
+            memory = ChromaMemory(cfg)
+            if init:
+                memory.clear()
 
     if memory is None:
         memory = LocalCache(cfg)
@@ -49,4 +63,5 @@ def get_supported_memory_backends():
     "LocalCache",
     "RedisMemory",
     "PineconeMemory",
+    "ChromaMemory"
 ]

scripts/memory/chroma.py

@@ -0,0 +1,48 @@
+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()