pull/709/head
harshalmore31 4 months ago
parent baa4215083
commit 5cf3219e8e

@ -5,22 +5,10 @@ import json
import time import time
import asyncio import asyncio
import gradio as gr import gradio as gr
import re
# Import necessary classes and functions from swarms library
from swarms.structs.agent import Agent from swarms.structs.agent import Agent
from swarms.structs.concurrent_workflow import ConcurrentWorkflow
from swarms.structs.mixture_of_agents import MixtureOfAgents
from swarms.structs.rearrange import AgentRearrange
from swarms.structs.sequential_workflow import SequentialWorkflow
from swarms.structs.spreadsheet_swarm import SpreadSheetSwarm
from swarms.structs.swarm_matcher import swarm_matcher, SwarmMatcher, SwarmMatcherConfig, initialize_swarm_types
from swarms.structs.swarm_router import SwarmRouter from swarms.structs.swarm_router import SwarmRouter
from swarms.utils.loguru_logger import initialize_logger from swarms.utils.loguru_logger import initialize_logger
from groq_model import OpenAIChat # Import OpenAIChat from the correct location from swarm_models import OpenAIChat
from swarms.utils.file_processing import create_file_in_folder
from doc_master import doc_master
# Initialize logger # Initialize logger
@ -33,6 +21,9 @@ load_dotenv()
# Get the OpenAI API key from the environment variable # Get the OpenAI API key from the environment variable
api_key = os.getenv("GROQ_API_KEY") api_key = os.getenv("GROQ_API_KEY")
# changed to swarms_models
# adding functionality to view other models of swarms models
# Model initialization # Model initialization
model = OpenAIChat( model = OpenAIChat(
openai_api_base="https://api.groq.com/openai/v1", openai_api_base="https://api.groq.com/openai/v1",
@ -42,21 +33,18 @@ model = OpenAIChat(
) )
# Define the path to agent_prompts.json # Define the path to agent_prompts.json
# locates the json file and then fetches the promopts
PROMPT_JSON_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "agent_prompts.json") PROMPT_JSON_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "agent_prompts.json")
logger.info(f"Loading prompts from: {PROMPT_JSON_PATH}") logger.info(f"Loading prompts from: {PROMPT_JSON_PATH}")
# Global log storage # Global log storage
execution_logs = [] execution_logs = []
def log_event(level: str, message: str, metadata: Optional[Dict] = None):
"""
Log an event and store it in the execution logs.
Args: def log_event(level: str, message: str, metadata: Optional[Dict] = None):
level: Log level (info, warning, error, etc.)
message: Log message
metadata: Optional metadata dictionary
"""
timestamp = time.strftime("%Y-%m-%d %H:%M:%S") timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
log_entry = { log_entry = {
"timestamp": timestamp, "timestamp": timestamp,
@ -65,25 +53,12 @@ def log_event(level: str, message: str, metadata: Optional[Dict] = None):
"metadata": metadata or {} "metadata": metadata or {}
} }
execution_logs.append(log_entry) execution_logs.append(log_entry)
# Also log to the logger
log_func = getattr(logger, level.lower(), logger.info) log_func = getattr(logger, level.lower(), logger.info)
log_func(message) log_func(message)
def get_logs(router: Optional['SwarmRouter'] = None) -> List[str]: def get_logs(router: Optional['SwarmRouter'] = None) -> List[str]:
"""
Get formatted logs from both the execution logs and router logs if available.
Args:
router: Optional SwarmRouter instance to get additional logs from
Returns:
List of formatted log strings
"""
formatted_logs = [] formatted_logs = []
# Add execution logs
for log in execution_logs: for log in execution_logs:
metadata_str = "" metadata_str = ""
if log["metadata"]: if log["metadata"]:
@ -107,12 +82,10 @@ def get_logs(router: Optional['SwarmRouter'] = None) -> List[str]:
def clear_logs(): def clear_logs():
"""Clear the execution logs."""
execution_logs.clear() execution_logs.clear()
def load_prompts_from_json() -> Dict[str, str]: def load_prompts_from_json() -> Dict[str, str]:
"""Robust prompt loading with comprehensive error handling."""
try: try:
if not os.path.exists(PROMPT_JSON_PATH): if not os.path.exists(PROMPT_JSON_PATH):
error_msg = f"Prompts file not found at: {PROMPT_JSON_PATH}" error_msg = f"Prompts file not found at: {PROMPT_JSON_PATH}"
@ -169,15 +142,10 @@ def load_prompts_from_json() -> Dict[str, str]:
} }
# Load prompts
AGENT_PROMPTS = load_prompts_from_json() AGENT_PROMPTS = load_prompts_from_json()
def initialize_agents(
data_temp: float, def initialize_agents(dynamic_temp: float, agent_keys: List[str]) -> List[Agent]:
sum_temp: float,
agent_keys: List[str]
) -> List[Agent]:
"""Enhanced agent initialization with more robust configuration."""
agents = [] agents = []
seen_names = set() seen_names = set()
for agent_key in agent_keys: for agent_key in agent_keys:
@ -207,17 +175,18 @@ def initialize_agents(
user_name="pe_firm", user_name="pe_firm",
retry_attempts=1, retry_attempts=1,
context_length=200000, context_length=200000,
output_type="string", output_type="string", # here is the output type which is string
temperature=data_temp, temperature=dynamic_temp,
) )
agents.append(agent) agents.append(agent)
return agents return agents
async def execute_task(task: str, max_loops: int, data_temp: float, sum_temp: float,
swarm_type: str, agent_keys: List[str], flow: str = None) -> Tuple[Dict[str, str], 'SwarmRouter', str]: async def execute_task(task: str, max_loops: int, dynamic_temp: float,
swarm_type: str, agent_keys: List[str], flow: str = None) -> Tuple[Any, 'SwarmRouter', str]:
""" """
Enhanced task execution with comprehensive error handling and result processing. Enhanced task execution with comprehensive error handling and raw result return.
""" """
start_time = time.time() start_time = time.time()
log_event("info", f"Starting task execution: {task}") log_event("info", f"Starting task execution: {task}")
@ -225,14 +194,14 @@ async def execute_task(task: str, max_loops: int, data_temp: float, sum_temp: fl
try: try:
# Initialize agents # Initialize agents
try: try:
agents = initialize_agents(data_temp, sum_temp, agent_keys) agents = initialize_agents(dynamic_temp, agent_keys)
log_event("info", f"Successfully initialized {len(agents)} agents") log_event("info", f"Successfully initialized {len(agents)} agents")
except Exception as e: except Exception as e:
error_msg = f"Agent initialization error: {str(e)}" error_msg = f"Agent initialization error: {str(e)}"
log_event("error", error_msg) log_event("error", error_msg)
return {}, None, error_msg return None, None, error_msg
# Create a SwarmRouter to manage the different swarm types # Swarm-specific configurations
router_kwargs = { router_kwargs = {
"name": "multi-agent-workflow", "name": "multi-agent-workflow",
"description": f"Executing {swarm_type} workflow", "description": f"Executing {swarm_type} workflow",
@ -240,11 +209,21 @@ async def execute_task(task: str, max_loops: int, data_temp: float, sum_temp: fl
"agents": agents, "agents": agents,
"autosave": True, "autosave": True,
"return_json": True, "return_json": True,
"output_type": "string" "output_type": "string",
"swarm_type": swarm_type, # Pass swarm_type here
} }
# Swarm-specific configurations if swarm_type == "AgentRearrange":
if not flow:
return None, None, "Flow configuration is required for AgentRearrange"
router_kwargs["flow"] = flow
if swarm_type == "MixtureOfAgents":
if len(agents) < 2:
return None, None, "MixtureOfAgents requires at least 2 agents"
if swarm_type == "SpreadSheetSwarm": if swarm_type == "SpreadSheetSwarm":
# spread sheet swarm needs specific setup
output_dir = "swarm_outputs" output_dir = "swarm_outputs"
os.makedirs(output_dir, exist_ok=True) os.makedirs(output_dir, exist_ok=True)
@ -262,139 +241,50 @@ async def execute_task(task: str, max_loops: int, data_temp: float, sum_temp: fl
except OSError as e: except OSError as e:
error_msg = f"Invalid output path: {str(e)}" error_msg = f"Invalid output path: {str(e)}"
log_event("error", error_msg) log_event("error", error_msg)
return {}, None, error_msg return None, None, error_msg
# Create and initialize SpreadSheetSwarm router_kwargs["output_path"] = output_path
# Create and execute SwarmRouter
try: try:
swarm = SpreadSheetSwarm( timeout = 450 if swarm_type != "SpreadSheetSwarm" else 900 # SpreadSheetSwarm will have different timeout.
agents=agents, router = SwarmRouter(**router_kwargs)
max_loops=max_loops,
name="spreadsheet-swarm", if swarm_type == "AgentRearrange":
description="SpreadSheet processing workflow", result = await asyncio.wait_for(
save_file_path=output_path, asyncio.to_thread(router._run, task),
workspace_dir=output_dir, timeout=timeout
llm=model,
autosave=True
) )
except Exception as e: return result, router, ""
error_msg = f"Failed to initialize SpreadSheetSwarm: {str(e)}"
log_event("error", error_msg)
return {}, None, error_msg
# Execute the swarm with proper error handling
try:
result = await asyncio.wait_for( result = await asyncio.wait_for(
asyncio.to_thread(lambda: swarm.run(task=task)), asyncio.to_thread(router.run, task),
timeout=900 # 15 minutes timeout timeout=timeout
) )
if swarm_type == "SpreadSheetSwarm":
# Verify the output file was created # Verify the output file was created
if not os.path.exists(output_path): if not os.path.exists(output_path):
error_msg = "Output file was not created" error_msg = "Output file was not created"
log_event("error", error_msg) log_event("error", error_msg)
return {}, None, error_msg return None, None, error_msg
return { return output_path, router, ""
"CSV File Path": output_path,
"Status": "Success",
"Message": "Spreadsheet processing completed successfully",
"Result": str(result)
}, swarm, ""
except asyncio.TimeoutError:
error_msg = "SpreadSheetSwarm execution timed out after 900 seconds"
log_event("error", error_msg)
return {}, None, error_msg
except Exception as e:
error_msg = f"SpreadSheetSwarm execution error: {str(e)}"
log_event("error", error_msg)
return {}, None, error_msg
# Rest of the swarm type handling...
elif swarm_type == "AgentRearrange":
if not flow:
return {}, None, "Flow configuration is required for AgentRearrange"
router_kwargs["flow"] = flow
elif swarm_type == "MixtureOfAgents":
if len(agents) < 2:
return {}, None, "MixtureOfAgents requires at least 2 agents"
router_kwargs.update({
"aggregator_agent": agents[-1],
"layers": max_loops
})
# Create router and execute task for non-SpreadSheetSwarm types
if swarm_type != "SpreadSheetSwarm":
try:
timeout = 450 # Default timeout
await asyncio.sleep(0.5)
router = SwarmRouter(**router_kwargs)
router.swarm_type = swarm_type
result = await asyncio.wait_for(
asyncio.to_thread(router.run, task=task),
timeout=timeout
)
# Process results return result, router, ""
def process_result(result):
"""Process and standardize the result output."""
if isinstance(result, str):
try:
# Attempt to parse as JSON
parsed_result = json.loads(result)
return parsed_result
except json.JSONDecodeError:
# Fallback to string result
return {"Final Output": result}
if isinstance(result, dict):
# Clean and standardize dictionary results
return {str(k): str(v) for k, v in result.items()}
return {"Final Output": str(result)}
responses = process_result(result)
return responses, router, ""
except asyncio.TimeoutError: except asyncio.TimeoutError:
error_msg = f"Task execution timed out after {timeout} seconds" error_msg = f"Task execution timed out after {timeout} seconds"
log_event("error", error_msg) log_event("error", error_msg)
return {}, None, error_msg return None, None, error_msg
except Exception as e: except Exception as e:
error_msg = f"Task execution error: {str(e)}" error_msg = f"Task execution error: {str(e)}"
log_event("error", error_msg) log_event("error", error_msg)
return {}, None, error_msg return None, None, error_msg
except Exception as e: except Exception as e:
error_msg = f"Unexpected error in task execution: {str(e)}" error_msg = f"Unexpected error in task execution: {str(e)}"
log_event("error", error_msg) log_event("error", error_msg)
return {}, None, error_msg return None, None, error_msg
def _extract_concurrent_responses(result: str, agents: List[Agent]) -> Dict[str, str]:
"""
Extract unique responses for each agent in a ConcurrentWorkflow.
Args:
result (str): Full output from SwarmRouter
agents (List[Agent]): List of agents used in the task
Returns:
Dict[str, str]: Unique responses for each agent
"""
agent_responses = {}
for agent in agents:
# Pattern to capture "Agent Name: ... Response: ... " format
pattern = rf"Agent Name:\s*{re.escape(agent.agent_name)}\s*Response:\s*(.+?)(?=Agent Name:|$)"
match = re.search(pattern, result, re.DOTALL | re.IGNORECASE | re.MULTILINE)
if match:
agent_responses[agent.agent_name] = match.group(1).strip()
else:
agent_responses[agent.agent_name] = "No response from the Agent"
return agent_responses
class UI: class UI:
@ -480,7 +370,7 @@ class UI:
def create_json_output(self, label, placeholder=""): def create_json_output(self, label, placeholder=""):
json_output = gr.JSON( json_output = gr.JSON(
label=label, label=label,
value = {}, value={},
elem_classes=["custom-output"], elem_classes=["custom-output"],
) )
self.components[f'json_output_{label}'] = json_output self.components[f'json_output_{label}'] = json_output
@ -545,6 +435,7 @@ class UI:
) )
return logs_display return logs_display
def create_app(): def create_app():
# Initialize UI # Initialize UI
theme = UI.create_ui_theme(primary_color="red") theme = UI.create_ui_theme(primary_color="red")
@ -677,7 +568,7 @@ def create_app():
f"Selected {swarm_type}" # For loading_status f"Selected {swarm_type}" # For loading_status
) )
async def run_task_wrapper(task, max_loops, data_temp, swarm_type, agent_prompt_selector, flow_text, sum_temp): async def run_task_wrapper(task, max_loops, dynamic_temp, swarm_type, agent_prompt_selector, flow_text):
"""Execute the task and update the UI with progress.""" """Execute the task and update the UI with progress."""
try: try:
if not task: if not task:
@ -702,11 +593,10 @@ def create_app():
flow = flow_text flow = flow_text
# Execute task # Execute task
responses, router, error = await execute_task( result, router, error = await execute_task(
task=task, task=task,
max_loops=max_loops, max_loops=max_loops,
data_temp=data_temp, dynamic_temp=dynamic_temp,
sum_temp=sum_temp,
swarm_type=swarm_type, swarm_type=swarm_type,
agent_keys=agent_prompt_selector, agent_keys=agent_prompt_selector,
flow=flow flow=flow
@ -718,28 +608,25 @@ def create_app():
# Format output based on swarm type # Format output based on swarm type
output_lines = [] output_lines = []
if router.swarm_type == "AgentRearrange": if swarm_type == "SpreadSheetSwarm":
for key, value in responses.items(): output_lines.append(f"### Spreadsheet Output ###\n{result}\n{'=' * 50}\n")
output_lines.append(f"### Step {key} ###\n{value}\n{'='*50}\n") elif isinstance(result, dict): # checking if result is dict or string.
elif router.swarm_type == "MixtureOfAgents": if swarm_type == "AgentRearrange":
for key, value in result.items():
output_lines.append(f"### Step {key} ###\n{value}\n{'=' * 50}\n")
elif swarm_type == "MixtureOfAgents":
# Add individual agent outputs # Add individual agent outputs
for key, value in responses.items(): for key, value in result.items():
if key != "Aggregated Summary": if key != "Aggregated Summary":
output_lines.append(f"### {key} ###\n{value}\n") output_lines.append(f"### {key} ###\n{value}\n")
# Add aggregated summary at the end # Add aggregated summary at the end
if "Aggregated Summary" in responses: if "Aggregated Summary" in result:
output_lines.append(f"\n### Aggregated Summary ###\n{responses['Aggregated Summary']}\n{'='*50}\n") output_lines.append(f"\n### Aggregated Summary ###\n{result['Aggregated Summary']}\n{'=' * 50}\n")
elif router.swarm_type == "SpreadSheetSwarm": else: # SequentialWorkflow, ConcurrentWorkflow, Auto
output_lines.append(f"### Spreadsheet Output ###\n{responses.get('CSV File Path', 'No file path provided')}\n{'='*50}\n") for key, value in result.items():
elif router.swarm_type == "ConcurrentWorkflow": output_lines.append(f"### {key} ###\n{value}\n{'=' * 50}\n")
for key, value in responses.items(): elif isinstance(result, str):
output_lines.append(f"### {key} ###\n{value}\n{'='*50}\n") output_lines.append(str(result))
else: # SequentialWorkflow, auto
if isinstance(responses, dict):
for key, value in responses.items():
output_lines.append(f"### {key} ###\n{value}\n{'='*50}\n")
else:
output_lines.append(str(responses))
yield "\n".join(output_lines), "Completed" yield "\n".join(output_lines), "Completed"
@ -782,16 +669,6 @@ def create_app():
cancels=run_event cancels=run_event
) )
# with gr.Column(scale=1): # Right column
# with gr.Tabs():
# with gr.Tab("Agent Details"):
# gr.Markdown("""
# ### Available Agent Types
# - **Data Extraction Agent**: Specialized in extracting relevant information
# - **Summary Agent**: Creates concise summaries of information
# - **Analysis Agent**: Performs detailed analysis of data
with gr.Column(scale=1): # Right column with gr.Column(scale=1): # Right column
with gr.Tabs(): with gr.Tabs():
with gr.Tab("Agent Details"): with gr.Tab("Agent Details"):
@ -799,6 +676,7 @@ def create_app():
with gr.Tab("Logs"): with gr.Tab("Logs"):
logs_display = ui.create_logs_tab() logs_display = ui.create_logs_tab()
def update_logs_display(): def update_logs_display():
"""Update logs display with current logs.""" """Update logs display with current logs."""
logs = get_logs() logs = get_logs()
@ -809,10 +687,8 @@ def create_app():
logs_tab = gr.Tab("Logs") logs_tab = gr.Tab("Logs")
logs_tab.select(fn=update_logs_display, inputs=None, outputs=[logs_display]) logs_tab.select(fn=update_logs_display, inputs=None, outputs=[logs_display])
return ui.build() return ui.build()
# Launch the app
if __name__ == "__main__": if __name__ == "__main__":
app = create_app() app = create_app()
app.launch() app.launch()
Loading…
Cancel
Save