Use Chroma to augment and make LocalCache memory/embeddings pluggable

README.md

@@ -166,7 +166,27 @@ export CUSTOM_SEARCH_ENGINE_ID="YOUR_CUSTOM_SEARCH_ENGINE_ID"
 
 ```
 
-## Redis Setup
+## Types of Memory
+
+AutoGPT comes with pluggable memory systems. By default it stores things in memory ("LocalCache") or disk (if a `memory_directory` is specified), but can be persisted to Pinecone (cloud hosted) or Redis (open source).
+
+### Local Setup
+
+By default, AutoGPT uses Chroma, which runs in-memory, with no setup needed.
+
+If you would like to persist Chroma to disk, you can specify your path to a memory directory (from root).
+
+```
+MEMORY_DIRECTORY=/path/to/persistence/directory
+```
+
+By default, AutoGPT also uses [Sentence Transformers](https://docs.trychroma.com/embeddings#default-sentence-transformers) to embed your data and queries (which doesn't cost anything), but you can also specify your choice of [OpenAI embedding model](https://platform.openai.com/docs/guides/embeddings/what-are-embeddings) (will use Azure's version instead if you set `USE_AZURE` env variable):
+
+```
+OPENAI_EMBEDDINGS_MODEL=text-embedding-ada-002
+```
+
+### Redis Setup
 
 Install docker desktop.
 
@@ -203,7 +223,7 @@ You can specify the memory index for redis using the following:
 MEMORY_INDEX=whatever
 ```
 
-## 🌲 Pinecone API Key Setup
+### Pinecone API Key Setup
 
 Pinecone enables the storage of vast amounts of vector-based memory, allowing for only relevant memories to be loaded for the agent at any given time.
 

requirements.txt

@@ -13,6 +13,7 @@ duckduckgo-search
 google-api-python-client #(https://developers.google.com/custom-search/v1/overview)
 pinecone-client==2.2.1
 redis
+chromadb
 orjson
 Pillow
 coverage

scripts/config.py

@@ -43,6 +43,7 @@ def __init__(self):
         self.smart_token_limit = int(os.getenv("SMART_TOKEN_LIMIT", 8000))
 
         self.openai_api_key = os.getenv("OPENAI_API_KEY")
+        self.openai_embeddings_model = os.getenv("OPENAI_EMBEDDINGS_MODEL")
         self.use_azure = False
         self.use_azure = os.getenv("USE_AZURE") == 'True'
         if self.use_azure:
@@ -65,6 +66,8 @@ def __init__(self):
         self.google_api_key = os.getenv("GOOGLE_API_KEY")
         self.custom_search_engine_id = os.getenv("CUSTOM_SEARCH_ENGINE_ID")
 
+        self.memory_directory = os.getenv("MEMORY_DIRECTORY")
+
         self.pinecone_api_key = os.getenv("PINECONE_API_KEY")
         self.pinecone_region = os.getenv("PINECONE_ENV")
 

scripts/memory/base.py

@@ -4,14 +4,12 @@
 import openai
 cfg = Config()
 
-
-def get_ada_embedding(text):
+def get_embedding(text, model_name="text-embedding-ada-002"):
     text = text.replace("\n", " ")
     if cfg.use_azure:
-        return openai.Embedding.create(input=[text], engine=cfg.azure_embeddigs_deployment_id, model="text-embedding-ada-002")["data"][0]["embedding"]
+        return openai.Embedding.create(input=[text], engine=cfg.azure_embeddigs_deployment_id, model=model_name)["data"][0]["embedding"]
     else:
-        return openai.Embedding.create(input=[text], model="text-embedding-ada-002")["data"][0]["embedding"]
-
+        return openai.Embedding.create(input=[text], model=model_name)["data"][0]["embedding"]
 
 class MemoryProviderSingleton(AbstractSingleton):
     @abc.abstractmethod

scripts/memory/local.py

@@ -1,90 +1,56 @@
-import dataclasses
-import orjson
 from typing import Any, List, Optional
-import numpy as np
 import os
-from memory.base import MemoryProviderSingleton, get_ada_embedding
-
-
-EMBED_DIM = 1536
-SAVE_OPTIONS = orjson.OPT_SERIALIZE_NUMPY | orjson.OPT_SERIALIZE_DATACLASS
-
-
-def create_default_embeddings():
-    return np.zeros((0, EMBED_DIM)).astype(np.float32)
-
-
-@dataclasses.dataclass
-class CacheContent:
-    texts: List[str] = dataclasses.field(default_factory=list)
-    embeddings: np.ndarray = dataclasses.field(
-        default_factory=create_default_embeddings
-    )
-
+import uuid
+import datetime
+import chromadb
+from memory.base import MemoryProviderSingleton, get_embedding
+from chromadb.errors import NoIndexException
 
 class LocalCache(MemoryProviderSingleton):
 
     # on load, load our database
     def __init__(self, cfg) -> None:
-        self.filename = f"{cfg.memory_index}.json"
-        if os.path.exists(self.filename):
-            try:
-                with open(self.filename, 'w+b') as f:
-                    file_content = f.read()
-                    if not file_content.strip():
-                        file_content = b'{}'
-                        f.write(file_content)
-
-                    loaded = orjson.loads(file_content)
-                    self.data = CacheContent(**loaded)
-            except orjson.JSONDecodeError:
-                print(f"Error: The file '{self.filename}' is not in JSON format.")
-                self.data = CacheContent()
+        self.persistence = cfg.memory_directory
+        if self.persistence is not None and os.path.exists(self.persistence):
+            self.chromaClient = chromadb.Client(Settings(
+                chroma_db_impl="duckdb+parquet", # duckdb+parquet = persisted, duckdb = in-memory
+                persist_directory=self.persistence
+            ))
         else:
-            print(f"Warning: The file '{self.filename}' does not exist. Local memory would not be saved to a file.")
-            self.data = CacheContent()
+            # in memory
+            self.chromaClient = chromadb.Client()
+        self.chromaCollection = self.chromaClient.create_collection(name="autogpt")
+        # we will key off of cfg.openai_embeddings_model to determine if using sentence transformers or openai embeddings
+        self.useOpenAIEmbeddings = True if (cfg.openai_embeddings_model) else False
 
     def add(self, text: str):
-        """
-        Add text to our list of texts, add embedding as row to our
-            embeddings-matrix
-
-        Args:
-            text: str
-
-        Returns: None
-        """
-        if 'Command Error:' in text:
-            return ""
-        self.data.texts.append(text)
-
-        embedding = get_ada_embedding(text)
-
-        vector = np.array(embedding).astype(np.float32)
-        vector = vector[np.newaxis, :]
-        self.data.embeddings = np.concatenate(
-            [
-                self.data.embeddings,
-                vector,
-            ],
-            axis=0,
-        )
-
-        with open(self.filename, 'wb') as f:
-            out = orjson.dumps(
-                self.data,
-                option=SAVE_OPTIONS
+        current_time = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M:%S.%f")
+        metadata = {"time_added": current_time}
+        if self.useOpenAIEmbeddings:
+            embeddings = get_embedding(text)
+            self.chromaCollection.add(
+                embeddings=[embeddings],
+                ids=[str(uuid.uuid4())],
+                metadatas=[metadata]
+            )
+        else:
+            self.chromaCollection.add(
+                documents=[text],
+                ids=[str(uuid.uuid4())],
+                metadatas=[metadata]
             )
-            f.write(out)
         return text
 
     def clear(self) -> str:
         """
-        Clears the redis server.
+        Resets the Chroma database.
 
-        Returns: A message indicating that the memory has been cleared.
+        Returns: A message indicating that the db has been cleared.
         """
-        self.data = CacheContent()
+
+        chroma_client = self.chromaClient
+        chroma_client.reset()
+        self.chromaCollection = chroma_client.create_collection(name="autogpt")
         return "Obliviated"
 
     def get(self, data: str) -> Optional[List[Any]]:
@@ -96,29 +62,41 @@ def get(self, data: str) -> Optional[List[Any]]:
 
         Returns: The most relevant data.
         """
-        return self.get_relevant(data, 1)
+        results = None
+        if self.useOpenAIEmbeddings:
+            embeddings = get_embedding(data)
+            results = self.chromaCollection.query(
+                query_embeddings=[data],
+                n_results=1
+            )
+        else:
+            results = self.chromaCollection.query(
+                query_texts=[data],
+                n_results=1
+            )
+        return results
 
     def get_relevant(self, text: str, k: int) -> List[Any]:
-        """"
-        matrix-vector mult to find score-for-each-row-of-matrix
-         get indices for top-k winning scores
-         return texts for those indices
-        Args:
-            text: str
-            k: int
-
-        Returns: List[str]
-        """
-        embedding = get_ada_embedding(text)
-
-        scores = np.dot(self.data.embeddings, embedding)
-
-        top_k_indices = np.argsort(scores)[-k:][::-1]
-
-        return [self.data.texts[i] for i in top_k_indices]
+        results = None
+        try: 
+            if self.useOpenAIEmbeddings:
+                embeddings = get_embedding(text)
+                results = self.chromaCollection.query(
+                    query_embeddings=[text],
+                    n_results=min(k, self.chromaCollection.count())
+                )
+            else:
+                results = self.chromaCollection.query(
+                    query_texts=[text],
+                    n_results=min(k, self.chromaCollection.count())
+                )
+        except NoIndexException:
+            # print("No index found - suppressed because this is a common issue for first-run users")
+            pass
+        return results
 
     def get_stats(self):
         """
-        Returns: The stats of the local cache.
+        Returns: The number of items that have been added to the Chroma collection
         """
-        return len(self.data.texts), self.data.embeddings.shape
+        return self.chromaCollection.count()

scripts/memory/pinecone.py

@@ -1,7 +1,7 @@
 
 import pinecone
 
-from memory.base import MemoryProviderSingleton, get_ada_embedding
+from memory.base import MemoryProviderSingleton, get_embedding
 
 
 class PineconeMemory(MemoryProviderSingleton):
@@ -22,7 +22,7 @@ def __init__(self, cfg):
         self.index = pinecone.Index(table_name)
 
     def add(self, data):
-        vector = get_ada_embedding(data)
+        vector = get_embedding(data)
         # no metadata here. We may wish to change that long term.
         resp = self.index.upsert([(str(self.vec_num), vector, {"raw_text": data})])
         _text = f"Inserting data into memory at index: {self.vec_num}:\n data: {data}"
@@ -42,7 +42,7 @@ def get_relevant(self, data, num_relevant=5):
         :param data: The data to compare to.
         :param num_relevant: The number of relevant data to return. Defaults to 5
         """
-        query_embedding = get_ada_embedding(data)
+        query_embedding = get_embedding(data)
         results = self.index.query(query_embedding, top_k=num_relevant, include_metadata=True)
         sorted_results = sorted(results.matches, key=lambda x: x.score)
         return [str(item['metadata']["raw_text"]) for item in sorted_results]

scripts/memory/redismem.py

@@ -6,7 +6,7 @@
 from redis.commands.search.indexDefinition import IndexDefinition, IndexType
 import numpy as np
 
-from memory.base import MemoryProviderSingleton, get_ada_embedding
+from memory.base import MemoryProviderSingleton, get_embedding
 
 
 SCHEMA = [
@@ -71,7 +71,7 @@ def add(self, data: str) -> str:
         """
         if 'Command Error:' in data:
             return ""
-        vector = get_ada_embedding(data)
+        vector = get_embedding(data)
         vector = np.array(vector).astype(np.float32).tobytes()
         data_dict = {
             b"data": data,
@@ -119,7 +119,7 @@ def get_relevant(
 
         Returns: A list of the most relevant data.
         """
-        query_embedding = get_ada_embedding(data)
+        query_embedding = get_embedding(data)
         base_query = f"*=>[KNN {num_relevant} @embedding $vector AS vector_score]"
         query = Query(base_query).return_fields(
             "data",