swarms 7.7.5__py3-none-any.whl → 7.7.7__py3-none-any.whl

swarms/structs/concurrent_workflow.py

@@ -1,70 +1,32 @@
 import os
-import uuid
+import time
 from concurrent.futures import ThreadPoolExecutor
-from datetime import datetime
+from functools import lru_cache
 from typing import Any, Callable, Dict, List, Optional, Union
 
-from pydantic import BaseModel, Field
+from tqdm import tqdm
 
 from swarms.structs.agent import Agent
 from swarms.structs.base_swarm import BaseSwarm
-from swarms.utils.file_processing import create_file_in_folder
-from swarms.utils.loguru_logger import initialize_logger
 from swarms.structs.conversation import Conversation
-from swarms.structs.swarm_id_generator import generate_swarm_id
-from swarms.structs.output_types import OutputType
+from swarms.utils.formatter import formatter
+from swarms.utils.history_output_formatter import (
+    history_output_formatter,
+)
+from swarms.utils.loguru_logger import initialize_logger
 
 logger = initialize_logger(log_folder="concurrent_workflow")
 
 
-class AgentOutputSchema(BaseModel):
-    run_id: Optional[str] = Field(
-        ..., description="Unique ID for the run"
-    )
-    agent_name: Optional[str] = Field(
-        ..., description="Name of the agent"
-    )
-    task: Optional[str] = Field(
-        ..., description="Task or query given to the agent"
-    )
-    output: Optional[str] = Field(
-        ..., description="Output generated by the agent"
-    )
-    start_time: Optional[datetime] = Field(
-        ..., description="Start time of the task"
-    )
-    end_time: Optional[datetime] = Field(
-        ..., description="End time of the task"
-    )
-    duration: Optional[float] = Field(
-        ...,
-        description="Duration taken to complete the task (in seconds)",
-    )
-
-
-class MetadataSchema(BaseModel):
-    swarm_id: Optional[str] = Field(
-        generate_swarm_id(), description="Unique ID for the run"
-    )
-    task: Optional[str] = Field(
-        ..., description="Task or query given to all agents"
-    )
-    description: Optional[str] = Field(
-        "Concurrent execution of multiple agents",
-        description="Description of the workflow",
-    )
-    agents: Optional[List[AgentOutputSchema]] = Field(
-        ..., description="List of agent outputs and metadata"
-    )
-    timestamp: Optional[datetime] = Field(
-        default_factory=datetime.now,
-        description="Timestamp of the workflow execution",
-    )
-
-
 class ConcurrentWorkflow(BaseSwarm):
     """
     Represents a concurrent workflow that executes multiple agents concurrently in a production-grade manner.
+    Features include:
+    - Interactive model support
+    - Caching for repeated prompts
+    - Optional progress tracking
+    - Enhanced error handling and retries
+    - Input validation
 
     Args:
         name (str): The name of the workflow. Defaults to "ConcurrentWorkflow".
@@ -72,11 +34,16 @@ class ConcurrentWorkflow(BaseSwarm):
         agents (List[Agent]): The list of agents to be executed concurrently. Defaults to an empty list.
         metadata_output_path (str): The path to save the metadata output. Defaults to "agent_metadata.json".
         auto_save (bool): Flag indicating whether to automatically save the metadata. Defaults to False.
-        output_schema (BaseModel): The output schema for the metadata. Defaults to MetadataSchema.
+        output_type (str): The type of output format. Defaults to "dict".
         max_loops (int): The maximum number of loops for each agent. Defaults to 1.
         return_str_on (bool): Flag indicating whether to return the output as a string. Defaults to False.
-        agent_responses (list): The list of agent responses. Defaults to an empty list.
         auto_generate_prompts (bool): Flag indicating whether to auto-generate prompts for agents. Defaults to False.
+        return_entire_history (bool): Flag indicating whether to return the entire conversation history. Defaults to False.
+        interactive (bool): Flag indicating whether to enable interactive mode. Defaults to False.
+        cache_size (int): The size of the cache. Defaults to 100.
+        max_retries (int): The maximum number of retry attempts. Defaults to 3.
+        retry_delay (float): The delay between retry attempts in seconds. Defaults to 1.0.
+        show_progress (bool): Flag indicating whether to show progress. Defaults to False.
 
     Raises:
         ValueError: If the list of agents is empty or if the description is empty.
@@ -87,13 +54,18 @@ class ConcurrentWorkflow(BaseSwarm):
         agents (List[Agent]): The list of agents to be executed concurrently.
         metadata_output_path (str): The path to save the metadata output.
         auto_save (bool): Flag indicating whether to automatically save the metadata.
-        output_schema (BaseModel): The output schema for the metadata.
+        output_type (str): The type of output format.
         max_loops (int): The maximum number of loops for each agent.
         return_str_on (bool): Flag indicating whether to return the output as a string.
-        agent_responses (list): The list of agent responses.
         auto_generate_prompts (bool): Flag indicating whether to auto-generate prompts for agents.
-        retry_attempts (int): The number of retry attempts for failed agent executions.
-        retry_wait_time (int): The initial wait time for retries in seconds.
+        return_entire_history (bool): Flag indicating whether to return the entire conversation history.
+        interactive (bool): Flag indicating whether to enable interactive mode.
+        cache_size (int): The size of the cache.
+        max_retries (int): The maximum number of retry attempts.
+        retry_delay (float): The delay between retry attempts in seconds.
+        show_progress (bool): Flag indicating whether to show progress.
+        _cache (dict): The cache for storing agent outputs.
+        _progress_bar (tqdm): The progress bar for tracking execution.
     """
 
     def __init__(
@@ -103,13 +75,16 @@ class ConcurrentWorkflow(BaseSwarm):
         agents: List[Union[Agent, Callable]] = [],
         metadata_output_path: str = "agent_metadata.json",
         auto_save: bool = True,
-        output_schema: BaseModel = MetadataSchema,
+        output_type: str = "dict-all-except-first",
         max_loops: int = 1,
         return_str_on: bool = False,
-        agent_responses: list = [],
         auto_generate_prompts: bool = False,
-        output_type: OutputType = "dict",
         return_entire_history: bool = False,
+        interactive: bool = False,
+        cache_size: int = 100,
+        max_retries: int = 3,
+        retry_delay: float = 1.0,
+        show_progress: bool = False,
         *args,
         **kwargs,
     ):
@@ -125,18 +100,22 @@ class ConcurrentWorkflow(BaseSwarm):
         self.agents = agents
         self.metadata_output_path = metadata_output_path
         self.auto_save = auto_save
-        self.output_schema = output_schema
         self.max_loops = max_loops
         self.return_str_on = return_str_on
-        self.agent_responses = agent_responses
         self.auto_generate_prompts = auto_generate_prompts
         self.max_workers = os.cpu_count()
         self.output_type = output_type
         self.return_entire_history = return_entire_history
         self.tasks = []  # Initialize tasks list
+        self.interactive = interactive
+        self.cache_size = cache_size
+        self.max_retries = max_retries
+        self.retry_delay = retry_delay
+        self.show_progress = show_progress
+        self._cache = {}
+        self._progress_bar = None
 
         self.reliability_check()
-
         self.conversation = Conversation()
 
     def disable_agent_prints(self):
@@ -145,29 +124,47 @@ class ConcurrentWorkflow(BaseSwarm):
 
     def reliability_check(self):
         try:
-            logger.info("Starting reliability checks")
+            formatter.print_panel(
+                content=f"\n 🏷️ Name: {self.name}\n 📝 Description: {self.description}\n 🤖 Agents: {len(self.agents)}\n 🔄 Max Loops: {self.max_loops}\n ",
+                title="⚙️ Concurrent Workflow Settings",
+                style="bold blue",
+            )
+            formatter.print_panel(
+                content="🔍 Starting reliability checks",
+                title="🔒 Reliability Checks",
+                style="bold blue",
+            )
 
             if self.name is None:
-                logger.error("A name is required for the swarm")
-                raise ValueError("A name is required for the swarm")
+                logger.error("❌ A name is required for the swarm")
+                raise ValueError(
+                    "❌ A name is required for the swarm"
+                )
 
-            if not self.agents:
-                logger.error("The list of agents must not be empty.")
+            if not self.agents or len(self.agents) <= 1:
+                logger.error(
+                    "❌ The list of agents must not be empty."
+                )
                 raise ValueError(
-                    "The list of agents must not be empty."
+                    "❌ The list of agents must not be empty."
                 )
 
             if not self.description:
-                logger.error("A description is required.")
-                raise ValueError("A description is required.")
+                logger.error("❌ A description is required.")
+                raise ValueError("❌ A description is required.")
+
+            formatter.print_panel(
+                content="✅ Reliability checks completed successfully",
+                title="🎉 Reliability Checks",
+                style="bold green",
+            )
 
-            logger.info("Reliability checks completed successfully")
         except ValueError as e:
-            logger.error(f"Reliability check failed: {e}")
+            logger.error(f"❌ Reliability check failed: {e}")
             raise
         except Exception as e:
             logger.error(
-                f"An unexpected error occurred during reliability checks: {e}"
+                f"💥 An unexpected error occurred during reliability checks: {e}"
             )
             raise
 
@@ -184,147 +181,179 @@ class ConcurrentWorkflow(BaseSwarm):
             for agent in self.agents:
                 agent.auto_generate_prompt = True
 
-    # @retry(wait=wait_exponential(min=2), stop=stop_after_attempt(3))
-
-    def transform_metadata_schema_to_str(
-        self, schema: MetadataSchema
-    ):
-        """
-        Converts the metadata swarm schema into a string format with the agent name, response, and time.
-
-        Args:
-            schema (MetadataSchema): The metadata schema to convert.
-
-        Returns:
-            str: The string representation of the metadata schema.
-
-        Example:
-            >>> metadata_schema = MetadataSchema()
-            >>> metadata_str = workflow.transform_metadata_schema_to_str(metadata_schema)
-            >>> print(metadata_str)
-        """
-        self.agent_responses = [
-            f"Agent Name: {agent.agent_name}\nResponse: {agent.output}\n\n"
-            for agent in schema.agents
-        ]
-
-        # Return the agent responses as a string
-        return "\n".join(self.agent_responses)
-
-    def save_metadata(self):
-        """
-        Saves the metadata to a JSON file based on the auto_save flag.
-
-        Example:
-            >>> workflow.save_metadata()
-            >>> # Metadata will be saved to the specified path if auto_save is True.
-        """
-        # Save metadata to a JSON file
-        if self.auto_save:
-            logger.info(
-                f"Saving metadata to {self.metadata_output_path}"
+    @lru_cache(maxsize=100)
+    def _cached_run(self, task: str, agent_id: int) -> Any:
+        """Cached version of agent execution to avoid redundant computations"""
+        return self.agents[agent_id].run(task=task)
+
+    def enable_progress_bar(self):
+        """Enable progress bar display"""
+        self.show_progress = True
+
+    def disable_progress_bar(self):
+        """Disable progress bar display"""
+        if self._progress_bar:
+            self._progress_bar.close()
+            self._progress_bar = None
+        self.show_progress = False
+
+    def _create_progress_bar(self, total: int):
+        """Create a progress bar for tracking execution"""
+        if self.show_progress:
+            try:
+                self._progress_bar = tqdm(
+                    total=total,
+                    desc="Processing tasks",
+                    unit="task",
+                    disable=not self.show_progress,
+                    ncols=100,
+                    bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]",
+                )
+            except Exception as e:
+                logger.warning(f"Failed to create progress bar: {e}")
+                self.show_progress = False
+                self._progress_bar = None
+        return self._progress_bar
+
+    def _update_progress(self, increment: int = 1):
+        """Update the progress bar"""
+        if self._progress_bar and self.show_progress:
+            try:
+                self._progress_bar.update(increment)
+            except Exception as e:
+                logger.warning(f"Failed to update progress bar: {e}")
+                self.disable_progress_bar()
+
+    def _validate_input(self, task: str) -> bool:
+        """Validate input task"""
+        if not isinstance(task, str):
+            raise ValueError("Task must be a string")
+        if not task.strip():
+            raise ValueError("Task cannot be empty")
+        return True
+
+    def _handle_interactive(self, task: str) -> str:
+        """Handle interactive mode for task input"""
+        if self.interactive:
+            from swarms.utils.formatter import formatter
+
+            # Display current task in a panel
+            formatter.print_panel(
+                content=f"Current task: {task}",
+                title="Task Status",
+                style="bold blue",
             )
-            create_file_in_folder(
-                os.getenv("WORKSPACE_DIR"),
-                self.metadata_output_path,
-                self.output_schema.model_dump_json(indent=4),
+
+            # Get user input with formatted prompt
+            formatter.print_panel(
+                content="Do you want to modify this task? (y/n/q to quit): ",
+                title="User Input",
+                style="bold green",
             )
+            response = input().lower()
+
+            if response == "q":
+                return None
+            elif response == "y":
+                formatter.print_panel(
+                    content="Enter new task: ",
+                    title="New Task Input",
+                    style="bold yellow",
+                )
+                new_task = input()
+                return new_task
+        return task
+
+    def _run_with_retry(
+        self, agent: Agent, task: str, img: str = None
+    ) -> Any:
+        """Run agent with retry mechanism"""
+        for attempt in range(self.max_retries):
+            try:
+                output = agent.run(task=task, img=img)
+                self.conversation.add(agent.agent_name, output)
+                return output
+            except Exception as e:
+                if attempt == self.max_retries - 1:
+                    logger.error(
+                        f"Error running agent {agent.agent_name} after {self.max_retries} attempts: {e}"
+                    )
+                    raise
+                logger.warning(
+                    f"Attempt {attempt + 1} failed for agent {agent.agent_name}: {e}"
+                )
+                time.sleep(
+                    self.retry_delay * (attempt + 1)
+                )  # Exponential backoff
 
     def _run(
         self, task: str, img: str = None, *args, **kwargs
     ) -> Union[Dict[str, Any], str]:
         """
-        Runs the workflow for the given task, executes agents concurrently using ThreadPoolExecutor, and saves metadata.
+        Enhanced run method with caching, progress tracking, and better error handling
+        """
 
-        Args:
-            task (str): The task or query to give to all agents.
-            img (str): The image to be processed by the agents.
+        # Validate and potentially modify task
+        self._validate_input(task)
+        task = self._handle_interactive(task)
 
-        Returns:
-            Dict[str, Any]: The final metadata as a dictionary.
-            str: The final metadata as a string if return_str_on is True.
+        # Add task to conversation
+        self.conversation.add("User", task)
 
-        Example:
-            >>> metadata = workflow.run(task="Example task", img="example.jpg")
-            >>> print(metadata)
-        """
-        logger.info(
-            f"Running concurrent workflow with {len(self.agents)} agents."
-        )
-
-        self.conversation.add(
-            "User",
-            task,
-        )
+        # Create progress bar if enabled
+        if self.show_progress:
+            self._create_progress_bar(len(self.agents))
 
         def run_agent(
             agent: Agent, task: str, img: str = None
-        ) -> AgentOutputSchema:
-            start_time = datetime.now()
+        ) -> Any:
             try:
-                output = agent.run(task=task)
-
-                self.conversation.add(
-                    agent.agent_name,
-                    output,
-                )
+                # Check cache first
+                cache_key = f"{task}_{agent.agent_name}"
+                if cache_key in self._cache:
+                    output = self._cache[cache_key]
+                else:
+                    output = self._run_with_retry(agent, task, img)
+                    # Update cache
+                    if len(self._cache) >= self.cache_size:
+                        self._cache.pop(next(iter(self._cache)))
+                    self._cache[cache_key] = output
+
+                self._update_progress()
+                return output
             except Exception as e:
                 logger.error(
                     f"Error running agent {agent.agent_name}: {e}"
                 )
+                self._update_progress()
                 raise
 
-            end_time = datetime.now()
-            duration = (end_time - start_time).total_seconds()
-
-            agent_output = AgentOutputSchema(
-                run_id=uuid.uuid4().hex,
-                agent_name=agent.agent_name,
-                task=task,
-                output=output,
-                start_time=start_time,
-                end_time=end_time,
-                duration=duration,
-            )
-
-            logger.info(
-                f"Agent {agent.agent_name} completed task: {task} in {duration:.2f} seconds."
-            )
-
-            return agent_output
-
-        with ThreadPoolExecutor(
-            max_workers=os.cpu_count()
-        ) as executor:
-            agent_outputs = list(
-                executor.map(
-                    lambda agent: run_agent(agent, task), self.agents
+        try:
+            with ThreadPoolExecutor(
+                max_workers=self.max_workers
+            ) as executor:
+                list(
+                    executor.map(
+                        lambda agent: run_agent(agent, task),
+                        self.agents,
+                    )
                 )
-            )
-
-        self.output_schema = MetadataSchema(
-            swarm_id=uuid.uuid4().hex,
-            task=task,
-            description=self.description,
-            agents=agent_outputs,
+        finally:
+            if self._progress_bar and self.show_progress:
+                try:
+                    self._progress_bar.close()
+                except Exception as e:
+                    logger.warning(
+                        f"Failed to close progress bar: {e}"
+                    )
+                finally:
+                    self._progress_bar = None
+
+        return history_output_formatter(
+            self.conversation,
+            type=self.output_type,
         )
 
-        self.save_metadata()
-
-        if self.return_str_on:
-            return self.transform_metadata_schema_to_str(
-                self.output_schema
-            )
-
-        elif self.return_entire_history:
-            return self.conversation.return_history_as_string()
-        elif self.output_type == "list":
-            return self.conversation.return_messages_as_list()
-        elif self.output_type == "dict":
-            return self.conversation.return_messages_as_dictionary()
-        else:
-            return self.output_schema.model_dump_json(indent=4)
-
     def run(
         self,
         task: Optional[str] = None,
@@ -333,9 +362,11 @@ class ConcurrentWorkflow(BaseSwarm):
         **kwargs,
     ) -> Any:
         """
-        Executes the agent's run method on a specified device.
+        Executes the agent's run method on a specified device with optional interactive mode.
 
-        This method attempts to execute the agent's run method on a specified device, either CPU or GPU. It logs the device selection and the number of cores or GPU ID used. If the device is set to CPU, it can use all available cores or a specific core specified by `device_id`. If the device is set to GPU, it uses the GPU specified by `device_id`.
+        This method attempts to execute the agent's run method on a specified device, either CPU or GPU.
+        It supports both standard execution and interactive mode where users can modify tasks and continue
+        the workflow interactively.
 
         Args:
             task (Optional[str], optional): The task to be executed. Defaults to None.
@@ -359,8 +390,73 @@ class ConcurrentWorkflow(BaseSwarm):
             self.tasks.append(task)
 
         try:
-            outputs = self._run(task, img, *args, **kwargs)
-            return outputs
+            # Handle interactive mode
+            if self.interactive:
+                current_task = task
+                loop_count = 0
+
+                while loop_count < self.max_loops:
+                    if (
+                        self.max_loops is not None
+                        and loop_count >= self.max_loops
+                    ):
+                        formatter.print_panel(
+                            content=f"Maximum number of loops ({self.max_loops}) reached.",
+                            title="Session Complete",
+                            style="bold red",
+                        )
+                        break
+
+                    if current_task is None:
+                        formatter.print_panel(
+                            content="Enter your task (or 'q' to quit): ",
+                            title="Task Input",
+                            style="bold blue",
+                        )
+                        current_task = input()
+                        if current_task.lower() == "q":
+                            break
+
+                    # Run the workflow with the current task
+                    try:
+                        outputs = self._run(
+                            current_task, img, *args, **kwargs
+                        )
+                        formatter.print_panel(
+                            content=str(outputs),
+                            title="Workflow Result",
+                            style="bold green",
+                        )
+                    except Exception as e:
+                        formatter.print_panel(
+                            content=f"Error: {str(e)}",
+                            title="Error",
+                            style="bold red",
+                        )
+
+                    # Ask if user wants to continue
+                    formatter.print_panel(
+                        content="Do you want to continue with a new task? (y/n): ",
+                        title="Continue Session",
+                        style="bold yellow",
+                    )
+                    if input().lower() != "y":
+                        break
+
+                    current_task = None
+                    loop_count += 1
+
+                formatter.print_panel(
+                    content="Interactive session ended.",
+                    title="Session Complete",
+                    style="bold blue",
+                )
+                return outputs
+            else:
+                # Standard non-interactive execution
+                outputs = self._run(task, img, *args, **kwargs)
+                return outputs
+
         except ValueError as e:
             logger.error(f"Invalid device specified: {e}")
             raise e
@@ -368,29 +464,48 @@ class ConcurrentWorkflow(BaseSwarm):
             logger.error(f"An error occurred during execution: {e}")
             raise e
 
-    def run_batched(
-        self, tasks: List[str]
-    ) -> List[Union[Dict[str, Any], str]]:
+    def run_batched(self, tasks: List[str]) -> Any:
         """
-        Runs the workflow for a batch of tasks, executes agents concurrently for each task, and saves metadata in a production-grade manner.
+        Enhanced batched execution with progress tracking
+        """
+        if not tasks:
+            raise ValueError("Tasks list cannot be empty")
 
-        Args:
-            tasks (List[str]): A list of tasks or queries to give to all agents.
+        results = []
 
-        Returns:
-            List[Union[Dict[str, Any], str]]: A list of final metadata for each task, either as a dictionary or a string.
+        # Create progress bar if enabled
+        if self.show_progress:
+            self._create_progress_bar(len(tasks))
+
+        try:
+            for task in tasks:
+                result = self.run(task)
+                results.append(result)
+                self._update_progress()
+        finally:
+            if self._progress_bar and self.show_progress:
+                try:
+                    self._progress_bar.close()
+                except Exception as e:
+                    logger.warning(
+                        f"Failed to close progress bar: {e}"
+                    )
+                finally:
+                    self._progress_bar = None
 
-        Example:
-            >>> tasks = ["Task 1", "Task 2"]
-            >>> results = workflow.run_batched(tasks)
-            >>> print(results)
-        """
-        results = []
-        for task in tasks:
-            result = self.run(task)
-            results.append(result)
         return results
 
+    def clear_cache(self):
+        """Clear the task cache"""
+        self._cache.clear()
+
+    def get_cache_stats(self) -> Dict[str, int]:
+        """Get cache statistics"""
+        return {
+            "cache_size": len(self._cache),
+            "max_cache_size": self.cache_size,
+        }
+
 
 # if __name__ == "__main__":
 #     # Assuming you've already initialized some agents outside of this class