diff --git a/requirements.txt b/requirements.txt
index ab46b6d2..042d9a53 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -109,4 +109,9 @@ huggingface_hub
fastapi_cache
fastapi-limiter
-sphinx_rtd_theme
\ No newline at end of file
+sphinx_rtd_theme
+pegasusx
+
+
+
+oceandb
diff --git a/setup.py b/setup.py
index 5d0b009f..3871d75b 100644
--- a/setup.py
+++ b/setup.py
@@ -27,6 +27,8 @@ setup(
'asyncio',
'nest_asyncio',
'bs4',
+ 'pegasusx',
+ 'oceandb',
'playwright',
'duckduckgo_search',
'faiss-cpu',
diff --git a/swarms/agents/memory/ocean.py b/swarms/agents/memory/ocean.py
index 744300bc..6b5587e5 100644
--- a/swarms/agents/memory/ocean.py
+++ b/swarms/agents/memory/ocean.py
@@ -1,3 +1,38 @@
#init ocean
# TODO upload ocean to pip and config it to the abstract class
+import logging
+from typing import Union, List
+import oceandb
+from oceandb.utils.embedding_function import MultiModalEmbeddingFunction
+
+class OceanDB:
+ def __init__(self):
+ try:
+ self.client = oceandb.Client()
+ print(self.client.heartbeat())
+ except Exception as e:
+ logging.error(f"Failed to initialize OceanDB client. Error: {e}")
+
+ def create_collection(self, collection_name: str, modality: str):
+ try:
+ embedding_function = MultiModalEmbeddingFunction(modality=modality)
+ collection = self.client.create_collection(collection_name, embedding_function=embedding_function)
+ return collection
+ except Exception as e:
+ logging.error(f"Failed to create collection. Error {e}")
+
+ def add_documents(self, collection, documents: List[str], ids: List[str]):
+ try:
+ return collection.add(documents=documents, ids=ids)
+ except Exception as e:
+ logging.error(f"Failed to add documents to collection. Error: {e}")
+ raise
+
+ def query(self, collection, query_texts: list[str], n_results: int):
+ try:
+ results = collection.query(query_texts=query_texts, n_results=n_results)
+ return results
+ except Exception as e:
+ logging.error(f"Failed to query the collection. Error {e}")
+ raise
\ No newline at end of file
diff --git a/swarms/swarms.py b/swarms/swarms.py
index 708e429b..892ca91f 100644
--- a/swarms/swarms.py
+++ b/swarms/swarms.py
@@ -21,9 +21,14 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(
# TODO: Add RLHF Data collection, ask user how the swarm is performing
class HierarchicalSwarm:
- def __init__(self, model_id: str = None, openai_api_key="",
- use_vectorstore=True, embedding_size: int = None, use_async=True,
- human_in_the_loop=True, model_type: str = None, boss_prompt: str = None):
+ def __init__(self, model_id: str = None,
+ openai_api_key="",
+ use_vectorstore=True, embedding_size: int = None, use_async=True,
+ human_in_the_loop=True, model_type: str = None, boss_prompt: str = None,
+ worker_prompt:str = None,
+ temperature=None,
+ max_iterations=None,
+ ):
#openai_api_key: the openai key. Default is empty
if not model_id:
logging.error("Model ID is not provided")
@@ -38,12 +43,18 @@ class HierarchicalSwarm:
self.use_async = use_async
self.human_in_the_loop = human_in_the_loop
+ self.model_type = model_type
+
self.embedding_size = embedding_size
-
self.boss_prompt = boss_prompt
+ self.worker_prompt = worker_prompt
+
+ self.temperature = temperature
+ self.max_iterations = max_iterations
- def initialize_llm(self, llm_class, temperature=0.5):
+
+ def initialize_llm(self, llm_class):
"""
Init LLM
@@ -54,9 +65,9 @@ class HierarchicalSwarm:
try:
# Initialize language model
if self.llm_class == 'openai' or OpenAI:
- return llm_class(openai_api_key=self.openai_api_key, temperature=temperature)
+ return llm_class(openai_api_key=self.openai_api_key, temperature=self.temperature)
elif self.model_type == "huggingface":
- return HuggingFaceLLM(model_id=self.model_id, temperature=temperature)
+ return HuggingFaceLLM(model_id=self.model_id, temperature=self.temperature)
except Exception as e:
logging.error(f"Failed to initialize language model: {e}")
@@ -97,9 +108,7 @@ class HierarchicalSwarm:
"""
Init vector store
"""
-
- try:
-
+ try:
embeddings_model = OpenAIEmbeddings(openai_api_key=self.openai_api_key)
embedding_size = self.embedding_size or 8192
index = faiss.IndexFlatL2(embedding_size)
@@ -119,21 +128,20 @@ class HierarchicalSwarm:
llm_class (class): The Language model class. Default is ChatOpenAI
ai_name (str): The AI name. Default is "Swarms worker AI assistant"
"""
-
- try:
-
+ try:
# Initialize worker node
llm = self.initialize_llm(ChatOpenAI)
worker_node = WorkerNodeInitializer(llm=llm, tools=worker_tools, vectorstore=vectorstore)
worker_node.create_agent(ai_name=ai_name, ai_role="Assistant", search_kwargs={}, human_in_the_loop=self.human_in_the_loop) # add search kwargs
+ worker_description = self.worker_prompt
- worker_node_tool = Tool(name="WorkerNode AI Agent", func=worker_node.run, description="Input: an objective with a todo list for that objective. Output: your task completed: Please be very clear what the objective and task instructions are. The Swarm worker agent is Useful for when you need to spawn an autonomous agent instance as a worker to accomplish any complex tasks, it can search the internet or write code or spawn child multi-modality models to process and generate images and text or audio and so on")
+ worker_node_tool = Tool(name="WorkerNode AI Agent", func=worker_node.run, description= worker_description or "Input: an objective with a todo list for that objective. Output: your task completed: Please be very clear what the objective and task instructions are. The Swarm worker agent is Useful for when you need to spawn an autonomous agent instance as a worker to accomplish any complex tasks, it can search the internet or write code or spawn child multi-modality models to process and generate images and text or audio and so on")
return worker_node_tool
except Exception as e:
logging.error(f"Failed to initialize worker node: {e}")
raise
- def initialize_boss_node(self, vectorstore, worker_node, llm_class=OpenAI, max_iterations=5, verbose=False):
+ def initialize_boss_node(self, vectorstore, worker_node, llm_class=OpenAI, max_iterations=None, verbose=False):
"""
Init BossNode
@@ -167,7 +175,7 @@ class HierarchicalSwarm:
agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=[tool.name for tool in tools])
agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=verbose)
- return BossNode(llm, vectorstore, agent_executor, max_iterations=max_iterations)
+ return BossNode(llm, vectorstore, agent_executor, max_iterations=self.max_iterations)
except Exception as e:
logging.error(f"Failed to initialize boss node: {e}")
raise
@@ -206,7 +214,7 @@ class HierarchicalSwarm:
return None
# usage-# usage-
-def swarm(api_key="", objective="", model_type=""):
+def swarm(api_key="", objective="", model_type="", model_id=""):
"""
Run the swarm with the given API key and objective.
@@ -225,7 +233,7 @@ def swarm(api_key="", objective="", model_type=""):
logging.error("Invalid objective")
raise ValueError("A valid objective is required")
try:
- swarms = HierarchicalSwarm(api_key, use_async=False, model_type=model_type) # Turn off async
+ swarms = HierarchicalSwarm(api_key, model_id, use_async=False, model_type=model_type) # Turn off async
result = swarms.run(objective)
if result is None:
logging.error("Failed to run swarms")
@@ -236,4 +244,3 @@ def swarm(api_key="", objective="", model_type=""):
logging.error(f"An error occured in swarm: {e}")
return None
-
diff --git a/swarms/utils/embeddings/pegasus.py b/swarms/utils/embeddings/pegasus.py
new file mode 100644
index 00000000..de3a4e36
--- /dev/null
+++ b/swarms/utils/embeddings/pegasus.py
@@ -0,0 +1,26 @@
+import logging
+from typing import Union
+from pegasus import Pegasus
+
+# import oceandb
+# from oceandb.utils.embedding_functions import MultiModalEmbeddingfunction
+
+
+class PegasusEmbedding:
+ def __init__(self, modality: str, multi_process: bool = False, n_processes: int = 4):
+ self.modality = modality
+ self.multi_process = multi_process
+ self.n_processes = n_processes
+ try:
+ self.pegasus = Pegasus(modality, multi_process, n_processes)
+ except Exception as e:
+ logging.error(f"Failed to initialize Pegasus with modality: {modality}: {e}")
+ raise
+
+ def embed(self, data: Union[str, list[str]]):
+ try:
+ return self.pegasus.embed(data)
+ except Exception as e:
+ logging.error(f"Failed to generate embeddings. Error: {e}")
+ raise
+