|
|
|
@ -44,35 +44,72 @@ import celery
|
|
|
|
|
from typing import List, Dict, Any
|
|
|
|
|
import numpy as np
|
|
|
|
|
|
|
|
|
|
from swarms.agents.
|
|
|
|
|
from swarms.agents.memory.ocean import OceanDB
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class Orchestrator(ABC):
|
|
|
|
|
def __init__(self, agent, agent_list: List[Any], task_queue: celery.Celery, vector_db: np.ndarray):
|
|
|
|
|
def __init__(self,
|
|
|
|
|
agent,
|
|
|
|
|
agent_list: List[Any],
|
|
|
|
|
task_queue: celery.Celery,
|
|
|
|
|
vector_db: OceanDB
|
|
|
|
|
):
|
|
|
|
|
self.agent = agent
|
|
|
|
|
self.agents = agent_list
|
|
|
|
|
self.agents = [agent_class() for _ in range(agent_list)]
|
|
|
|
|
self.task_queue = task_queue
|
|
|
|
|
self.vector_db = vector_db
|
|
|
|
|
self.current_tasks = {}
|
|
|
|
|
self.lock = Lock()
|
|
|
|
|
|
|
|
|
|
@abstractmethod
|
|
|
|
|
def assign_task(self, agent_id: int, task: Dict[str, Any]) -> None:
|
|
|
|
|
"""Assign a task to a specific agent"""
|
|
|
|
|
pass
|
|
|
|
|
with self.lock:
|
|
|
|
|
if self.task_queue:
|
|
|
|
|
#get and agent and a task
|
|
|
|
|
agent = self.agents.pop(0)
|
|
|
|
|
task = self.task_queue.popleft()
|
|
|
|
|
|
|
|
|
|
#process the task and get result and vector representation
|
|
|
|
|
result, vector_representation = agent.process_task()
|
|
|
|
|
|
|
|
|
|
#store the vector representation in the database
|
|
|
|
|
self.vector_db.add_documents([vector_representation],[str(id(task))])
|
|
|
|
|
|
|
|
|
|
#put the agent back to agent slist
|
|
|
|
|
self.agents.append(agent)
|
|
|
|
|
|
|
|
|
|
logging.info(f"Task {id(str)} has been processed by agent {id(agent)} ")
|
|
|
|
|
|
|
|
|
|
return result
|
|
|
|
|
else:
|
|
|
|
|
logging.error("Task queue is empty")
|
|
|
|
|
|
|
|
|
|
@abstractmethod
|
|
|
|
|
def retrieve_results(self, agent_id: int) -> Any:
|
|
|
|
|
"""Retrieve results from a specific agent"""
|
|
|
|
|
pass
|
|
|
|
|
try:
|
|
|
|
|
#Query the vector database for documents created by the agents
|
|
|
|
|
results = self.vector_db.query(query_texts=[str(agent_id)], n_results=10)
|
|
|
|
|
return results
|
|
|
|
|
except Exception as e:
|
|
|
|
|
logging.error(f"Failed to retrieve results from agent {agent_id}. Error {e}")
|
|
|
|
|
raise
|
|
|
|
|
|
|
|
|
|
@abstractmethod
|
|
|
|
|
def update_vector_db(self, data: np.ndarray) -> None:
|
|
|
|
|
"""Update the vector database"""
|
|
|
|
|
pass
|
|
|
|
|
try:
|
|
|
|
|
self.vector_db.add_documents([data['vector']], [str(data['task_id'])])
|
|
|
|
|
except Exception as e:
|
|
|
|
|
logging.error(f"Failed to update the vector database. Error: {e}")
|
|
|
|
|
raise
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@abstractmethod
|
|
|
|
|
def get_vector_db(self) -> np.ndarray:
|
|
|
|
|
"""Retrieve the vector database"""
|
|
|
|
|
pass
|
|
|
|
|
return self.vector_db
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|