From 982601614e1071a105088d8dbbbb9800ec3cc2d3 Mon Sep 17 00:00:00 2001 From: Kye Gomez Date: Thu, 20 Feb 2025 07:48:01 -0800 Subject: [PATCH] ui docs fix --- cuda_swarm.py | 359 +++++++++++++++++++++ docs/swarms/ui/main.md | 4 - duo_agent.py | 49 +++ loguru_test.py | 506 +++++++++++++++++++++++++++++ open_scientist.py | 541 ++++++++++++++++++++++++++++++++ swarms/structs/rearrange.py | 10 +- swarms/tools/tool_parse_exec.py | 96 +++++- 7 files changed, 1554 insertions(+), 11 deletions(-) create mode 100644 cuda_swarm.py create mode 100644 duo_agent.py create mode 100644 loguru_test.py create mode 100644 open_scientist.py diff --git a/cuda_swarm.py b/cuda_swarm.py new file mode 100644 index 00000000..64ae7159 --- /dev/null +++ b/cuda_swarm.py @@ -0,0 +1,359 @@ +import os +import re + +import litellm + +from swarms import Agent + +litellm.drop_params = True + + +def extract_code_blocks(text: str) -> str: + """ + Extracts code blocks enclosed in triple backticks from the given text. + + Args: + text (str): The input text containing code blocks. + + Returns: + str: The extracted code blocks joined together as a single string. + """ + # Regular expression to match code blocks enclosed in triple backticks + code_block_pattern = re.compile( + r"```(?:[a-zA-Z]*)\n(.*?)```", re.DOTALL + ) + + # Find all matches and join them into a single string + matches = code_block_pattern.findall(text) + return "\n".join(matches) + + +# Prompt #1: Translate PyTorch code into CUDA (extensive instructions) +translate_pytorch_to_cuda_prompt = """ +You are an AI agent specialized in converting PyTorch code into efficient, +well-structured, and properly functioning CUDA code. Your role is to transform +the given PyTorch code (which may include Python-based tensor operations, +layers, and training loops) into a CUDA codebase capable of running directly on +NVIDIA GPUs without relying on the high-level abstractions of PyTorch. + +Detailed Instructions: +1. Read and thoroughly understand every line of the provided PyTorch code, + including model definitions, forward passes, loss functions, backward passes, + and optimization steps. + +2. Rewrite the code in pure CUDA C/C++: + - Convert high-level tensor operations to raw GPU kernel operations and + memory management. + - Allocate and deallocate memory on the GPU using CUDA APIs like cudaMalloc + and cudaFree. + - Handle data transfers between the host (CPU) and the device (GPU) + appropriately (e.g., cudaMemcpy). + - Convert PyTorch's autograd-based gradient calculations into manual CUDA + kernel operations or unify them within your own backward pass + implementation if needed. + - Replace Pythonic loops and array operations with explicit CUDA kernels. + - Implement or port any custom CUDA kernels for special functionality. + +3. Optimize the CUDA code for parallel execution: + - Use grid and block dimensions that take advantage of the target GPU's + compute capabilities and memory bandwidth. + - Leverage shared memory, constant memory, or registers when beneficial. + - Unroll loops and reduce warp divergence where possible. + - Use efficient memory access patterns (e.g., coalesced memory access). + +4. Preserve the overall logic and functionality: + - The final CUDA code must produce equivalent numerical results as the + original PyTorch code for the same inputs and initialization. + - Keep the data types, precision, and numerical stability in mind while + performing computations in CUDA. + +5. Ensure readability and maintainability: + - Add comments where necessary to explain tricky parts of the CUDA kernels + and memory management. + - Use clear function or file organization to separate different components + (e.g., forward pass, backward pass, kernel definitions, helper utilities). + +6. Remember that the goal is to create an extensive and fully functional .cu + file or files that can be compiled with NVCC or integrated into a larger + project. + +7. Ensure PyTorch Integration: + - Implement proper PyTorch C++ extension interfaces using torch::extension + - Include necessary PyTorch headers and macros for tensor operations + - Create Python bindings using pybind11 to expose CUDA functions to PyTorch + - Handle PyTorch tensor types and conversions appropriately + - Ensure compatibility with PyTorch's autograd system + - Follow PyTorch's extension conventions for error handling and memory management + - Make the code loadable as a PyTorch extension module + +Output Requirements: +- Provide only the final CUDA code. +- Do not include any explanatory text outside of comments in the code itself. +- The CUDA code should stand on its own as a complete solution that can be + compiled or integrated without additional references. +- The code must be fully compatible with PyTorch's extension system and + able to be imported and used seamlessly within PyTorch models. +- Ensure the code is fully compatible with PyTorch's extension system. +- And, the code should be very long and extensive +""" + +# Prompt #2: Make the CUDA code super reliable and fast (extensive instructions) +make_cuda_super_reliable_and_fast_prompt = """ +You are an advanced AI agent whose goal is to take a fully functional CUDA codebase +and optimize it to be extraordinarily robust, reliable, and efficient. You must +focus on performance improvements at both the kernel level and architectural +level, ensuring that the code is streamlined and built to handle potential edge +cases gracefully. + +Detailed Instructions: +1. Dive deeply into the CUDA kernels and identify performance bottlenecks: + - Look for any uncoalesced memory accesses, unnecessary global memory reads, + or writes that could be optimized. + - Identify any opportunities to reduce shared memory usage, if that memory + usage does not yield significant benefit, or to increase it if it can help + reduce global memory accesses. + +2. Implement advanced optimization techniques: + - Use loop unrolling where beneficial to reduce overhead. + - Employ occupancy analysis to find the ideal block size and grid size for + maximum parallelism. + - Consider the use of constant memory for frequently accessed read-only data. + - Evaluate the benefits of pinned (page-locked) memory for host-device + transfers. + +3. Improve reliability and error handling: + - Insert meaningful checks for CUDA API function calls (e.g., cudaMalloc, + cudaMemcpy, cudaFree, kernel launches). Ensure proper cleanup if any call + fails or returns an error code. + - Handle edge cases where input sizes might be too large or too small, + preventing invalid memory accesses or out-of-bounds kernel launches. + - Ensure that any macros, compilation flags, or library calls align with + best practices for portability across different GPU architectures. + +4. Document any advanced tricks or strategies used: + - In-line commentary is crucial. Briefly explain why a specific optimization + was chosen and how it impacts performance. + +5. Maintain numerical equivalence: + - All transformations must preserve the original numerical outcomes within + reasonable floating-point precision limits. Do not alter the algorithmic + results unexpectedly. + +6. Provide a clean final output: + - The output must be only the improved and optimized CUDA source code. + - All extraneous explanation beyond code comments should be avoided. + - The code must be very long and extensive, containing all necessary functionality. + - Output only the complete code file with no additional text. + +Goal: +- Deliver a high-performance CUDA code that is not only efficient in running + time but is also resilient, capable of handling unexpected conditions, and + thoroughly commented to allow easy maintenance and future modifications. +- Ensure the output is a complete, extensive codebase with all functionality included. +""" + +# Prompt #3: Cleanup errors and add extensive documentation (extensive instructions) +cleanup_and_document_cuda_code_prompt = """ +You are a specialized AI agent focused on thoroughly debugging, cleaning up, and +documenting a CUDA codebase. Your mission is to ensure the code compiles cleanly, +handles errors gracefully, follows best coding practices, and includes detailed +documentation for maintainability and educational purposes. + +Detailed Instructions: +1. Debug and error cleanup: + - Identify any compilation or runtime errors and fix them. + - Check for mismatched types, undeclared variables, improper memory usage, or + incorrect kernel configurations. + - Make sure that each kernel launch has correct grid and block dimensions and + that indexing inside kernels is handled properly (avoid out-of-bounds + threads). + +2. Strengthen error handling: + - Wrap all CUDA library calls (cudaMalloc, cudaMemcpy, kernel launches, etc.) + with macros or functions that check for and report errors in a consistent + manner (e.g., using cudaGetErrorString). + - Ensure resources are deallocated in case of failure and that the program + can exit safely without memory leaks or GPU lockups. + +3. Add thorough documentation: + - Provide high-level explanations near the top of the file describing the + overall code structure and flow. + - Within each function, write docstrings or block comments to explain the + function's purpose, inputs, outputs, and major steps. + - For each kernel, add comments describing how threads are mapped to data, + how memory is accessed, and what the main loop or computational logic + accomplishes. + +4. Check performance remains robust: + - Ensure that none of the debugging or cleanup processes introduce unnecessary + slowdowns. Where possible, maintain or improve previous optimizations. + +5. Provide final cleaned, well-documented CUDA source code: + - The output must contain only the updated source code, with no additional + explanation outside of code comments and docstrings. + - The code must be ready to compile, fully functional, and in a polished + state that one can confidently integrate into a production environment. + - The code must be very long and extensive, containing all necessary functionality. + - Output only the complete code file with no additional text. + +Goal: +- Deliver a thoroughly cleaned, expertly documented CUDA code file. The + readability, reliability, and educational clarity of the code should be + enhanced without sacrificing computational performance. +- Ensure the output is a complete, extensive codebase with all functionality included. +""" + +# Prompt #4: Produce a final, extensive CUDA file (extensive instructions) +produce_extensive_cuda_file_prompt = """ +You are an AI agent tasked with producing a final, extensive .cu or .cuh file +containing all functionality needed to run the code originally derived from +PyTorch. This final file should encapsulate: +- The core CUDA kernels. +- The initialization and teardown logic for GPU resources. +- Optimized computations and any specialized routines (e.g., custom backward + passes, advanced math operations). +- Comprehensive yet concise in-code documentation. + +Detailed Instructions: +1. Merge all relevant kernels, utility functions, and structures into a cohesive + single file or a well-structured set of files. +2. Ensure you apply all previous optimizations and cleanup efforts, reflecting + them in this combined final output. +3. Use robust function prototypes for any utility routines, paying attention to + scope and avoiding unnecessary global variables. +4. Keep the code easily navigable with clear sections: + - Data structures and utility definitions. + - Kernel definitions. + - Host functions for kernel launches. + - Initialization and cleanup logic. + - Document each section thoroughly for ease of reference. +5. Ensure it compiles error-free under standard NVCC compilation. +6. Provide only the CUDA code, including all documentation within comments: + - No additional external explanation should be outside these comments. + - The code must be very long and extensive, containing all necessary functionality. + - Output only the complete code file with no additional text. + +Goal: +- A single or modular set of .cu/.cuh files that stand as the definitive version + of the CUDA codebase, balancing performance, reliability, and maintainability. +- Ensure the output is a complete, extensive codebase with all functionality included. +""" + + +# Now create one agent for each prompt, similar to the example with the Financial-Analysis-Agent. +translate_agent = Agent( + agent_name="Translate-PyTorch-To-CUDA-Agent", + system_prompt=translate_pytorch_to_cuda_prompt, + model_name="openai/o1", + max_loops=1, + max_tokens=10000, + output_type="str", + temperature=0, +) + +super_fast_agent = Agent( + agent_name="Make-CUDA-Code-Super-Fast-Agent", + system_prompt=make_cuda_super_reliable_and_fast_prompt, + model_name="openai/o1", + max_loops=1, + max_tokens=10000, + output_type="str", + temperature=0, +) + +cleanup_agent = Agent( + agent_name="Cleanup-and-Document-CUDA-Agent", + system_prompt=cleanup_and_document_cuda_code_prompt, + model_name="openai/o1", + max_loops=1, + max_tokens=10000, + output_type="str", + temperature=0, +) + +final_cuda_agent = Agent( + agent_name="Produce-Final-Extensive-CUDA-File-Agent", + system_prompt=produce_extensive_cuda_file_prompt, + model_name="openai/o1", + max_loops=1, + max_tokens=10000, + output_type="str", + temperature=0, +) + + +class CudaSwarm: + def __init__( + self, + name: str = "CudaSwarm", + description: str = "A swarm of agents that convert PyTorch code into CUDA code", + agents: list[Agent] = [ + translate_agent, + super_fast_agent, + cleanup_agent, + final_cuda_agent, + ], + max_loops: int = 1, + file_path: str = "cuda_code.cu", + ): + self.name = name + self.description = description + self.agents = agents + self.max_loops = max_loops + self.file_path = file_path + + def write_file(self, content: str = "") -> None: + """ + Creates a new CUDA file or overwrites an existing one with the specified content. + + Args: + content (str): The content to write into the file. Defaults to an empty string. + """ + # Ensure the file has a .cu extension + if not self.file_path.endswith(".cu"): + self.file_path += ".cu" + + # Create the directory if it doesn't exist + directory = os.path.dirname(self.file_path) + if directory and not os.path.exists(directory): + os.makedirs(directory) + + # Use write mode to overwrite the file + mode = "w" + + try: + # Write content to the file + with open(self.file_path, mode) as file: + file.write(content) + print(f"Content successfully written to {self.file_path}") + except IOError as e: + print(f"An error occurred while writing to the file: {e}") + + def run(self, task: str): + """ + Runs the swarm of agents to complete the task. + """ + first_iteration = self.agents[0].run(task) + first_iteration = extract_code_blocks(first_iteration) + self.write_file(first_iteration) + + # second_iteration = self.agents[1].run(task) + # second_iteration = extract_code_blocks(second_iteration) + # self.write_file(second_iteration) + + # third_iteration = self.agents[2].run(task) + # third_iteration = extract_code_blocks(third_iteration) + # self.write_file(third_iteration) + + # final_iteration = self.agents[3].run(task) + # final_iteration = extract_code_blocks(final_iteration) + # self.write_file(final_iteration) + + return first_iteration + + +if __name__ == "__main__": + swarm = CudaSwarm(file_path="cuda_code.cu") + swarm.run( + "Create the cuda code for a highly optimized liquid continous learner ssm model" + ) diff --git a/docs/swarms/ui/main.md b/docs/swarms/ui/main.md index 582c09b6..640e96b3 100644 --- a/docs/swarms/ui/main.md +++ b/docs/swarms/ui/main.md @@ -1,7 +1,3 @@ -Below is a revised version of the Swarms Chat UI documentation. It improves the flow, adds detailed examples for every use case, and clarifies the installation and usage instructions for a more reliable experience. - ---- - # Swarms Chat UI Documentation The Swarms Chat interface provides a customizable, multi-agent chat experience using Gradio. It supports various specialized AI agents—from finance to healthcare and news analysis—by leveraging Swarms models. diff --git a/duo_agent.py b/duo_agent.py new file mode 100644 index 00000000..d54df64c --- /dev/null +++ b/duo_agent.py @@ -0,0 +1,49 @@ +from swarms import Agent +from swarms.prompts.finance_agent_sys_prompt import ( + FINANCIAL_AGENT_SYS_PROMPT, +) + +# Initialize the equity analyst agents +equity_analyst_1 = Agent( + agent_name="Equity-Analyst-1", + agent_description="Equity research analyst focused on fundamental analysis", + system_prompt=FINANCIAL_AGENT_SYS_PROMPT, + max_loops=1, + model_name="gpt-4o", + dynamic_temperature_enabled=True, + user_name="swarms_corp", + retry_attempts=3, + context_length=8192, + return_step_meta=False, + output_type="str", + auto_generate_prompt=False, + max_tokens=4000, + saved_state_path="equity_analyst_1.json", + interactive=False, + roles="analyst" +) + +equity_analyst_2 = Agent( + agent_name="Equity-Analyst-2", + agent_description="Equity research analyst focused on technical analysis", + system_prompt=FINANCIAL_AGENT_SYS_PROMPT, + max_loops=1, + model_name="gpt-4o", + dynamic_temperature_enabled=True, + user_name="swarms_corp", + retry_attempts=3, + context_length=8192, + return_step_meta=False, + output_type="str", + auto_generate_prompt=False, + max_tokens=4000, + saved_state_path="equity_analyst_2.json", + interactive=False, + roles="analyst" +) + +# Run analysis with both analysts +equity_analyst_1.talk_to( + equity_analyst_2, + "Analyze high growth tech stocks focusing on fundamentals like revenue growth, margins, and market position. Create a detailed analysis table in markdown." +) diff --git a/loguru_test.py b/loguru_test.py new file mode 100644 index 00000000..f72ca938 --- /dev/null +++ b/loguru_test.py @@ -0,0 +1,506 @@ +from typing import Callable, Any +import psutil +import asyncio +from rich.console import Console +from rich.table import Table +from rich.panel import Panel +from rich.layout import Layout +from rich.tree import Tree +from rich.progress import ( + Progress, + TimeElapsedColumn, + TextColumn, + SpinnerColumn, +) +from rich.live import Live +from rich.text import Text +from typing import Dict, Optional, List +from dataclasses import dataclass +from queue import Queue +from datetime import datetime + +try: + import pynvml + + pynvml.nvmlInit() + GPU_ENABLED = True +except ImportError: + GPU_ENABLED = False + + +@dataclass +class SwarmMetadata: + name: str + description: str + version: str + type: str # hierarchical, parallel, sequential + created_at: datetime + author: str + tags: List[str] + primary_objective: str + secondary_objectives: List[str] + + +@dataclass +class Agent: + name: str + role: str + description: str + agent_type: str # e.g., "LLM", "Neural", "Rule-based" + capabilities: List[str] + parameters: Dict[str, any] + metadata: Dict[str, str] + children: List["Agent"] = None + parent: Optional["Agent"] = None + output_stream: Queue = None + + def __post_init__(self): + self.children = self.children or [] + self.output_stream = Queue() + + @property + def hierarchy_level(self) -> int: + level = 0 + current = self + while current.parent: + level += 1 + current = current.parent + return level + + +class SwarmVisualizationRich: + def __init__( + self, + swarm_metadata: SwarmMetadata, + root_agent: Agent, + update_resources: bool = True, + refresh_rate: float = 0.1, + ): + self.swarm_metadata = swarm_metadata + self.root_agent = root_agent + self.update_resources = update_resources + self.refresh_rate = refresh_rate + self.console = Console() + self.live = None + self.output_history = {} + + # System monitoring + self.cores_available = 0 + self.memory_usage = "N/A" + self.gpu_power = "N/A" + self.start_time = datetime.now() + + if self.update_resources: + self._update_resource_stats() + + def _format_uptime(self) -> str: + """Formats the swarm's uptime.""" + delta = datetime.now() - self.start_time + hours, remainder = divmod(delta.seconds, 3600) + minutes, seconds = divmod(remainder, 60) + return f"{hours:02d}:{minutes:02d}:{seconds:02d}" + + def _build_agent_tree( + self, agent: Agent, tree: Optional[Tree] = None + ) -> Tree: + """Builds a detailed tree visualization of the agent hierarchy.""" + agent_info = [ + f"[bold cyan]{agent.name}[/bold cyan]", + f"[yellow]Role:[/yellow] {agent.role}", + f"[green]Type:[/green] {agent.agent_type}", + f"[blue]Level:[/blue] {agent.hierarchy_level}", + f"[magenta]Capabilities:[/magenta] {', '.join(agent.capabilities)}", + ] + + # Add any custom metadata + for key, value in agent.metadata.items(): + agent_info.append(f"[white]{key}:[/white] {value}") + + # Parameters summary + param_summary = ", ".join( + f"{k}: {v}" for k, v in agent.parameters.items() + ) + agent_info.append( + f"[white]Parameters:[/white] {param_summary}" + ) + + node_text = "\n".join(agent_info) + + if tree is None: + tree = Tree(node_text) + else: + branch = tree.add(node_text) + tree = branch + + for child in agent.children: + self._build_agent_tree(child, tree) + + return tree + + def _count_agents(self, agent: Agent) -> int: + """Recursively counts total number of agents in the swarm.""" + count = 1 # Count current agent + for child in agent.children or []: + count += self._count_agents(child) + return count + + def _create_unified_info_panel(self) -> Panel: + """Creates a unified panel showing both swarm metadata and architecture.""" + # Create the main container + info_layout = Layout() + info_layout.split_column( + Layout(name="metadata", size=13), + Layout(name="architecture"), + ) + + # Calculate total agents + total_agents = self._count_agents(self.root_agent) + + # Metadata section + metadata_table = Table.grid(padding=1, expand=True) + metadata_table.add_column("Label", style="bold cyan") + metadata_table.add_column("Value", style="white") + + # System resources + if self.update_resources: + self._update_resource_stats() + + # Add description with proper wrapping + description_text = Text( + self.swarm_metadata.description, style="italic" + ) + description_text.wrap(self.console, width=60, overflow="fold") + + metadata_table.add_row("Swarm Name", self.swarm_metadata.name) + metadata_table.add_row("Description", description_text) + metadata_table.add_row("Version", self.swarm_metadata.version) + metadata_table.add_row("Total Agents", str(total_agents)) + metadata_table.add_row("Author", self.swarm_metadata.author) + metadata_table.add_row( + "System", + f"CPU: {self.cores_available} cores | Memory: {self.memory_usage}", + ) + metadata_table.add_row( + "Primary Objective", self.swarm_metadata.primary_objective + ) + + info_layout["metadata"].update(metadata_table) + + info_layout["metadata"].update(metadata_table) + + # Architecture section with tree visualization + architecture_tree = self._build_agent_tree(self.root_agent) + info_layout["architecture"].update(architecture_tree) + + return Panel( + info_layout, + title="[bold]Swarm Information & Architecture[/bold]", + ) + + def _create_outputs_panel(self) -> Panel: + """Creates a panel that displays stacked message history for all agents.""" + # Create a container for all messages across all agents + all_messages = [] + + def collect_agent_messages(agent: Agent): + """Recursively collect messages from all agents.""" + messages = self.output_history.get(agent.name, []) + for msg in messages: + all_messages.append( + { + "agent": agent.name, + "time": msg["time"], + "content": msg["content"], + "style": msg["style"], + } + ) + for child in agent.children: + collect_agent_messages(child) + + # Collect all messages + collect_agent_messages(self.root_agent) + + # Sort messages by timestamp + all_messages.sort(key=lambda x: x["time"]) + + # Create the stacked message display + Layout() + messages_container = [] + + for msg in all_messages: + # Create a panel for each message + message_text = Text() + message_text.append(f"[{msg['time']}] ", style="dim") + message_text.append( + f"{msg['agent']}: ", style="bold cyan" + ) + message_text.append(msg["content"], style=msg["style"]) + + messages_container.append(message_text) + + # Join all messages with line breaks + if messages_container: + final_text = Text("\n").join(messages_container) + else: + final_text = Text("No messages yet...", style="dim") + + # Create scrollable panel for all messages + return Panel( + final_text, + title="[bold]Agent Communication Log[/bold]", + border_style="green", + padding=(1, 2), + ) + + def _update_resource_stats(self): + """Updates system resource statistics.""" + self.cores_available = psutil.cpu_count(logical=True) + mem_info = psutil.virtual_memory() + total_gb = mem_info.total / (1024**3) + used_gb = mem_info.used / (1024**3) + self.memory_usage = f"{used_gb:.1f}GB / {total_gb:.1f}GB ({mem_info.percent}%)" + + if GPU_ENABLED: + try: + device_count = pynvml.nvmlDeviceGetCount() + gpu_info = [] + for i in range(device_count): + handle = pynvml.nvmlDeviceGetHandleByIndex(i) + name = pynvml.nvmlDeviceGetName(handle).decode() + mem = pynvml.nvmlDeviceGetMemoryInfo(handle) + usage = (mem.used / mem.total) * 100 + gpu_info.append(f"{name}: {usage:.1f}%") + self.gpu_power = " | ".join(gpu_info) + except Exception as e: + self.gpu_power = f"GPU Error: {str(e)}" + else: + self.gpu_power = "No GPU detected" + + async def stream_output( + self, + agent: Agent, + text: str, + title: Optional[str] = None, + style: str = "bold cyan", + delay: float = 0.05, + by_word: bool = False, + ): + """ + Streams output for a specific agent with sophisticated token-by-token animation. + + Args: + agent (Agent): The agent whose output is being streamed + text (str): The text to stream + title (Optional[str]): Custom title for the output panel + style (str): Style for the output text + delay (float): Delay between tokens + by_word (bool): If True, streams word by word instead of character by character + """ + display_text = Text(style=style) + current_output = "" + + # Split into words or characters + tokens = text.split() if by_word else text + + # Create a panel for this agent's output + title = title or f"{agent.name} Output" + + for token in tokens: + # Add appropriate spacing + token_with_space = token + (" " if by_word else "") + current_output += token_with_space + display_text.append(token_with_space) + + # Initialize history list if it doesn't exist + if agent.name not in self.output_history: + self.output_history[agent.name] = [] + + # Store the complete message when finished + if token == tokens[-1]: + timestamp = datetime.now().strftime("%H:%M:%S") + self.output_history[agent.name].append( + { + "time": timestamp, + "content": current_output, + "style": style, + } + ) + + # Update live display if active + if self.live: + self.live.update(self._create_layout()) + await asyncio.sleep(delay) + + async def print_progress( + self, + description: str, + task_fn: Callable, + *args: Any, + **kwargs: Any, + ) -> Any: + """ + Displays a progress spinner while executing a task. + + Args: + description (str): Task description + task_fn (Callable): Function to execute + *args (Any): Arguments for task_fn + **kwargs (Any): Keyword arguments for task_fn + + Returns: + Any: Result of task_fn + """ + progress = Progress( + SpinnerColumn(), + TextColumn("[progress.description]{task.description}"), + TimeElapsedColumn(), + ) + + try: + with progress: + task = progress.add_task(description, total=None) + result = await task_fn(*args, **kwargs) + progress.update(task, completed=True) + return result + except Exception as e: + progress.stop() + raise e + + def _create_layout(self) -> Layout: + """Creates the main visualization layout.""" + layout = Layout() + layout.split_row( + Layout(name="info", ratio=2), + Layout(name="outputs", ratio=3), + ) + + layout["info"].update(self._create_unified_info_panel()) + layout["outputs"].update(self._create_outputs_panel()) + + return layout + + async def start(self): + """Starts the visualization with live updates.""" + with Live( + self._create_layout(), + refresh_per_second=int(1 / self.refresh_rate), + ) as self.live: + while True: + + def process_agent_streams(agent: Agent): + while not agent.output_stream.empty(): + new_output = agent.output_stream.get() + asyncio.create_task( + self.stream_output(agent, new_output) + ) + for child in agent.children: + process_agent_streams(child) + + process_agent_streams(self.root_agent) + await asyncio.sleep(self.refresh_rate) + + +# Example usage +if __name__ == "__main__": + # Create swarm metadata + swarm_metadata = SwarmMetadata( + name="Financial Advisory Swarm", + description="Intelligent swarm for financial analysis and advisory", + version="1.0.0", + type="hierarchical", + created_at=datetime.now(), + author="AI Research Team", + # tags=["finance", "analysis", "advisory"], + primary_objective="Provide comprehensive financial analysis and recommendations", + secondary_objectives=[ + "Monitor market trends", + "Analyze competitor behavior", + "Generate investment strategies", + ], + ) + + # Create agent hierarchy with detailed parameters + analyst = Agent( + name="Financial Analyst", + role="Analysis", + description="Analyzes financial data and market trends", + agent_type="LLM", + capabilities=[ + "data analysis", + "trend detection", + "risk assessment", + ], + parameters={"model": "gpt-4", "temperature": 0.7}, + metadata={ + "specialty": "Market Analysis", + "confidence_threshold": "0.85", + }, + ) + + researcher = Agent( + name="Market Researcher", + role="Research", + description="Conducts market research and competitor analysis", + agent_type="Neural", + capabilities=[ + "competitor analysis", + "market sentiment", + "trend forecasting", + ], + parameters={"batch_size": 32, "learning_rate": 0.001}, + metadata={ + "data_sources": "Bloomberg, Reuters", + "update_frequency": "1h", + }, + ) + + advisor = Agent( + name="Investment Advisor", + role="Advisory", + description="Provides investment recommendations", + agent_type="Hybrid", + capabilities=[ + "portfolio optimization", + "risk management", + "strategy generation", + ], + parameters={ + "risk_tolerance": "moderate", + "time_horizon": "long", + }, + metadata={ + "certification": "CFA Level 3", + "specialization": "Equity", + }, + children=[analyst, researcher], + ) + + # Create visualization + viz = SwarmVisualizationRich( + swarm_metadata=swarm_metadata, + root_agent=advisor, + refresh_rate=0.1, + ) + + # Example of streaming output simulation + async def simulate_outputs(): + await viz.stream_output( + advisor, + "Analyzing market conditions...\nGenerating investment advice...", + ) + await viz.stream_output( + analyst, + "Processing financial data...\nIdentifying trends...", + ) + await viz.stream_output( + researcher, + "Researching competitor movements...\nAnalyzing market share...", + ) + + # Run the visualization + async def main(): + viz_task = asyncio.create_task(viz.start()) + await simulate_outputs() + await viz_task + + asyncio.run(main()) diff --git a/open_scientist.py b/open_scientist.py new file mode 100644 index 00000000..a6fbebdd --- /dev/null +++ b/open_scientist.py @@ -0,0 +1,541 @@ +import json +import os +from datetime import datetime +from typing import Any, Dict + +import requests +from dotenv import load_dotenv +from rich.console import Console + +from swarms import Agent, SequentialWorkflow + +console = Console() +load_dotenv() + + +############################################################################### +# 1. System Prompts for Each Scientist Agent +############################################################################### + + +def format_exa_results(json_data: Dict[str, Any]) -> str: + """Formats Exa.ai search results into structured text""" + formatted_text = [] + + if "error" in json_data: + return f"### Error\n{json_data['error']}\n" + + # Extract search metadata + search_params = json_data.get("effectiveFilters", {}) + query = search_params.get("query", "General web search") + formatted_text.append( + f"### Exa Search Results for: '{query}'\n\n---\n" + ) + + # Process results + results = json_data.get("results", []) + + if not results: + formatted_text.append("No results found.\n") + else: + for i, result in enumerate(results, 1): + title = result.get("title", "No title") + url = result.get("url", result.get("id", "No URL")) + published_date = result.get("publishedDate", "") + + # Handle highlights + highlights = result.get("highlights", []) + highlight_text = ( + "\n".join( + [ + ( + h.get("text", h) + if isinstance(h, dict) + else str(h) + ) + for h in highlights[:3] + ] + ) + if highlights + else "No summary available" + ) + + formatted_text.extend( + [ + f"{i}. **{title}**\n", + f" - URL: {url}\n", + f" - Published: {published_date.split('T')[0] if published_date else 'Date unknown'}\n", + f" - Key Points:\n {highlight_text}\n\n", + ] + ) + + return "".join(formatted_text) + + +def exa_search(query: str, **kwargs: Any) -> str: + """Performs web search using Exa.ai API""" + api_url = "https://api.exa.ai/search" + headers = { + "x-api-key": os.getenv("EXA_API_KEY"), + "Content-Type": "application/json", + } + + payload = { + "query": query, + "useAutoprompt": True, + "numResults": kwargs.get("num_results", 10), + "contents": { + "text": True, + "highlights": {"numSentences": 2}, + }, + **kwargs, + } + + try: + response = requests.post( + api_url, json=payload, headers=headers + ) + response.raise_for_status() + response_json = response.json() + + console.print("\n[bold]Exa Raw Response:[/bold]") + console.print(json.dumps(response_json, indent=2)) + + formatted_text = format_exa_results( + response_json + ) # Correct function call + + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + filename = f"exa_search_results_{timestamp}.txt" + with open(filename, "w", encoding="utf-8") as f: + f.write(formatted_text) + + return formatted_text + + except requests.exceptions.RequestException as e: + error_msg = f"Exa search request failed: {str(e)}" + console.print(f"[bold red]{error_msg}[/bold red]") + return error_msg + except json.JSONDecodeError as e: + error_msg = f"Invalid Exa response: {str(e)}" + console.print(f"[bold red]{error_msg}[/bold red]") + return error_msg + except Exception as e: + error_msg = f"Unexpected error: {str(e)}" + console.print(f"[bold red]{error_msg}[/bold red]") + return error_msg + + +# if __name__ == "__main__": +# console.print("\n[bold]Example Exa.ai Search:[/bold]") +# results = exa_search("Deepseek news") +# console.print("\n[bold green]Formatted Exa Results:[/bold green]") +# console.print(results) + + +SUPERVISOR_AGENT_SYS_PROMPT = """ +You are the SUPERVISOR AGENT, the central coordinator of a multi-agent research system. +Your responsibilities include: +1. **Overall Orchestration**: You manage and delegate tasks to specialized scientist agents: + - Generation Agent + - Review Agent + - Ranking Agent + - Evolution Agent + - Proximity Agent + - Meta-Review Agent +2. **Workflow Guidance**: You oversee the flow of information between agents, ensuring + that each agent receives the necessary data to fulfill its role effectively. +3. **Quality Assurance**: You evaluate and monitor the outputs of all agents to confirm + that they align with the user’s overarching objectives and constraints. +4. **Conflict Resolution**: If there are conflicting viewpoints or recommendations among + the agents, you facilitate a resolution by encouraging deeper analysis or gathering + more information as needed. +5. **User-Facing Summary**: You compile the important findings, recommendations, or + next-step instructions from each agent into a cohesive overview for the user. + +**KEY POINTS FOR THE SUPERVISOR AGENT**: +- Maintain clarity: explicitly direct each agent in how to proceed. +- Provide necessary context to each agent when you pass tasks along, ensuring they have + relevant details to operate effectively. +- Continuously keep track of each agent’s outcomes: generation, review, ranking, + evolutionary refinement, proximity assessment, and meta-review summaries. +- Strive to deliver consolidated, actionable suggestions or final decisions that directly + reflect the user's aims. + +Your tone should be **methodical, organized, and firm**. Always ensure that each subsequent +step is clear to both the agents and the user. If something is ambiguous, request +clarifications or deeper analysis from the appropriate agent. +""" + +GENERATION_AGENT_SYS_PROMPT = """ +You are the GENERATION AGENT, responsible for ideation and initial proposal creation +within a multi-agent research system. Your duties include: +1. **Idea Synthesis**: Transform the user’s research goal, along with any guidance from + the Supervisor Agent, into multiple innovative, feasible proposals. +2. **Conceptual Breadth**: Provide a variety of approaches to ensure that the system can + explore multiple avenues. Avoid fixating on a single concept. +3. **Rationale and Context**: For each proposed idea, supply clear reasoning, referencing + relevant background knowledge, possible methods, or prior art when useful. +4. **Clarity and Organization**: Present your ideas so that the next agents (Review, + Ranking, Evolution, etc.) can parse and evaluate them effectively. + +**GUIDELINES**: +- Remain imaginative yet grounded in the practical constraints provided by the user + (e.g., limited compute resources, timeline constraints). +- Where beneficial, highlight potential trade-offs or known challenges, but do not let + these limit your creativity; note them as points the Review Agent might scrutinize. +- Provide enough detail to enable further refinement, but do not overload the system with + excessive complexity at this stage. + +Your tone should be **inquisitive, creative, and detailed**. Aim to generate ideas +that could stimulate rigorous evaluation by the other agents in the system. +""" + +REVIEW_AGENT_SYS_PROMPT = """ +You are the REVIEW AGENT, tasked with critically examining the ideas generated by +the Generation Agent. Your primary roles: +1. **Critical Analysis**: Evaluate each idea's feasibility, potential impact, strengths, + and weaknesses. Consider real-world constraints such as data availability, compute + limitations, complexity of implementation, or state-of-the-art performance standards. +2. **Constructive Feedback**: Offer specific suggestions for how an idea might be + improved, streamlined, or combined with other approaches. Clearly highlight what is + missing, needs more elaboration, or risks failure. +3. **Consistency and Credibility**: Check that each idea adheres to the user’s + overarching goals and does not conflict with known facts or constraints. If an idea + is overly ambitious or deviates from the domain’s realities, note this and propose + ways to mitigate. +4. **Readiness Level**: Provide a sense of whether each idea is ready for immediate + testing or if it requires further planning. Label each as "promising and nearly + ready," "promising but needs more work," or "high risk / questionable feasibility." + +**EXPECTED STYLE**: +- Write with **thoroughness and clarity**. You serve as an internal “peer reviewer.” +- Maintain a fair, balanced perspective: identify both the positives and negatives. +- If multiple ideas overlap or can be combined, advise on how that might yield + stronger outcomes. + +Remember, your assessment informs the Ranking, Evolution, and other agents, so +be both structured and concise where possible. +""" + +RANKING_AGENT_SYS_PROMPT = """ +You are the RANKING AGENT, responsible for organizing the ideas (and their accompanying +reviews) by importance, feasibility, or impact. Your specific aims: +1. **Assessment of Criteria**: Leverage the user’s primary objectives (e.g., performance, + resource constraints, novelty) and the Review Agent’s feedback to determine a robust + ranking of ideas from most to least promising. +2. **Transparent Rationale**: Explicitly justify your ranking methodology. Discuss why + one idea outranks another (e.g., higher potential for success, better alignment with + constraints, synergy with existing research). +3. **Concise Format**: Provide a clear, itemized hierarchy (1st, 2nd, 3rd, etc.) so that + it can be easily read and acted upon by the Supervisor Agent or other agents. + +**THINGS TO CONSIDER**: +- The user’s stated constraints (like limited GPU resources, time to market, etc.). +- The Review Agent’s critiques (strengths, weaknesses, suggestions). +- Potential synergy or overlaps between ideas—sometimes a combined approach can rank + higher if it builds on strengths. + +**WRITING STYLE**: +- Keep it **orderly and succinct**. State the ranking results plainly. +- If two ideas are effectively tied, you can mention that, but make a clear call + regarding priority. + +Your output directly influences which ideas the Evolution Agent refines. +""" + +EVOLUTION_AGENT_SYS_PROMPT = """ +You are the EVOLUTION AGENT, focusing on refining, adapting, or "mutating" the top-ranked +ideas. Your objectives: +1. **Targeted Improvements**: Based on the Ranking Agent’s selection of top concepts + and the Review Agent’s critiques, systematically enhance each idea. Incorporate + suggestions from the reviews and address potential weaknesses or limitations. +2. **Novelty and Iteration**: Where beneficial, propose new variations or sub-concepts + that push the idea’s capabilities further while staying within the user’s constraints. +3. **Implementation Detailing**: Provide more detail on how the evolved ideas might be + implemented in practical scenarios—what frameworks, data, or training procedures + might be used, and how to measure success. +4. **Ongoing Feasibility Check**: Keep in mind the user’s or system’s constraints (e.g., + limited GPU hours, certain hardware specs). Make sure your evolutions do not + inadvertently break these constraints. + +**RESPONSE GUIDELINES**: +- Offer a brief summary of how you modified each idea. +- Explain **why** those changes address prior weaknesses or align more closely with + the user's needs. +- If you propose multiple variants of the same idea, clarify how they differ and + the pros/cons of each. + +Your style should be **solution-oriented, thoughtful,** and mindful of resource usage +and domain constraints. +""" + +PROXIMITY_AGENT_SYS_PROMPT = """ +You are the PROXIMITY AGENT, responsible for evaluating how closely each refined idea +meets the user’s declared research goal. Your role: +1. **Alignment Check**: Compare the current iteration of ideas to the original + requirements (e.g., a specific performance threshold, limited computational budget). +2. **Scoring or Rating**: Provide a numeric or qualitative measure of "distance" from + full satisfaction of the goal. If possible, highlight the factors that keep the idea + from meeting the goal (e.g., insufficient training efficiency, unclear data sources). +3. **Gap Analysis**: Suggest which aspects need further refinement, or note if an idea + is essentially ready for practical deployment or testing. + +**COMMUNICATION STYLE**: +- Be **direct and evidence-based**: reference the user’s constraints or the system’s + specs as found in the other agents’ discussions. +- Provide a clear, easily interpretable scoring metric (e.g., a scale of 0–100, or + a textual label like “very close,” “moderately aligned,” “far from ready”). +- Keep it succinct so subsequent steps can parse your results quickly. + +Your evaluation will help the Supervisor Agent and the Meta-Review Agent see how +much remains to be done for each idea. +""" + +META_REVIEW_AGENT_SYS_PROMPT = """ +You are the META-REVIEW AGENT, tasked with synthesizing input from the Review Agent, +Ranking Agent, Evolution Agent, and Proximity Agent (and any other relevant feedback). +Your jobs: +1. **Holistic Synthesis**: Provide an overarching analysis that captures the major + pros, cons, controversies, and consensus points across all agent reports. +2. **Actionable Summary**: Summarize the final or near-final recommendations, offering + a bottom-line statement on which idea(s) stand out and how ready they are. +3. **Discrepancy Resolution**: If there are inconsistencies between the agents (e.g., + the Ranking Agent has a different view than the Proximity Agent), address them and + either reconcile or highlight them for the Supervisor Agent to decide. +4. **Roadmap**: Propose next steps (e.g., further refinement, additional data + collection, experiments) if the ideas are not yet fully converged. + +**EXPECTED OUTPUT STYLE**: +- Provide a **concise but comprehensive** overview. Do not drown the user in repetition. +- Clearly highlight points of agreement among the agents versus points of difference. +- End with a recommendation for how to finalize or further develop the top ideas, + taking into account the user’s research goal. + +Your tone should be **balanced and authoritative**, reflecting the aggregated wisdom +of the entire system. +""" + +RESEARCH_AGENT_SYS_PROMPT = """ +You are the RESEARCH AGENT, tasked with formulating search queries and gathering +relevant information on any given topic. Your responsibilities include: +1. **Query Formulation**: Create effective search queries that can be used to find + information on the topic provided by the user or other agents. +2. **Information Gathering**: Use the formulated queries to collect data, articles, + papers, or any other relevant information that can aid in understanding the topic. +3. **Summarization**: Provide a concise summary of the gathered information, highlighting + key points, trends, or insights that are relevant to the research goal. + +**GUIDELINES**: +- Ensure that the search queries are specific enough to yield useful results but broad + enough to cover different aspects of the topic. +- Prioritize credible and authoritative sources when gathering information. +- Present the information in a clear and organized manner, making it easy for other + agents or the user to understand and utilize. + +Your tone should be **informative and precise**. Aim to provide comprehensive insights +that can support further exploration or decision-making by the system. + + +""" + + +PROTEIN_GENERATION_AGENT_SYS_PROMPT = """ +You are the PROTEIN GENERATION AGENT, responsible for generating a protein sequence that can be used to create a drug for alzheimer's disease. + +Output only the protein sequence, nothing else. + + +""" + + +############################################################################### +# 2. Instantiate Each Agent +############################################################################### + +# Supervisor Agent +supervisor_agent = Agent( + agent_name="Supervisor-Agent", + system_prompt=SUPERVISOR_AGENT_SYS_PROMPT, + model_name="gpt-4o-mini", # Example placeholder + max_loops=1, +) + +research_agent = Agent( + agent_name="Research-Agent", + system_prompt=RESEARCH_AGENT_SYS_PROMPT, + model_name="gpt-4o-mini", + max_loops=1, + tools=[exa_search], +) + + +# Generation Agent +generation_agent = Agent( + agent_name="Generation-Agent", + system_prompt=GENERATION_AGENT_SYS_PROMPT, + model_name="gpt-4o-mini", + max_loops=1, +) + +# Review Agent +review_agent = Agent( + agent_name="Review-Agent", + system_prompt=REVIEW_AGENT_SYS_PROMPT, + model_name="gpt-4o-mini", + max_loops=1, +) + +# Ranking Agent +ranking_agent = Agent( + agent_name="Ranking-Agent", + system_prompt=RANKING_AGENT_SYS_PROMPT, + model_name="gpt-4o-mini", + max_loops=1, +) + +# Evolution Agent +evolution_agent = Agent( + agent_name="Evolution-Agent", + system_prompt=EVOLUTION_AGENT_SYS_PROMPT, + model_name="gpt-4o-mini", + max_loops=1, +) + +# Proximity Agent +proximity_agent = Agent( + agent_name="Proximity-Agent", + system_prompt=PROXIMITY_AGENT_SYS_PROMPT, + model_name="gpt-4o-mini", + max_loops=1, +) + +# Meta-Review Agent +meta_review_agent = Agent( + agent_name="Meta-Review-Agent", + system_prompt=META_REVIEW_AGENT_SYS_PROMPT, + model_name="gpt-4o-mini", + max_loops=1, +) + + +protein_generation_agent = Agent( + agent_name="Protein-Generation-Agent", + system_prompt=PROTEIN_GENERATION_AGENT_SYS_PROMPT, + model_name="gpt-4o-mini", + max_loops=1, +) + +############################################################################### +# 3. Example Usage (Supervisor-Orchestrated Workflow) +############################################################################### + +# def run_example_flow(): +# """ +# Demonstrates an example of how the Supervisor Agent might orchestrate +# the multi-agent system to explore a user’s research goal. +# """ + +# # A sample user-defined research goal +# research_goal = ( +# "Design a novel approach to training a neural architecture that can " +# "outperform state-of-the-art models in image classification " +# "with limited GPU resources available." +# ) + +# # ------------------------------------------------------------------------- +# # Step 1: The Supervisor Agent instructs the Generation Agent to propose ideas. +# # ------------------------------------------------------------------------- +# generated_ideas = generation_agent.run( +# f"Please propose 3 distinctive ideas to achieve this goal:\n\n{research_goal}" +# ) + +# # ------------------------------------------------------------------------- +# # Step 2: The Supervisor Agent sends the generated ideas to the Review Agent. +# # ------------------------------------------------------------------------- +# reviewed_ideas = review_agent.run( +# f"Here are the generated ideas:\n{generated_ideas}\n\n" +# "Please critique each idea in detail and suggest improvements." +# ) + +# # ------------------------------------------------------------------------- +# # Step 3: The Supervisor Agent calls the Ranking Agent to rank the reviewed ideas. +# # ------------------------------------------------------------------------- +# ranked_ideas = ranking_agent.run( +# f"The Review Agent offered these critiques:\n{reviewed_ideas}\n\n" +# "Please provide a ranked list of these ideas from most to least promising, " +# "with brief justifications." +# ) + +# # ------------------------------------------------------------------------- +# # Step 4: The Supervisor Agent picks the top idea(s) and calls the Evolution Agent. +# # ------------------------------------------------------------------------- +# evolved_ideas = evolution_agent.run( +# f"Top-ranked concept(s):\n{ranked_ideas}\n\n" +# "Based on the feedback above, evolve or refine the best ideas. Please provide " +# "detailed modifications or new variants that address the critiques." +# ) + +# # ------------------------------------------------------------------------- +# # Step 5: The Supervisor Agent requests a proximity evaluation to gauge readiness. +# # ------------------------------------------------------------------------- +# proximity_feedback = proximity_agent.run( +# f"User goal:\n{research_goal}\n\n" +# f"Refined ideas:\n{evolved_ideas}\n\n" +# "How close are these ideas to achieving the stated goal? Provide a proximity " +# "metric or rating and justify your reasoning." +# ) + +# # ------------------------------------------------------------------------- +# # Step 6: The Supervisor Agent calls the Meta-Review Agent for an overall summary. +# # ------------------------------------------------------------------------- +# meta_review = meta_review_agent.run( +# f"Review Feedback:\n{reviewed_ideas}\n\n" +# f"Ranking:\n{ranked_ideas}\n\n" +# f"Evolved Ideas:\n{evolved_ideas}\n\n" +# f"Proximity Feedback:\n{proximity_feedback}\n\n" +# "Please synthesize all of this feedback into a final meta-review. Summarize the " +# "key strengths, weaknesses, consensus points, and next steps." +# ) + +# # ------------------------------------------------------------------------- +# # Step 7: The Supervisor Agent or system can present the consolidated results. +# # ------------------------------------------------------------------------- +# print("=== Generated Ideas ===") +# print(generated_ideas, "\n") +# print("=== Review Feedback ===") +# print(reviewed_ideas, "\n") +# print("=== Ranking ===") +# print(ranked_ideas, "\n") +# print("=== Evolved Ideas ===") +# print(evolved_ideas, "\n") +# print("=== Proximity Feedback ===") +# print(proximity_feedback, "\n") +# print("=== Meta-Review ===") +# print(meta_review, "\n") + + +# if __name__ == "__main__": +# # Example run to demonstrate the multi-agent workflow +# run_example_flow() + + +swarm = SequentialWorkflow( + agents=[ + generation_agent, + review_agent, + # research_agent, + ranking_agent, + # generation_agent, + # supervisor_agent, + # evolution_agent, + # proximity_agent, + meta_review_agent, + # protein_generation_agent + ], + name="Open Scientist", + description="This is a swarm that uses a multi-agent system to explore a user's research goal.", +) + + +swarm.run("Let's create a drug for alzheimer's disease.") diff --git a/swarms/structs/rearrange.py b/swarms/structs/rearrange.py index b604eb59..ff15993b 100644 --- a/swarms/structs/rearrange.py +++ b/swarms/structs/rearrange.py @@ -263,11 +263,11 @@ class AgentRearrange(BaseSwarm): ) agents_in_flow.append(agent_name) - # If the length of the agents does not equal the length of the agents in flow - if len(set(agents_in_flow)) != len(agents_in_flow): - raise ValueError( - "Duplicate agent names in the flow are not allowed." - ) + # # If the length of the agents does not equal the length of the agents in flow + # if len(set(agents_in_flow)) != len(agents_in_flow): + # raise ValueError( + # "Duplicate agent names in the flow are not allowed." + # ) logger.info(f"Flow: {self.flow} is valid.") return True diff --git a/swarms/tools/tool_parse_exec.py b/swarms/tools/tool_parse_exec.py index f118319d..f09b7d7d 100644 --- a/swarms/tools/tool_parse_exec.py +++ b/swarms/tools/tool_parse_exec.py @@ -1,5 +1,6 @@ import json from typing import List, Any, Callable +import re from swarms.utils.parse_code import extract_code_from_markdown from swarms.utils.loguru_logger import initialize_logger @@ -33,9 +34,100 @@ def parse_and_execute_json( if parse_md: try: - json_string = extract_code_from_markdown(json_string) + code_blocks = re.findall( + r"```(?:json)?\s*([\s\S]*?)```", json_string + ) + if code_blocks and len(code_blocks) > 1: + function_dict = { + func.__name__: func for func in functions + } + + def process_json_block(json_block: str) -> dict: + try: + json_block = json_block.strip() + if not json_block: + raise ValueError("JSON block is empty") + data = json.loads(json_block) + function_list = [] + if "functions" in data: + function_list = data["functions"] + elif "function" in data: + function_list = [data["function"]] + else: + function_list = [data] + if isinstance(function_list, dict): + function_list = [function_list] + function_list = [ + f for f in function_list if f + ] + + block_results = {} + for function_data in function_list: + function_name = function_data.get("name") + parameters = function_data.get( + "parameters", {} + ) + + if not function_name: + logger.warning( + "Function data missing 'name' field" + ) + continue + + if function_name not in function_dict: + logger.warning( + f"Function '{function_name}' not found" + ) + block_results[function_name] = ( + "Error: Function not found" + ) + continue + + for attempt in range(max_retries): + try: + result = function_dict[ + function_name + ](**parameters) + block_results[function_name] = ( + str(result) + ) + logger.info( + f"Result for {function_name}: {result}" + ) + break + except Exception as e: + logger.error( + f"Attempt {attempt + 1} failed for {function_name}: {e}" + ) + if attempt == max_retries - 1: + block_results[ + function_name + ] = f"Error after {max_retries} attempts: {str(e)}" + return block_results + except Exception as e: + logger.error( + f"Failed to process JSON block: {e}" + ) + return { + "error": f"Failed to process block: {str(e)}" + } + + combined_results = {} + for idx, block in enumerate(code_blocks, start=1): + combined_results[f"block_{idx}"] = ( + process_json_block(block) + ) + return json.dumps( + {"results": combined_results}, indent=4 + ) + elif code_blocks: + json_string = code_blocks[0] + else: + json_string = extract_code_from_markdown(json_string) except Exception as e: - logger.error(f"Error extracting code from Markdown: {e}") + logger.error( + f"Error extracting code blocks from Markdown: {e}" + ) return {"error": f"Markdown parsing failed: {str(e)}"} try: