diff --git a/forge/autogpt/sdk/memory/memstore.py b/forge/autogpt/sdk/memory/memstore.py index b53b02b5..a19118b2 100644 --- a/forge/autogpt/sdk/memory/memstore.py +++ b/forge/autogpt/sdk/memory/memstore.py @@ -1,10 +1,47 @@ +import abc import hashlib import chromadb from chromadb.config import Settings -class MemStore: +class MemStore(abc.ABC): + """ + An abstract class that represents a Memory Store + """ + + @abc.abstractmethod + def __init__(self, store_path: str): + pass + + @abc.abstractmethod + def add(self, task_id: str, document: str, metadatas: dict) -> None: + pass + + @abc.abstractmethod + def query( + self, + task_id: str, + query: str, + filters: dict = None, + document_search: dict = None, + ) -> dict: + pass + + @abc.abstractmethod + def get(self, task_id: str, doc_ids: list = None, filters: dict = None) -> dict: + pass + + @abc.abstractmethod + def update(self, task_id: str, doc_ids: list, documents: list, metadatas: list): + pass + + @abc.abstractmethod + def delete(self, task_id: str, doc_id: str): + pass + + +class ChromaMemStore(MemStore): """ A class used to represent a Memory Store """ diff --git a/forge/autogpt/sdk/memory/memstore_test.py b/forge/autogpt/sdk/memory/memstore_test.py index 925982dc..5ade2d46 100644 --- a/forge/autogpt/sdk/memory/memstore_test.py +++ b/forge/autogpt/sdk/memory/memstore_test.py @@ -3,12 +3,12 @@ import shutil import pytest -from autogpt.sdk.memory.memstore import MemStore +from autogpt.sdk.memory.memstore import ChromaMemStore @pytest.fixture def memstore(): - mem = MemStore(".test_mem_store") + mem = ChromaMemStore(".test_mem_store") yield mem shutil.rmtree(".test_mem_store")