|
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 |
|
|
