coding-alt/AutoGPT

1from __future__ import annotations2 3import dataclasses4import os5from typing import Any, List6 7import numpy as np8import orjson9 10from autogpt.llm_utils import create_embedding_with_ada11from autogpt.memory.base import MemoryProviderSingleton12 13EMBED_DIM = 153614SAVE_OPTIONS = orjson.OPT_SERIALIZE_NUMPY | orjson.OPT_SERIALIZE_DATACLASS15 16 17def create_default_embeddings():18    return np.zeros((0, EMBED_DIM)).astype(np.float32)19 20 21@dataclasses.dataclass22class CacheContent:23    texts: List[str] = dataclasses.field(default_factory=list)24    embeddings: np.ndarray = dataclasses.field(25        default_factory=create_default_embeddings26    )27 28 29class LocalCache(MemoryProviderSingleton):30    """A class that stores the memory in a local file"""31 32    def __init__(self, cfg) -> None:33        """Initialize a class instance34 35        Args:36            cfg: Config object37 38        Returns:39            None40        """41        self.filename = f"{cfg.memory_index}.json"42        if os.path.exists(self.filename):43            try:44                with open(self.filename, "w+b") as f:45                    file_content = f.read()46                    if not file_content.strip():47                        file_content = b"{}"48                        f.write(file_content)49 50                    loaded = orjson.loads(file_content)51                    self.data = CacheContent(**loaded)52            except orjson.JSONDecodeError:53                print(f"Error: The file '{self.filename}' is not in JSON format.")54                self.data = CacheContent()55        else:56            print(57                f"Warning: The file '{self.filename}' does not exist. "58                "Local memory would not be saved to a file."59            )60            self.data = CacheContent()61 62    def add(self, text: str):63        """64        Add text to our list of texts, add embedding as row to our65            embeddings-matrix66 67        Args:68            text: str69 70        Returns: None71        """72        if "Command Error:" in text:73            return ""74        self.data.texts.append(text)75 76        embedding = create_embedding_with_ada(text)77 78        vector = np.array(embedding).astype(np.float32)79        vector = vector[np.newaxis, :]80        self.data.embeddings = np.concatenate(81            [82                self.data.embeddings,83                vector,84            ],85            axis=0,86        )87 88        with open(self.filename, "wb") as f:89            out = orjson.dumps(self.data, option=SAVE_OPTIONS)90            f.write(out)91        return text92 93    def clear(self) -> str:94        """95        Clears the redis server.96 97        Returns: A message indicating that the memory has been cleared.98        """99        self.data = CacheContent()100        return "Obliviated"101 102    def get(self, data: str) -> list[Any] | None:103        """104        Gets the data from the memory that is most relevant to the given data.105 106        Args:107            data: The data to compare to.108 109        Returns: The most relevant data.110        """111        return self.get_relevant(data, 1)112 113    def get_relevant(self, text: str, k: int) -> list[Any]:114        """ "115        matrix-vector mult to find score-for-each-row-of-matrix116         get indices for top-k winning scores117         return texts for those indices118        Args:119            text: str120            k: int121 122        Returns: List[str]123        """124        embedding = create_embedding_with_ada(text)125 126        scores = np.dot(self.data.embeddings, embedding)127 128        top_k_indices = np.argsort(scores)[-k:][::-1]129 130        return [self.data.texts[i] for i in top_k_indices]131 132    def get_stats(self) -> tuple[int, tuple[int, ...]]:133        """134        Returns: The stats of the local cache.135        """136        return len(self.data.texts), self.data.embeddings.shape137