1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
|
import abc
import hashlib
import chromadb
from chromadb.config import Settings
class MemStore(abc.ABC):
"""
An abstract class that represents a Memory Store
"""
@abc.abstractmethod
def __init__(self, store_path: str):
"""
Initialize the MemStore with a given store path.
Args:
store_path (str): The path to the store.
"""
pass
@abc.abstractmethod
def add_task_memory(self, task_id: str, document: str, metadatas: dict) -> None:
"""
Add a document to the current tasks MemStore.
This function calls the base version with the task_id as the collection_name.
Args:
task_id (str): The ID of the task.
document (str): The document to be added.
metadatas (dict): The metadata of the document.
"""
self.add(collection_name=task_id, document=document, metadatas=metadatas)
@abc.abstractmethod
def query_task_memory(
self,
task_id: str,
query: str,
filters: dict = None,
document_search: dict = None,
) -> dict:
"""
Query the current tasks MemStore.
This function calls the base version with the task_id as the collection_name.
Args:
task_id (str): The ID of the task.
query (str): The query string.
filters (dict, optional): The filters to be applied. Defaults to None.
document_search (dict, optional): The search string. Defaults to None.
Returns:
dict: The query results.
"""
return self.query(
collection_name=task_id,
query=query,
filters=filters,
document_search=document_search,
)
@abc.abstractmethod
def get_task_memory(
self, task_id: str, doc_ids: list = None, filters: dict = None
) -> dict:
"""
Get documents from the current tasks MemStore.
This function calls the base version with the task_id as the collection_name.
Args:
task_id (str): The ID of the task.
doc_ids (list, optional): The IDs of the documents to be retrieved. Defaults to None.
filters (dict, optional): The filters to be applied. Defaults to None.
Returns:
dict: The retrieved documents.
"""
return self.get(collection_name=task_id, doc_ids=doc_ids, filters=filters)
@abc.abstractmethod
def update_task_memory(
self, task_id: str, doc_ids: list, documents: list, metadatas: list
):
"""
Update documents in the current tasks MemStore.
This function calls the base version with the task_id as the collection_name.
Args:
task_id (str): The ID of the task.
doc_ids (list): The IDs of the documents to be updated.
documents (list): The updated documents.
metadatas (list): The updated metadata.
"""
self.update(
collection_name=task_id,
doc_ids=doc_ids,
documents=documents,
metadatas=metadatas,
)
@abc.abstractmethod
def delete_task_memory(self, task_id: str, doc_id: str):
"""
Delete a document from the current tasks MemStore.
This function calls the base version with the task_id as the collection_name.
Args:
task_id (str): The ID of the task.
doc_id (str): The ID of the document to be deleted.
"""
self.delete(collection_name=task_id, doc_id=doc_id)
@abc.abstractmethod
def add(self, collection_name: str, document: str, metadatas: dict) -> None:
"""
Add a document to the current collection's MemStore.
Args:
collection_name (str): The name of the collection.
document (str): The document to be added.
metadatas (dict): The metadata of the document.
"""
pass
@abc.abstractmethod
def query(
self,
collection_name: str,
query: str,
filters: dict = None,
document_search: dict = None,
) -> dict:
pass
@abc.abstractmethod
def get(
self, collection_name: str, doc_ids: list = None, filters: dict = None
) -> dict:
pass
@abc.abstractmethod
def update(
self, collection_name: str, doc_ids: list, documents: list, metadatas: list
):
pass
@abc.abstractmethod
def delete(self, collection_name: str, doc_id: str):
pass
|
