Spaces:
Sleeping
Sleeping
File size: 7,665 Bytes
d853cbf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 | """
Master Agent - coordinates ExcelAnalysisAgent and WebSearchAgent as tools
"""
import logging
import os
import asyncio
import json
from typing import Dict, Any, Optional
from agents import Agent, Runner
from agents.mcp import MCPServerStdio, create_static_tool_filter
from .excel_agent import create_excel_agent
from .web_agent import create_web_search_agent
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
MASTER_AGENT_PROMPT = """
You are the orchestrator of a multi-agent system. Your task is to take the user's query and the file path and pass it to the appropriate agent tool.
Available agent tools:
- excel_analysis_agent: Executes Python code for data analysis and visualization using pandas and matplotlib.
When calling this tool, you MUST pass the complete user query and the file path so it can execute the correct analysis.
- web_search_agent: Searches the web for documentation, examples, and solutions.
Your strategy:
1. First, try to use the excel_analysis_agent to directly answer the user's query using the file path.
IMPORTANT: When calling excel_analysis_agent, include the FULL user query in your message to the tool.
2. If the analysis fails or needs additional context, use the web_search_agent to find relevant information.
3. Use the web search results to guide a retry with the excel_analysis_agent.
Always provide clear, actionable results to the user.
"""
class MasterAgent:
"""
Master agent that coordinates ExcelAnalysisAgent and WebSearchAgent as tools.
"""
def __init__(self, api_key: str, model: str = "gpt-4o-mini"):
if api_key:
os.environ["OPENAI_API_KEY"] = api_key
self.model = model
def analyze(self, user_query: str, file_path: str) -> Dict[str, Any]:
"""
Coordinate the two agents to get the best possible result
"""
async def _arun():
# Create MCP servers
python_server = MCPServerStdio(
name="excel-tools-python",
params={"command": "python", "args": ["-m", "app_agents.mcp_server"]},
cache_tools_list=True,
use_structured_content=True,
tool_filter=create_static_tool_filter(allowed_tool_names=["execute_python_code"]),
)
web_server = MCPServerStdio(
name="excel-tools-web",
params={"command": "python", "args": ["-m", "app_agents.mcp_server"]},
cache_tools_list=True,
tool_filter=create_static_tool_filter(allowed_tool_names=["search_web"]),
)
# Connect servers
await python_server.connect()
await web_server.connect()
try:
# Create specialized agents using functions from their respective modules
excel_agent = create_excel_agent(mcp_server=python_server, model=self.model)
web_agent = create_web_search_agent(mcp_server=web_server, model=self.model)
# Create orchestrator agent with other agents as tools
orchestrator = Agent(
name="MasterAgent",
model=self.model,
instructions=MASTER_AGENT_PROMPT,
tools=[
excel_agent.as_tool(
tool_name="excel_analysis_agent",
tool_description="Execute Python code to analyze Excel/CSV files and create visualizations. The agent receives the user query and file path and must execute the exact analysis requested."
),
web_agent.as_tool(
tool_name="web_search_agent",
tool_description="Search the web for up-to-date information, documentation, and code examples"
),
],
)
# Prepare user message with file path
user_msg = (
f"User query: {user_query}\n"
f"File path: {file_path}\n\n"
f"Call the excel_analysis_agent tool with this exact message:\n"
f"'Analyze this request: {user_query}\\n\\nThe file is located at: {file_path}\\n\\n"
f"Write Python code and call execute_python_code with that code and the same file_path.'\n\n"
f"Make sure to pass the complete user query to the excel_analysis_agent so it can perform the correct analysis."
)
# Run orchestrator
result = await Runner.run(orchestrator, user_msg, max_turns=20)
return result
finally:
# Clean up servers
for server in [python_server, web_server]:
close_fn = getattr(server, "close", None) or getattr(server, "aclose", None)
if close_fn:
res = close_fn()
if hasattr(res, "__await__"):
await res
try:
loop = asyncio.new_event_loop()
try:
asyncio.set_event_loop(loop)
result = loop.run_until_complete(_arun())
finally:
loop.close()
asyncio.set_event_loop(None)
raw_output = result.final_output or ""
# Extract dataframe and images from tool output
extracted_df = None
extracted_images = []
final_text = raw_output
# Extract from result.new_items - Item 1 (ToolCallOutputItem) contains the JSON
for item in result.new_items:
if hasattr(item, 'output') and isinstance(item.output, str):
# Extract JSON from markdown code blocks if present
json_str = item.output
if "```json" in item.output:
parts = item.output.split("```json")
if len(parts) > 1:
json_str = parts[1].split("```")[0].strip()
try:
tool_result = json.loads(json_str)
if isinstance(tool_result, dict) and "success" in tool_result:
# Extract dataframe and images from tool result
if isinstance(tool_result.get("dataframe"), list) and tool_result.get("dataframe"):
extracted_df = tool_result.get("dataframe")
if isinstance(tool_result.get("images"), list) and tool_result.get("images"):
extracted_images = tool_result.get("images")
break # Found the JSON, no need to continue
except (json.JSONDecodeError, ValueError):
continue
return {
'success': True,
'output': final_text,
'dataframe': extracted_df,
'images': extracted_images,
'code': None,
'error': None
}
except Exception as e:
err = f"MasterAgent error: {e}"
logger.error(err)
return {
"success": False,
"output": None,
"dataframe": None,
"images": [],
"code": None,
"error": err,
}
|