diff --git a/README.md b/README.md index 5dc477e31f34..760e62cf5962 100644 --- a/README.md +++ b/README.md @@ -141,6 +141,39 @@ export CUSTOM_SEARCH_ENGINE_ID="YOUR_CUSTOM_SEARCH_ENGINE_ID" ``` +## Redis Setup + +Install docker desktop. + +Run: +``` +docker run -d --name redis-stack-server -p 6379:6379 redis/redis-stack-server:latest +``` + +Set the following environment variables: +``` +MEMORY_BACKEND=redis +REDIS_HOST=localhost +REDIS_PORT=6379 +REDIS_PASSWORD= +``` + +Note that this is not intended to be run facing the internet and is not secure, do not expose redis to the internet without a password or at all really. + +You can optionally set + +``` +WIPE_REDIS_ON_START=False +``` + +To persist memory stored in Redis. + +You can specify the memory index for redis using the following: + +```` +MEMORY_INDEX=whatever +```` + ## 🌲 Pinecone API Key Setup Pinecone enable a vector based memory so a vast memory can be stored and only relevant memories diff --git a/requirements.txt b/requirements.txt index 7b1040401bb3..6a9ba6433004 100644 --- a/requirements.txt +++ b/requirements.txt @@ -12,4 +12,6 @@ docker duckduckgo-search google-api-python-client #(https://developers.google.com/custom-search/v1/overview) pinecone-client==2.2.1 +redis +orjson Pillow diff --git a/scripts/commands.py b/scripts/commands.py index a45fb8963fcd..1f255751c05e 100644 --- a/scripts/commands.py +++ b/scripts/commands.py @@ -1,6 +1,6 @@ import browse import json -from memory import PineconeMemory +from memory import get_memory import datetime import agent_manager as agents import speak @@ -53,10 +53,11 @@ def get_command(response): def execute_command(command_name, arguments): - memory = PineconeMemory() + memory = get_memory(cfg) + try: if command_name == "google": - + # Check if the Google API key is set and use the official search method # If the API key is not set or has only whitespaces, use the unofficial search method if cfg.google_api_key and (cfg.google_api_key.strip() if cfg.google_api_key else None): diff --git a/scripts/config.py b/scripts/config.py index 959c3eb226c9..d5f1a3f0660f 100644 --- a/scripts/config.py +++ b/scripts/config.py @@ -1,3 +1,4 @@ +import abc import os import openai from dotenv import load_dotenv @@ -5,7 +6,7 @@ load_dotenv() -class Singleton(type): +class Singleton(abc.ABCMeta, type): """ Singleton metaclass for ensuring only one instance of a class. """ @@ -20,6 +21,10 @@ def __call__(cls, *args, **kwargs): return cls._instances[cls] +class AbstractSingleton(abc.ABC, metaclass=Singleton): + pass + + class Config(metaclass=Singleton): """ Configuration class to store the state of bools for different scripts access. @@ -59,7 +64,14 @@ def __init__(self): # User agent headers to use when browsing web # Some websites might just completely deny request with an error code if no user agent was found. self.user_agent_header = {"User-Agent":"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_4) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/83.0.4103.97 Safari/537.36"} - + self.redis_host = os.getenv("REDIS_HOST", "localhost") + self.redis_port = os.getenv("REDIS_PORT", "6379") + self.redis_password = os.getenv("REDIS_PASSWORD", "") + self.wipe_redis_on_start = os.getenv("WIPE_REDIS_ON_START", "True") == 'True' + self.memory_index = os.getenv("MEMORY_INDEX", 'auto-gpt') + # Note that indexes must be created on db 0 in redis, this is not configureable. + + self.memory_backend = os.getenv("MEMORY_BACKEND", 'local') # Initialize the OpenAI API client openai.api_key = self.openai_api_key diff --git a/scripts/main.py b/scripts/main.py index 17385bf33938..10f9d0dcaa0b 100644 --- a/scripts/main.py +++ b/scripts/main.py @@ -1,7 +1,7 @@ import json import random import commands as cmd -from memory import PineconeMemory +from memory import get_memory import data import chat from colorama import Fore, Style @@ -281,12 +281,9 @@ def parse_arguments(): # Make a constant: user_input = "Determine which next command to use, and respond using the format specified above:" -# raise an exception if pinecone_api_key or region is not provided -if not cfg.pinecone_api_key or not cfg.pinecone_region: raise Exception("Please provide pinecone_api_key and pinecone_region") # Initialize memory and make sure it is empty. # this is particularly important for indexing and referencing pinecone memory -memory = PineconeMemory() -memory.clear() +memory = get_memory(cfg, init=True) print('Using memory of type: ' + memory.__class__.__name__) # Interaction Loop diff --git a/scripts/memory/__init__.py b/scripts/memory/__init__.py new file mode 100644 index 000000000000..a441a46aa94e --- /dev/null +++ b/scripts/memory/__init__.py @@ -0,0 +1,44 @@ +from memory.local import LocalCache +try: + from memory.redismem import RedisMemory +except ImportError: + print("Redis not installed. Skipping import.") + RedisMemory = None + +try: + from memory.pinecone import PineconeMemory +except ImportError: + print("Pinecone not installed. Skipping import.") + PineconeMemory = None + + +def get_memory(cfg, init=False): + memory = None + if cfg.memory_backend == "pinecone": + if not PineconeMemory: + print("Error: Pinecone is not installed. Please install pinecone" + " to use Pinecone as a memory backend.") + else: + memory = PineconeMemory(cfg) + if init: + memory.clear() + elif cfg.memory_backend == "redis": + if not RedisMemory: + print("Error: Redis is not installed. Please install redis-py to" + " use Redis as a memory backend.") + else: + memory = RedisMemory(cfg) + + if memory is None: + memory = LocalCache(cfg) + if init: + memory.clear() + return memory + + +__all__ = [ + "get_memory", + "LocalCache", + "RedisMemory", + "PineconeMemory", +] diff --git a/scripts/memory/base.py b/scripts/memory/base.py new file mode 100644 index 000000000000..d7ab7fcf1f55 --- /dev/null +++ b/scripts/memory/base.py @@ -0,0 +1,31 @@ +"""Base class for memory providers.""" +import abc +from config import AbstractSingleton +import openai + + +def get_ada_embedding(text): + text = text.replace("\n", " ") + return openai.Embedding.create(input=[text], model="text-embedding-ada-002")["data"][0]["embedding"] + + +class MemoryProviderSingleton(AbstractSingleton): + @abc.abstractmethod + def add(self, data): + pass + + @abc.abstractmethod + def get(self, data): + pass + + @abc.abstractmethod + def clear(self): + pass + + @abc.abstractmethod + def get_relevant(self, data, num_relevant=5): + pass + + @abc.abstractmethod + def get_stats(self): + pass diff --git a/scripts/memory/local.py b/scripts/memory/local.py new file mode 100644 index 000000000000..8dc90021ff6e --- /dev/null +++ b/scripts/memory/local.py @@ -0,0 +1,114 @@ +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 + ) + + +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): + with open(self.filename, 'rb') as f: + loaded = orjson.loads(f.read()) + self.data = CacheContent(**loaded) + else: + self.data = CacheContent() + + 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( + [ + vector, + self.data.embeddings, + ], + axis=0, + ) + + with open(self.filename, 'wb') as f: + out = orjson.dumps( + self.data, + option=SAVE_OPTIONS + ) + f.write(out) + return text + + def clear(self) -> str: + """ + Clears the redis server. + + Returns: A message indicating that the memory has been cleared. + """ + self.data = CacheContent() + return "Obliviated" + + def get(self, data: str) -> Optional[List[Any]]: + """ + Gets the data from the memory that is most relevant to the given data. + + Args: + data: The data to compare to. + + Returns: The most relevant data. + """ + return self.get_relevant(data, 1) + + 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] + + def get_stats(self): + """ + Returns: The stats of the local cache. + """ + return len(self.data.texts), self.data.embeddings.shape diff --git a/scripts/memory.py b/scripts/memory/pinecone.py similarity index 80% rename from scripts/memory.py rename to scripts/memory/pinecone.py index 0d265a31d8f4..8e1eaa570fee 100644 --- a/scripts/memory.py +++ b/scripts/memory/pinecone.py @@ -1,21 +1,11 @@ -from config import Config, Singleton -import pinecone -import openai - -cfg = Config() - - -def get_ada_embedding(text): - text = text.replace("\n", " ") - return openai.Embedding.create(input=[text], model="text-embedding-ada-002")["data"][0]["embedding"] +import pinecone -def get_text_from_embedding(embedding): - return openai.Embedding.retrieve(embedding, model="text-embedding-ada-002")["data"][0]["text"] +from memory.base import MemoryProviderSingleton, get_ada_embedding -class PineconeMemory(metaclass=Singleton): - def __init__(self): +class PineconeMemory(MemoryProviderSingleton): + def __init__(self, cfg): pinecone_api_key = cfg.pinecone_api_key pinecone_region = cfg.pinecone_region pinecone.init(api_key=pinecone_api_key, environment=pinecone_region) diff --git a/scripts/memory/redismem.py b/scripts/memory/redismem.py new file mode 100644 index 000000000000..2082fe588764 --- /dev/null +++ b/scripts/memory/redismem.py @@ -0,0 +1,143 @@ +"""Redis memory provider.""" +from typing import Any, List, Optional +import redis +from redis.commands.search.field import VectorField, TextField +from redis.commands.search.query import Query +from redis.commands.search.indexDefinition import IndexDefinition, IndexType +import numpy as np + +from memory.base import MemoryProviderSingleton, get_ada_embedding + + +SCHEMA = [ + TextField("data"), + VectorField( + "embedding", + "HNSW", + { + "TYPE": "FLOAT32", + "DIM": 1536, + "DISTANCE_METRIC": "COSINE" + } + ), +] + + +class RedisMemory(MemoryProviderSingleton): + def __init__(self, cfg): + """ + Initializes the Redis memory provider. + + Args: + cfg: The config object. + + Returns: None + """ + redis_host = cfg.redis_host + redis_port = cfg.redis_port + redis_password = cfg.redis_password + self.dimension = 1536 + self.redis = redis.Redis( + host=redis_host, + port=redis_port, + password=redis_password, + db=0 # Cannot be changed + ) + self.cfg = cfg + if cfg.wipe_redis_on_start: + self.redis.flushall() + try: + self.redis.ft(f"{cfg.memory_index}").create_index( + fields=SCHEMA, + definition=IndexDefinition( + prefix=[f"{cfg.memory_index}:"], + index_type=IndexType.HASH + ) + ) + except Exception as e: + print("Error creating Redis search index: ", e) + existing_vec_num = self.redis.get(f'{cfg.memory_index}-vec_num') + self.vec_num = int(existing_vec_num.decode('utf-8')) if\ + existing_vec_num else 0 + + def add(self, data: str) -> str: + """ + Adds a data point to the memory. + + Args: + data: The data to add. + + Returns: Message indicating that the data has been added. + """ + if 'Command Error:' in data: + return "" + vector = get_ada_embedding(data) + vector = np.array(vector).astype(np.float32).tobytes() + data_dict = { + b"data": data, + "embedding": vector + } + pipe = self.redis.pipeline() + pipe.hset(f"{self.cfg.memory_index}:{self.vec_num}", mapping=data_dict) + _text = f"Inserting data into memory at index: {self.vec_num}:\n"\ + f"data: {data}" + self.vec_num += 1 + pipe.set(f'{self.cfg.memory_index}-vec_num', self.vec_num) + pipe.execute() + return _text + + def get(self, data: str) -> Optional[List[Any]]: + """ + Gets the data from the memory that is most relevant to the given data. + + Args: + data: The data to compare to. + + Returns: The most relevant data. + """ + return self.get_relevant(data, 1) + + def clear(self) -> str: + """ + Clears the redis server. + + Returns: A message indicating that the memory has been cleared. + """ + self.redis.flushall() + return "Obliviated" + + def get_relevant( + self, + data: str, + num_relevant: int = 5 + ) -> Optional[List[Any]]: + """ + Returns all the data in the memory that is relevant to the given data. + Args: + data: The data to compare to. + num_relevant: The number of relevant data to return. + + Returns: A list of the most relevant data. + """ + query_embedding = get_ada_embedding(data) + base_query = f"*=>[KNN {num_relevant} @embedding $vector AS vector_score]" + query = Query(base_query).return_fields( + "data", + "vector_score" + ).sort_by("vector_score").dialect(2) + query_vector = np.array(query_embedding).astype(np.float32).tobytes() + + try: + results = self.redis.ft(f"{self.cfg.memory_index}").search( + query, query_params={"vector": query_vector} + ) + except Exception as e: + print("Error calling Redis search: ", e) + return None + return [result.data for result in results.docs] + + def get_stats(self): + """ + Returns: The stats of the memory index. + """ + return self.redis.ft(f"{self.cfg.memory_index}").info()