| | from __future__ import annotations |
| |
|
| | import dataclasses |
| | import os |
| | from typing import Any, List |
| |
|
| | import numpy as np |
| | import orjson |
| |
|
| | from autogpt.llm_utils import create_embedding_with_ada |
| | from autogpt.memory.base import MemoryProviderSingleton |
| |
|
| | 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): |
| | """A class that stores the memory in a local file""" |
| |
|
| | def __init__(self, cfg) -> None: |
| | """Initialize a class instance |
| | |
| | Args: |
| | cfg: Config object |
| | |
| | Returns: |
| | 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() |
| | else: |
| | print( |
| | f"Warning: The file '{self.filename}' does not exist. " |
| | "Local memory would not be saved to a file." |
| | ) |
| | 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 = create_embedding_with_ada(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) |
| | 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) -> list[Any] | None: |
| | """ |
| | 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 = create_embedding_with_ada(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) -> tuple[int, tuple[int, ...]]: |
| | """ |
| | Returns: The stats of the local cache. |
| | """ |
| | return len(self.data.texts), self.data.embeddings.shape |
| |
|
