enum

swarms/swarms/orchestrate.py

@@ -4,11 +4,18 @@ import threading
 # from abc import ABC, abstractmethod
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any, Dict, List
+from enum import Enum
 
 import chromadb
 from chromadb.utils import embedding_functions
 
 
+class TaskStatus(Enum):
+    QUEUED = 1
+    RUNNING = 2
+    COMPLETED = 3
+    FAILED = 4
+
 class Orchestrator:
     """
     The Orchestrator takes in an agent, worker, or boss as input 
@@ -71,7 +78,8 @@ class Orchestrator:
         task_queue: List[Any], 
         collection_name: str = "swarm",
         api_key: str = None,
-        model_name: str = None
+        model_name: str = None,
+        embed_func = None
     ):
         self.agent = agent
         self.agents = queue.Queue()
@@ -93,6 +101,8 @@ class Orchestrator:
         self.condition = threading.Condition(self.lock)
         self.executor = ThreadPoolExecutor(max_workers=len(agent_list))
         
+        self.embed_func = embed_func if embed_func else self.embed
+
         
     # @abstractmethod
     def assign_task(