From 072f129d8597a4cff0e22f4c31da9a784b574467 Mon Sep 17 00:00:00 2001
From: Kye <kye@apacmediasolutions.com>
Date: Thu, 14 Dec 2023 11:14:22 -0800
Subject: [PATCH] [CLEANUP OP]

---
 swarms/memory/chroma.py                  |  0
 swarms/memory/chroma_db.py               | 10 +++++++++-
 swarms/memory/vector_store_retrieveer.py |  4 ----
 3 files changed, 9 insertions(+), 5 deletions(-)
 delete mode 100644 swarms/memory/chroma.py
 delete mode 100644 swarms/memory/vector_store_retrieveer.py

diff --git a/swarms/memory/chroma.py b/swarms/memory/chroma.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/swarms/memory/chroma_db.py b/swarms/memory/chroma_db.py
index 52dd7bed..189d6ec2 100644
--- a/swarms/memory/chroma_db.py
+++ b/swarms/memory/chroma_db.py
@@ -61,6 +61,8 @@ class ChromaDB:
         openai_api_key: str = OPENAI_API_KEY,
         top_results_num: int = 3,
         limit_tokens: Optional[int] = 1000,
+        *args,
+        **kwargs
     ):
         self.metric = metric
         self.RESULTS_STORE_NAME = RESULTS_STORE_NAME
@@ -91,7 +93,7 @@ class ChromaDB:
             embedding_function=embedding_function,
         )
 
-    def add(self, task: Dict, result: str, result_id: str):
+    def add(self, task: Dict, result: str, result_id: str, *args, **kwargs):
         """Adds a result to the ChromaDB collection
 
         Args:
@@ -137,6 +139,8 @@ class ChromaDB:
                         "task": task["task_name"],
                         "result": result,
                     },
+                    *args,
+                    **kwargs
                 )
         except Exception as error:
             print(
@@ -146,6 +150,8 @@ class ChromaDB:
     def query(
         self,
         query: str,
+        *args,
+        **kwargs
     ) -> List[dict]:
         """Queries the ChromaDB collection with a query for the top results
 
@@ -164,6 +170,8 @@ class ChromaDB:
                 query_texts=query,
                 n_results=min(self.top_results_num, count),
                 include=["metadatas"],
+                *args,
+                **kwargs
             )
             out = [item["task"] for item in results["metadatas"][0]]
             out = limit_tokens_from_string(
diff --git a/swarms/memory/vector_store_retrieveer.py b/swarms/memory/vector_store_retrieveer.py
deleted file mode 100644
index 5171eb29..00000000
--- a/swarms/memory/vector_store_retrieveer.py
+++ /dev/null
@@ -1,4 +0,0 @@
-"""
-Implement retreiever for vector store
-
-"""