From 0a48a97bf611106650649a9a77d0c2ee6b0af299 Mon Sep 17 00:00:00 2001
From: Kye <kye@apacmediasolutions.com>
Date: Tue, 26 Mar 2024 11:30:32 -0700
Subject: [PATCH] [CLEANUP]

---
 swarms/agents/omni_modal_agent.py | 3 ++-
 swarms/agents/worker_agent.py     | 9 ++-------
 swarms/memory/chroma_db.py        | 3 ++-
 swarms/memory/lanchain_chroma.py  | 6 +++---
 swarms/memory/pg.py               | 7 ++++---
 swarms/memory/qdrant.py           | 3 ++-
 swarms/structs/auto_swarm.py      | 5 ++++-
 7 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/swarms/agents/omni_modal_agent.py b/swarms/agents/omni_modal_agent.py
index 4af03906..90d0ce28 100644
--- a/swarms/agents/omni_modal_agent.py
+++ b/swarms/agents/omni_modal_agent.py
@@ -9,8 +9,9 @@ from langchain_experimental.autonomous_agents.hugginggpt.task_planner import (
     load_chat_planner,
 )
 from transformers import load_tool
-from swarms.utils.loguru_logger import logger
+
 from swarms.structs.agent import Agent
+from swarms.utils.loguru_logger import logger
 
 
 class OmniModalAgent(Agent):
diff --git a/swarms/agents/worker_agent.py b/swarms/agents/worker_agent.py
index 12760318..f9ba2094 100644
--- a/swarms/agents/worker_agent.py
+++ b/swarms/agents/worker_agent.py
@@ -1,18 +1,13 @@
-import os
 from typing import List
 
-# import faiss
-from langchain.docstore import InMemoryDocstore
-from langchain.embeddings import OpenAIEmbeddings
-
-# from langchain.vectorstores import FAISS
 from langchain_experimental.autonomous_agents import AutoGPT
 
+from swarms.structs.agent import Agent
 from swarms.tools.tool import BaseTool
 from swarms.utils.decorators import error_decorator, timing_decorator
 
 
-class Worker:
+class Worker(Agent):
     """
     The Worker class represents an autonomous agent that can perform tassks through
     function calls or by running a chat.
diff --git a/swarms/memory/chroma_db.py b/swarms/memory/chroma_db.py
index 0ef34286..033be6f6 100644
--- a/swarms/memory/chroma_db.py
+++ b/swarms/memory/chroma_db.py
@@ -9,13 +9,14 @@ from dotenv import load_dotenv
 
 from swarms.utils.data_to_text import data_to_text
 from swarms.utils.markdown_message import display_markdown_message
+from swarms.memory.base_vectordb import AbstractVectorDatabase
 
 # Load environment variables
 load_dotenv()
 
 
 # Results storage using local ChromaDB
-class ChromaDB:
+class ChromaDB(AbstractVectorDatabase):
     """
 
     ChromaDB database
diff --git a/swarms/memory/lanchain_chroma.py b/swarms/memory/lanchain_chroma.py
index 95a2e9e3..cd5d832a 100644
--- a/swarms/memory/lanchain_chroma.py
+++ b/swarms/memory/lanchain_chroma.py
@@ -6,8 +6,8 @@ from langchain.chains.question_answering import load_qa_chain
 from langchain.embeddings.openai import OpenAIEmbeddings
 from langchain.text_splitter import CharacterTextSplitter
 from langchain.vectorstores import Chroma
-
-from swarms.models.openai_models import OpenAIChat
+from swarms.models.popular_llms import OpenAIChat
+from swarms.memory.base_vectordb import AbstractVectorDatabase
 
 
 def synchronized_mem(method):
@@ -31,7 +31,7 @@ def synchronized_mem(method):
     return wrapper
 
 
-class LangchainChromaVectorMemory:
+class LangchainChromaVectorMemory(AbstractVectorDatabase):
     """
     A class representing a vector memory for storing and retrieving text entries.
 
diff --git a/swarms/memory/pg.py b/swarms/memory/pg.py
index b04beacf..e0bc72d2 100644
--- a/swarms/memory/pg.py
+++ b/swarms/memory/pg.py
@@ -5,9 +5,10 @@ from sqlalchemy import JSON, Column, String, create_engine
 from sqlalchemy.dialects.postgresql import UUID
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import Session
+from swarms.memory.base_vectordb import AbstractVectorDatabase
 
 
-class PostgresDB:
+class PostgresDB(AbstractVectorDatabase):
     """
     A class representing a Postgres database.
 
@@ -67,7 +68,7 @@ class PostgresDB:
 
         return VectorModel
 
-    def add_or_update_vector(
+    def add(
         self,
         vector: str,
         vector_id: Optional[str] = None,
@@ -97,7 +98,7 @@ class PostgresDB:
         except Exception as e:
             print(f"Error adding or updating vector: {e}")
 
-    def query_vectors(
+    def query(
         self, query: Any, namespace: Optional[str] = None
     ) -> List[Any]:
         """
diff --git a/swarms/memory/qdrant.py b/swarms/memory/qdrant.py
index 0a553a16..49c5ef62 100644
--- a/swarms/memory/qdrant.py
+++ b/swarms/memory/qdrant.py
@@ -1,6 +1,7 @@
 from typing import List
 
 from httpx import RequestError
+from swarms.memory.base_vectordb import AbstractVectorDatabase
 
 try:
     from sentence_transformers import SentenceTransformer
@@ -20,7 +21,7 @@ except ImportError:
     print("pip install qdrant-client")
 
 
-class Qdrant:
+class Qdrant(AbstractVectorDatabase):
     """
     Qdrant class for managing collections and performing vector operations using QdrantClient.
 
diff --git a/swarms/structs/auto_swarm.py b/swarms/structs/auto_swarm.py
index a0d0f05c..f1c61f0f 100644
--- a/swarms/structs/auto_swarm.py
+++ b/swarms/structs/auto_swarm.py
@@ -25,7 +25,6 @@ class SequentialAccountingSwarm(AbstractSwarm):
             Run the swarm simulation.
 
     """
-
     def __init__(
         self,
         name: Optional[str] = "kyegomez/sequential-accounting-swarm",
@@ -35,6 +34,8 @@ class SequentialAccountingSwarm(AbstractSwarm):
         iters: Optional[int] = 100,
         max_agents: Optional[int] = 100,
         agents: Sequence[AbstractLLM] = None,
+        *args,
+        **kwargs,
     ):
         super().__init__()
         self.name = name
@@ -98,6 +99,8 @@ class AutoSwarmRouter(AbstractSwarm):
         custom_preprocess: Optional[Callable] = None,
         custom_postprocess: Optional[Callable] = None,
         custom_router: Optional[Callable] = None,
+        *args,
+        **kwargs,
     ):
         super().__init__()
         self.name = name