coding-alt/AutoGPT

1"""Redis memory provider."""2from __future__ import annotations3 4from typing import Any5 6import numpy as np7import redis8from colorama import Fore, Style9from redis.commands.search.field import TextField, VectorField10from redis.commands.search.indexDefinition import IndexDefinition, IndexType11from redis.commands.search.query import Query12 13from autogpt.llm_utils import create_embedding_with_ada14from autogpt.logs import logger15from autogpt.memory.base import MemoryProviderSingleton16 17SCHEMA = [18    TextField("data"),19    VectorField(20        "embedding",21        "HNSW",22        {"TYPE": "FLOAT32", "DIM": 1536, "DISTANCE_METRIC": "COSINE"},23    ),24]25 26 27class RedisMemory(MemoryProviderSingleton):28    def __init__(self, cfg):29        """30        Initializes the Redis memory provider.31 32        Args:33            cfg: The config object.34 35        Returns: None36        """37        redis_host = cfg.redis_host38        redis_port = cfg.redis_port39        redis_password = cfg.redis_password40        self.dimension = 153641        self.redis = redis.Redis(42            host=redis_host,43            port=redis_port,44            password=redis_password,45            db=0,  # Cannot be changed46        )47        self.cfg = cfg48 49        # Check redis connection50        try:51            self.redis.ping()52        except redis.ConnectionError as e:53            logger.typewriter_log(54                "FAILED TO CONNECT TO REDIS",55                Fore.RED,56                Style.BRIGHT + str(e) + Style.RESET_ALL,57            )58            logger.double_check(59                "Please ensure you have setup and configured Redis properly for use. "60                + f"You can check out {Fore.CYAN + Style.BRIGHT}"61                f"https://github.com/Torantulino/Auto-GPT#redis-setup{Style.RESET_ALL}"62                " to ensure you've set up everything correctly."63            )64            exit(1)65 66        if cfg.wipe_redis_on_start:67            self.redis.flushall()68        try:69            self.redis.ft(f"{cfg.memory_index}").create_index(70                fields=SCHEMA,71                definition=IndexDefinition(72                    prefix=[f"{cfg.memory_index}:"], index_type=IndexType.HASH73                ),74            )75        except Exception as e:76            print("Error creating Redis search index: ", e)77        existing_vec_num = self.redis.get(f"{cfg.memory_index}-vec_num")78        self.vec_num = int(existing_vec_num.decode("utf-8")) if existing_vec_num else 079 80    def add(self, data: str) -> str:81        """82        Adds a data point to the memory.83 84        Args:85            data: The data to add.86 87        Returns: Message indicating that the data has been added.88        """89        if "Command Error:" in data:90            return ""91        vector = create_embedding_with_ada(data)92        vector = np.array(vector).astype(np.float32).tobytes()93        data_dict = {b"data": data, "embedding": vector}94        pipe = self.redis.pipeline()95        pipe.hset(f"{self.cfg.memory_index}:{self.vec_num}", mapping=data_dict)96        _text = (97            f"Inserting data into memory at index: {self.vec_num}:\n" f"data: {data}"98        )99        self.vec_num += 1100        pipe.set(f"{self.cfg.memory_index}-vec_num", self.vec_num)101        pipe.execute()102        return _text103 104    def get(self, data: str) -> list[Any] | None:105        """106        Gets the data from the memory that is most relevant to the given data.107 108        Args:109            data: The data to compare to.110 111        Returns: The most relevant data.112        """113        return self.get_relevant(data, 1)114 115    def clear(self) -> str:116        """117        Clears the redis server.118 119        Returns: A message indicating that the memory has been cleared.120        """121        self.redis.flushall()122        return "Obliviated"123 124    def get_relevant(self, data: str, num_relevant: int = 5) -> list[Any] | None:125        """126        Returns all the data in the memory that is relevant to the given data.127        Args:128            data: The data to compare to.129            num_relevant: The number of relevant data to return.130 131        Returns: A list of the most relevant data.132        """133        query_embedding = create_embedding_with_ada(data)134        base_query = f"*=>[KNN {num_relevant} @embedding $vector AS vector_score]"135        query = (136            Query(base_query)137            .return_fields("data", "vector_score")138            .sort_by("vector_score")139            .dialect(2)140        )141        query_vector = np.array(query_embedding).astype(np.float32).tobytes()142 143        try:144            results = self.redis.ft(f"{self.cfg.memory_index}").search(145                query, query_params={"vector": query_vector}146            )147        except Exception as e:148            print("Error calling Redis search: ", e)149            return None150        return [result.data for result in results.docs]151 152    def get_stats(self):153        """154        Returns: The stats of the memory index.155        """156        return self.redis.ft(f"{self.cfg.memory_index}").info()157