ai-agent-book 精选快照(<2MB 代码与文档,来自 github.com/bojieli/ai-agent-book)
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s

This commit is contained in:
2026-08-20 13:12:50 +00:00
commit b119135836
10275 changed files with 3284984 additions and 0 deletions
@@ -0,0 +1,76 @@
FROM python:3.11-slim-bookworm AS base
## Basis ##
ENV ENV=prod \
PORT=9099
# Install GCC and build tools.
# These are kept in the final image to enable installing packages on the fly.
RUN apt-get update && \
apt-get install -y gcc build-essential curl git && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
RUN apt-get update && apt-get install -y wget unzip openssh-client procps nodejs npm
ARG PIP_OPTIONS='-i https://mirrors.aliyun.com/pypi/simple/'
ARG ENABLE_OSS_MOUNT=""
RUN pip install -U pip pysocks ${PIP_OPTIONS}
# Install Chrome Driver
RUN mkdir /app
RUN cd /app/ && \
wget https://storage.googleapis.com/chrome-for-testing-public/136.0.7103.92/linux64/chromedriver-linux64.zip && \
unzip chromedriver-linux64.zip && \
rm chromedriver-linux64.zip
ENV CHROME_DRIVER_PATH=/app/chromedriver-linux64/chromedriver
RUN if [ "$ENABLE_OSS_MOUNT" = "true" ]; then \
apt-get update && \
apt-get install -y gdebi-core mime-support && \
cd /tmp && \
wget https://gosspublic.alicdn.com/ossfs/ossfs_1.91.6_ubuntu22.04_amd64.deb && \
gdebi ossfs_1.91.6_ubuntu22.04_amd64.deb -n && \
rm ossfs_1.91.6_ubuntu22.04_amd64.deb && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*; \
fi
FROM base as runner
WORKDIR /app
# Install Python dependencies
COPY ./requirements.txt .
RUN pip3 install uv ${PIP_OPTIONS}
RUN uv pip install --system -r requirements.txt --no-cache-dir ${PIP_OPTIONS}
# Copy the application code
RUN echo "start install aworld"
RUN mkdir -p /app/lib
RUN cd /app/lib && git clone https://github.com/inclusionAI/AWorld.git
RUN cd /app/lib/AWorld && git checkout framework_upgrade_aworldserver_gaia && pip install -r aworld/requirements.txt ${PIP_OPTIONS} && python setup.py install
RUN npx playwright install chrome --with-deps --no-shell
RUN cd /app
# Layer on for other components
FROM runner AS app
WORKDIR /app
COPY . .
# Expose the port
ENV HOST="0.0.0.0"
ENV PORT="9099"
# if we already installed the requirements on build, we can skip this step on run
ENTRYPOINT [ "bash", "start.sh" ]
@@ -0,0 +1,158 @@
# AworldServer
AworldServer is an execution environment for the Aworld framework that integrates MCP LLM models. It supports distributed deployment and dynamic scaling.
![img.png](img.png)
The system features:
- Distributed Architecture: Supports multi-server deployment with load balancing
- Dynamic Scaling: Ability to adjust server capacity based on demand
- LLM Integration: Built-in MCP LLM model support
- Asynchronous Processing: Uses asynchronous programming patterns for improved performance
- Containerized Deployment: Docker containerization support for easy environment management
## 🚀 Quick Start
1. Start services using Docker Compose:
```sh
docker build --build-arg MINIMUM_BUILD=true -f Dockerfile --progress=plain -t aworldserver:main .
docker compose up -d
```
2. Configure the number of server instances:
You can modify the `docker-compose.yaml` file to adjust the number of server instances. The default configuration includes 3 instances:
3. Usage Methods:
a. OpenWebUI Integration:
- Configure external link in OpenWebUI settings
- Add AworldServer endpoints to the configuration
- Set up API key authentication
b. Python Client Usage:
```python
# Initialize AworldTaskClient with server endpoints
AWORLD_TASK_CLIENT = AworldTaskClient(
know_hosts=["localhost:9299", "localhost:9399", "localhost:9499"]
)
async def _run_gaia_task(gaia_question_id: str) -> None:
"""Run a single Gaia task with the given question ID.
Args:
gaia_question_id: The ID of the question to process
"""
global AWORLD_TASK_CLIENT
task_id = str(uuid.uuid4())
# Submit task to Aworld server
await AWORLD_TASK_CLIENT.submit_task(
AworldTask(
task_id=task_id,
agent_id="gaia_agent",
agent_input=gaia_question_id,
session_id="session_id",
user_id="SYSTEM"
)
)
# Get and print task result
task_result = await AWORLD_TASK_CLIENT.get_task_state(task_id=task_id)
print(task_result)
async def _batch_run_gaia_task(start_i: int, end_i: int) -> None:
"""Run multiple Gaia tasks in parallel.
Args:
start_i: Starting question ID
end_i: Ending question ID
"""
tasks = [
_run_gaia_task(str(i))
for i in range(start_i, end_i + 1)
]
await asyncio.gather(*tasks)
if __name__ == '__main__':
# Run batch processing for questions 1-5
asyncio.run(_batch_run_gaia_task(1, 5))
```
c. user curl
```shell
curl http://localhost:9299/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer 0p3n-w3bu!" \
-d '{
"model": "gaia_agent",
"messages": [
{
"role": "user",
"content": [
{
"type": "text",
"text": "5"
}
]
}
]
}'
```
## 🔑 Key Features
- **Distributed Task Processing System**
- Multi-server load balancing
- Round-robin task distribution
- Asynchronous task processing
- **Docker Containerization**
- Multi-instance deployment
- Environment variable configuration
- Auto-restart mechanism
- **API Services**
- FastAPI framework support
- RESTful API design
- Asynchronous request handling
- **Development Tools**
- Debug mode support
- Batch task processing
- Task state tracking
- **Security Features**
- API key authentication
- Session management
- User authentication
## 📦 Installation and Setup
Get started with aworldserver in a few easy steps:
1. **Ensure Python 3.11 is installed.**
2. **Install the required dependencies:**
```sh
pip install -r requirements-minimux.txt
```
3. **Start the aworld server:**
```sh
sh ./start.sh
```
### Custom debug
please run `debug_run.py`
##
@@ -0,0 +1,59 @@
import logging
import os
from aworld.config import ModelConfig
from aworld.config.conf import AgentConfig, ClientType
from pydantic import BaseModel
from aworldspace.base_agent import AworldBaseAgent
from aworldspace.utils.mcp_utils import load_all_mcp_config
SYSTEM_PROMPT = f"""You are an helpful AI assistant, aimed at solving any task presented by the user. """
class Pipeline(AworldBaseAgent):
class Valves(BaseModel):
pass
def __init__(self):
self.valves = self.Valves()
logging.info("default init success")
async def get_agent_config(self, body):
default_llm_provider = os.environ.get("LLM_PROVIDER")
llm_model_name = os.environ.get("LLM_MODEL_NAME")
llm_api_key = os.environ.get("LLM_API_KEY")
llm_base_url = os.environ.get("LLM_BASE_URL")
task = await self.get_task_from_body(body)
logging.info(f"task llm config is: {task.llm_provider}, {task.llm_model_name},{task.llm_base_url}")
llm_config = ModelConfig(
llm_provider=task.llm_provider if task and task.llm_provider else default_llm_provider,
llm_model_name=task.llm_model_name if task and task.llm_model_name else llm_model_name,
llm_api_key=task.llm_api_key if task and task.llm_api_key else llm_api_key,
llm_base_url=task.llm_base_url if task and task.llm_base_url else llm_base_url,
max_retries=task.max_retries if task and task.max_retries else 3
)
return AgentConfig(
name=self.agent_name(),
llm_config=llm_config,
system_prompt=task.task_system_prompt if task and task.task_system_prompt else SYSTEM_PROMPT
)
def agent_name(self) -> str:
return "DefaultAgent"
async def get_mcp_servers(self, body) -> list[str]:
task = await self.get_task_from_body(body)
if task.mcp_servers:
logging.info(f"mcp_servers from task: {task.mcp_servers}")
return task.mcp_servers
return [
"ms-playwright"
]
async def load_mcp_config(self) -> dict:
return load_all_mcp_config()
@@ -0,0 +1,264 @@
import logging
import os
import re
from pathlib import Path
from typing import Dict, Any, List, Optional
from aworld.config import ModelConfig
from aworld.config.conf import AgentConfig, TaskConfig, ClientType
from aworld.core.task import Task
from aworld.output import Outputs, Output, StreamingOutputs
from aworld.utils.common import get_local_ip
from datasets import load_dataset, concatenate_datasets
from pydantic import BaseModel, Field
from aworldspace.base_agent import AworldBaseAgent
from aworldspace.utils.mcp_utils import load_all_mcp_config
from aworldspace.utils.utils import question_scorer
GAIA_SYSTEM_PROMPT = f"""You are an all-capable AI assistant, aimed at solving any task presented by the user. You have various tools at your disposal that you can call upon to efficiently complete complex requests. Whether it's programming, information retrieval, file processing, or web browsing, you can handle it all.
Please note that the task may be complex. Do not attempt to solve it all at once. You should break the task down and use different tools step by step to solve it. After using each tool, clearly explain the execution results and suggest the next steps.
Please utilize appropriate tools for the task, analyze the results obtained from these tools, and provide your reasoning. Always use available tools such as browser, calcutor, etc. to verify correctness rather than relying on your internal knowledge.
If you believe the problem has been solved, please output the `final answer`. The `final answer` should be given in <answer></answer> format, while your other thought process should be output in <think></think> tags.
Your `final answer` should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
Here are some tips to help you give better instructions:
<tips>
1. Do not use any tools outside of the provided tools list.
2. Even if the task is complex, there is always a solution. If you cant find the answer using one method, try another approach or use different tools to find the solution.
3. When using browser `playwright_click` tool, you need to check if the element exists and is clickable before clicking it.
4. Before providing the `final answer`, carefully reflect on whether the task has been fully solved. If you have not solved the task, please provide your reasoning and suggest the next steps.
5. Due to context length limitations, always try to complete browser-based tasks with the minimal number of steps possible.
6. When providing the `final answer`, answer the user's question directly and precisely. For example, if asked "what animal is x?" and x is a monkey, simply answer "monkey" rather than "x is a monkey".
7. When you need to process excel file, prioritize using the `excel` tool instead of writing custom code with `terminal-controller` tool.
8. If you need to download a file, please use the `terminal-controller` tool to download the file and save it to the specified path.
9. The browser doesn't support direct searching on www.google.com. Use the `google-search` to get the relevant website URLs or contents instead of `ms-playwright` directly.
10. Always use only one tool at a time in each step of your execution.
11. Using `mcp__ms-playwright__browser_pdf_save` tool to save the pdf file of URLs to the specified path.
12. Using `mcp__terminal-controller__execute_command` tool to set the timeout to 300 seconds when downloading large files such as pdf.
13. Using `mcp__ms-playwright__browser_take_screenshot` tool to save the screenshot of URLs to the specified path when you need to understand the gif / jpg of the URLs.
14. When there are questions related to YouTube video comprehension, use tools in `youtube_download_server` and `video_server` to analyze the video content by the given question.
</tips>
Now, here is the task. Stay focused and complete it carefully using the appropriate tools!
"""
class Pipeline(AworldBaseAgent):
class Valves(BaseModel):
llm_provider: Optional[str] = Field(default=None, description="llm_model_name")
llm_model_name: Optional[str] = Field(default=None, description="llm_model_name")
llm_base_url: Optional[str] = Field(default=None,description="llm_base_urly")
llm_api_key: Optional[str] = Field(default=None,description="llm api key" )
system_prompt: str = Field(default=GAIA_SYSTEM_PROMPT,description="system_prompt")
history_messages: int = Field(default=100, description="rounds of history messages")
def __init__(self):
self.valves = self.Valves()
self.gaia_files = os.path.abspath(os.path.join(os.path.curdir, "aworldspace", "datasets", "gaia_dataset"))
logging.info(f"gaia_files path {self.gaia_files}")
self.full_dataset = load_dataset(
os.path.join(self.gaia_files, "GAIA.py"),
name="2023_all",
trust_remote_code=True
)
self.full_dataset = concatenate_datasets([self.full_dataset['validation'], self.full_dataset['test']])
# Create task_id to index mapping for improved lookup performance
self.task_id_to_index = {}
for i, task in enumerate(self.full_dataset):
self.task_id_to_index[task['task_id']] = i
logging.info(f"Loaded {len(self.full_dataset)} tasks, created task_id mapping")
logging.info("gaia_agent init success")
async def get_custom_input(self, user_message: str, model_id: str, messages: List[dict], body: dict) -> Any:
task = await self.get_gaia_task(user_message)
logging.info(f"🌈 -----------------------------------------------")
logging.info(f"🚀 Start to process: gaia_task_{task['task_id']}")
logging.info(f"📝 Detail: {task}")
logging.info(f"❓ Question: {task['Question']}")
logging.info(f"⭐ Level: {task['Level']}")
logging.info(f"🛠️ Tools: {task['Annotator Metadata']['Tools']}")
logging.info(f"🌈 -----------------------------------------------")
return task['Question']
async def get_agent_config(self, body):
default_llm_provider = self.valves.llm_provider if self.valves.llm_provider else os.environ.get("LLM_PROVIDER")
llm_model_name = self.valves.llm_model_name if self.valves.llm_model_name else os.environ.get("LLM_MODEL_NAME")
llm_api_key = self.valves.llm_api_key if self.valves.llm_api_key else os.environ.get("LLM_API_KEY")
llm_base_url = self.valves.llm_base_url if self.valves.llm_base_url else os.environ.get("LLM_BASE_URL")
system_prompt = self.valves.system_prompt if self.valves.system_prompt else GAIA_SYSTEM_PROMPT
task = await self.get_task_from_body(body)
if task:
logging.info(f"task llm config is: {task.llm_provider}, {task.llm_model_name}, {task.llm_api_key}, {task.llm_base_url}")
llm_config = ModelConfig(
llm_provider=task.llm_provider if task and task.llm_provider else default_llm_provider,
llm_model_name=task.llm_model_name if task and task.llm_model_name else llm_model_name,
llm_api_key=task.llm_api_key if task and task.llm_api_key else llm_api_key,
llm_base_url=task.llm_base_url if task and task.llm_base_url else llm_base_url,
max_retries=task.max_retries if task and task.max_retries else 3
)
return AgentConfig(
name=self.agent_name(),
llm_config=llm_config,
system_prompt=task.task_system_prompt if task and task.task_system_prompt else system_prompt
)
def agent_name(self) -> str:
return "GaiaAgent"
async def get_mcp_servers(self, body) -> list[str]:
task = await self.get_task_from_body(body)
if task and task.mcp_servers:
logging.info(f"mcp_servers from task: {task.mcp_servers}")
return task.mcp_servers
return [
"e2b-server",
"terminal-controller",
"excel",
"calculator",
"ms-playwright",
"audio_server",
"image_server",
"video_server",
"search_server",
"download_server",
"document_server",
"youtube_server",
"reasoning_server",
]
async def get_gaia_task(self, task_id: str) -> dict:
"""
Get GAIA task by task_id
Args:
task_id: Unique identifier of the task
Returns:
Corresponding task dictionary
"""
# Search by task_id
if task_id in self.task_id_to_index:
index = self.task_id_to_index[task_id]
gaia_task = self.full_dataset[index]
else:
raise ValueError(f"Task with task_id '{task_id}' not found in dataset")
return self.add_file_path(gaia_task)
def get_all_task_ids(self) -> List[str]:
"""
Get list of all available task_ids
Returns:
List of all task_ids
"""
return list(self.task_id_to_index.keys())
def get_task_count(self) -> int:
"""
Get total number of tasks
Returns:
Total task count
"""
return len(self.full_dataset)
def get_task_index_by_id(self, task_id: str) -> int:
"""
Get task index in dataset by task_id
Args:
task_id: Unique identifier of the task
Returns:
Index of the task in the dataset
"""
if task_id in self.task_id_to_index:
return self.task_id_to_index[task_id]
else:
raise ValueError(f"Task with task_id '{task_id}' not found in dataset")
async def custom_output_before_task(self, outputs: Outputs, chat_id: str, task: Task) -> None:
task_config:TaskConfig = task.conf
gaia_task = await self.get_gaia_task(task_config.ext['origin_message'])
result = f"\n\n`{get_local_ip()}` execute `GAIA TASK#{task_config.ext['origin_message']}`:\n\n---\n\n"
result += f"**Question**: {gaia_task['Question']}\n"
result += f"**Answer**: {gaia_task['Final answer']}\n"
result += f"**Level**: {gaia_task['Level']}\n"
result += f"**Tools**: \n {gaia_task['Annotator Metadata']['Tools']}\n"
result += f"\n\n-----\n\n"
await outputs.add_output(Output(data = result))
async def custom_output_after_task(self, outputs: Outputs, chat_id: str, task: Task):
"""
check gaia task output
Args:
outputs:
chat_id:
task:
Returns:
"""
task_config: TaskConfig = task.conf
gaia_task_id = task_config['ext']['origin_message']
gaia_task = await self.get_gaia_task(gaia_task_id)
agent_result = ""
if isinstance(outputs, StreamingOutputs):
agent_result = await outputs._visited_outputs[-2].get_finished_response() # read llm result
match = re.search(r"<answer>(.*?)</answer>", agent_result)
answer = agent_result
if match:
answer = match.group(1)
logging.info(f"🤖 Agent answer: {answer}")
logging.info(f"👨‍🏫 Correct answer: {gaia_task['Final answer']}")
is_correct = question_scorer(answer, gaia_task["Final answer"])
if is_correct:
logging.info(f"📝Question {gaia_task_id} Correct! 🎉")
result = f"\n\n📝 **Question: {gaia_task_id} -> Agent Answer:[{answer}] is `Correct`**"
else:
logging.info(f"📝Question {gaia_task_id} Incorrect! ❌")
result = f"\n\n📝 **Question: {gaia_task_id} -> Agent Answer:`{answer}` != Correct answer: `{gaia_task['Final answer']}` is `Incorrect` ❌**"
metadata = await outputs.get_metadata()
if not metadata:
await outputs.set_metadata({})
metadata = await outputs.get_metadata()
metadata['gaia_correct'] = is_correct
metadata['gaia_result'] = result
metadata['agent_answer'] = answer
metadata['correct_answer'] = gaia_task['Final answer']
return result
def add_file_path(self, task: Dict[str, Any]
):
split = "validation" if task["Annotator Metadata"]["Steps"] != "" else "test"
if task["file_name"]:
file_path = Path(f"{self.gaia_files}/2023/{split}/" + task["file_name"])
if file_path.suffix in [".pdf", ".docx", ".doc", ".txt"]:
task["Question"] += f" Here are the necessary document files: {file_path}"
elif file_path.suffix in [".jpg", ".jpeg", ".png"]:
task["Question"] += f" Here are the necessary image files: {file_path}"
elif file_path.suffix in [".xlsx", "xls", ".csv"]:
task[
"Question"
] += f" Here are the necessary table files: {file_path}, for processing excel file, you can use the excel tool or write python code to process the file step-by-step and get the information."
elif file_path.suffix in [".py"]:
task["Question"] += f" Here are the necessary python files: {file_path}"
else:
task["Question"] += f" Here are the necessary files: {file_path}"
return task
async def load_mcp_config(self) -> dict:
return load_all_mcp_config()
@@ -0,0 +1,794 @@
import logging
import os
import traceback
from typing import Dict, Any, List, Union
from typing import Optional
from aworld.core.event.base import Message
from aworldspace.base_agent import AworldBaseAgent
from pydantic import BaseModel, Field
import aworld.trace as trace
from aworld.config.conf import AgentConfig, ConfigDict
from aworld.config.conf import TaskConfig
from aworld.agents.llm_agent import Agent
from aworld.core.common import Observation, ActionModel
from aworld.core.memory import MemoryItem
from aworld.core.task import Task
from aworld.logs.util import logger
from aworld.models.llm import acall_llm_model
from aworld.models.model_response import ToolCall, Function
from aworld.output import Output, StreamingOutputs
from aworld.output import Outputs
from aworld.output.base import MessageOutput
from aworld.utils.common import sync_exec
BROWSER_SYSTEM_PROMPT = """You are a GUI agent. You are given a task and your action history, with screenshots. You need to perform the next action to complete the task.
## Output Format
```
Thought: ...
Action: ...
```
## Action Space
navigate(website='xxx') #Open the target website, usually the first action to open browser.
click(start_box='[x1, y1, x2, y2]')
left_double(start_box='[x1, y1, x2, y2]')
right_single(start_box='[x1, y1, x2, y2]')
drag(start_box='[x1, y1, x2, y2]', end_box='[x3, y3, x4, y4]')
hotkey(key='')
type(content='') #If you want to submit your input, use "\n" at the end of `content`.
scroll(direction='down or up or right or left')
wait() #Sleep for 5s and take a screenshot to check for any changes.
finished(content='xxx') # Use escape characters \\', \\", and \\n in content part to ensure we can parse the content in normal python string format.
## Note
- only one action per step.
- Use Chinese in `Thought` part.
- Write a small plan and finally summarize your next action (with its target element) in one sentence in `Thought` part.
## User Instruction
"""
import json
import re
MAX_IMAGE = 50
def parse_action_output(output_text):
# 提取Thought部分
logger.info(f"{output_text=}")
thought_match = re.search(r'Thought:(.*?)\nAction:', output_text, re.DOTALL)
thought = thought_match.group(1).strip() if thought_match else ""
# 提取Action部分
action_match = re.search(r'Action:(.*?)(?:\n|$)', output_text, re.DOTALL)
action_text = action_match.group(1).strip() if action_match else ""
# 初始化结果字典
result = {
"thought": thought,
"action": "",
"key": None,
"content": None,
"start_box": None,
"end_box": None,
"direction": None,
"website": None,
}
if not action_text:
return json.dumps(result, ensure_ascii=False)
# tmp 兼容ui-tars1.5-7b
action_text = action_text.replace("'(","'[").replace(")'","]'")
# 解析action类型
action_parts = action_text.split('(')
action_type = action_parts[0]
result["action"] = action_type
# 解析参数
if len(action_parts) > 1:
params_text = action_parts[1].rstrip(')')
params = {}
# gpt-4o兼容
if 'start_box' in params_text:
params_text = params_text.replace(", ", " ").replace(",", " ")
if 'end_box' in params_text:
params_text = params_text.replace(" end_box", ", end_box")
# 处理键值对参数
for param in params_text.split(','):
param = param.strip()
if '=' in param:
key, value = param.split('=', 1)
key = key.strip()
value = value.strip().strip('\'"')
# 处理bbox格式
if 'box' in key:
print(value)
# 提取坐标数字
numbers = re.findall(r'\d+', value)
print(numbers)
if numbers:
coords = [int(num) for num in numbers]
if len(coords) == 4:
if key == 'start_box':
result["start_box"] = coords
elif key == 'end_box':
result["end_box"] = coords
if len(coords) == 2:
if key == 'start_box':
result["start_box"] = [coords[0], coords[1], coords[0], coords[1]]
elif key == 'end_box':
result["end_box"] = [coords[0], coords[1], coords[0], coords[1]]
elif key == 'key':
result["key"] = value.replace("pagedown", "PageDown").replace("pageup", "PageUp").replace("enter","Enter")
elif key == 'content':
# 处理转义字符
value = value.replace('\\n', '\n').replace('\\"', '"').replace("\\'", "'")
result["content"] = value
elif key == 'website':
result["website"] = value
elif key == 'direction':
result["direction"] = value
return result, thought, action_text
def parse_tool_call(line):
# 提取 Action和param
result, thought, action_text = parse_action_output(line)
action = result['action']
# 映射到实际函数名和参数
if action == 'navigate':
func_name = 'mcp__ms-playwright__browser_navigate'
content = {'url': result['website']}
elif action == 'click':
func_name = 'mcp__ms-playwright__browser_screen_click'
x = int((result["start_box"][0] + result["start_box"][2]) / 2)
y = int((result["start_box"][1] + result["start_box"][3]) / 2)
content = {'element': '', 'x': x, 'y': y}
elif action == 'right_single':
func_name = 'mcp__ms-playwright__browser_screen_click'
x = int((result["start_box"][0] + result["start_box"][2]) / 2)
y = int((result["start_box"][1] + result["start_box"][3]) / 2)
content = {'element': 'right click target', 'x': x, 'y': y, 'button': 'right'}
elif action == 'drag':
func_name = 'mcp__ms-playwright__browser_screen_drag'
x1 = int((result["start_box"][0] + result["start_box"][2]) / 2)
y1 = int((result["start_box"][1] + result["start_box"][3]) / 2)
x2 = int((result["end_box"][0] + result["end_box"][2]) / 2)
y2 = int((result["end_box"][1] + result["end_box"][3]) / 2)
content = {
'element': f'drag from [{x1},{y1}] to [{x2},{y2}]',
'startX': x1,
'startY': y1,
'endX': x2,
'endY': y2
}
elif action == 'hotkey':
func_name = 'mcp__ms-playwright__browser_press_key'
content = {'key': result["key"]}
elif action == 'type':
func_name = 'mcp__ms-playwright__browser_screen_type'
content = {'text': result['content']}
elif action == 'scroll':
# 暂时使用presskey代替scroll
func_name = 'mcp__ms-playwright__browser_press_key'
direction = result['direction']
key_map = {
'up': 'PageUp',
'down': 'PageDown',
'left': 'ArrowLeft',
'right': 'ArrowRight'
}
key = key_map.get(direction, 'ArrowDown')
content = {'key': key}
elif action == 'wait':
func_name = 'mcp__ms-playwright__browser_wait_for'
content = {'time': 5}
elif action == 'finished':
func_name = "finished"
content = result['content']
else:
return ""
return Function(name=func_name, arguments=json.dumps(content)), thought, action_text, result
# eval code start
def identify_key_points(task):
system_msg = """You are an expert tasked with analyzing a given task to identify the key points explicitly stated in the task description.
**Objective**: Carefully analyze the task description and extract the critical elements explicitly mentioned in the task for achieving its goal.
**Instructions**:
1. Read the task description carefully.
2. Identify and extract **key points** directly stated in the task description.
- A **key point** is a critical element, condition, or step explicitly mentioned in the task description.
- Do not infer or add any unstated elements.
- Words such as "best," "highest," "cheapest," "latest," "most recent," "lowest," "closest," "highest-rated," "largest," and "newest" must go through the sort function(e.g., the key point should be "Filter by highest").
**Respond with**:
- **Key Points**: A numbered list of the explicit key points for completing this task, one per line, without explanations or additional details."""
prompt = """Task: {task}"""
text = prompt.format(task=task)
messages = [
{"role": "system", "content": system_msg},
{
"role": "user",
"content": [
{"type": "text", "text": text}
],
}
]
return messages
def judge_image(task, image_path, key_points):
system_msg = """You are an expert evaluator tasked with determining whether an image contains information about the necessary steps to complete a task.
**Objective**: Analyze the provided image and decide if it shows essential steps or evidence required for completing the task. Use your reasoning to explain your decision before assigning a score.
**Instructions**:
1. Provide a detailed description of the image, including its contents, visible elements, text (if any), and any notable features.
2. Carefully examine the image and evaluate whether it contains necessary steps or evidence crucial to task completion:
- Identify key points that could be relevant to task completion, such as actions, progress indicators, tool usage, applied filters, or step-by-step instructions.
- Does the image show actions, progress indicators, or critical information directly related to completing the task?
- Is this information indispensable for understanding or ensuring task success?
- If the image contains partial but relevant information, consider its usefulness rather than dismissing it outright.
3. Provide your response in the following format:
- **Reasoning**: Explain your thought process and observations. Mention specific elements in the image that indicate necessary steps, evidence, or lack thereof.
- **Score**: Assign a score based on the reasoning, using the following scale:
- **1**: The image does not contain any necessary steps or relevant information.
- **2**: The image contains minimal or ambiguous information, unlikely to be essential.
- **3**: The image includes some relevant steps or hints but lacks clarity or completeness.
- **4**: The image contains important steps or evidence that are highly relevant but not fully comprehensive.
- **5**: The image clearly displays necessary steps or evidence crucial for completing the task.
Respond with:
1. **Reasoning**: [Your explanation]
2. **Score**: [1-5]"""
# jpg_base64_str = encode_image(Image.open(image_path))
prompt = """**Task**: {task}
**Key Points for Task Completion**: {key_points}
The snapshot of the web page is shown in the image."""
text = prompt.format(task=task, key_points=key_points)
messages = [
{"role": "system", "content": system_msg},
{
"role": "user",
"content": [
{"type": "text", "text": text},
{
"type": "image_url",
"image_url": {"url": image_path, "detail": "high"},
},
],
}
]
return messages
def WebJudge_Online_Mind2Web_eval(task, last_actions, images_path, image_responses, key_points, score_threshold):
system_msg = """You are an expert in evaluating the performance of a web navigation agent. The agent is designed to help a human user navigate a website to complete a task. Given the user's task, the agent's action history, key points for task completion, some potentially important web pages in the agent's trajectory and their reasons, your goal is to determine whether the agent has completed the task and achieved all requirements.
Your response must strictly follow the following evaluation criteria!
*Important Evaluation Criteria*:
1: The filtered results must be displayed correctly. If filters were not properly applied (i.e., missing selection, missing confirmation, or no visible effect in results), the task is not considered successful.
2: You must carefully check whether these snapshots and action history meet these key points. Ensure that specific filter conditions, such as "best," "highest," "cheapest," "latest," "most recent," "lowest," "closest," "highest-rated," "largest," and "newest" are correctly applied using the filter function(e.g., sort function).
3: Certain key points or requirements should be applied by the filter. Otherwise, a search with all requirements as input will be deemed a failure since it cannot guarantee that all results meet the requirements!
4: If the task requires filtering by a specific range of money, years, or the number of beds and bathrooms, the applied filter must exactly match the given requirement. Any deviation results in failure. To ensure the task is successful, the applied filter must precisely match the specified range without being too broad or too narrow.
Examples of Failure Cases:
- If the requirement is less than $50, but the applied filter is less than $25, it is a failure.
- If the requirement is $1500-$2500, but the applied filter is $2000-$2500, it is a failure.
- If the requirement is $25-$200, but the applied filter is $0-$200, it is a failure.
- If the required years are 2004-2012, but the filter applied is 2001-2012, it is a failure.
- If the required years are before 2015, but the applied filter is 2000-2014, it is a failure.
- If the task requires exactly 2 beds, but the filter applied is 2+ beds, it is a failure.
5: Some tasks require a submission action or a display of results to be considered successful.
6: If the retrieved information is invalid or empty(e.g., No match was found), but the agent has correctly performed the required action, it should still be considered successful.
7: If the current page already displays all available items, then applying a filter is not necessary. As long as the agent selects items that meet the requirements (e.g., the cheapest or lowest price), the task is still considered successful.
*IMPORTANT*
Format your response into two lines as shown below:
Thoughts: <your thoughts and reasoning process based on double-checking each key points and the evaluation criteria>
Status: "success" or "failure"
"""
prompt = """User Task: {task}
Key Points: {key_points}
Action History:
{last_actions}
The potentially important snapshots of the webpage in the agent's trajectory and their reasons:
{thoughts}"""
whole_content_img = []
whole_thoughts = []
record = []
pattern = r"[1-5]"
for response, image_path in zip(image_responses, images_path):
try:
score_text = response.split("Score")[1]
thought = response.split("**Reasoning**:")[-1].strip().lstrip("\n").split("\n\n")[0].replace('\n', ' ')
score = re.findall(pattern, score_text)[0]
record.append({"Response": response, "Score": int(score)})
except Exception as e:
print(f"Error processing response: {e}")
score = 0
record.append({"Response": response, "Score": 0})
if int(score) >= score_threshold:
# jpg_base64_str = encode_image(Image.open(image_path))
whole_content_img.append(
{
'type': 'image_url',
"image_url": {"url": image_path, "detail": "high"},
}
)
if thought != "":
whole_thoughts.append(thought)
whole_content_img = whole_content_img[:MAX_IMAGE]
whole_thoughts = whole_thoughts[:MAX_IMAGE]
if len(whole_content_img) == 0:
prompt = """User Task: {task}
Key Points: {key_points}
Action History:
{last_actions}"""
text = prompt.format(task=task,
last_actions="\n".join(f"{i + 1}. {action}" for i, action in enumerate(last_actions)),
key_points=key_points,
thoughts="\n".join(f"{i + 1}. {thought}" for i, thought in enumerate(whole_thoughts)))
messages = [
{"role": "system", "content": system_msg},
{
"role": "user",
"content": [
{"type": "text", "text": text}]
+ whole_content_img
}
]
return messages, text, system_msg, record
# eval code end
class PlayWrightAgent(Agent):
def __init__(self, conf: Union[Dict[str, Any], ConfigDict, AgentConfig], **kwargs):
self.screen_capture = True
self.step_images = []
self.step_thoughts = []
self.step_actions = []
self.step_results = []
self.success = False
super().__init__(conf, **kwargs)
async def async_policy(self, observation: Observation, info: Dict[str, Any] = {}, message: Message = None,
**kwargs) -> Union[
List[ActionModel], None]:
"""The strategy of an agent can be to decide which tools to use in the environment, or to delegate tasks to other agents.
Args:
observation: The state observed from tools in the environment.
info: Extended information is used to assist the agent to decide a policy.
Returns:
ActionModel sequence from agent policy
"""
outputs = None
if kwargs.get("outputs") and isinstance(kwargs.get("outputs"), Outputs):
outputs = kwargs.get("outputs")
# Get current step information for trace recording
step = kwargs.get("step", 0)
exp_id = kwargs.get("exp_id", None)
source_span = trace.get_current_span()
if hasattr(observation, 'context') and observation.context:
self.task_histories = observation.context
self._finished = False
await self.async_desc_transform(message.context)
self.tools = None
if "data:image/jpeg;base64," in observation.content:
logger.info("transfer base64 content to image")
observation.image = observation.content
observation.content = "observation:"
self.step_images.append(observation.image)
images = observation.images if self.conf.use_vision else None
if self.conf.use_vision and not images and observation.image:
images = [observation.image]
messages = self.messages_transform(content=observation.content,
image_urls=images,
sys_prompt=self.system_prompt,
agent_prompt=self.agent_prompt)
self._log_messages(messages)
if isinstance(messages[-1]['content'], list):
messages[-1]['role'] = 'user' # 有image的话必须使用user请求,而且不写入历史对话
# self.memory.add(MemoryItem(
# content=messages[-1]['content'],
# metadata={
# "role": messages[-1]['role'],
# "agent_name": self.name(),
# }
# ))
else:
self.memory.add(MemoryItem(
content=messages[-1]['content'],
metadata={
"role": messages[-1]['role'],
"agent_name": self.name(),
}
))
llm_response = None
span_name = f"llm_call_{exp_id}"
with trace.span(span_name) as llm_span:
llm_span.set_attributes({
"exp_id": exp_id,
"step": step,
"messages": json.dumps([str(m) for m in messages], ensure_ascii=False)
})
if source_span:
source_span.set_attribute("messages", json.dumps([str(m) for m in messages], ensure_ascii=False))
try:
llm_response = await acall_llm_model(
self.llm,
messages=messages,
model=self.model_name,
# temperature=self.conf.llm_config.llm_temperature,
temperature=0.0,
tools=self.tools if not self.use_tools_in_prompt and self.tools else None,
stream=kwargs.get("stream", False)
)
# Record LLM response
llm_span.set_attributes({
"llm_response": json.dumps(llm_response.to_dict(), ensure_ascii=False),
"tool_calls": json.dumps([tool_call.model_dump() for tool_call in
llm_response.tool_calls] if llm_response.tool_calls else [],
ensure_ascii=False),
"error": llm_response.error if llm_response.error else ""
})
except Exception as e:
logger.warn(traceback.format_exc())
llm_span.set_attribute("error", str(e))
raise e
finally:
if llm_response:
use_tools = self.use_tool_list(llm_response)
is_use_tool_prompt = len(use_tools) > 0
if llm_response.error:
logger.info(f"llm result error: {llm_response.error}")
else:
self.memory.add(MemoryItem(
content=llm_response.content,
metadata={
"role": "assistant",
"agent_name": self.name(),
"tool_calls": llm_response.tool_calls if not self.use_tools_in_prompt else use_tools,
"is_use_tool_prompt": is_use_tool_prompt if not self.use_tools_in_prompt else False
}
))
function, origin_thought, origin_action, origin_result = parse_tool_call(
llm_response.message['content'])
self.step_thoughts.append(origin_thought)
self.step_actions.append(origin_action)
self.step_results.append(origin_result)
if function.name == "finished":
self._finished = True
llm_response.content = "<answer>" + llm_response.content + "</answer>"
llm_response.tool_calls = None
else:
llm_response.content = None
tool_call = ToolCall(
id="tooluse_mock",
type="function",
function=function,
)
screen_capture = ToolCall(
id="screen_capture",
type="function",
function=Function(
name="mcp__ms-playwright__browser_screen_capture",
arguments="{}"
)
)
llm_response.tool_calls = [tool_call, screen_capture]
else:
logger.error(f"{self.name()} failed to get LLM response")
raise RuntimeError(f"{self.name()} failed to get LLM response")
if outputs and isinstance(outputs, Outputs):
await outputs.add_output(MessageOutput(source=llm_response, json_parse=False))
agent_result = await self.model_output_parser.parse(llm_response, agent_id=self.id())
if not agent_result.is_call_tool:
self._finished = True
logger.info(self.step_thoughts)
logger.info(self.step_actions)
# now is eval code:
logger.info(f"step:{step}")
if self.finished or step >= 20: # 暂时写死,这里应该是max_step
task = self.task.split("Please first navigate to the target")[0]
key_points_messages = identify_key_points(task)
# eval_model_name = "shangshu.gpt-4o"
eval_model_name = self.model_name
tmp_llm_response = await acall_llm_model(
self.llm,
messages=key_points_messages,
model=eval_model_name,
temperature=0
)
key_points = tmp_llm_response.content
key_points = key_points.replace("\n\n", "\n")
try:
key_points = key_points.split("**Key Points**:")[1]
key_points = "\n".join(line.lstrip() for line in key_points.splitlines())
except:
key_points = key_points.split("Key Points:")[-1]
key_points = "\n".join(line.lstrip() for line in key_points.splitlines())
logger.info(f"key_points: {key_points}")
tasks_messages = [judge_image(task, image_path, key_points) for image_path in self.step_images]
# 这里暂时使用串行执行的写法
image_responses = []
for task_messages in tasks_messages:
logger.info(task_messages)
image_response = await acall_llm_model(
self.llm, # 假设这是你传给函数的第一个参数
messages=task_messages, # 每个请求的消息内容
model=eval_model_name, # 模型名称
temperature=0 # 温度参数
)
image_responses.append(image_response)
image_responses = [i.content for i in image_responses]
logger.info(f"image_responses: {image_responses}")
eval_messages, text, system_msg, record = WebJudge_Online_Mind2Web_eval(
self.task, self.step_actions, self.step_images, image_responses, key_points, 3)
response = await acall_llm_model(
self.llm,
messages=eval_messages,
model=eval_model_name,
temperature=0
)
eval_response = response.content
logger.info(f"eval_response: {eval_response}")
if "success" in eval_response.lower().split('status:')[1]:
self.success = True
# now is saving code:
result_dict = {
'task': task,
'images': self.step_images,
'actions': self.step_actions,
'thoughts': self.step_thoughts,
'results': self.step_results,
'success': self.success,
'final_answer': llm_response.content,
'eval_response': eval_response,
'is_done': self.finished,
'done_step': step,
}
result_dict = json.dumps(result_dict, ensure_ascii=False)
agent_result.actions[0].policy_info = result_dict
agent_result.actions[0].tool_name = None
agent_result.actions[0].action_name = None
agent_result.actions[0].agent_name = self.name()
# saving is over...
return agent_result.actions
class Pipeline(AworldBaseAgent):
class Valves(BaseModel):
llm_provider: Optional[str] = Field(default=None, description="llm_model_name")
llm_model_name: Optional[str] = Field(default=None, description="llm_model_name")
llm_base_url: Optional[str] = Field(default=None, description="llm_base_urly")
llm_api_key: Optional[str] = Field(default=None, description="llm api key")
system_prompt: str = Field(default=BROWSER_SYSTEM_PROMPT, description="system_prompt")
history_messages: int = Field(default=100, description="rounds of history messages")
def __init__(self):
self.valves = self.Valves()
self.agent_config = AgentConfig(
name=self.agent_name(),
llm_provider=self.valves.llm_provider if self.valves.llm_provider else os.environ.get("LLM_PROVIDER"),
llm_model_name=self.valves.llm_model_name if self.valves.llm_model_name else os.environ.get(
"LLM_MODEL_NAME"),
llm_api_key=self.valves.llm_api_key if self.valves.llm_api_key else os.environ.get("LLM_API_KEY"),
llm_base_url=self.valves.llm_base_url if self.valves.llm_base_url else os.environ.get("LLM_BASE_URL"),
system_prompt=self.valves.system_prompt if self.valves.system_prompt else BROWSER_SYSTEM_PROMPT
)
self.m2w_files = os.path.abspath(os.path.join(os.path.curdir, "aworldspace", "datasets", "online-mind2web"))
logging.info(f"m2w_files path {self.m2w_files}")
file_path = os.path.join(self.m2w_files, "Online_Mind2Web.json")
with open(file_path, 'r') as file:
self.full_dataset = json.load(file)
logging.info("playwright_agent init success")
# 重写build_agent
async def build_agent(self, body: dict):
agent_config = await self.get_agent_config(body)
mcp_servers = await self.get_mcp_servers(body)
agent = PlayWrightAgent(
conf=agent_config,
name=agent_config.name,
system_prompt=agent_config.system_prompt,
mcp_servers=mcp_servers,
mcp_config=await self.load_mcp_config(),
history_messages=await self.get_history_messages(body)
)
return agent
async def get_custom_input(self, user_message: str, model_id: str, messages: List[dict], body: dict) -> Any:
task = await self.get_m2w_task(int(user_message))
return task['Task']
async def get_agent_config(self, body):
default_llm_provider = self.valves.llm_provider if self.valves.llm_provider else os.environ.get("LLM_PROVIDER")
llm_model_name = self.valves.llm_model_name if self.valves.llm_model_name else os.environ.get("LLM_MODEL_NAME")
llm_api_key = self.valves.llm_api_key if self.valves.llm_api_key else os.environ.get("LLM_API_KEY")
llm_base_url = self.valves.llm_base_url if self.valves.llm_base_url else os.environ.get("LLM_BASE_URL")
system_prompt = self.valves.system_prompt if self.valves.system_prompt else BROWSER_SYSTEM_PROMPT
task = await self.get_task_from_body(body)
logging.info(
f"task llm config is: {task.llm_provider}, {task.llm_model_name}, {task.llm_api_key}, {task.llm_base_url}")
return AgentConfig(
name=self.agent_name(),
llm_provider=task.llm_provider if task and task.llm_provider else default_llm_provider,
llm_model_name=task.llm_model_name if task and task.llm_model_name else llm_model_name,
llm_api_key=task.llm_api_key if task and task.llm_api_key else llm_api_key,
llm_base_url=task.llm_base_url if task and task.llm_base_url else llm_base_url,
system_prompt=task.task_system_prompt if task and task.task_system_prompt else system_prompt
)
def agent_name(self) -> str:
return "PlaywrightAgent"
async def get_mcp_servers(self, body) -> list[str]:
task = await self.get_task_from_body(body)
if task.mcp_servers:
logging.info(f"mcp_servers from task: {task.mcp_servers}")
return task.mcp_servers
return [
"ms-playwright"
]
async def get_m2w_task(self, index) -> dict:
logging.info(f"Start to process: m2w_task_{index}")
m2w_task = self.full_dataset[index]
logging.info(f"Detail: {m2w_task}")
logging.info(f"Task: {m2w_task['confirmed_task']}")
logging.info(f"Level: {m2w_task['level']}")
logging.info(f"Website: {m2w_task['website']}")
return self.add_file_path(m2w_task)
async def custom_output_before_task(self, outputs: Outputs, chat_id: str, task: Task) -> None:
task_config: TaskConfig = task.conf
m2w_task = await self.get_m2w_task(int(task_config.ext['origin_message']))
result = f"\n\n`Web TASK#{task_config.ext['origin_message']}`\n\n---\n\n"
result += f"**Task**: {m2w_task['Task']}\n"
result += f"**Level**: {m2w_task['level']}\n"
result += f"**Website**: \n {m2w_task['website']}\n"
result += f"\n\n-----\n\n"
await outputs.add_output(Output(data=result))
async def custom_output_after_task(self, outputs: Outputs, chat_id: str, task: Task):
"""
check gaia task output
Args:
outputs:
chat_id:
task:
Returns:
"""
task_config: TaskConfig = task.conf
web_task_id = int(task_config['ext']['origin_message'])
web_task = await self.get_m2w_task(web_task_id)
agent_result = ""
if isinstance(outputs, StreamingOutputs):
agent_result = await outputs._visited_outputs[-2].get_finished_response() # read llm result
# match = re.search(r"<answer>(.*?)</answer>", agent_result)
result = ""
# if match:
# answer = match.group(1)
logging.info(f"Agent answer: {agent_result}")
metadata = await outputs.get_metadata()
if not metadata:
await outputs.set_metadata({})
metadata = await outputs.get_metadata()
metadata['web_task'] = web_task
return result
def add_file_path(self, task: Dict[str, Any]
):
task["Task"] = "Task: " + task['confirmed_task'] + '\n' + "Please first navigate to the target " + "Website: " + \
task['website']
return task
async def load_mcp_config(self) -> dict:
return {
"mcpServers": {
"ms-playwright": {
"command": "npx",
"args": [
"@playwright/mcp@0.0.27",
"--vision",
"--no-sandbox",
"--headless",
"--isolated"
],
"env": {
"PLAYWRIGHT_TIMEOUT": "120000",
"SESSION_REQUEST_CONNECT_TIMEOUT": "120"
}
}
}
}
@@ -0,0 +1,32 @@
from typing import Optional
from pydantic import BaseModel, Field
from aworldspace.base_agent import AworldBaseAgent
"""
Agent Space
"""
class AgentMeta(BaseModel):
name: str = None
desc: str = None
class AgentSpace(BaseModel):
agent_modules: Optional[dict] = Field(default_factory=dict, description="agent module")
agents_meta: Optional[dict] = Field(default_factory=dict, description="agents meta")
def register(self, agent_name: str, agent_instance: AworldBaseAgent, metadata: dict=None):
# Register agent metadata and instance
self.agent_modules[agent_name] = agent_instance
async def get_agent_modules(self):
return self.agent_modules
async def get_agents_meta(self):
return self.agents_meta
AGENT_SPACE = AgentSpace()
@@ -0,0 +1,242 @@
import json
import logging
import os
import traceback
import uuid
from abc import abstractmethod
from typing import List, AsyncGenerator, Any
from aworld.config import AgentConfig, TaskConfig, ContextRuleConfig, OptimizationConfig
from aworld.agents.llm_agent import Agent
from aworld.core.task import Task
from aworld.output import WorkSpace, AworldUI, Outputs
from aworld.output.ui.markdown_aworld_ui import MarkdownAworldUI
from aworld.output.utils import load_workspace
from aworld.runner import Runners
from client.aworld_client import AworldTask
class AworldBaseAgent:
def pipes(self) -> list[dict]:
return [{"id": self.agent_name(), "name": self.agent_name()}]
@abstractmethod
def agent_name(self) -> str:
pass
async def pipe(
self,
user_message: str,
model_id: str,
messages: List[dict],
body: dict
):
try:
logging.info(f"🤖{self.agent_name()} received user_message is {user_message}, form-data = {body}")
task = await self.get_task_from_body(body)
if task:
logging.info(f"🤖{self.agent_name()} received task is {task.task_id}_{task.client_id}_{task.user_id}")
task_id = task.task_id
else:
task_id = str(uuid.uuid4())
session_id = task_id
if body.get('metadata'):
# user_id = body.get('metadata').get('user_id')
session_id = body.get('metadata').get('chat_id', task_id)
task_id = body.get('metadata').get('message_id', task_id)
user_input = await self.get_custom_input(user_message, model_id, messages, body)
if task and task.llm_custom_input:
user_input = task.llm_custom_input
logging.info(f"🤖{self.agent_name()} call llm input is [{user_input}]")
# build agent task read from config
swarm = await self.build_swarm(body=body)
agent = None
if not swarm:
# build single agent task read from config
agent = await self.build_agent(body=body)
logging.info(f"🤖{self.agent_name()} build agent finished")
# return task
task = await self.build_task(agent=agent, task_id=task_id, user_input=user_input, user_message=user_message, body=body)
logging.info(f"🤖{self.agent_name()} build task finished, task_id is {task_id}")
workspace_type = os.environ.get("WORKSPACE_TYPE", "local")
workspace_path = os.environ.get("WORKSPACE_PATH", "./data/workspaces")
workspace = await load_workspace(session_id, workspace_type, workspace_path)
# render output
async_generator = await self.parse_task_output(session_id, task, workspace)
return async_generator()
except Exception as e:
return await self._format_exception(e)
async def _format_exception(self, e: Exception) -> str:
traceback.print_exc()
# tb_lines = traceback.format_exception(type(e), e, e.__traceback__)
# detailed_error = "".join(tb_lines)
# logging.error(e)
# return json.dumps({"error": detailed_error}, ensure_ascii=False)
return "💥💥💥process failed💥💥💥"
async def _format_error(self, status_code: int, error: bytes) -> str:
if isinstance(error, str):
error_str = error
else:
error_str = error.decode(errors="ignore")
try:
err_msg = json.loads(error_str).get("message", error_str)[:200]
except Exception:
err_msg = error_str[:200]
return json.dumps(
{"error": f"HTTP {status_code}: {err_msg}"}, ensure_ascii=False
)
async def get_custom_input(self, user_message: str,
model_id: str,
messages: List[dict],
body: dict) -> Any:
user_input = body["messages"][-1]["content"]
return user_input
@abstractmethod
async def get_history_messages(self, body) -> int:
task = await self.get_task_from_body(body)
if task:
return task.history_messages
return 100
@abstractmethod
async def get_agent_config(self, body) -> AgentConfig:
pass
@abstractmethod
async def get_mcp_servers(self, body) -> list[str]:
pass
async def build_agent(self, body: dict):
agent_config =await self.get_agent_config(body)
mcp_servers = await self.get_mcp_servers(body)
agent = Agent(
conf=agent_config,
name=agent_config.name,
system_prompt=agent_config.system_prompt,
mcp_servers=mcp_servers,
mcp_config=await self.load_mcp_config(),
history_messages=await self.get_history_messages(body),
context_rule=ContextRuleConfig(
optimization_config=OptimizationConfig(
enabled=False,
)
)
)
return agent
async def build_task(self, agent, task_id, user_input, user_message, body):
aworld_task = await self.get_task_from_body(body)
task = Task(
id=task_id,
name=task_id,
input=user_input,
agent=agent,
conf=TaskConfig(
task_id=task_id,
stream=False,
ext={
"origin_message": user_message
},
max_steps=aworld_task.max_steps if aworld_task else 100
)
)
return task
async def parse_task_output(self, chat_id, task: Task, workspace: WorkSpace):
_SENTINEL = object()
async def async_generator():
from asyncio import Queue
queue = Queue()
async def consume_all():
openwebui_ui = MarkdownAworldUI(
session_id=chat_id,
workspace=workspace
)
# get outputs
outputs = Runners.streamed_run_task(task)
# output hooks
await self.custom_output_before_task(outputs, chat_id, task)
# render output
try:
async for output in outputs.stream_events():
res = await AworldUI.parse_output(output, openwebui_ui)
if res:
if isinstance(res, AsyncGenerator):
async for item in res:
await queue.put(item)
else:
await queue.put(res)
custom_output = await self.custom_output_after_task(outputs, chat_id, task)
if custom_output:
await queue.put(custom_output)
await queue.put(task)
finally:
await queue.put(_SENTINEL)
# Start the consumer in the background
import asyncio
consumer_task = asyncio.create_task(consume_all())
while True:
item = await queue.get()
if item is _SENTINEL:
break
yield item
await consumer_task
logging.info(f"🤖{self.agent_name()} task#{task.id} output finished🔚🔚🔚")
return async_generator
async def custom_output_before_task(self, outputs: Outputs, chat_id: str, task: Task) -> str | None:
return None
async def custom_output_after_task(self, outputs: Outputs, chat_id: str, task: Task):
pass
async def get_task_from_body(self, body: dict) -> AworldTask | None:
try:
if not body.get("user") or not body.get("user").get("aworld_task"):
return None
return AworldTask.model_validate_json(body.get("user").get("aworld_task"))
except Exception as err:
logging.error(f"Error parsing AworldTask: {err}; data: {body.get('user_message')}")
traceback.print_exc()
return None
@abstractmethod
async def load_mcp_config(self) -> dict:
pass
async def build_swarm(self, body):
return None
@@ -0,0 +1,210 @@
from abc import ABC, abstractmethod
from datetime import datetime
from typing import Optional
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from base import AworldTask, AworldTaskResult
from aworldspace.db.models import (
Base, AworldTaskModel, AworldTaskResultModel,
orm_to_pydantic_task, pydantic_to_orm_task,
orm_to_pydantic_result, pydantic_to_orm_result
)
class AworldTaskDB(ABC):
@abstractmethod
async def query_task_by_id(self, task_id: str) -> AworldTask:
pass
@abstractmethod
async def query_latest_task_result_by_id(self, task_id: str) -> Optional[AworldTaskResult]:
pass
@abstractmethod
async def insert_task(self, task: AworldTask):
pass
@abstractmethod
async def query_tasks_by_status(self, status: str, nums: int) -> list[AworldTask]:
pass
@abstractmethod
async def update_task(self, task: AworldTask):
pass
@abstractmethod
async def page_query_tasks(self, filter: dict, page_size: int, page_num: int) -> dict:
pass
@abstractmethod
async def save_task_result(self, result: AworldTaskResult):
pass
class SqliteTaskDB(AworldTaskDB):
def __init__(self, db_path: str):
self.engine = create_engine(db_path, echo=False, future=True)
Base.metadata.create_all(self.engine)
self.Session = sessionmaker(bind=self.engine, expire_on_commit=False)
async def query_task_by_id(self, task_id: str) -> Optional[AworldTask]:
with self.Session() as session:
orm_task = session.query(AworldTaskModel).filter_by(task_id=task_id).first()
return orm_to_pydantic_task(orm_task) if orm_task else None
async def query_latest_task_result_by_id(self, task_id: str) -> Optional[AworldTaskResult]:
with self.Session() as session:
orm_result = (
session.query(AworldTaskResultModel)
.filter_by(task_id=task_id)
.order_by(AworldTaskResultModel.created_at.desc())
.first()
)
return orm_to_pydantic_result(orm_result) if orm_result else None
async def insert_task(self, task: AworldTask):
with self.Session() as session:
orm_task = pydantic_to_orm_task(task)
session.add(orm_task)
session.commit()
async def query_tasks_by_status(self, status: str, nums: int) -> list[AworldTask]:
with self.Session() as session:
orm_tasks = (
session.query(AworldTaskModel)
.filter_by(status=status)
.limit(nums)
.all()
)
return [orm_to_pydantic_task(t) for t in orm_tasks]
async def update_task(self, task: AworldTask):
with self.Session() as session:
orm_task = session.query(AworldTaskModel).filter_by(task_id=task.task_id).first()
if orm_task:
for k, v in task.model_dump().items():
setattr(orm_task, k, v)
orm_task.updated_at = datetime.utcnow()
session.commit()
async def save_task_result(self, result: AworldTaskResult):
with self.Session() as session:
orm_task = pydantic_to_orm_result(result)
session.add(orm_task)
session.commit()
async def page_query_tasks(self, filter: dict, page_size: int, page_num: int) -> dict:
with self.Session() as session:
query = session.query(AworldTaskModel)
# Handle special filters for time ranges
start_time = filter.pop('start_time', None)
end_time = filter.pop('end_time', None)
# Apply regular filters
for k, v in filter.items():
if hasattr(AworldTaskModel, k):
query = query.filter(getattr(AworldTaskModel, k) == v)
# Apply time range filters
if start_time:
query = query.filter(AworldTaskModel.created_at >= start_time)
if end_time:
query = query.filter(AworldTaskModel.created_at <= end_time)
total = query.count()
orm_tasks = query.offset((page_num - 1) * page_size).limit(page_size).all()
items = [orm_to_pydantic_task(t) for t in orm_tasks]
return {
"total": total,
"page_num": page_num,
"page_size": page_size,
"items": items
}
class PostgresTaskDB(AworldTaskDB):
def __init__(self, db_url: str):
# db_url example: 'postgresql+psycopg2://user:password@host:port/dbname'
self.engine = create_engine(db_url, echo=False, future=True)
Base.metadata.create_all(self.engine)
self.Session = sessionmaker(bind=self.engine, expire_on_commit=False)
async def query_task_by_id(self, task_id: str) -> Optional[AworldTask]:
with self.Session() as session:
orm_task = session.query(AworldTaskModel).filter_by(task_id=task_id).first()
return orm_to_pydantic_task(orm_task) if orm_task else None
async def query_latest_task_result_by_id(self, task_id: str) -> Optional[AworldTaskResult]:
with self.Session() as session:
orm_result = (
session.query(AworldTaskResultModel)
.filter_by(task_id=task_id)
.order_by(AworldTaskResultModel.created_at.desc())
.first()
)
return orm_to_pydantic_result(orm_result) if orm_result else None
async def insert_task(self, task: AworldTask):
with self.Session() as session:
orm_task = pydantic_to_orm_task(task)
session.add(orm_task)
session.commit()
async def query_tasks_by_status(self, status: str, nums: int) -> list[AworldTask]:
with self.Session() as session:
orm_tasks = (
session.query(AworldTaskModel)
.filter_by(status=status)
.limit(nums)
.all()
)
return [orm_to_pydantic_task(t) for t in orm_tasks]
async def update_task(self, task: AworldTask):
with self.Session() as session:
orm_task = session.query(AworldTaskModel).filter_by(task_id=task.task_id).first()
if orm_task:
for k, v in task.model_dump().items():
setattr(orm_task, k, v)
orm_task.updated_at = datetime.utcnow()
session.commit()
async def save_task_result(self, result: AworldTaskResult):
with self.Session() as session:
orm_task = pydantic_to_orm_result(result)
session.add(orm_task)
session.commit()
async def page_query_tasks(self, filter: dict, page_size: int, page_num: int) -> dict:
with self.Session() as session:
query = session.query(AworldTaskModel)
# Handle special filters for time ranges
start_time = filter.pop('start_time', None)
end_time = filter.pop('end_time', None)
# Apply regular filters
for k, v in filter.items():
if hasattr(AworldTaskModel, k):
query = query.filter(getattr(AworldTaskModel, k) == v)
# Apply time range filters
if start_time:
query = query.filter(AworldTaskModel.created_at >= start_time)
if end_time:
query = query.filter(AworldTaskModel.created_at <= end_time)
total = query.count()
orm_tasks = query.offset((page_num - 1) * page_size).limit(page_size).all()
items = [orm_to_pydantic_task(t) for t in orm_tasks]
return {
"total": total,
"page_num": page_num,
"page_size": page_size,
"items": items
}
@@ -0,0 +1,66 @@
from sqlalchemy import Column, String, Integer, Text, DateTime, JSON, create_engine
from sqlalchemy.orm import declarative_base
from datetime import datetime
from typing import Optional
from base import AworldTask, AworldTaskResult
Base = declarative_base()
class AworldTaskModel(Base):
__tablename__ = 'aworld_tasks'
task_id = Column(String, primary_key=True)
agent_id = Column(String)
agent_input = Column(Text)
session_id = Column(String)
user_id = Column(String)
llm_provider = Column(String)
llm_model_name = Column(String)
llm_api_key = Column(String)
llm_base_url = Column(String)
llm_custom_input = Column(Text)
task_system_prompt = Column(Text)
mcp_servers = Column(JSON)
node_id = Column(String)
client_id = Column(String)
status = Column(String, default='INIT')
history_messages = Column(Integer, default=100)
max_steps = Column(Integer, default=100)
max_retries = Column(Integer, default=5)
ext_info = Column(JSON, default=dict)
created_at = Column(DateTime, default=datetime.utcnow)
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
class AworldTaskResultModel(Base):
__tablename__ = 'aworld_tasks_results'
task_result_id = Column(Integer, primary_key=True, autoincrement=True)
task_id = Column(String)
server_host = Column(String)
data = Column(JSON)
created_at = Column(DateTime, default=datetime.utcnow)
def orm_to_pydantic_task(orm_obj: AworldTaskModel) -> AworldTask:
return AworldTask(**{c.name: getattr(orm_obj, c.name) for c in orm_obj.__table__.columns})
def pydantic_to_orm_task(pydantic_obj: AworldTask) -> AworldTaskModel:
return AworldTaskModel(**pydantic_obj.model_dump())
def orm_to_pydantic_result(orm_obj: AworldTaskResultModel) -> AworldTaskResult:
return AworldTaskResult(
server_host=orm_obj.server_host,
data=orm_obj.data
)
def pydantic_to_orm_result(pydantic_obj: AworldTaskResult) -> AworldTaskResultModel:
return AworldTaskResultModel(
task_id=pydantic_obj.task.task_id if pydantic_obj.task else None,
server_host=pydantic_obj.server_host,
data=pydantic_obj.data
)
@@ -0,0 +1,439 @@
import json
import os
import time
from datetime import datetime
from typing import AsyncGenerator, Optional, List
from aworld.utils.common import get_local_ip
from fastapi import APIRouter, Query, Response
from fastapi.responses import StreamingResponse
import logging
import traceback
from asyncio import Queue
import asyncio
from aworld.models.model_response import ModelResponse
from pydantic import BaseModel, Field, PrivateAttr
from aworldspace.db.db import AworldTaskDB, SqliteTaskDB, PostgresTaskDB
from aworldspace.utils.job import generate_openai_chat_completion, call_pipeline
from aworldspace.utils.log import task_logger
from base import AworldTask, AworldTaskResult, OpenAIChatCompletionForm, OpenAIChatMessage, AworldTaskForm
from config import ROOT_DIR
__STOP_TASK__ = object()
class AworldTaskExecutor(BaseModel):
"""
task executor
- load task from db and execute task in a loop
- use semaphore to limit concurrent tasks
"""
_task_db: AworldTaskDB = PrivateAttr()
_tasks: Queue = PrivateAttr()
max_concurrent: int = Field(default=os.environ.get("AWORLD_MAX_CONCURRENT_TASKS", 2), description="max concurrent tasks")
def __init__(self, task_db: AworldTaskDB):
super().__init__()
self._task_db = task_db
self._tasks = Queue()
self._semaphore = asyncio.BoundedSemaphore(self.max_concurrent)
async def start(self):
"""
execute task in a loop
"""
await asyncio.sleep(5)
logging.info(f"🚀[task executor] start, max concurrent is {self.max_concurrent}")
while True:
# load task if queue is empty and semaphore is not full
if self._tasks.empty():
await self.load_task()
task = await self._tasks.get()
if not task:
logging.info("task is none")
continue
if task == __STOP_TASK__:
logging.info("✅[task executor] stop, all tasks finished")
break
# acquire semaphore
await self._semaphore.acquire()
asyncio.create_task(self._run_task_and_release_semaphore(task))
async def stop(self):
logging.info("🛑 task executor stop, wait for all tasks to finish")
await self._tasks.put(__STOP_TASK__)
async def _run_task_and_release_semaphore(self, task: AworldTask):
"""
execute task and release semaphore when done
"""
start_time = time.time()
logging.info(f"🚀[task executor] execute task#{task.task_id} start, lock acquired")
try:
await self.execute_task(task)
finally:
# release semaphore
self._semaphore.release()
logging.info(f"✅[task executor] execute task#{task.task_id} success, use time {time.time() - start_time:.2f}s")
async def load_task(self):
interval = os.environ.get("AWORLD_TASK_LOAD_INTERVAL", 10)
# calculate the number of tasks to load
need_load = self._semaphore._value
if need_load <= 0:
logging.info(f"🔍[task executor] runner is busy, wait {interval}s and retry")
await asyncio.sleep(interval)
return await self.load_task()
tasks = await self._task_db.query_tasks_by_status(status="INIT", nums=need_load)
logging.info(f"🔍[task executor] load {len(tasks)} tasks from db (need {need_load})")
if not tasks or len(tasks) == 0:
logging.info(f"🔍[task executor] no task to load, wait {interval}s and retry")
await asyncio.sleep(interval)
return await self.load_task()
for task in tasks:
task.mark_running()
await self._task_db.update_task(task)
await self._tasks.put(task)
return True
async def execute_task(self, task: AworldTask):
"""
execute task
"""
try:
result = await self._execute_task(task)
task.mark_success()
await self._task_db.update_task(task)
await self._task_db.save_task_result(result)
task_logger.log_task_submission(task, "execute_finished", task_result=result)
except Exception as err:
task.mark_failed()
await self._task_db.update_task(task)
traceback.print_exc()
task_logger.log_task_submission(task, "execute_failed", details=f"err is {err}")
async def _execute_task(self, task: AworldTask):
# build params
messages = [
OpenAIChatMessage(role="user", content=task.agent_input)
]
# call_llm_model
form_data = OpenAIChatCompletionForm(
model=task.agent_id,
messages=messages,
stream=True,
user={
"user_id": task.user_id,
"session_id": task.session_id,
"task_id": task.task_id,
"aworld_task": task.model_dump_json()
}
)
data = await generate_openai_chat_completion(form_data)
task_result = {}
task.node_id = get_local_ip()
items = []
md_file = ""
if data.body_iterator:
if isinstance(data.body_iterator, AsyncGenerator):
async for item_content in data.body_iterator:
async def parse_item(_item_content) -> Optional[ModelResponse]:
if item_content == "data: [DONE]":
return None
return ModelResponse.from_openai_stream_chunk(json.loads(item_content.replace("data:", "")))
# if isinstance(item, ModelResponse)
item = await parse_item(item_content)
items.append(item)
if not item:
continue
if item.content:
md_file = task_logger.log_task_result(task, item)
logging.info(f"task#{task.task_id} response data chunk is: {item}"[:500])
if item.raw_response and item.raw_response and isinstance(item.raw_response, dict) and item.raw_response.get('task_output_meta'):
task_result = item.raw_response.get('task_output_meta')
data = {
"task_result": task_result,
"md_file": md_file,
"replays_file": f"trace_data/{datetime.now().strftime('%Y%m%d')}/{get_local_ip()}/replays/task_replay_{task.task_id}.json"
}
result = AworldTaskResult(task=task, server_host=get_local_ip(), data=data)
return result
class AworldTaskManager(BaseModel):
_task_db: AworldTaskDB = PrivateAttr()
_task_executor: AworldTaskExecutor = PrivateAttr()
def __init__(self, task_db: AworldTaskDB):
super().__init__()
self._task_db = task_db
self._task_executor = AworldTaskExecutor(task_db=self._task_db)
async def start_task_executor(self):
asyncio.create_task(self._task_executor.start())
async def stop_task_executor(self):
self._task_executor.tasks.put_nowait(None)
async def submit_task(self, task: AworldTask):
# save to db
await self._task_db.insert_task(task)
# log it
task_logger.log_task_submission(task, status="init")
return AworldTaskResult(task = task)
async def load_one_unfinished_task(self) -> Optional[AworldTask]:
tasks = await self._task_db.query_tasks_by_status(status="INIT", nums=1)
if not tasks or len(tasks) == 0:
return None
cur_task = tasks[0]
cur_task.mark_running()
await self._task_db.update_task(cur_task)
# from db load one task by locked and mark task running
return cur_task
async def get_task_result(self, task_id: str) -> Optional[AworldTaskResult]:
task = await self._task_db.query_task_by_id(task_id)
if task:
task_result = await self._task_db.query_latest_task_result_by_id(task_id)
if task_result:
return task_result
return AworldTaskResult(task=task)
async def get_batch_task_results(self, task_ids: List[str]) -> List[dict]:
"""
Batch retrieve task results, returns dictionary format
Each dict contains: task (required) and task_result (may be None)
"""
results = []
for task_id in task_ids:
task = await self._task_db.query_task_by_id(task_id)
if task:
task_result = await self._task_db.query_latest_task_result_by_id(task_id)
result_dict = {
"task": task,
"task_result": task_result # May be None
}
results.append(result_dict)
return results
async def query_and_download_task_results(
self,
start_time: Optional[datetime] = None,
end_time: Optional[datetime] = None,
task_id: Optional[str] = None,
page_size: int = 100
) -> List[dict]:
"""
Query tasks and get results, support time range and task_id filtering
"""
all_results = []
page_num = 1
while True:
# Build query filter conditions
filter_dict = {}
if start_time:
filter_dict['start_time'] = start_time
if end_time:
filter_dict['end_time'] = end_time
if task_id:
filter_dict['task_id'] = task_id
# Page query tasks
page_result = await self._task_db.page_query_tasks(
filter=filter_dict,
page_size=page_size,
page_num=page_num
)
if not page_result['items']:
break
tasks = page_result['items']
for task in tasks:
# Only query task_result (may not exist)
task_result = await self._task_db.query_latest_task_result_by_id(task.task_id)
# Use task information to build results
result_data = {
"task_id": task.task_id,
"agent_id": task.agent_id,
"status": task.status,
"created_at": task.created_at.isoformat() if task.created_at else None,
"updated_at": task.updated_at.isoformat() if task.updated_at else None,
"user_id": task.user_id,
"session_id": task.session_id,
"node_id": task.node_id,
"client_id": task.client_id,
"task_data": task.model_dump(mode='json'),
"has_result": task_result is not None,
"server_host": task_result.server_host if task_result else None,
"result_data": task_result.data if task_result else None,
}
all_results.append(result_data)
if len(page_result['items']) < page_size:
break
page_num += 1
return all_results
########################################################################################
########################### API
########################################################################################
router = APIRouter()
task_db_path = os.environ.get("AWORLD_TASK_DB_PATH", f"sqlite:///{ROOT_DIR}/db/aworld.db")
if task_db_path.startswith("sqlite://"):
task_db = SqliteTaskDB(db_path = task_db_path)
elif task_db_path.startswith("mysql://"):
task_db = None # todo: add mysql task db
elif task_db_path.startswith("postgresql://") or task_db_path.startswith("postgresql+"):
task_db = PostgresTaskDB(db_url=task_db_path)
else:
raise ValueError("❌ task_db_path is not a valid sqlite, mysql or postgresql path")
task_manager = AworldTaskManager(task_db)
@router.post("/submit_task")
async def submit_task(form_data: AworldTaskForm) -> Optional[AworldTaskResult]:
logging.info(f"🚀 submit task#{form_data.task.task_id} start")
if not form_data.task:
raise ValueError("task is empty")
try:
task_result = await task_manager.submit_task(form_data.task)
logging.info(f"✅ submit task#{form_data.task.task_id} success")
return task_result
except Exception as err:
traceback.print_exc()
logging.error(f"❌ submit task#{form_data.task.task_id} failed, err is {err}")
raise ValueError("❌ submit task failed, please see logs for details")
@router.get("/task_result")
async def get_task_result(task_id) -> Optional[AworldTaskResult]:
if not task_id:
raise ValueError("❌ task_id is empty")
logging.info(f"🚀 get task result#{task_id} start")
try:
task_result = await task_manager.get_task_result(task_id)
logging.info(f"✅ get task result#{task_id} success, task result is {task_result}")
return task_result
except Exception as err:
traceback.print_exc()
logging.error(f"❌ get task result#{task_id} failed, err is {err}")
raise ValueError("❌ get task result failed, please see logs for details")
@router.post("/get_batch_task_results")
async def get_batch_task_results(task_ids: List[str]) -> List[dict]:
if not task_ids or len(task_ids) == 0:
raise ValueError("❌ task_ids is empty")
logging.info(f"🚀 get batch task results start, task_ids: {task_ids}")
try:
batch_results = await task_manager.get_batch_task_results(task_ids)
logging.info(f"✅ get batch task results success, found {len(batch_results)} results")
return batch_results
except Exception as err:
traceback.print_exc()
logging.error(f"❌ get batch task results failed, err is {err}")
raise ValueError("❌ get batch task results failed, please see logs for details")
@router.get("/download_task_results")
async def download_task_results(
start_time: Optional[str] = Query(None, description="Start time, format: YYYY-MM-DD HH:MM:SS"),
end_time: Optional[str] = Query(None, description="End time, format: YYYY-MM-DD HH:MM:SS"),
task_id: Optional[str] = Query(None, description="Task ID"),
page_size: int = Query(100, description="Page size, ge=1, le=1000")
) -> StreamingResponse:
"""
Download task results, generate jsonl format file
Query parameters support: time range (based on creation time), task_id
"""
logging.info(f"🚀 download task results start, start_time: {start_time}, end_time: {end_time}, task_id: {task_id}")
try:
start_datetime = None
end_datetime = None
if start_time:
try:
start_datetime = datetime.strptime(start_time, "%Y-%m-%d %H:%M:%S")
except ValueError:
raise ValueError("❌ start_time格式错误,请使用 YYYY-MM-DD HH:MM:SS 格式")
if end_time:
try:
end_datetime = datetime.strptime(end_time, "%Y-%m-%d %H:%M:%S")
except ValueError:
raise ValueError("❌ end_time格式错误,请使用 YYYY-MM-DD HH:MM:SS 格式")
results = await task_manager.query_and_download_task_results(
start_time=start_datetime,
end_time=end_datetime,
task_id=task_id,
page_size=page_size
)
if not results:
logging.info("📄 no task results found")
def generate_empty():
yield ""
return StreamingResponse(
generate_empty(),
media_type="application/jsonl",
headers={"Content-Disposition": "attachment; filename=task_results_empty.jsonl"}
)
# Generate jsonl content
def generate_jsonl():
for result in results:
yield json.dumps(result, ensure_ascii=False) + "\n"
# Generate file name
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
filename = f"task_results_{timestamp}.jsonl"
logging.info(f"✅ download task results success, total: {len(results)} results")
return StreamingResponse(
generate_jsonl(),
media_type="application/jsonl",
headers={"Content-Disposition": f"attachment; filename={filename}"}
)
except Exception as err:
traceback.print_exc()
logging.error(f"❌ download task results failed, err is {err}")
raise ValueError(f"❌ download task results failed: {str(err)}")
@@ -0,0 +1,259 @@
import inspect
import json
import inspect
import json
import logging
import time
import uuid
from typing import Generator, Iterator, AsyncGenerator, Optional
from aworld.core.task import Task
from aworld.utils.common import get_local_ip
from fastapi import status, HTTPException
from fastapi.concurrency import run_in_threadpool
from pydantic import BaseModel
from starlette.responses import StreamingResponse
from aworldspace.base import AGENT_SPACE
from aworldspace.utils.utils import get_last_user_message
from base import OpenAIChatCompletionForm
async def generate_openai_chat_completion(form_data: OpenAIChatCompletionForm):
messages = [message.model_dump() for message in form_data.messages]
user_message = get_last_user_message(messages)
PIPELINES = await AGENT_SPACE.get_agents_meta()
PIPELINE_MODULES = await AGENT_SPACE.get_agent_modules()
if (
form_data.model not in PIPELINES
or PIPELINES[form_data.model]["type"] == "filter"
):
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"Pipeline {form_data.model} not found",
)
def job():
pipeline = PIPELINES[form_data.model]
pipeline_id = form_data.model
if pipeline["type"] == "manifold":
manifold_id, pipeline_id = pipeline_id.split(".", 1)
pipe = PIPELINE_MODULES[manifold_id].pipe
else:
pipe = PIPELINE_MODULES[pipeline_id].pipe
def process_line(model, line):
if isinstance(line, Task):
task_output_meta = line.outputs._metadata
line = openai_chat_chunk_message_template(model, "", task_output_meta=task_output_meta)
return f"data: {json.dumps(line)}\n\n"
if isinstance(line, BaseModel):
line = line.model_dump_json()
line = f"data: {line}"
if isinstance(line, dict):
line = f"data: {json.dumps(line)}"
try:
line = line.decode("utf-8")
except Exception:
pass
if line.startswith("data:"):
return f"{line}\n\n"
else:
line = openai_chat_chunk_message_template(model, line)
return f"data: {json.dumps(line)}\n\n"
if form_data.stream:
async def stream_content():
async def execute_pipe(_pipe):
if inspect.iscoroutinefunction(_pipe):
return await _pipe(user_message=user_message,
model_id=pipeline_id,
messages=messages,
body=form_data.model_dump())
else:
return _pipe(user_message=user_message,
model_id=pipeline_id,
messages=messages,
body=form_data.model_dump())
try:
res = await execute_pipe(pipe)
# Directly return if the response is a StreamingResponse
if isinstance(res, StreamingResponse):
async for data in res.body_iterator:
yield data
return
if isinstance(res, dict):
yield f"data: {json.dumps(res)}\n\n"
return
except Exception as e:
logging.error(f"Error: {e}")
import traceback
traceback.print_exc()
yield f"data: {json.dumps({'error': {'detail': str(e)}})}\n\n"
return
if isinstance(res, str):
message = openai_chat_chunk_message_template(form_data.model, res)
yield f"data: {json.dumps(message)}\n\n"
if isinstance(res, Iterator):
for line in res:
yield process_line(form_data.model, line)
if isinstance(res, AsyncGenerator):
async for line in res:
yield process_line(form_data.model, line)
logging.info(f"AsyncGenerator end...")
if isinstance(res, str) or isinstance(res, Generator) or isinstance(res, AsyncGenerator):
finish_message = openai_chat_chunk_message_template(
form_data.model, ""
)
finish_message["choices"][0]["finish_reason"] = "stop"
print(f"Pipe-Dataline:::: DONE")
yield f"data: {json.dumps(finish_message)}\n\n"
yield "data: [DONE]"
return StreamingResponse(stream_content(), media_type="text/event-stream")
else:
res = pipe(
user_message=user_message,
model_id=pipeline_id,
messages=messages,
body=form_data.model_dump(),
)
logging.info(f"stream:false:{res}")
if isinstance(res, dict):
return res
elif isinstance(res, BaseModel):
return res.model_dump()
else:
message = ""
if isinstance(res, str):
message = res
if isinstance(res, Generator):
for stream in res:
message = f"{message}{stream}"
logging.info(f"stream:false:{message}")
return {
"id": f"{form_data.model}-{str(uuid.uuid4())}",
"object": "chat.completion",
"created": int(time.time()),
"model": form_data.model,
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": message,
},
"logprobs": None,
"finish_reason": "stop",
}
],
}
return await run_in_threadpool(job)
async def call_pipeline(form_data: OpenAIChatCompletionForm):
messages = [message.model_dump() for message in form_data.messages]
user_message = get_last_user_message(messages)
PIPELINES = await AGENT_SPACE.get_agents_meta()
PIPELINE_MODULES = await AGENT_SPACE.get_agent_modules()
if (
form_data.model not in PIPELINES
or PIPELINES[form_data.model]["type"] == "filter"
):
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"Pipeline {form_data.model} not found",
)
pipeline = PIPELINES[form_data.model]
pipeline_id = form_data.model
if pipeline["type"] == "manifold":
manifold_id, pipeline_id = pipeline_id.split(".", 1)
pipe = PIPELINE_MODULES[manifold_id].pipe
else:
pipe = PIPELINE_MODULES[pipeline_id].pipe
if form_data.stream:
async def execute_pipe(_pipe):
if inspect.iscoroutinefunction(_pipe):
return await _pipe(user_message=user_message,
model_id=pipeline_id,
messages=messages,
body=form_data.model_dump())
else:
return _pipe(user_message=user_message,
model_id=pipeline_id,
messages=messages,
body=form_data.model_dump())
res = await execute_pipe(pipe)
return res
else:
if not inspect.iscoroutinefunction(pipe):
return await run_in_threadpool(
pipe,
user_message=user_message,
model_id=pipeline_id,
messages=messages,
body=form_data.model_dump()
)
else:
return await pipe(
user_message=user_message,
model_id=pipeline_id,
messages=messages,
body=form_data.model_dump()
)
def openai_chat_chunk_message_template(
model: str,
content: Optional[str] = None,
tool_calls: Optional[list[dict]] = None,
usage: Optional[dict] = None,
**kwargs
) -> dict:
template = openai_chat_message_template(model, **kwargs)
template["object"] = "chat.completion.chunk"
template["choices"][0]["index"] = 0
template["choices"][0]["delta"] = {}
if content:
template["choices"][0]["delta"]["content"] = content
if tool_calls:
template["choices"][0]["delta"]["tool_calls"] = tool_calls
if not content and not tool_calls:
template["choices"][0]["finish_reason"] = "stop"
if usage:
template["usage"] = usage
return template
def openai_chat_message_template(model: str, **kwargs):
return {
"id": f"{model}-{str(uuid.uuid4())}",
"created": int(time.time()),
"model": model,
"node_id": get_local_ip(),
"task_output_meta": kwargs.get("task_output_meta"),
"choices": [{"index": 0, "logprobs": None, "finish_reason": None}],
}
@@ -0,0 +1,197 @@
import importlib.util
import json
import logging
import os
import subprocess
import sys
import traceback
from aworldspace.base import AGENT_SPACE
import aworld.trace as trace # noqa
from config import AGENTS_DIR
if not os.path.exists(AGENTS_DIR):
os.makedirs(AGENTS_DIR)
PIPELINES = {}
PIPELINE_MODULES = {}
def get_all_pipelines():
pipelines = {}
for pipeline_id in PIPELINE_MODULES.keys():
pipeline = PIPELINE_MODULES[pipeline_id]
if hasattr(pipeline, "type"):
if pipeline.type == "manifold":
manifold_pipelines = []
# Check if pipelines is a function or a list
if callable(pipeline.pipelines):
manifold_pipelines = pipeline.pipelines()
else:
manifold_pipelines = pipeline.pipelines
for p in manifold_pipelines:
manifold_pipeline_id = f'{pipeline_id}.{p["id"]}'
manifold_pipeline_name = p["name"]
if hasattr(pipeline, "name"):
manifold_pipeline_name = (
f"{pipeline.name}{manifold_pipeline_name}"
)
pipelines[manifold_pipeline_id] = {
"module": pipeline_id,
"type": pipeline.type if hasattr(pipeline, "type") else "pipe",
"id": manifold_pipeline_id,
"name": manifold_pipeline_name,
"valves": (
pipeline.valves if hasattr(pipeline, "valves") else None
),
}
if pipeline.type == "filter":
pipelines[pipeline_id] = {
"module": pipeline_id,
"type": (pipeline.type if hasattr(pipeline, "type") else "pipe"),
"id": pipeline_id,
"name": (
pipeline.name if hasattr(pipeline, "name") else pipeline_id
),
"pipelines": (
pipeline.valves.pipelines
if hasattr(pipeline, "valves")
and hasattr(pipeline.valves, "pipelines")
else []
),
"priority": (
pipeline.valves.priority
if hasattr(pipeline, "valves")
and hasattr(pipeline.valves, "priority")
else 0
),
"valves": pipeline.valves if hasattr(pipeline, "valves") else None,
}
else:
pipelines[pipeline_id] = {
"module": pipeline_id,
"type": (pipeline.type if hasattr(pipeline, "type") else "pipe"),
"id": pipeline_id,
"name": (pipeline.name if hasattr(pipeline, "name") else pipeline_id),
"valves": pipeline.valves if hasattr(pipeline, "valves") else None,
}
return pipelines
def parse_frontmatter(content):
frontmatter = {}
for line in content.split("\n"):
if ":" in line:
key, value = line.split(":", 1)
frontmatter[key.strip().lower()] = value.strip()
return frontmatter
def install_frontmatter_requirements(requirements):
if requirements:
req_list = [req.strip() for req in requirements.split(",")]
for req in req_list:
print(f"Installing requirement: {req}")
subprocess.check_call([sys.executable, "-m", "pip", "install", req])
else:
print("No requirements found in frontmatter.")
async def load_module_from_path(module_name, module_path):
try:
# Read the module content
with open(module_path, "r") as file:
content = file.read()
# Parse frontmatter
frontmatter = {}
if content.startswith('"""'):
end = content.find('"""', 3)
if end != -1:
frontmatter_content = content[3:end]
frontmatter = parse_frontmatter(frontmatter_content)
# Install requirements if specified
if "requirements" in frontmatter:
install_frontmatter_requirements(frontmatter["requirements"])
# Load the module
spec = importlib.util.spec_from_file_location(module_name, module_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
logging.info(f"Loaded module start: {module.__name__}")
if hasattr(module, "Pipeline"):
return module.Pipeline()
else:
logging.info(f"Loaded module failed: {module.__name__ } No Pipeline class found")
raise Exception("No Pipeline class found")
except Exception as e:
logging.info(f"Error loading module: {module_name}, error is {e}")
traceback.print_exc()
# Move the file to the error folder
failed_pipelines_folder = os.path.join(AGENTS_DIR, "failed")
if not os.path.exists(failed_pipelines_folder):
os.makedirs(failed_pipelines_folder)
# failed_file_path = os.path.join(failed_pipelines_folder, f"{module_name}.py")
# if module_path.__contains__(PIPELINES_DIR):
# os.rename(module_path, failed_file_path)
print(e)
return None
async def load_modules_from_directory(directory):
logging.info(f"load_modules_from_directory: {directory}")
global PIPELINE_MODULES
for filename in os.listdir(directory):
if filename.endswith(".py"):
module_name = filename[:-3] # Remove the .py extension
module_path = os.path.join(directory, filename)
# Create subfolder matching the filename without the .py extension
subfolder_path = os.path.join(directory, module_name)
if not os.path.exists(subfolder_path):
os.makedirs(subfolder_path)
logging.info(f"Created subfolder: {subfolder_path}")
# Create a valves.json file if it doesn't exist
valves_json_path = os.path.join(subfolder_path, "valves.json")
if not os.path.exists(valves_json_path):
with open(valves_json_path, "w") as f:
json.dump({}, f)
logging.info(f"Created valves.json in: {subfolder_path}")
pipeline = await load_module_from_path(module_name, module_path)
if pipeline:
# Overwrite pipeline.valves with values from valves.json
if os.path.exists(valves_json_path):
with open(valves_json_path, "r") as f:
valves_json = json.load(f)
if hasattr(pipeline, "valves"):
ValvesModel = pipeline.valves.__class__
# Create a ValvesModel instance using default values and overwrite with valves_json
combined_valves = {
**pipeline.valves.model_dump(),
**valves_json,
}
valves = ValvesModel(**combined_valves)
pipeline.valves = valves
logging.info(f"Updated valves for module: {module_name}")
pipeline_id = pipeline.id if hasattr(pipeline, "id") else module_name
PIPELINE_MODULES[pipeline_id] = pipeline
logging.info(f"Loaded module success: {module_name}")
else:
logging.warning(f"No Pipeline class found in {module_name}")
AGENT_SPACE.agent_modules = PIPELINE_MODULES
AGENT_SPACE.agents_meta = get_all_pipelines()
@@ -0,0 +1,75 @@
import logging
import os
from datetime import datetime
from aworld.models.model_response import ModelResponse
from base import AworldTask, AworldTaskResult
from config import ROOT_LOG
class TaskLogger:
"""任务提交日志记录器"""
def __init__(self, log_file: str = "aworld_task_submissions.log"):
self.log_file = os.path.join(ROOT_LOG, 'task_logs' , log_file)
self._ensure_log_file_exists()
def _ensure_log_file_exists(self):
"""确保日志文件存在"""
if not os.path.exists(self.log_file):
os.makedirs(os.path.dirname(self.log_file), exist_ok=True)
with open(self.log_file, 'w', encoding='utf-8') as f:
f.write("# Aworld Task Submission Log\n")
f.write(
"# Format: [timestamp] task_id | agent_id | server | status | agent_answer | correct_answer | is_correct | details\n\n")
def log_task_submission(self, task: AworldTask, status: str, details: str = "",
task_result: AworldTaskResult = None):
"""记录任务提交日志"""
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
log_entry = f"[{timestamp}] {task.task_id} | {task.agent_id} | {task.node_id} | {status} | {task_result.data.get('agent_answer') if task_result and task_result.data else None} | {task_result.data.get('correct_answer') if task_result and task_result.data else None} | {task_result.data.get('gaia_correct') if task_result and task_result.data else None} |{details}\n"
try:
with open(self.log_file, 'a', encoding='utf-8') as f:
f.write(log_entry)
except Exception as e:
logging.error(f"Failed to write task submission log: {e}")
def log_task_result(self, task: AworldTask, result: ModelResponse):
try:
date_str = datetime.now().strftime("%Y%m%d")
result_dir = os.path.join(ROOT_LOG, 'task_logs', 'result', date_str)
os.makedirs(result_dir, exist_ok=True)
md_file = f"{result_dir}/{task.task_id}.md"
content_parts = []
if hasattr(result, 'content') and result.content:
if isinstance(result.content, list):
content_parts.extend(result.content)
else:
content_parts.append(str(result.content))
file_exists = os.path.exists(md_file)
with open(md_file, 'a', encoding='utf-8') as f:
if not file_exists:
f.write(f"# Task Result: {task.task_id}\n\n")
f.write(f"**Agent ID:** {task.agent_id}\n\n")
f.write(f"**Timestamp:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
f.write("## Content\n\n")
if content_parts:
for i, content in enumerate(content_parts, 1):
f.write(f"{content}\n\n")
else:
f.write("No content available.\n\n")
return md_file
except Exception as e:
logging.error(f"Failed to write task result log: {e}")
return None
task_logger = TaskLogger(log_file=f"aworld_task_submissions_{datetime.now().strftime('%Y%m%d')}.log")
@@ -0,0 +1,199 @@
import os
def load_all_mcp_config():
return {
"mcpServers": {
"e2b-server": {
"command": "npx",
"args": [
"-y",
"@e2b/mcp-server"
],
"env": {
"E2B_API_KEY": os.environ["E2B_API_KEY"],
"SESSION_REQUEST_CONNECT_TIMEOUT": "60"
}
},
"filesystem": {
"command": "npx",
"args": [
"-y",
"@modelcontextprotocol/server-filesystem",
"${FILESYSTEM_SERVER_WORKDIR}"
]
},
"terminal-controller": {
"command": "python",
"args": [
"-m",
"terminal_controller"
],
"env": {
"SESSION_REQUEST_CONNECT_TIMEOUT": "300"
}
},
"calculator": {
"command": "python",
"args": [
"-m",
"mcp_server_calculator"
],
"env": {
"SESSION_REQUEST_CONNECT_TIMEOUT": "20"
}
},
"excel": {
"command": "uvx",
"args": ["excel-mcp-server", "stdio"],
"env": {
"EXCEL_MCP_PAGING_CELLS_LIMIT": "4000",
"SESSION_REQUEST_CONNECT_TIMEOUT": "120"
}
},
"google-search": {
"command": "npx",
"args": [
"-y",
"@adenot/mcp-google-search"
],
"env": {
"GOOGLE_API_KEY": os.environ["GOOGLE_API_KEY"],
"GOOGLE_SEARCH_ENGINE_ID": os.environ["GOOGLE_CSE_ID"],
"SESSION_REQUEST_CONNECT_TIMEOUT": "60"
}
},
"ms-playwright": {
"command": "npx",
"args": [
"@playwright/mcp@latest",
"--no-sandbox",
"--headless",
"--isolated"
],
"env": {
"PLAYWRIGHT_TIMEOUT": "120000",
"SESSION_REQUEST_CONNECT_TIMEOUT": "120"
}
},
"audio_server": {
"command": "python",
"args": [
"-m",
"mcp_servers.audio_server"
],
"env": {
"AUDIO_LLM_API_KEY": os.environ["AUDIO_LLM_API_KEY"],
"AUDIO_LLM_BASE_URL": os.environ["AUDIO_LLM_BASE_URL"],
"AUDIO_LLM_MODEL_NAME": os.environ["AUDIO_LLM_MODEL_NAME"],
"SESSION_REQUEST_CONNECT_TIMEOUT": "60"
}
},
"image_server": {
"command": "python",
"args": [
"-m",
"mcp_servers.image_server"
],
"env": {
"LLM_API_KEY": os.environ.get("LLM_API_KEY"),
"LLM_MODEL_NAME": os.environ.get("LLM_MODEL_NAME"),
"LLM_BASE_URL": os.environ.get("LLM_BASE_URL"),
"SESSION_REQUEST_CONNECT_TIMEOUT": "60"
}
},
"youtube_server": {
"command": "python",
"args": [
"-m",
"mcp_servers.youtube_server"
],
"env": {
"CHROME_DRIVER_PATH": os.environ['CHROME_DRIVER_PATH'],
"SESSION_REQUEST_CONNECT_TIMEOUT": "120"
}
},
"video_server": {
"command": "python",
"args": [
"-m",
"mcp_servers.video_server"
],
"env": {
"LLM_API_KEY": os.environ.get("LLM_API_KEY"),
"LLM_MODEL_NAME": os.environ.get("LLM_MODEL_NAME"),
"LLM_BASE_URL": os.environ.get("LLM_BASE_URL"),
"SESSION_REQUEST_CONNECT_TIMEOUT": "60"
}
},
"search_server": {
"command": "python",
"args": [
"-m",
"mcp_servers.search_server"
],
"env": {
"GOOGLE_API_KEY": os.environ["GOOGLE_API_KEY"],
"GOOGLE_CSE_ID": os.environ["GOOGLE_CSE_ID"],
"SESSION_REQUEST_CONNECT_TIMEOUT": "60"
}
},
"download_server": {
"command": "python",
"args": [
"-m",
"mcp_servers.download_server"
],
"env": {
"SESSION_REQUEST_CONNECT_TIMEOUT": "120"
}
},
"document_server": {
"command": "python",
"args": [
"-m",
"mcp_servers.document_server"
],
"env": {
"SESSION_REQUEST_CONNECT_TIMEOUT": "120"
}
},
"browser_server": {
"command": "python",
"args": [
"-m",
"mcp_servers.browser_server"
],
"env": {
"LLM_API_KEY": os.environ.get("LLM_API_KEY"),
"LLM_MODEL_NAME": os.environ.get("LLM_MODEL_NAME"),
"LLM_BASE_URL": os.environ.get("LLM_BASE_URL"),
"SESSION_REQUEST_CONNECT_TIMEOUT": "120"
}
},
"reasoning_server": {
"command": "python",
"args": [
"-m",
"mcp_servers.reasoning_server"
],
"env": {
"LLM_API_KEY": os.environ.get("LLM_API_KEY"),
"LLM_MODEL_NAME": os.environ.get("LLM_MODEL_NAME"),
"LLM_BASE_URL": os.environ.get("LLM_BASE_URL"),
"SESSION_REQUEST_CONNECT_TIMEOUT": "120"
}
},
"e2b-code-server": {
"command": "python",
"args": [
"-m",
"mcp_servers.e2b_code_server"
],
"env": {
"E2B_API_KEY": os.environ["E2B_API_KEY"],
"SESSION_REQUEST_CONNECT_TIMEOUT": "120"
}
},
}
}
@@ -0,0 +1,344 @@
import json
import re
import string
from pathlib import Path
from typing import Any, Dict, List, Optional
from loguru import logger
from tabulate import tabulate
def normalize_str(input_str, remove_punct=True) -> str:
no_spaces = re.sub(r"\s", "", input_str)
if remove_punct:
translator = str.maketrans("", "", string.punctuation)
return no_spaces.lower().translate(translator)
else:
return no_spaces.lower()
def split_string(s: str, char_list: Optional[List[str]] = None) -> list[str]:
if char_list is None:
char_list = [",", ";"]
pattern = f"[{''.join(char_list)}]"
return re.split(pattern, s)
def normalize_number_str(number_str: str) -> float:
for char in ["$", "%", ","]:
number_str = number_str.replace(char, "")
try:
return float(number_str)
except ValueError:
logger.error(f"String {number_str} cannot be normalized to number str.")
return float("inf")
def question_scorer(model_answer: str, ground_truth: str) -> bool:
def is_float(element: Any) -> bool:
try:
float(element)
return True
except ValueError:
return False
try:
if is_float(ground_truth):
logger.info(f"Evaluating {model_answer} as a number.")
normalized_answer = normalize_number_str(model_answer)
return normalized_answer == float(ground_truth)
elif any(char in ground_truth for char in [",", ";"]):
logger.info(f"Evaluating {model_answer} as a comma separated list.")
gt_elems = split_string(ground_truth)
ma_elems = split_string(model_answer)
if len(gt_elems) != len(ma_elems):
logger.warning("Answer lists have different lengths, returning False.")
return False
comparisons = []
for ma_elem, gt_elem in zip(ma_elems, gt_elems):
if is_float(gt_elem):
normalized_ma_elem = normalize_number_str(ma_elem)
comparisons.append(normalized_ma_elem == float(gt_elem))
else:
ma_elem = normalize_str(ma_elem, remove_punct=False)
gt_elem = normalize_str(gt_elem, remove_punct=False)
comparisons.append(ma_elem == gt_elem)
return all(comparisons)
else:
logger.info(f"Evaluating {model_answer} as a string.")
ma_elem = normalize_str(model_answer)
gt_elem = normalize_str(ground_truth)
return ma_elem == gt_elem
except Exception as e:
logger.error(f"Error during evaluation: {e}")
return False
def load_dataset_meta(path: str, split: str = "validation"):
data_dir = Path(path) / split
dataset = []
with open(data_dir / "metadata.jsonl", "r", encoding="utf-8") as metaf:
lines = metaf.readlines()
for line in lines:
data = json.loads(line)
if data["task_id"] == "0-0-0-0-0":
continue
if data["file_name"]:
data["file_name"] = data_dir / data["file_name"]
dataset.append(data)
return dataset
def load_dataset_meta_dict(path: str, split: str = "validation"):
data_dir = Path(path) / split
dataset = {}
with open(data_dir / "metadata.jsonl", "r", encoding="utf-8") as metaf:
lines = metaf.readlines()
for line in lines:
data = json.loads(line)
if data["task_id"] == "0-0-0-0-0":
continue
if data["file_name"]:
data["file_name"] = data_dir / data["file_name"]
dataset[data["task_id"]] = data
return dataset
def add_file_path(
task: Dict[str, Any], file_path: str = "./gaia_dataset", split: str = "validation"
):
if task["file_name"]:
file_path = Path(f"{file_path}/{split}") / task["file_name"]
if file_path.suffix in [".pdf", ".docx", ".doc", ".txt"]:
task["Question"] += f" Here are the necessary document files: {file_path}"
elif file_path.suffix in [".jpg", ".jpeg", ".png"]:
task["Question"] += f" Here are the necessary image files: {file_path}"
elif file_path.suffix in [".xlsx", "xls", ".csv"]:
task["Question"] += (
f" Here are the necessary table files: {file_path}, for processing excel file,"
" you can use the excel tool or write python code to process the file"
" step-by-step and get the information."
)
elif file_path.suffix in [".py"]:
task["Question"] += f" Here are the necessary python files: {file_path}"
else:
task["Question"] += f" Here are the necessary files: {file_path}"
return task
def report_results(entries):
# Initialize counters
total_entries = len(entries)
total_correct = 0
# Initialize level statistics
level_stats = {}
# Process each entry
for entry in entries:
level = entry.get("level")
is_correct = entry.get("is_correct", False)
# Initialize level stats if not already present
if level not in level_stats:
level_stats[level] = {"total": 0, "correct": 0, "accuracy": 0}
# Update counters
level_stats[level]["total"] += 1
if is_correct:
total_correct += 1
level_stats[level]["correct"] += 1
# Calculate accuracy for each level
for level, stats in level_stats.items():
if stats["total"] > 0:
stats["accuracy"] = (stats["correct"] / stats["total"]) * 100
# Print overall statistics with colorful logging
logger.info("Overall Statistics:")
overall_accuracy = (total_correct / total_entries) * 100
# Create overall statistics table
overall_table = [
["Total Entries", total_entries],
["Total Correct", total_correct],
["Overall Accuracy", f"{overall_accuracy:.2f}%"],
]
logger.success(tabulate(overall_table, tablefmt="grid"))
logger.info("")
# Create level statistics table
logger.info("Statistics by Level:")
level_table = []
headers = ["Level", "Total Entries", "Correct Answers", "Accuracy"]
for level in sorted(level_stats.keys()):
stats = level_stats[level]
level_table.append(
[level, stats["total"], stats["correct"], f"{stats['accuracy']:.2f}%"]
)
logger.success(tabulate(level_table, headers=headers, tablefmt="grid"))
import uuid
import time
from typing import List
import inspect
from typing import get_type_hints, Tuple
def stream_message_template(model: str, message: str):
return {
"id": f"{model}-{str(uuid.uuid4())}",
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": model,
"choices": [
{
"index": 0,
"delta": {"content": message},
"logprobs": None,
"finish_reason": None,
}
],
}
def get_last_user_message(messages: List[dict]) -> str:
for message in reversed(messages):
if message["role"] == "user":
if isinstance(message["content"], list):
for item in message["content"]:
if item["type"] == "text":
return item["text"]
return message["content"]
return None
def get_last_assistant_message(messages: List[dict]) -> str:
for message in reversed(messages):
if message["role"] == "assistant":
if isinstance(message["content"], list):
for item in message["content"]:
if item["type"] == "text":
return item["text"]
return message["content"]
return None
def get_system_message(messages: List[dict]) -> dict:
for message in messages:
if message["role"] == "system":
return message
return None
def remove_system_message(messages: List[dict]) -> List[dict]:
return [message for message in messages if message["role"] != "system"]
def pop_system_message(messages: List[dict]) -> Tuple[dict, List[dict]]:
return get_system_message(messages), remove_system_message(messages)
def add_or_update_system_message(content: str, messages: List[dict]) -> List[dict]:
"""
Adds a new system message at the beginning of the messages list
or updates the existing system message at the beginning.
:param msg: The message to be added or appended.
:param messages: The list of message dictionaries.
:return: The updated list of message dictionaries.
"""
if messages and messages[0].get("role") == "system":
messages[0]["content"] += f"{content}\n{messages[0]['content']}"
else:
# Insert at the beginning
messages.insert(0, {"role": "system", "content": content})
return messages
def doc_to_dict(docstring):
lines = docstring.split("\n")
description = lines[1].strip()
param_dict = {}
for line in lines:
if ":param" in line:
line = line.replace(":param", "").strip()
param, desc = line.split(":", 1)
param_dict[param.strip()] = desc.strip()
ret_dict = {"description": description, "params": param_dict}
return ret_dict
def get_tools_specs(tools) -> List[dict]:
function_list = [
{"name": func, "function": getattr(tools, func)}
for func in dir(tools)
if callable(getattr(tools, func)) and not func.startswith("__")
]
specs = []
for function_item in function_list:
function_name = function_item["name"]
function = function_item["function"]
function_doc = doc_to_dict(function.__doc__ or function_name)
specs.append(
{
"name": function_name,
# TODO: multi-line desc?
"description": function_doc.get("description", function_name),
"parameters": {
"type": "object",
"properties": {
param_name: {
"type": param_annotation.__name__.lower(),
**(
{
"enum": (
param_annotation.__args__
if hasattr(param_annotation, "__args__")
else None
)
}
if hasattr(param_annotation, "__args__")
else {}
),
"description": function_doc.get("params", {}).get(
param_name, param_name
),
}
for param_name, param_annotation in get_type_hints(
function
).items()
if param_name != "return"
},
"required": [
name
for name, param in inspect.signature(
function
).parameters.items()
if param.default is param.empty
],
},
}
)
return specs
@@ -0,0 +1,72 @@
import uuid
from typing import Any
from typing import List, Optional
from datetime import datetime
from pydantic import BaseModel, ConfigDict
from pydantic import Field
class AworldTask(BaseModel):
task_id: str = Field(default=None, description="task id")
agent_id: str = Field(default=None, description="agent id")
agent_input: str = Field(default=None, description="agent input")
session_id: Optional[str] = Field(default=None, description="session id")
user_id: Optional[str] = Field(default=None, description="user id")
llm_provider: Optional[str] = Field(default=None, description="llm provider")
llm_model_name: Optional[str] = Field(default=None, description="llm model name")
llm_api_key: Optional[str] = Field(default=None, description="llm api key")
llm_base_url: Optional[str] = Field(default=None, description="llm base url")
llm_custom_input: Optional[str] = Field(default=None, description="custom_input")
task_system_prompt: Optional[str] = Field(default=None, description="task_system_prompt")
mcp_servers: Optional[list[str]] = Field(default=None, description="mcp_servers")
node_id: Optional[str] = Field(default=None, description="execute task node_id")
client_id: Optional[str] = Field(default=None, description="submit client ip")
status: Optional[str] = Field(default="INIT", description="submitted/running/execute_failed/execute_success")
history_messages: Optional[int] = Field(default=100, description="history_message")
max_steps: Optional[int] = Field(default=100, description="max_steps")
max_retries: Optional[int] = Field(default=5, description="max_retries use Exponential backoff with jitter")
ext_info: Optional[dict] = Field(default_factory=dict, description="custom")
created_at: Optional[datetime] = Field(default=None, description="created time")
updated_at: Optional[datetime] = Field(default=None, description="updated time")
def mark_running(self):
self.status = 'RUNNING'
def mark_failed(self):
self.status = 'FAILED'
def mark_success(self):
self.status = 'SUCCESS'
class AworldTaskResult(BaseModel):
task: AworldTask = Field(default=None, description="task")
server_host: Optional[str] = Field(default=None, description="aworld server id")
data: Any = Field(default=None, description="result data")
class AworldTaskForm(BaseModel):
batch_id: str = Field(default=str(uuid.uuid4()), description="batch_id")
task: Optional[AworldTask] = Field(default=None, description="task")
user_id: Optional[str] = Field(default=None, description="user id")
client_id: Optional[str] = Field(default=None, description="submit client ip")
class OpenAIChatMessage(BaseModel):
role: str
content: str | List
model_config = ConfigDict(extra="allow")
class OpenAIChatCompletionForm(BaseModel):
stream: bool = True
model: str
messages: List[OpenAIChatMessage]
model_config = ConfigDict(extra="allow")
class FilterForm(BaseModel):
body: dict
user: Optional[dict] = None
model_config = ConfigDict(extra="allow")
@@ -0,0 +1,369 @@
import asyncio
import logging
import os
import traceback
from datetime import datetime
from typing import AsyncGenerator
from aworld.models.llm import acall_llm_model, get_llm_model, acall_llm_model_stream
from aworld.models.model_response import ModelResponse, LLMResponseError
from pydantic import BaseModel, Field
from base import AworldTask, AworldTaskResult, AworldTaskForm
class TaskLogger:
"""Task submission logger"""
def __init__(self, log_file: str = "aworld_task_submissions.log"):
self.log_file = 'task_logs/' + log_file
self._ensure_log_file_exists()
def _ensure_log_file_exists(self):
"""ensure log file exists"""
if not os.path.exists(self.log_file):
os.makedirs(os.path.dirname(self.log_file), exist_ok=True)
with open(self.log_file, 'w', encoding='utf-8') as f:
f.write("# Aworld Task Submission Log\n")
f.write("# Format: [timestamp] task_id | agent_id | server | status | agent_answer | correct_answer | is_correct | details\n\n")
def log_task_submission(self, task: AworldTask, server: str, status: str, details: str = "", task_result: AworldTaskResult = None):
"""log task submission"""
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
log_entry = f"[{timestamp}] {task.task_id} | {task.agent_id} | {task.node_id} | {status} | { task_result.data.get('agent_answer') if task_result and task_result.data else None } | {task_result.data.get('correct_answer') if task_result and task_result.data else None} | {task_result.data.get('gaia_correct') if task_result and task_result.data else None} |{details}\n"
try:
with open(self.log_file, 'a', encoding='utf-8') as f:
f.write(log_entry)
except Exception as e:
logging.error(f"Failed to write task submission log: {e}")
def log_task_result(self, task: AworldTask, result: ModelResponse):
"""log task result to markdown file"""
try:
# create result directory
date_str = datetime.now().strftime("%Y%m%d")
result_dir = f"task_logs/result/{date_str}"
os.makedirs(result_dir, exist_ok=True)
# create markdown file
md_file = f"{result_dir}/{task.task_id}.md"
# concat content
content_parts = []
if hasattr(result, 'content') and result.content:
if isinstance(result.content, list):
content_parts.extend(result.content)
else:
content_parts.append(str(result.content))
# write to markdown file
file_exists = os.path.exists(md_file)
with open(md_file, 'a', encoding='utf-8') as f:
# only write title info when file not exists
if not file_exists:
f.write(f"# Task Result: {task.task_id}\n\n")
f.write(f"**Agent ID:** {task.agent_id}\n\n")
f.write(f"**Timestamp:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
f.write("## Content\n\n")
# write content parts
if content_parts:
for i, content in enumerate(content_parts, 1):
f.write(f"{content}\n\n")
else:
f.write("No content available.\n\n")
except Exception as e:
logging.error(f"Failed to write task result log: {e}")
task_logger = TaskLogger(log_file=f"aworld_task_submissions_{datetime.now().strftime('%Y%m%d')}.log")
class AworldTaskClient(BaseModel):
"""
AworldTaskClient
"""
know_hosts: list[str] = Field(default_factory=list, description="aworldserver list")
tasks: list[AworldTask] = Field(default_factory=list, description="submitted task list")
task_states: dict[str, AworldTaskResult] = Field(default_factory=dict, description="task_states")
async def submit_task(self, task: AworldTask, background: bool = True):
if not self.know_hosts:
raise ValueError("No aworld server hosts configured.")
# 1. select aworld server from know_hosts using round-robin
if not hasattr(self, '_current_server_index'):
self._current_server_index = 0
aworld_server = self.know_hosts[self._current_server_index]
if not aworld_server.startswith("http"):
aworld_server = "http://" + aworld_server
self._current_server_index = (self._current_server_index + 1) % len(self.know_hosts)
# 2. call _submit_task
result = await self._submit_task(aworld_server, task, background)
# 3. update task_states
self.task_states[task.task_id] = result
async def _submit_task(self, aworld_server, task: AworldTask, background: bool = True):
try:
logging.info(f"submit task#{task.task_id} to cluster#[{aworld_server}]")
if not background:
task_result = await self._submit_task_to_server(aworld_server, task)
else:
task_result = await self._async_submit_task_to_server(aworld_server, task)
return task_result
except Exception as e:
if isinstance(e, LLMResponseError):
if e.message and 'peer closed connection without sending complete message body (incomplete chunked read)' == e.message:
task_logger.log_task_submission(task, aworld_server, "server_close_connection", str(e))
logging.error(f"execute task to {task.node_id} server_close_connection: [{e}], please see replays wait a moment")
return
traceback.print_exc()
logging.error(f"execute task to {task.node_id} execute_failed: [{e}], please see logs from server ")
task_logger.log_task_submission(task, aworld_server, "execute_failed", str(e))
async def _async_submit_task_to_server(self, aworld_server, task: AworldTask):
import httpx
from base import AworldTaskForm, AworldTaskResult
# 构建 AworldTaskForm
form_data = AworldTaskForm(task=task)
async with httpx.AsyncClient() as client:
resp = await client.post(f"{aworld_server}/api/v1/tasks/submit_task", json=form_data.model_dump())
resp.raise_for_status()
data = resp.json()
task_logger.log_task_submission(task, aworld_server, "submitted")
return AworldTaskResult(**data)
async def _submit_task_to_server(self, aworld_server, task: AworldTask):
# build params
llm_model = get_llm_model(
llm_provider="openai",
model_name=task.agent_id,
base_url=f"{aworld_server}/v1",
api_key="0p3n-w3bu!"
)
messages = [
{"role": "user", "content": task.agent_input}
]
#call_llm_model
data = acall_llm_model_stream(llm_model, messages, stream=True, user={
"user_id": task.user_id,
"session_id": task.session_id,
"task_id": task.task_id,
"aworld_task": task.model_dump_json()
})
items = []
task_result = {}
if isinstance(data, AsyncGenerator):
async for item in data:
items.append(item)
if item.raw_response and item.raw_response.model_extra and item.raw_response.model_extra.get('node_id'):
if not task.node_id:
logging.info(f"submit task#{task.task_id} success. execute pod ip is [{item.raw_response.model_extra.get('node_id')}]")
task.node_id = item.raw_response.model_extra.get('node_id')
task_logger.log_task_submission(task, aworld_server, "submitted")
if item.content:
task_logger.log_task_result(task, item)
logging.info(f"task#{task.task_id} response data chunk is: {item}"[:500])
if item.raw_response and item.raw_response.model_extra and item.raw_response.model_extra.get(
'task_output_meta'):
task_result = item.raw_response.model_extra.get('task_output_meta')
elif isinstance(data, ModelResponse):
if data.raw_response and data.raw_response.model_extra and data.raw_response.model_extra.get('node_id'):
if not task.node_id:
logging.info(f"submit task#{task.task_id} success. execute pod ip is [{data.raw_response.model_extra.get('node_id')}]")
task.node_id = data.raw_response.model_extra.get('node_id')
logging.info(f"task#{task.task_id} response data is: {data}")
task_logger.log_task_result(task, data)
if data.raw_response and data.raw_response.model_extra and data.raw_response.model_extra.get('task_output_meta'):
task_result = data.raw_response.model_extra.get('task_output_meta')
result = AworldTaskResult(task=task, server_host=aworld_server, data=task_result)
task_logger.log_task_submission(task, aworld_server, "execute_finished", task_result=result)
return result
async def get_task_state(self, task_id: str):
if not isinstance(self.task_states, dict):
self.task_states = dict(self.task_states)
return self.task_states.get(task_id, None)
async def download_task_results(
self,
start_time: str = None,
end_time: str = None,
task_id: str = None,
page_size: int = 100,
save_path: str = None
) -> str:
"""
Download task results and generate a JSONL format file
Args:
start_time: Start time, format: YYYY-MM-DD HH:MM:SS
end_time: End time, format: YYYY-MM-DD HH:MM:SS
task_id: Task ID
page_size: Page size
save_path: Save path, if not specified, it will be generated automatically
Returns:
str: Save path
"""
if not self.know_hosts:
raise ValueError("No aworld server hosts configured.")
# select server
if not hasattr(self, '_current_server_index'):
self._current_server_index = 0
aworld_server = self.know_hosts[self._current_server_index]
logging.info(f"🚀 downloading task results from server: {aworld_server}")
try:
import httpx
# build query params
params = {"page_size": page_size}
if start_time:
params["start_time"] = start_time
if end_time:
params["end_time"] = end_time
if task_id:
params["task_id"] = task_id
# send download request
async with httpx.AsyncClient(timeout=300.0) as client: # 5分钟超时
response = await client.get(
f"{aworld_server}/api/v1/tasks/download_task_results",
params=params
)
response.raise_for_status()
# if not specified save path, generate automatically
if not save_path:
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
save_path = f"task_results_{timestamp}.jsonl"
# ensure directory exists
save_dir = os.path.dirname(save_path) if os.path.dirname(save_path) else "."
os.makedirs(save_dir, exist_ok=True)
with open(save_path, 'wb') as f:
for chunk in response.iter_bytes():
f.write(chunk)
# calculate file size
file_size = os.path.getsize(save_path)
logging.info(f"✅ task results downloaded successfully, file: {save_path}, size: {file_size} bytes")
return save_path
except Exception as e:
logging.error(f"❌ download task results failed: {e}")
raise ValueError(f"❌ download task results failed: {str(e)}")
async def download_task_results_to_memory(
self,
start_time: str = None,
end_time: str = None,
task_id: str = None,
page_size: int = 100
) -> list:
"""
Download task results to memory, return parsed data list
Args:
start_time: Start time, format: YYYY-MM-DD HH:MM:SS
end_time: End time, format: YYYY-MM-DD HH:MM:SS
task_id: Task ID
page_size: Page size
Returns:
list: Task results data list
"""
if not self.know_hosts:
raise ValueError("No aworld server hosts configured.")
# select server
if not hasattr(self, '_current_server_index'):
self._current_server_index = 0
aworld_server = self.know_hosts[self._current_server_index]
logging.info(f"🚀 downloading task results to memory from server: {aworld_server}")
try:
import httpx
import json
# build query params
params = {"page_size": page_size}
if start_time:
params["start_time"] = start_time
if end_time:
params["end_time"] = end_time
if task_id:
params["task_id"] = task_id
# send download request
async with httpx.AsyncClient(timeout=300.0) as client: # 5分钟超时
response = await client.get(
f"{aworld_server}/api/v1/tasks/download_task_results",
params=params
)
response.raise_for_status()
# parse jsonl content
results = []
content = response.text
if content.strip(): # check content is not empty
for line in content.strip().split('\n'):
if line.strip(): # skip empty line
try:
result_data = json.loads(line)
results.append(result_data)
except json.JSONDecodeError as e:
logging.warning(f"Failed to parse line: {line}, error: {e}")
logging.info(f"✅ task results downloaded to memory successfully, total: {len(results)} records")
return results
except Exception as e:
logging.error(f"❌ download task results to memory failed: {e}")
raise ValueError(f"❌ download task results to memory failed: {str(e)}")
def parse_task_results_file(self, file_path: str) -> list:
"""
Parse local task results jsonl file
Args:
file_path: jsonl file path
Returns:
list: Parsed task results list
"""
import json
results = []
try:
with open(file_path, 'r', encoding='utf-8') as f:
for line_num, line in enumerate(f, 1):
line = line.strip()
if line: # 跳过空行
try:
result_data = json.loads(line)
results.append(result_data)
except json.JSONDecodeError as e:
logging.warning(f"Failed to parse line {line_num} in {file_path}: {e}")
logging.info(f"✅ parsed {len(results)} task results from {file_path}")
return results
except Exception as e:
logging.error(f"❌ failed to parse task results file {file_path}: {e}")
raise ValueError(f"❌ failed to parse task results file: {str(e)}")
@@ -0,0 +1,91 @@
# Initialize AworldTaskClient with server endpoints
import asyncio
import logging
from datetime import datetime
import os
import random
import uuid
from aworld.utils.common import get_local_ip
from client.aworld_client import AworldTask, AworldTaskClient
AWORLD_TASK_CLIENT = AworldTaskClient(
know_hosts = ["localhost:9999"]
)
async def _run_gaia_task(gaia_task: AworldTask, delay: int, background: bool = False) -> None:
"""Run a single Gaia task with the given question ID.
Args:
gaia_task_id: The ID of the question to process
"""
global AWORLD_TASK_CLIENT
await asyncio.sleep(delay)
# Submit task to Aworld server
await AWORLD_TASK_CLIENT.submit_task(gaia_task, background=background)
# Get and print task result
task_result = await AWORLD_TASK_CLIENT.get_task_state(task_id=gaia_task.task_id)
if not background:
logging.info(f"execute task_result#{gaia_task.task_id} is {task_result.data if task_result else None}")
else:
logging.info(f"submit task_result#{gaia_task.task_id} background success, please use task_id get task_result await a moment")
async def _batch_run_gaia_task(gaia_tasks: list[AworldTask]) -> None:
"""Run multiple Gaia tasks in parallel.
"""
tasks = [
_run_gaia_task(gaia_task, index * 3, background=True)
for index, gaia_task in enumerate(gaia_tasks)
]
await asyncio.gather(*tasks)
CUSTOM_SYSTEM_PROMPT = f""" **PLEASE CUSTOM IT **"""
if __name__ == '__main__':
gaia_task_ids = ['c61d22de-5f6c-4958-a7f6-5e9707bd3466']
gaia_tasks = []
custom_mcp_servers = [
# "e2b-server",
"e2b-code-server",
"terminal-controller",
"excel",
# "filesystem",
"calculator",
"ms-playwright",
"audio_server",
"image_server",
"google-search",
# "video_server",
# "search_server",
# "download_server",
# "document_server",
# "youtube_server",
# "reasoning_server",
]
for gaia_task_id in gaia_task_ids:
task_id = datetime.now().strftime("%Y%m%d%H%M%S") + "_" + gaia_task_id + "_" + str(uuid.uuid4())
gaia_tasks.append(
AworldTask(
task_id=task_id,
agent_id="gaia_agent",
agent_input=gaia_task_id,
session_id="session_id",
user_id=os.getenv("USER", "SYSTEM"),
client_id=get_local_ip(),
mcp_servers=custom_mcp_servers,
max_retries=5,
llm_custom_input="你好"
# llm_model_name="gpt-4o",
# task_system_prompt=CUSTOM_SYSTEM_PROMPT
)
)
asyncio.run(_batch_run_gaia_task(gaia_tasks))
@@ -0,0 +1,55 @@
# Initialize AworldTaskClient with server endpoints
import asyncio
import random
import uuid
from base import AworldTask
from client.aworld_client import AworldTaskClient
AWORLD_TASK_CLIENT = AworldTaskClient(
know_hosts=["localhost:9999"]
)
async def _run_web_task(web_question_id: str) -> None:
"""Run a single Web task with the given question ID.
Args:
web_question_id: The ID of the question to process
"""
global AWORLD_TASK_CLIENT
task_id = str(uuid.uuid4())
# Submit task to Aworld server
await AWORLD_TASK_CLIENT.submit_task(
AworldTask(
task_id=task_id,
agent_id="playwright_agent",
agent_input=web_question_id,
session_id="session_id",
user_id="SYSTEM"
)
)
# Get and print task result
task_result = await AWORLD_TASK_CLIENT.get_task_state(task_id=task_id)
print(task_result)
async def _batch_run_web_task(start_i: int, end_i: int) -> None:
"""Run multiple Web tasks in parallel.
Args:
start_i: Starting question ID
end_i: Ending question ID
"""
tasks = [
_run_web_task(str(i))
for i in range(start_i, end_i + 1)
]
await asyncio.gather(*tasks)
if __name__ == '__main__':
# Run batch processing for questions 1-5
asyncio.run(_batch_run_web_task(25, 25))
@@ -0,0 +1,28 @@
import asyncio
from client.aworld_client import AworldTaskClient
async def download_with_timerange(know_hosts: list[str], start_time, end_time, save_path):
# create client
client = AworldTaskClient(know_hosts = know_hosts)
# 1. download task results to file
file_path = await client.download_task_results(
start_time=start_time,
end_time=end_time,
save_path=save_path
)
# 2. parse local jsonl file
local_results = client.parse_task_results_file(save_path)
# 3. analyze results data
for result in local_results:
print(f"Submit User ID: {result['user_id']}, Task ID: {result['task_id']},Status: {result['status']}, Replays: {result['result_data']['replays_file'] if result['result_data'] else ''}")
if __name__ == '__main__':
asyncio.run(download_with_timerange(know_hosts= ["http://localhost:9999"],
start_time="2025-06-12 00:00:00",
end_time="2025-06-12 23:59:59",
save_path="results/january_tasks.jsonl"))
@@ -0,0 +1,30 @@
import os
import logging
from pathlib import Path
from aworld.utils.common import get_local_ip
####################################
# Load .env file
####################################
try:
from dotenv import load_dotenv, find_dotenv
load_dotenv(find_dotenv("./.env"))
except ImportError:
print("dotenv not installed, skipping...")
# Define log levels dictionary
LOG_LEVELS = {
'DEBUG': logging.DEBUG,
'INFO': logging.INFO,
'WARNING': logging.WARNING,
'ERROR': logging.ERROR,
'CRITICAL': logging.CRITICAL
}
ROOT_DIR = Path(__file__).parent # the path containing this file
AGENTS_DIR = os.getenv("AGENTS_DIR", "./aworldspace/agents")
ROOT_LOG = os.path.join(os.getenv("LOG_DIR_PATH", "logs") , get_local_ip())
WORKSPACE_TYPE = os.environ.get("WORKSPACE_TYPE", "local")
WORKSPACE_PATH = os.environ.get("WORKSPACE_PATH", "./data/workspaces")
@@ -0,0 +1,7 @@
import uvicorn
from dotenv import load_dotenv
if __name__ == "__main__":
load_dotenv()
import main
uvicorn.run(main.app, host="0.0.0.0", port=9999)
@@ -0,0 +1,22 @@
services:
aworldserver-1:
image: aworldserver:main
volumes:
- ./.env:/app/.env
ports:
- "9299:9099"
restart: always
aworldserver-2:
image: aworldserver:main
volumes:
- ./.env:/app/.env
ports:
- "9399:9099"
restart: always
aworldserver-3:
image: aworldserver:main
volumes:
- ./.env:/app/.env
ports:
- "9499:9099"
restart: always
Binary file not shown.

After

Width:  |  Height:  |  Size: 550 KiB

@@ -0,0 +1,143 @@
import logging
import os
import time
from contextlib import asynccontextmanager
from logging.handlers import TimedRotatingFileHandler
import uvicorn
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from aworld.cmd.web.routers import workspaces
from aworldspace.routes import tasks
from aworldspace.utils.job import generate_openai_chat_completion
from aworldspace.utils.loader import load_modules_from_directory, PIPELINE_MODULES, PIPELINES
from base import OpenAIChatCompletionForm
from config import AGENTS_DIR, LOG_LEVELS, ROOT_LOG
if not os.path.exists(AGENTS_DIR):
os.makedirs(AGENTS_DIR)
# Add GLOBAL_LOG_LEVEL for Pipeplines
log_level = os.getenv("GLOBAL_LOG_LEVEL", "INFO").upper()
logging.basicConfig(level=LOG_LEVELS[log_level])
def setup_logging():
logger = logging.getLogger()
logger.setLevel(logging.INFO)
log_dir = ROOT_LOG
if not os.path.exists(log_dir):
os.makedirs(log_dir)
log_path = os.path.join(log_dir, "aworldserver.log")
file_handler = TimedRotatingFileHandler(log_path, when='H', interval=1, backupCount=24)
file_handler.setLevel(logging.INFO)
formatter = logging.Formatter(
"%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
file_handler.setFormatter(formatter)
error_log_path = os.path.join(log_dir, "aworldserver_error.log")
error_file_handler = TimedRotatingFileHandler(error_log_path, when='D', interval=1, backupCount=24)
error_file_handler.setLevel(logging.WARNING)
error_file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
logger.addHandler(error_file_handler)
setup_logging()
async def on_startup():
await load_modules_from_directory(AGENTS_DIR)
await tasks.task_manager.start_task_executor()
for module in PIPELINE_MODULES.values():
if hasattr(module, "on_startup"):
await module.on_startup()
async def on_shutdown():
for module in PIPELINE_MODULES.values():
if hasattr(module, "on_shutdown"):
await module.on_shutdown()
async def reload():
await on_shutdown()
# Clear existing pipelines
PIPELINES.clear()
PIPELINE_MODULES.clear()
# Load pipelines afresh
await on_startup()
@asynccontextmanager
async def lifespan(app: FastAPI):
await on_startup()
yield
await on_shutdown()
app = FastAPI(docs_url="/docs", redoc_url=None, lifespan=lifespan)
origins = ["*"]
app.add_middleware(
CORSMiddleware,
allow_origins=origins,
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
app.include_router(tasks.router, prefix="/api/v1/tasks", tags=["tasks"])
app.include_router(workspaces.router, prefix="/api/v1/workspaces", tags=["workspace"])
@app.middleware("http")
async def check_url(request: Request, call_next):
start_time = int(time.time())
response = await call_next(request)
process_time = int(time.time()) - start_time
response.headers["X-Process-Time"] = str(process_time)
return response
@app.get("/v1")
@app.get("/")
async def get_status():
return {"status": True}
@app.post("/v1/chat/completions")
@app.post("/chat/completions")
async def chat_completion(form_data: OpenAIChatCompletionForm, request: Request
):
# Extract headers into a dict
headers = request.headers
if headers.get("x-aworld-session-id"):
metadata = {
"user_id": headers.get("x-aworld-user-id"),
"chat_id": headers.get("x-aworld-session-id"),
"message_id": headers.get("x-aworld-message-id")
}
# Add metadata to form_data
form_data.metadata = metadata
return await generate_openai_chat_completion(form_data)
@app.get("/health")
async def healthcheck():
return {
"status": True
}
if __name__ == "__main__":
uvicorn.run(
"aworlddistributed.main:app",
host="0.0.0.0",
port=8088,
reload=True,
)
@@ -0,0 +1,149 @@
import base64
import json
import os
import traceback
from typing import List
from mcp.server.fastmcp import FastMCP
from openai import OpenAI
from pydantic import Field
from aworld.logs.util import logger
from mcp_servers.utils import get_file_from_source
# Initialize MCP server
mcp = FastMCP("audio-server")
client = OpenAI(
api_key=os.getenv("AUDIO_LLM_API_KEY"), base_url=os.getenv("AUDIO_LLM_BASE_URL")
)
AUDIO_TRANSCRIBE = (
"Input is a base64 encoded audio. Transcribe the audio content. "
"Return a json string with the following format: "
'{"audio_text": "transcribed text from audio"}'
)
def encode_audio(audio_source: str, with_header: bool = True) -> str:
"""
Encode audio to base64 format with robust file handling
Args:
audio_source: URL or local file path of the audio
with_header: Whether to include MIME type header
Returns:
str: Base64 encoded audio string, with MIME type prefix if with_header is True
Raises:
ValueError: When audio source is invalid or audio format is not supported
IOError: When audio file cannot be read
"""
if not audio_source:
raise ValueError("Audio source cannot be empty")
try:
# Get file with validation (only audio files allowed)
file_path, mime_type, content = get_file_from_source(
audio_source,
allowed_mime_prefixes=["audio/"],
max_size_mb=200.0, # 200MB limit for audio files
type="audio", # Specify type as audio to handle audio files
)
# Encode to base64
audio_base64 = base64.b64encode(content).decode()
# Format with header if requested
final_audio = (
f"data:{mime_type};base64,{audio_base64}" if with_header else audio_base64
)
# Clean up temporary file if it was created for a URL
if file_path != os.path.abspath(audio_source) and os.path.exists(file_path):
os.unlink(file_path)
return final_audio
except Exception:
logger.error(
f"Error encoding audio from {audio_source}: {traceback.format_exc()}"
)
raise
@mcp.tool(description="Transcribe the given audio in a list of filepaths or urls.")
async def mcp_transcribe_audio(
audio_urls: List[str] = Field(
description="The input audio in given a list of filepaths or urls."
),
) -> str:
"""
Transcribe the given audio in a list of filepaths or urls.
Args:
audio_urls: List of audio file paths or URLs
Returns:
str: JSON string containing transcriptions
"""
transcriptions = []
for audio_url in audio_urls:
try:
# Get file with validation (only audio files allowed)
file_path, _, _ = get_file_from_source(
audio_url,
allowed_mime_prefixes=["audio/"],
max_size_mb=200.0, # 200MB limit for audio files
type="audio", # Specify type as audio to handle audio files
)
# Use the file for transcription
with open(file_path, "rb") as audio_file:
transcription = client.audio.transcriptions.create(
file=audio_file,
model=os.getenv("AUDIO_LLM_MODEL_NAME"),
response_format="text",
)
transcriptions.append(transcription)
# Clean up temporary file if it was created for a URL
if file_path != os.path.abspath(audio_url) and os.path.exists(file_path):
os.unlink(file_path)
except Exception as e:
logger.error(f"Error transcribing {audio_url}: {traceback.format_exc()}")
transcriptions.append(f"Error: {str(e)}")
logger.info(f"---get_text_by_transcribe-transcription:{transcriptions}")
return json.dumps(transcriptions, ensure_ascii=False)
def main():
from dotenv import load_dotenv
load_dotenv()
print("Starting Audio MCP Server...", file=sys.stderr)
mcp.run(transport="stdio")
# Make the module callable
def __call__():
"""
Make the module callable for uvx.
This function is called when the module is executed directly.
"""
main()
# Add this for compatibility with uvx
import sys
sys.modules[__name__].__call__ = __call__
# Run the server when the script is executed directly
if __name__ == "__main__":
main()
@@ -0,0 +1,236 @@
import asyncio
import json
import logging
import os
import sys
from typing import List, Dict, Any, Optional, Union
import aiohttp
from mcp.server import FastMCP
from mcp.types import TextContent
from pydantic import Field
mcp = FastMCP("aworldsearch-server")
async def search_single(query: str, num: int = 5) -> Optional[Dict[str, Any]]:
"""Execute a single search query, returns None on error"""
try:
url = os.getenv('AWORLD_SEARCH_URL')
searchMode = os.getenv('AWORLD_SEARCH_SEARCHMODE')
source = os.getenv('AWORLD_SEARCH_SOURCE')
domain = os.getenv('AWORLD_SEARCH_DOMAIN')
uid = os.getenv('AWORLD_SEARCH_UID')
if not url or not searchMode or not source or not domain:
logging.warning(f"Query failed: url, searchMode, source, domain parameters incomplete")
return None
headers = {
'Content-Type': 'application/json'
}
data = {
"domain": domain,
"extParams": {},
"page": 0,
"pageSize": num,
"query": query,
"searchMode": searchMode,
"source": source,
"userId": uid
}
async with aiohttp.ClientSession() as session:
try:
async with session.post(url, headers=headers, json=data) as response:
if response.status != 200:
logging.warning(f"Query failed: {query}, status code: {response.status}")
return None
result = await response.json()
return result
except aiohttp.ClientError:
logging.warning(f"Request error: {query}")
return None
except Exception:
logging.warning(f"Query exception: {query}")
return None
def filter_valid_docs(result: Optional[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""Filter valid document results, returns empty list if input is None"""
if result is None:
return []
try:
valid_docs = []
# Check success field
if not result.get("success"):
return valid_docs
# Check searchDocs field
search_docs = result.get("searchDocs", [])
if not search_docs:
return valid_docs
# Extract required fields
required_fields = ["title", "docAbstract", "url", "doc"]
for doc in search_docs:
# Check if all required fields exist and are not empty
is_valid = True
for field in required_fields:
if field not in doc or not doc[field]:
is_valid = False
break
if is_valid:
# Keep only required fields
filtered_doc = {field: doc[field] for field in required_fields}
valid_docs.append(filtered_doc)
return valid_docs
except Exception:
return []
@mcp.tool(description="Search based on the user's input query list")
async def search(
query_list: List[str] = Field(
description="List format, queries to search for"
),
num: int = Field(
5,
description="Maximum number of results per query, default is 5, please keep the total results within 15"
)
) -> Union[str, TextContent]:
"""Execute search main function, supports single query or query list"""
try:
# Get configuration from environment variables
env_total_num = os.getenv('AWORLD_SEARCH_TOTAL_NUM')
if env_total_num and env_total_num.isdigit():
# Force override input num parameter with environment variable
num = int(env_total_num)
# If no queries provided, return empty list
if not query_list:
# Initialize TextContent with additional parameters
return TextContent(
type="text",
text="", # Empty string instead of None
**{"metadata": {}} # Pass as additional fields
)
# When query count is >= 3 or slice_num is set, use corresponding value
slice_num = os.getenv('AWORLD_SEARCH_SLICE_NUM')
if slice_num and slice_num.isdigit():
actual_num = int(slice_num)
else:
actual_num = 2 if len(query_list) >= 3 else num
# Execute all queries in parallel
tasks = [search_single(q, actual_num) for q in query_list]
raw_results = await asyncio.gather(*tasks)
# Filter and merge results
all_valid_docs = []
for result in raw_results:
valid_docs = filter_valid_docs(result)
all_valid_docs.extend(valid_docs)
# If no valid results found, return empty list
if not all_valid_docs:
# Initialize TextContent with additional parameters
return TextContent(
type="text",
text="", # Empty string instead of None
**{"metadata": {}} # Pass as additional fields
)
# Format results as JSON
result_json = json.dumps(all_valid_docs, ensure_ascii=False)
# Create dictionary structure directly
combined_query = ",".join(query_list)
search_items = []
# Use a dictionary to deduplicate by URL
url_dict = {}
for doc in all_valid_docs:
url = doc.get("url", "")
if url not in url_dict:
url_dict[url] = {
"title": doc.get("title", ""),
"url": url,
"snippet": doc.get("doc", "")[:100] + "..." if len(doc.get("doc", "")) > 100 else doc.get("doc",
""),
"content": doc.get("doc", "") # Map doc field to content
}
# Convert dictionary values to list
search_items = list(url_dict.values())
search_output_dict = {
"artifact_type": "WEB_PAGES",
"artifact_data": {
"query": combined_query,
"results": search_items
}
}
# Log results
logging.info(f"Completed {len(query_list)} queries, found {len(all_valid_docs)} valid documents")
# Initialize TextContent with additional parameters
return TextContent(
type="text",
text=result_json,
**{"metadata": search_output_dict} # Pass processed data as metadata
)
except Exception as e:
# Handle errors
logging.error(f"Search error: {e}")
# Initialize TextContent with additional parameters
return TextContent(
type="text",
text="", # Empty string instead of None
**{"metadata": {}} # Pass as additional fields
)
def main():
from dotenv import load_dotenv
load_dotenv(override=True)
print("Starting Audio MCP aworldsearch-server...", file=sys.stderr)
mcp.run(transport="stdio")
# Make the module callable
def __call__():
"""
Make the module callable for uvx.
This function is called when the module is executed directly.
"""
main()
sys.modules[__name__].__call__ = __call__
if __name__ == "__main__":
main()
# if __name__ == "__main__":
# # Configure logging
# logging.basicConfig(
# level=logging.INFO,
# format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
# )
#
#
# # Test single query
# # asyncio.run(search("Alibaba financial report"))
#
# # Test multiple queries
# test_queries = ["Alibaba financial report", "Tencent financial report", "Baidu financial report"]
# asyncio.run(search(query_list=test_queries))
@@ -0,0 +1,149 @@
"""
Browser MCP Server
This module provides MCP server functionality for browser automation and interaction.
It handles tasks such as web scraping, form submission, and automated browsing.
Main functions:
- browse_url: Opens a URL and performs specified actions
- submit_form: Fills and submits forms on web pages
"""
import json
import os
import sys
import traceback
from browser_use import Agent
from browser_use.agent.views import AgentHistoryList
from browser_use.browser.browser import Browser, BrowserConfig
from browser_use.browser.context import BrowserContext, BrowserContextConfig
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from mcp.server.fastmcp import FastMCP
from pydantic import Field
from aworld.logs.util import logger
mcp = FastMCP("browser-server")
browser_system_prompt = """
===== NAVIGATION STRATEGY =====
1. START: Navigate to the most authoritative source for this information
- For general queries: Use Google with specific search terms
- For known sources: Go directly to the relevant website
2. EVALUATE: Assess each page methodically
- Scan headings and highlighted text first
- Look for data tables, charts, or official statistics
- Check publication dates for timeliness
3. EXTRACT: Capture exactly what's needed
- Take screenshots of visual evidence (charts, tables, etc.)
- Copy precise text that answers the query
- Note source URLs for citation
4. DOWNLOAD: Save the most relevant file to local path for further processing
- Save the text if possible for futher text reading and analysis
- Save the image if possible for futher image reasoning analysis
- Save the pdf if possible for futher pdf reading and analysis
5. ROBOT DETECTION:
- If the page is a robot detection page, abort immediately
- Navigate to the most authoritative source for similar information instead
===== EFFICIENCY GUIDELINES =====
- Use specific search queries with key terms from the task
- Avoid getting distracted by tangential information
- If blocked by paywalls, try archive.org or similar alternatives
- Document each significant finding clearly and concisely
Your goal is to extract precisely the information needed with minimal browsing steps.
"""
@mcp.tool(description="Perform browser actions using the browser-use package.")
async def browser_use(
task: str = Field(description="The task to perform using the browser."),
) -> str:
"""
Perform browser actions using the browser-use package.
Args:
task (str): The task to perform using the browser.
Returns:
str: The result of the browser actions.
"""
browser = Browser(
config=BrowserConfig(
headless=False,
new_context_config=BrowserContextConfig(
disable_security=True,
user_agent="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
minimum_wait_page_load_time=10,
maximum_wait_page_load_time=30,
),
)
)
browser_context = BrowserContext(
config=BrowserContextConfig(
trace_path=os.getenv("LOG_FILE_PATH" + "/browser_trace.log")
),
browser=browser,
)
agent = Agent(
task=task,
llm=ChatOpenAI(
model=os.getenv("LLM_MODEL_NAME"),
api_key=os.getenv("LLM_API_KEY"),
base_url=os.getenv("LLM_BASE_URL"),
model_name=os.getenv("LLM_MODEL_NAME"),
openai_api_base=os.getenv("LLM_BASE_URL"),
openai_api_key=os.getenv("LLM_API_KEY"),
temperature=1.0,
),
browser_context=browser_context,
extend_system_message=browser_system_prompt,
)
try:
browser_execution: AgentHistoryList = await agent.run(max_steps=50)
if (
browser_execution is not None
and browser_execution.is_done()
and browser_execution.is_successful()
):
exec_trace = browser_execution.extracted_content()
logger.info(
">>> 🌏 Browse Execution Succeed!\n"
f">>> 💡 Result: {json.dumps(exec_trace, ensure_ascii=False, indent=4)}\n"
">>> 🌏 Browse Execution Succeed!\n"
)
return browser_execution.final_result()
else:
return f"Browser execution failed for task: {task}"
except Exception as e:
logger.error(f"Browser execution failed: {traceback.format_exc()}")
return f"Browser execution failed for task: {task} due to {str(e)}"
finally:
await browser.close()
logger.info("Browser Closed!")
def main():
load_dotenv()
print("Starting Browser MCP Server...", file=sys.stderr)
mcp.run(transport="stdio")
# Make the module callable
def __call__():
"""
Make the module callable for uvx.
This function is called when the module is executed directly.
"""
main()
sys.modules[__name__].__call__ = __call__
# Run the server when the script is executed directly
if __name__ == "__main__":
main()
@@ -0,0 +1,998 @@
"""
Document MCP Server
This module provides MCP server functionality for document processing and analysis.
It handles various document formats including:
- Text files
- PDF documents
- Word documents (DOCX)
- Excel spreadsheets
- PowerPoint presentations
- JSON and XML files
- Source code files
Each document type has specialized processing functions that extract content,
structure, and metadata. The server focuses on local file processing with
appropriate validation and error handling.
Main functions:
- mcpreadtext: Reads plain text files
- mcpreadpdf: Reads PDF files with optional image extraction
- mcpreaddocx: Reads Word documents
- mcpreadexcel: Reads Excel spreadsheets
- mcpreadpptx: Reads PowerPoint presentations
- mcpreadjson: Reads and parses JSON/JSONL files
- mcpreadxml: Reads and parses XML files
- mcpreadsourcecode: Reads and analyzes source code files
"""
import io
import json
import os
import sys
import tempfile
import traceback
from datetime import date, datetime
from typing import Any, Dict, List, Optional
import fitz
import html2text
import pandas as pd
import xmltodict
from bs4 import BeautifulSoup
from docx2markdown._docx_to_markdown import docx_to_markdown
from dotenv import load_dotenv
from mcp.server.fastmcp import FastMCP
from PIL import Image, ImageDraw, ImageFont
from pptx import Presentation
from pydantic import BaseModel, Field
from PyPDF2 import PdfReader
from tabulate import tabulate
from xls2xlsx import XLS2XLSX
from aworld.logs.util import logger
from aworld.utils import import_package
from mcp_servers.image_server import encode_images
mcp = FastMCP("document-server")
# Define model classes for different document types
class TextDocument(BaseModel):
"""Model representing a text document"""
content: str
file_path: str
file_name: str
file_size: int
last_modified: str
class HtmlDocument(BaseModel):
"""Model representing an HTML document"""
content: str # Extracted text content
html_content: str # Original HTML content
file_path: str
file_name: str
file_size: int
last_modified: str
title: Optional[str] = None
links: Optional[List[Dict[str, str]]] = None
images: Optional[List[Dict[str, str]]] = None
tables: Optional[List[str]] = None
markdown: Optional[str] = None # HTML converted to Markdown format
class JsonDocument(BaseModel):
"""Model representing a JSON document"""
format: str # "json" or "jsonl"
type: Optional[str] = None # "array" or "object" for standard JSON
count: Optional[int] = None
keys: Optional[List[str]] = None
data: Any
file_path: str
file_name: str
class XmlDocument(BaseModel):
"""Model representing an XML document"""
content: Dict
file_path: str
file_name: str
class PdfImage(BaseModel):
"""Model representing an image extracted from a PDF"""
page: int
format: str
width: int
height: int
path: str
class PdfDocument(BaseModel):
"""Model representing a PDF document"""
content: str
file_path: str
file_name: str
page_count: int
images: Optional[List[PdfImage]] = None
image_count: Optional[int] = None
image_dir: Optional[str] = None
error: Optional[str] = None
class PdfResult(BaseModel):
"""Model representing results from processing multiple PDF documents"""
total_files: int
success_count: int
failed_count: int
results: List[PdfDocument]
class DocxDocument(BaseModel):
"""Model representing a Word document"""
content: str
file_path: str
file_name: str
class ExcelSheet(BaseModel):
"""Model representing a sheet in an Excel file"""
name: str
data: List[Dict[str, Any]]
markdown_table: str
row_count: int
column_count: int
class ExcelDocument(BaseModel):
"""Model representing an Excel document"""
file_name: str
file_path: str
processed_path: Optional[str] = None
file_type: str
sheet_count: int
sheet_names: List[str]
sheets: List[ExcelSheet]
success: bool = True
error: Optional[str] = None
class ExcelResult(BaseModel):
"""Model representing results from processing multiple Excel documents"""
total_files: int
success_count: int
failed_count: int
results: List[ExcelDocument]
class PowerPointSlide(BaseModel):
"""Model representing a slide in a PowerPoint presentation"""
slide_number: int
image: str # Base64 encoded image
class PowerPointDocument(BaseModel):
"""Model representing a PowerPoint document"""
file_path: str
file_name: str
slide_count: int
slides: List[PowerPointSlide]
class SourceCodeDocument(BaseModel):
"""Model representing a source code document"""
content: str
file_type: str
file_path: str
file_name: str
line_count: int
size_bytes: int
last_modified: str
classes: Optional[List[str]] = None
functions: Optional[List[str]] = None
imports: Optional[List[str]] = None
package: Optional[List[str]] = None
methods: Optional[List[str]] = None
includes: Optional[List[str]] = None
class DocumentError(BaseModel):
"""Model representing an error in document processing"""
error: str
file_path: Optional[str] = None
file_name: Optional[str] = None
class ComplexEncoder(json.JSONEncoder):
def default(self, o):
if isinstance(o, datetime):
return o.strftime("%Y-%m-%d %H:%M:%S")
elif isinstance(o, date):
return o.strftime("%Y-%m-%d")
else:
return json.JSONEncoder.default(self, o)
def handle_error(e: Exception, error_type: str, file_path: Optional[str] = None) -> str:
"""Unified error handling and return standard format error message"""
error_msg = f"{error_type} error: {str(e)}"
logger.error(traceback.format_exc())
error = DocumentError(
error=error_msg,
file_path=file_path,
file_name=os.path.basename(file_path) if file_path else None,
)
return error.model_dump_json()
def check_file_readable(document_path: str) -> str:
"""Check if file exists and is readable, return error message or None"""
if not os.path.exists(document_path):
return f"File does not exist: {document_path}"
if not os.access(document_path, os.R_OK):
return f"File is not readable: {document_path}"
return None
@mcp.tool(
description="Read and return content from local text file. Cannot process https://URLs files."
)
def mcpreadtext(
document_path: str = Field(description="The input local text file path."),
) -> str:
"""Read and return content from local text file. Cannot process https://URLs files."""
error = check_file_readable(document_path)
if error:
return DocumentError(error=error, file_path=document_path).model_dump_json()
try:
with open(document_path, "r", encoding="utf-8") as f:
content = f.read()
result = TextDocument(
content=content,
file_path=document_path,
file_name=os.path.basename(document_path),
file_size=os.path.getsize(document_path),
last_modified=datetime.fromtimestamp(
os.path.getmtime(document_path)
).strftime("%Y-%m-%d %H:%M:%S"),
)
return result.model_dump_json()
except Exception as e:
return handle_error(e, "Text file reading", document_path)
@mcp.tool(
description="Read and parse JSON or JSONL file, return the parsed content. Cannot process https://URLs files."
)
def mcpreadjson(
document_path: str = Field(description="Local path to JSON or JSONL file"),
is_jsonl: bool = Field(
default=False,
description="Whether the file is in JSONL format (one JSON object per line)",
),
) -> str:
"""Read and parse JSON or JSONL file, return the parsed content. Cannot process https://URLs files."""
error = check_file_readable(document_path)
if error:
return DocumentError(error=error, file_path=document_path).model_dump_json()
try:
# Choose processing method based on file type
if is_jsonl:
# Process JSONL file (one JSON object per line)
results = []
with open(document_path, "r", encoding="utf-8") as f:
for line_num, line in enumerate(f, 1):
line = line.strip()
if not line:
continue
try:
json_obj = json.loads(line)
results.append(json_obj)
except json.JSONDecodeError as e:
logger.warning(
f"JSON parsing error at line {line_num}: {str(e)}"
)
# Create result model
result = JsonDocument(
format="jsonl",
count=len(results),
data=results,
file_path=document_path,
file_name=os.path.basename(document_path),
)
else:
# Process standard JSON file
with open(document_path, "r", encoding="utf-8") as f:
data = json.load(f)
# Create result model based on data type
if isinstance(data, list):
result = JsonDocument(
format="json",
type="array",
count=len(data),
data=data,
file_path=document_path,
file_name=os.path.basename(document_path),
)
else:
result = JsonDocument(
format="json",
type="object",
keys=list(data.keys()) if isinstance(data, dict) else [],
data=data,
file_path=document_path,
file_name=os.path.basename(document_path),
)
return result.model_dump_json()
except json.JSONDecodeError as e:
return handle_error(e, "JSON parsing", document_path)
except Exception as e:
return handle_error(e, "JSON file reading", document_path)
@mcp.tool(
description="Read and return content from XML file. return the parsed content. Cannot process https://URLs files."
)
def mcpreadxml(
document_path: str = Field(description="The local input XML file path."),
) -> str:
"""Read and return content from XML file. Cannot process https://URLs files."""
error = check_file_readable(document_path)
if error:
return DocumentError(error=error, file_path=document_path).model_dump_json()
try:
with open(document_path, "r", encoding="utf-8") as f:
data = f.read()
result = XmlDocument(
content=xmltodict.parse(data),
file_path=document_path,
file_name=os.path.basename(document_path),
)
return result.model_dump_json()
except Exception as e:
return handle_error(e, "XML file reading", document_path)
@mcp.tool(
description="Read and return content from PDF file with optional image extraction. return the parsed content. Cannot process https://URLs files."
)
def mcpreadpdf(
document_paths: List[str] = Field(description="The local input PDF file paths."),
extract_images: bool = Field(
default=False, description="Whether to extract images from PDF (default: False)"
),
) -> str:
"""Read and return content from PDF file with optional image extraction. Cannot process https://URLs files."""
try:
results = []
success_count = 0
failed_count = 0
for document_path in document_paths:
error = check_file_readable(document_path)
if error:
results.append(
PdfDocument(
content="",
file_path=document_path,
file_name=os.path.basename(document_path),
page_count=0,
error=error,
)
)
failed_count += 1
continue
try:
with open(document_path, "rb") as f:
reader = PdfReader(f)
content = " ".join(page.extract_text() for page in reader.pages)
page_count = len(reader.pages)
pdf_result = PdfDocument(
content=content,
file_path=document_path,
file_name=os.path.basename(document_path),
page_count=page_count,
)
# Extract images if requested
if extract_images:
images_data = []
# Use /tmp directory for storing images
output_dir = "/tmp/pdf_images"
# Create output directory if it doesn't exist
os.makedirs(output_dir, exist_ok=True)
# Generate a unique subfolder based on filename to avoid conflicts
pdf_name = os.path.splitext(os.path.basename(document_path))[0]
timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
image_dir = os.path.join(output_dir, f"{pdf_name}_{timestamp}")
os.makedirs(image_dir, exist_ok=True)
try:
# Open PDF with PyMuPDF
pdf_document = fitz.open(document_path)
# Iterate through each page
for page_index in range(len(pdf_document)):
page = pdf_document[page_index]
# Get image list
image_list = page.get_images(full=True)
# Process each image
for img_index, img in enumerate(image_list):
# Extract image information
xref = img[0]
base_image = pdf_document.extract_image(xref)
image_bytes = base_image["image"]
image_ext = base_image["ext"]
# Save image to file in /tmp directory
img_filename = f"pdf_image_p{page_index+1}_{img_index+1}.{image_ext}"
img_path = os.path.join(image_dir, img_filename)
with open(img_path, "wb") as img_file:
img_file.write(image_bytes)
logger.success(f"Image saved: {img_path}")
# Get image dimensions
with Image.open(img_path) as img:
width, height = img.size
# Add to results with file path instead of base64
images_data.append(
PdfImage(
page=page_index + 1,
format=image_ext,
width=width,
height=height,
path=img_path,
)
)
pdf_result.images = images_data
pdf_result.image_count = len(images_data)
pdf_result.image_dir = image_dir
except Exception as img_error:
logger.error(f"Error extracting images: {str(img_error)}")
# Don't clean up on error so we can keep any successfully extracted images
pdf_result.error = str(img_error)
results.append(pdf_result)
success_count += 1
except Exception as e:
results.append(
PdfDocument(
content="",
file_path=document_path,
file_name=os.path.basename(document_path),
page_count=0,
error=str(e),
)
)
failed_count += 1
# Create final result
pdf_result = PdfResult(
total_files=len(document_paths),
success_count=success_count,
failed_count=failed_count,
results=results,
)
return pdf_result.model_dump_json()
except Exception as e:
return handle_error(e, "PDF file reading")
@mcp.tool(
description="Read and return content from Word file. return the parsed content. Cannot process https://URLs files."
)
def mcpreaddocx(
document_path: str = Field(description="The local input Word file path."),
) -> str:
"""Read and return content from Word file. Cannot process https://URLs files."""
error = check_file_readable(document_path)
if error:
return DocumentError(error=error, file_path=document_path).model_dump_json()
try:
file_name = os.path.basename(document_path)
md_file_path = f"{file_name}.md"
docx_to_markdown(document_path, md_file_path)
with open(md_file_path, "r", encoding="utf-8") as f:
content = f.read()
os.remove(md_file_path)
result = DocxDocument(
content=content, file_path=document_path, file_name=file_name
)
return result.model_dump_json()
except Exception as e:
return handle_error(e, "Word file reading", document_path)
@mcp.tool(
description="Read multiple Excel/CSV files and convert sheets to Markdown tables. return the parsed content. Cannot process https://URLs files."
)
def mcpreadexcel(
document_paths: List[str] = Field(
description="List of local input Excel/CSV file paths."
),
max_rows: int = Field(
1000, description="Maximum number of rows to read per sheet (default: 1000)"
),
convert_xls_to_xlsx: bool = Field(
False,
description="Whether to convert XLS files to XLSX format (default: False)",
),
) -> str:
"""Read multiple Excel/CSV files and convert sheets to Markdown tables. Cannot process https://URLs files."""
try:
# Import required packages
import_package("tabulate")
# Import xls2xlsx package if conversion is requested
if convert_xls_to_xlsx:
import_package("xls2xlsx")
all_results = []
temp_files = [] # Track temporary files for cleanup
success_count = 0
failed_count = 0
# Process each file
for document_path in document_paths:
# Check if file exists and is readable
error = check_file_readable(document_path)
if error:
all_results.append(
ExcelDocument(
file_name=os.path.basename(document_path),
file_path=document_path,
file_type="UNKNOWN",
sheet_count=0,
sheet_names=[],
sheets=[],
success=False,
error=error,
)
)
failed_count += 1
continue
try:
# Check file extension
file_ext = os.path.splitext(document_path)[1].lower()
# Validate file type
if file_ext not in [".csv", ".xls", ".xlsx", ".xlsm"]:
error_msg = f"Unsupported file format: {file_ext}. Only CSV, XLS, XLSX, and XLSM formats are supported."
all_results.append(
ExcelDocument(
file_name=os.path.basename(document_path),
file_path=document_path,
file_type=file_ext.replace(".", "").upper(),
sheet_count=0,
sheet_names=[],
sheets=[],
success=False,
error=error_msg,
)
)
failed_count += 1
continue
# Convert XLS to XLSX if requested and file is XLS
processed_path = document_path
if convert_xls_to_xlsx and file_ext == ".xls":
try:
logger.info(f"Converting XLS to XLSX: {document_path}")
converter = XLS2XLSX(document_path)
# Create temp file with xlsx extension
xlsx_path = (
os.path.splitext(document_path)[0] + "_converted.xlsx"
)
converter.to_xlsx(xlsx_path)
processed_path = xlsx_path
temp_files.append(xlsx_path) # Track for cleanup
logger.success(f"Converted XLS to XLSX: {xlsx_path}")
except Exception as conv_error:
logger.error(f"XLS to XLSX conversion error: {str(conv_error)}")
# Continue with original file if conversion fails
excel_sheets = []
sheet_names = []
# Handle CSV files differently
if file_ext == ".csv":
# For CSV files, create a single sheet with the file name
sheet_name = os.path.basename(document_path).replace(".csv", "")
df = pd.read_csv(processed_path, nrows=max_rows)
# Create markdown table
markdown_table = "*Empty table*"
if not df.empty:
headers = df.columns.tolist()
table_data = df.values.tolist()
markdown_table = tabulate(
table_data, headers=headers, tablefmt="pipe"
)
if len(df) >= max_rows:
markdown_table += (
f"\n\n*Note: Table truncated to {max_rows} rows*"
)
# Create sheet model
excel_sheets.append(
ExcelSheet(
name=sheet_name,
data=df.to_dict(orient="records"),
markdown_table=markdown_table,
row_count=len(df),
column_count=len(df.columns),
)
)
sheet_names = [sheet_name]
else:
# For Excel files, process all sheets
with pd.ExcelFile(processed_path) as xls:
sheet_names = xls.sheet_names
for sheet_name in sheet_names:
# Read Excel sheet into DataFrame with row limit
df = pd.read_excel(
xls, sheet_name=sheet_name, nrows=max_rows
)
# Create markdown table
markdown_table = "*Empty table*"
if not df.empty:
headers = df.columns.tolist()
table_data = df.values.tolist()
markdown_table = tabulate(
table_data, headers=headers, tablefmt="pipe"
)
if len(df) >= max_rows:
markdown_table += f"\n\n*Note: Table truncated to {max_rows} rows*"
# Create sheet model
excel_sheets.append(
ExcelSheet(
name=sheet_name,
data=df.to_dict(orient="records"),
markdown_table=markdown_table,
row_count=len(df),
column_count=len(df.columns),
)
)
# Create result for this file
file_result = ExcelDocument(
file_name=os.path.basename(document_path),
file_path=document_path,
processed_path=(
processed_path if processed_path != document_path else None
),
file_type=file_ext.replace(".", "").upper(),
sheet_count=len(sheet_names),
sheet_names=sheet_names,
sheets=excel_sheets,
success=True,
)
all_results.append(file_result)
success_count += 1
except Exception as file_error:
# Handle errors for individual files
error_msg = str(file_error)
logger.error(f"File reading error for {document_path}: {error_msg}")
all_results.append(
ExcelDocument(
file_name=os.path.basename(document_path),
file_path=document_path,
file_type=os.path.splitext(document_path)[1]
.replace(".", "")
.upper(),
sheet_count=0,
sheet_names=[],
sheets=[],
success=False,
error=error_msg,
)
)
failed_count += 1
# Clean up temporary files
for temp_file in temp_files:
try:
if os.path.exists(temp_file):
os.remove(temp_file)
logger.info(f"Removed temporary file: {temp_file}")
except Exception as cleanup_error:
logger.warning(
f"Error cleaning up temporary file {temp_file}: {str(cleanup_error)}"
)
# Create final result
excel_result = ExcelResult(
total_files=len(document_paths),
success_count=success_count,
failed_count=failed_count,
results=all_results,
)
return excel_result.model_dump_json()
except Exception as e:
return handle_error(e, "Excel/CSV files processing")
@mcp.tool(
description="Read and convert PowerPoint slides to base64 encoded images. return the parsed content. Cannot process https://URLs files."
)
def mcpreadpptx(
document_path: str = Field(description="The local input PowerPoint file path."),
) -> str:
"""Read and convert PowerPoint slides to base64 encoded images. Cannot process https://URLs files."""
error = check_file_readable(document_path)
if error:
return DocumentError(error=error, file_path=document_path).model_dump_json()
# Create temporary directory
temp_dir = tempfile.mkdtemp()
slides_data = []
try:
presentation = Presentation(document_path)
total_slides = len(presentation.slides)
if total_slides == 0:
raise ValueError("PPT file does not contain any slides")
# Process each slide
for i, slide in enumerate(presentation.slides):
# Set slide dimensions
slide_width_px = 1920 # 16:9 ratio
slide_height_px = 1080
# Create blank image
slide_img = Image.new("RGB", (slide_width_px, slide_height_px), "white")
draw = ImageDraw.Draw(slide_img)
font = ImageFont.load_default()
# Draw slide number
draw.text((20, 20), f"Slide {i+1}/{total_slides}", fill="black", font=font)
# Process shapes in the slide
for shape in slide.shapes:
try:
# Process images
if hasattr(shape, "image") and shape.image:
image_stream = io.BytesIO(shape.image.blob)
img = Image.open(image_stream)
left = int(
shape.left * slide_width_px / presentation.slide_width
)
top = int(
shape.top * slide_height_px / presentation.slide_height
)
slide_img.paste(img, (left, top))
# Process text
elif hasattr(shape, "text") and shape.text:
text_left = int(
shape.left * slide_width_px / presentation.slide_width
)
text_top = int(
shape.top * slide_height_px / presentation.slide_height
)
draw.text(
(text_left, text_top),
shape.text,
fill="black",
font=font,
)
except Exception as shape_error:
logger.warning(
f"Error processing shape in slide {i+1}: {str(shape_error)}"
)
# Save slide image
img_path = os.path.join(temp_dir, f"slide_{i+1}.jpg")
slide_img.save(img_path, "JPEG")
# Convert to base64
base64_image = encode_images(img_path)
slides_data.append(
PowerPointSlide(
slide_number=i + 1, image=f"data:image/jpeg;base64,{base64_image}"
)
)
# Create result
result = PowerPointDocument(
file_path=document_path,
file_name=os.path.basename(document_path),
slide_count=total_slides,
slides=slides_data,
)
return result.model_dump_json()
except Exception as e:
return handle_error(e, "PowerPoint processing", document_path)
finally:
# Clean up temporary files
try:
for file in os.listdir(temp_dir):
os.remove(os.path.join(temp_dir, file))
os.rmdir(temp_dir)
except Exception as cleanup_error:
logger.warning(f"Error cleaning up temporary files: {str(cleanup_error)}")
@mcp.tool(
description="Read HTML file and extract text content, optionally extract links, images, and table information, and convert to Markdown format."
)
def mcpreadhtmltext(
document_path: str = Field(description="Local HTML file path or Web URL."),
extract_links: bool = Field(
default=True, description="Whether to extract link information"
),
extract_images: bool = Field(
default=True, description="Whether to extract image information"
),
extract_tables: bool = Field(
default=True, description="Whether to extract table information"
),
convert_to_markdown: bool = Field(
default=True, description="Whether to convert HTML to Markdown format"
),
) -> str:
"""Read HTML file and extract text content, optionally extract links, images, and table information, and convert to Markdown format."""
error = check_file_readable(document_path)
if error:
return DocumentError(error=error, file_path=document_path).model_dump_json()
try:
# Read HTML file
with open(document_path, "r", encoding="utf-8") as f:
html_content = f.read()
# Parse HTML using BeautifulSoup
soup = BeautifulSoup(html_content, "html.parser")
# Extract text content (remove script and style content)
for script in soup(["script", "style"]):
script.extract()
text_content = soup.get_text(separator="\n", strip=True)
# Extract title
title = soup.title.string if soup.title else None
# Initialize result object
result = HtmlDocument(
content=text_content,
html_content=html_content,
file_path=document_path,
file_name=os.path.basename(document_path),
file_size=os.path.getsize(document_path),
last_modified=datetime.fromtimestamp(
os.path.getmtime(document_path)
).strftime("%Y-%m-%d %H:%M:%S"),
title=title,
)
# Extract links
if extract_links:
links = []
for link in soup.find_all("a"):
href = link.get("href")
text = link.get_text(strip=True)
if href:
links.append({"url": href, "text": text})
result.links = links
# Extract images
if extract_images:
images = []
for img in soup.find_all("img"):
src = img.get("src")
alt = img.get("alt", "")
if src:
images.append({"src": src, "alt": alt})
result.images = images
# Extract tables
if extract_tables:
tables = []
for table in soup.find_all("table"):
tables.append(str(table))
result.tables = tables
# Convert to Markdown
if convert_to_markdown:
h = html2text.HTML2Text()
h.ignore_links = False
h.ignore_images = False
h.ignore_tables = False
markdown_content = h.handle(html_content)
result.markdown = markdown_content
return result.model_dump_json()
except Exception as e:
return handle_error(e, "HTML file reading", document_path)
def main():
load_dotenv()
print("Starting Document MCP Server...", file=sys.stderr)
mcp.run(transport="stdio")
# Make the module callable
def __call__():
"""
Make the module callable for uvx.
This function is called when the module is executed directly.
"""
main()
sys.modules[__name__].__call__ = __call__
# Run the server when the script is executed directly
if __name__ == "__main__":
main()
@@ -0,0 +1,199 @@
"""
Download MCP Server
This module provides MCP server functionality for downloading files from URLs.
It handles various download scenarios with proper validation, error handling,
and progress tracking.
Key features:
- File downloading from HTTP/HTTPS URLs
- Download progress tracking
- File validation
- Safe file saving
Main functions:
- mcpdownload: Downloads files from URLs to local filesystem
"""
import os
import sys
import traceback
import urllib.parse
from pathlib import Path
from typing import List, Optional
import requests
from dotenv import load_dotenv
from mcp.server.fastmcp import FastMCP
from pydantic import BaseModel, Field
from aworld.logs.util import logger
mcp = FastMCP("download-server")
class DownloadResult(BaseModel):
"""Download result model with file information"""
file_path: str
file_name: str
file_size: int
content_type: Optional[str] = None
success: bool
error: Optional[str] = None
class DownloadResults(BaseModel):
"""Download results model for multiple files"""
results: List[DownloadResult]
success_count: int
failed_count: int
@mcp.tool(description="Download files from URLs and save to the local filesystem.")
def mcpdownloadfiles(
urls: List[str] = Field(
..., description="The URLs of the files to download. Must be a list of URLs."
),
output_dir: str = Field(
"/tmp/mcp_downloads",
description="Directory to save the downloaded files (default: /tmp/mcp_downloads).",
),
timeout: int = Field(60, description="Download timeout in seconds (default: 60)."),
) -> str:
"""Download files from URLs and save to the local filesystem.
Args:
urls: The URLs of the files to download, must be a list of URLs
output_dir: Directory to save the downloaded files
timeout: Download timeout in seconds
Returns:
JSON string with download results information
"""
results = []
success_count = 0
failed_count = 0
for single_url in urls:
result_json = _download_single_file(single_url, output_dir, "", timeout)
result = DownloadResult.model_validate_json(result_json)
results.append(result)
if result.success:
success_count += 1
else:
failed_count += 1
batch_results = DownloadResults(
results=results, success_count=success_count, failed_count=failed_count
)
return batch_results.model_dump_json()
def _download_single_file(
url: str, output_dir: str, filename: str, timeout: int
) -> str:
"""Download a single file from URL and save it to the local filesystem."""
try:
# Validate URL
if not url.startswith(("http://", "https://")):
raise ValueError(
"Invalid URL format. URL must start with http:// or https://"
)
# Create output directory if it doesn't exist
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
# Determine filename if not provided
if not filename:
filename = os.path.basename(urllib.parse.urlparse(url).path)
if not filename:
filename = "downloaded_file"
# Full path to save the file
file_path = os.path.join(output_path, filename)
logger.info(f"Downloading file from {url} to {file_path}")
# Download the file with progress tracking
headers = {
"User-Agent": (
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
"AWorld/1.0 (https://github.com/inclusionAI/AWorld; qintong.wqt@antgroup.com) "
"Python/requests "
),
"Accept": "text/html,application/xhtml+xml,application/xml,application/pdf;q=0.9,image/webp,*/*;q=0.8",
"Accept-Language": "en-US,en;q=0.5",
"Accept-Encoding": "gzip, deflate, br",
"Connection": "keep-alive",
}
response = requests.get(url, headers=headers, stream=True, timeout=timeout)
response.raise_for_status()
# Get content type and size
content_type = response.headers.get("Content-Type")
# Save the file
with open(file_path, "wb") as f:
for chunk in response.iter_content(chunk_size=8192):
if chunk:
f.write(chunk)
# Get actual file size
actual_size = os.path.getsize(file_path)
logger.info(f"File downloaded successfully to {file_path}")
# Create result
result = DownloadResult(
file_path=file_path,
file_name=filename,
file_size=actual_size,
content_type=content_type,
success=True,
error=None,
)
return result.model_dump_json()
except Exception as e:
error_msg = str(e)
logger.error(f"Download error: {traceback.format_exc()}")
result = DownloadResult(
file_path="",
file_name="",
file_size=0,
content_type=None,
success=False,
error=error_msg,
)
return result.model_dump_json()
def main():
load_dotenv()
print("Starting Download MCP Server...", file=sys.stderr)
mcp.run(transport="stdio")
# Make the module callable
def __call__():
"""
Make the module callable for uvx.
This function is called when the module is executed directly.
"""
main()
sys.modules[__name__].__call__ = __call__
# Run the server when the script is executed directly
if __name__ == "__main__":
main()
@@ -0,0 +1,96 @@
from e2b_code_interpreter import Sandbox
from pydantic import Field
from mcp.server.fastmcp import FastMCP
import os
# Initialize MCP server
mcp = FastMCP("e2b-code-server")
@mcp.tool(description="Upload local file to e2b sandbox.")
async def e2b_upload_file(
path: str = Field(
description="The local file path to upload."
)
) -> str:
"""
Upload local file to e2b sandbox.
Args:
path (str): The local file path to upload.
Returns:
str: E2b file path and sandbox_id.
"""
try:
os.environ["E2B_API_KEY"] = os.getenv("E2B_API_KEY")
sbx = Sandbox()
local_file_name = os.path.basename(path)
e2b_file_path = f"/home/user/{local_file_name}"
# Read local file relative to the current working directory
with open(path, "rb") as file:
# Upload file to the sandbox to absolute path
sbx.files.write(e2b_file_path, file)
return f"{e2b_file_path}, {sbx.sandbox_id}"
except Exception as e:
return f"Upload failed. Error: {str(e)}"
@mcp.tool(description="Run code in a specified e2b sandbox.")
async def e2b_run_code(
sandbox_id: str = Field(
default=None,
description="The sandbox id to run code in, if you have uploaded a file, you should use the sandbox_id returned by the e2b_upload_file function."
),
code_block: str = Field(
default=None,
description="The code block to run in e2b sandbox."
),
) -> str:
"""
Run code in a specified e2b sandbox.
Args:
sandbox_id (str): The sandbox id to run code in.
code_block (str): The code block to run in e2b sandbox.
Returns:
str: The result of running the code block.
"""
try:
os.environ["E2B_API_KEY"] = os.getenv("E2B_API_KEY")
sbx = Sandbox(
sandbox_id=sandbox_id,
)
execution = sbx.run_code(code_block)
return execution.logs
except Exception as e:
return f"Run code failed. Error: {str(e)}"
def main():
from dotenv import load_dotenv
load_dotenv()
import sys
print("Starting E2b Code MCP Server...", file=sys.stderr)
mcp.run(transport='stdio')
# Make the module callable
def __call__():
"""
Make the module callable for uvx.
This function is called when the module is executed directly.
"""
main()
# Add this for compatibility with uvx
import sys
sys.modules[__name__].__call__ = __call__
# Run the server when the script is executed directly
if __name__ == "__main__":
main()
@@ -0,0 +1,197 @@
import os
import time
import json
import requests
import sys
import hashlib
from dotenv import load_dotenv
from mcp.server import FastMCP
from pydantic import Field
from typing_extensions import Any
from aworld.logs.util import logger
mcp = FastMCP("gen-audio-server")
def calculate_sha256(plain_text):
"""
Calculate SHA-256 digest of a string.
Args:
plain_text (str): The text to digest
Returns:
str: Hexadecimal representation of the digest
"""
try:
# Create SHA-256 hash object
sha256 = hashlib.sha256()
# Update with the bytes of the plain text (UTF-8 encoded)
sha256.update(plain_text.encode('utf-8'))
# Get the digest in bytes
digest_bytes = sha256.digest()
# Convert each byte to hexadecimal and join
hex_digest = ''.join([f'{b:02x}' for b in digest_bytes])
return hex_digest
except Exception as e:
logger.warning(f"Error calculating SHA-256 digest: {e}")
return ""
def generate_headers(app_key, secret):
"""Generate headers with fresh timestamp and digest"""
timestamp = str(int(time.time() * 1000))
plain_text = f"{app_key}_{secret}_{timestamp}"
digest = calculate_sha256(plain_text)
return {
'Content-Type': 'application/json',
'Alipay-Mf-Appkey': app_key,
'Alipay-Mf-Digest': digest,
'Alipay-Mf-Timestamp': timestamp
}
@mcp.tool(description="Generate audio from text content")
def gen_audio(content: str = Field(description="The text content to convert to audio")) -> Any:
"""Generate audio from text content using TTS service"""
task_url = os.getenv('AUDIO_TASK_URL')
query_url = os.getenv('AUDIO_QUERY_URL')
app_key = os.getenv('AUDIO_APP_KEY')
secret = os.getenv('AUDIO_SECRET')
if not (task_url and query_url and app_key and secret):
logger.warning(f"Query failed: task_url, query_url, app_key, secret parameters incomplete")
return None
# Generate initial headers
headers = generate_headers(app_key, secret)
sample_rate = os.getenv('AUDIO_SAMPLE_RATE', '16000')
audio_format = os.getenv('AUDIO_AUDIO_FORMAT', 'wav')
tts_voice = os.getenv('AUDIO_TTS_VOICE', 'DBCNF245')
tts_speech_rate = os.getenv('AUDIO_TTS_SPEECH_RATE', '0')
tts_volume = os.getenv('AUDIO_TTS_VOLUME', '50')
tts_pitch = os.getenv('AUDIO_TTS_PITCH', '0')
voice_type = os.getenv('AUDIO_VOICE_TYPE', 'VOICE_CLONE_LAM')
# task_data
task_data = {
"sample_rate": sample_rate,
"audio_format": audio_format,
"tts_voice": tts_voice,
"tts_speech_rate": tts_speech_rate,
"tts_volume": tts_volume,
"tts_pitch": tts_pitch,
"tts_text": content,
"voice_type": voice_type,
}
try:
# Step 1: Submit task to generate audio
response = requests.post(task_url, headers=headers, json=task_data)
if response.status_code != 200:
return None
result = response.json()
# Check if task was successfully submitted
if not result.get("success"):
return None
# Extract task ID
task_id = result.get("data")
if not task_id:
return None
logger.info(f"Task submitted successfully. Task ID: {task_id}")
# Step 2: Poll for results
max_attempts = int(os.getenv('AUDIO_RETRY_TIMES', 10))
wait_time = int(os.getenv('AUDIO_SLEEP_TIME', 5))
query_url = query_url + f"?async_task_id={task_id}"
for attempt in range(max_attempts):
# Wait before polling
time.sleep(wait_time)
logger.info(f"Polling attempt {attempt + 1}/{max_attempts}...")
# Generate fresh headers for each poll request
query_headers = generate_headers(app_key, secret)
# Poll for results
query_response = requests.post(query_url, headers=query_headers)
if query_response.status_code != 200:
logger.info(f"Poll request failed with status code {query_response.status_code}")
continue
try:
query_result = query_response.json()
except json.JSONDecodeError as e:
logger.warning(f"Failed to parse response as JSON: {e}")
continue
# Check if processing is complete
if query_result.get("success") and query_result.get("data", {}).get("status") == "ST_SUCCESS":
# Extract audio URL based on the correct JSON structure
# Navigate through the nested structure: data -> result -> result -> audioUrl
audio_url = query_result.get("data", {}).get("result", {}).get("result", {}).get("audioUrl")
if audio_url:
return json.dumps({"audio_data": audio_url})
else:
logger.info("Audio URL not found in the response")
return None
elif query_result.get("success") and query_result.get("data", {}).get("status") == "ST_RUNNING":
# If still running, continue to next polling attempt
logger.info("Task still running, continuing to next poll...")
continue
else:
# Any other status, return None
logger.warning(f"Unexpected status: {query_result.get('data', {}).get('status')}")
return None
# If we get here, polling timed out
logger.warning("Polling timed out after maximum attempts")
return None
except Exception as e:
import traceback
logger.warning(f"Exception occurred: {e}")
return None
def main():
from dotenv import load_dotenv
load_dotenv()
print("Starting Audio MCP gen-audio-server...", file=sys.stderr)
mcp.run(transport="stdio")
# Make the module callable
def __call__():
"""
Make the module callable for uvx.
This function is called when the module is executed directly.
"""
main()
sys.modules[__name__].__call__ = __call__
if __name__ == "__main__":
main()
# For testing without MCP
# result = gen_audio("hello ,this is test")
# print("\nFinal Result:")
# print(result)
@@ -0,0 +1,166 @@
import os
import time
import json
import requests
import sys
from dotenv import load_dotenv
from mcp.server import FastMCP
from pydantic import Field
from typing_extensions import Any
from aworld.logs.util import logger
mcp = FastMCP("gen-pic-server")
@mcp.tool(description="Generate picture from text content")
def gen_picture(prompt: str = Field(description="The text prompt to generate an image"),
num: int = Field(0,
description="Number of images to generate, 0 means use environment variable")) -> Any:
"""Generate picture from text prompt"""
api_key = os.getenv('DASHSCOPE_API_KEY')
submit_url = os.getenv('DASHSCOPE_SUBMIT_URL', '')
query_base_url = os.getenv('DASHSCOPE_QUERY_BASE_URL', '')
if not api_key or not submit_url or not query_base_url:
logger.warning(
"Query failed: DASHSCOPE_API_KEY,DASHSCOPE_SUBMIT_URL,DASHSCOPE_QUERY_BASE_URL environment variable is not set")
return None
headers = {
'X-DashScope-Async': 'enable',
'Authorization': f'Bearer {api_key}',
'Content-Type': 'application/json'
}
# Get parameters from environment variables or use defaults
model = os.getenv('DASHSCOPE_MODEL', 'wanx2.1-t2i-turbo')
size = os.getenv('DASHSCOPE_SIZE', '1024*1024')
# Use num parameter if provided (>0), otherwise use environment variable
n = num if num > 0 else int(os.getenv('DASHSCOPE_N', '1'))
task_data = {
"model": model,
"input": {
"prompt": prompt
},
"parameters": {
"size": size,
"n": n
}
}
try:
# Step 1: Submit task to generate image
logger.info("Submitting task to generate image...")
response = requests.post(submit_url, headers=headers, json=task_data)
if response.status_code != 200:
logger.warning(f"Task submission failed with status code {response.status_code}")
return None
result = response.json()
# Check if task was successfully submitted
if not result.get("output") or not result.get("output").get("task_id"):
logger.warning("Failed to get task_id from response")
return None
# Extract task ID
task_id = result.get("output").get("task_id")
logger.info(f"Task submitted successfully. Task ID: {task_id}")
# Step 2: Poll for results
max_attempts = int(os.getenv('DASHSCOPE_RETRY_TIMES', 10))
wait_time = int(os.getenv('DASHSCOPE_SLEEP_TIME', 5))
query_url = f"{query_base_url}{task_id}"
for attempt in range(max_attempts):
# Wait before polling
time.sleep(wait_time)
logger.info(f"Polling attempt {attempt + 1}/{max_attempts}...")
# Poll for results
query_response = requests.get(query_url, headers={'Authorization': f'Bearer {api_key}'})
if query_response.status_code != 200:
logger.info(f"Poll request failed with status code {query_response.status_code}")
continue
try:
query_result = query_response.json()
except json.JSONDecodeError as e:
logger.warning(f"Failed to parse response as JSON: {e}")
continue
# Check task status
task_status = query_result.get("output", {}).get("task_status")
if task_status == "SUCCEEDED":
# Extract image URLs
results = query_result.get("output", {}).get("results", [])
if results:
# Create a simple array of objects with image_url
image_urls = []
for result in results:
if "url" in result:
image_urls.append({"image_url": result["url"]})
if image_urls:
return json.dumps(image_urls)
else:
logger.info("No valid image URLs found in the response")
return None
else:
logger.info("No results found in the response")
return None
elif task_status in ["PENDING", "RUNNING"]:
# If still running, continue to next polling attempt
logger.info(f"Task status: {task_status}, continuing to next poll...")
continue
elif task_status == "FAILED":
logger.warning("Task failed")
return None
else:
# Any other status, return None
logger.warning(f"Unexpected status: {task_status}")
return None
# If we get here, polling timed out
logger.warning("Polling timed out after maximum attempts")
return None
except Exception as e:
logger.warning(f"Exception occurred: {e}")
return None
def main():
from dotenv import load_dotenv
load_dotenv()
print("Starting MCP gen-pic-server...", file=sys.stderr)
mcp.run(transport="stdio")
# Make the module callable
def __call__():
"""
Make the module callable for uvx.
This function is called when the module is executed directly.
"""
main()
sys.modules[__name__].__call__ = __call__
if __name__ == "__main__":
main()
# For testing without MCP
# result = gen_picture("sunflower", 2)
# print("\nFinal Result:")
# print(result)
@@ -0,0 +1,153 @@
import os
import time
import json
import requests
import sys
from dotenv import load_dotenv
from mcp.server import FastMCP
from pydantic import Field
from typing_extensions import Any
from aworld.logs.util import logger
mcp = FastMCP("gen-video-server")
@mcp.tool(description="Generate video from text content")
def gen_video(prompt: str = Field(description="The text prompt to generate a video")) -> Any:
"""Generate video from text prompt"""
api_key = os.getenv('DASHSCOPE_API_KEY')
submit_url = os.getenv('DASHSCOPE_VIDEO_SUBMIT_URL', '')
query_base_url = os.getenv('DASHSCOPE_QUERY_BASE_URL', '')
if not api_key or not submit_url or not query_base_url:
logger.warning("Query failed: DASHSCOPE_API_KEY, DASHSCOPE_VIDEO_SUBMIT_URL, DASHSCOPE_QUERY_BASE_URL environment variables are not set")
return None
headers = {
'X-DashScope-Async': 'enable',
'Authorization': f'Bearer {api_key}',
'Content-Type': 'application/json'
}
# Get parameters from environment variables or use defaults
model = os.getenv('DASHSCOPE_VIDEO_MODEL', 'wanx2.1-t2v-turbo')
size = os.getenv('DASHSCOPE_VIDEO_SIZE', '1280*720')
# Note: Currently the API only supports generating one video at a time
# But we keep the num parameter for API compatibility
task_data = {
"model": model,
"input": {
"prompt": prompt
},
"parameters": {
"size": size
}
}
try:
# Step 1: Submit task to generate video
logger.info("Submitting task to generate video...")
response = requests.post(submit_url, headers=headers, json=task_data)
if response.status_code != 200:
logger.warning(f"Task submission failed with status code {response.status_code}")
return None
result = response.json()
# Check if task was successfully submitted
if not result.get("output") or not result.get("output").get("task_id"):
logger.warning("Failed to get task_id from response")
return None
# Extract task ID
task_id = result.get("output").get("task_id")
logger.info(f"Task submitted successfully. Task ID: {task_id}")
# Step 2: Poll for results
max_attempts = int(os.getenv('DASHSCOPE_VIDEO_RETRY_TIMES', 10)) # Increased default retries for video
wait_time = int(os.getenv('DASHSCOPE_VIDEO_SLEEP_TIME', 5)) # Increased default wait time for video
query_url = f"{query_base_url}{task_id}"
for attempt in range(max_attempts):
# Wait before polling
time.sleep(wait_time)
logger.info(f"Polling attempt {attempt + 1}/{max_attempts}...")
# Poll for results
query_response = requests.get(query_url, headers={'Authorization': f'Bearer {api_key}'})
if query_response.status_code != 200:
logger.info(f"Poll request failed with status code {query_response.status_code}")
continue
try:
query_result = query_response.json()
except json.JSONDecodeError as e:
logger.warning(f"Failed to parse response as JSON: {e}")
continue
# Check task status
task_status = query_result.get("output", {}).get("task_status")
if task_status == "SUCCEEDED":
# Extract video URL
video_url = query_result.get("output", {}).get("video_url")
if video_url:
# Return as array of objects with video_url for consistency with image API
return json.dumps({"video_url": video_url})
else:
logger.info("Video URL not found in the response")
return None
elif task_status in ["PENDING", "RUNNING"]:
# If still running, continue to next polling attempt
logger.info(f"Task status: {task_status}, continuing to next poll...")
continue
elif task_status == "FAILED":
logger.warning("Task failed")
return None
else:
# Any other status, return None
logger.warning(f"Unexpected status: {task_status}")
return None
# If we get here, polling timed out
logger.warning("Polling timed out after maximum attempts")
return None
except Exception as e:
logger.warning(f"Exception occurred: {e}")
return None
def main():
from dotenv import load_dotenv
load_dotenv()
print("Starting MCP gen-video-server...", file=sys.stderr)
mcp.run(transport="stdio")
# Make the module callable
def __call__():
"""
Make the module callable for uvx.
This function is called when the module is executed directly.
"""
main()
sys.modules[__name__].__call__ = __call__
if __name__ == "__main__":
main()
# For testing without MCP
# result = gen_video("A cat running under moonlight")
# print("\nFinal Result:")
# print(result)
@@ -0,0 +1,240 @@
"""
Image MCP Server
This module provides MCP server functionality for image processing and analysis.
It handles image encoding, optimization, and various image analysis tasks such as
OCR (Optical Character Recognition) and visual reasoning.
The server supports both local image files and remote image URLs with proper validation
and handles various image formats including JPEG, PNG, GIF, and others.
Main functions:
- encode_images: Encodes images to base64 format with optimization
- optimize_image: Resizes and optimizes images for better performance
- Various MCP tools for image analysis and processing
"""
# import asyncio
import base64
import os
from io import BytesIO
from typing import Any, Dict, List
from PIL import Image
from pydantic import Field
from aworld.logs.util import logger
from mcp_servers.utils import get_file_from_source
from mcp.server.fastmcp import FastMCP
from openai import OpenAI
# Initialize MCP server
mcp = FastMCP("image-server")
IMAGE_OCR = (
"Input is a base64 encoded image. Read text from image if present. "
"Return a json string with the following format: "
'{"image_text": "text from image"}'
)
IMAGE_REASONING = (
"Input is a base64 encoded image. Given user's task: {task}, "
"solve it following the guide line:\n"
"1. Careful visual inspection\n"
"2. Contextual reasoning\n"
"3. Text transcription where relevant\n"
"4. Logical deduction from visual evidence\n"
"Return a json string with the following format: "
'{"image_reasoning_result": "reasoning result given task and image"}'
)
def optimize_image(image_data: bytes, max_size: int = 1024) -> bytes:
"""
Optimize image by resizing if needed
Args:
image_data: Raw image data
max_size: Maximum dimension size in pixels
Returns:
bytes: Optimized image data
Raises:
ValueError: When image cannot be processed
"""
try:
image = Image.open(BytesIO(image_data))
# Resize if image is too large
if max(image.size) > max_size:
ratio = max_size / max(image.size)
new_size = (int(image.size[0] * ratio), int(image.size[1] * ratio))
image = image.resize(new_size, Image.Resampling.LANCZOS)
# Save to buffer
buffered = BytesIO()
image_format = image.format if image.format else "JPEG"
image.save(buffered, format=image_format)
return buffered.getvalue()
except Exception as e:
logger.warning(f"Failed to optimize image: {str(e)}")
return image_data # Return original data if optimization fails
def encode_images(image_sources: List[str], with_header: bool = True) -> List[str]:
"""
Encode images to base64 format with robust file handling
Args:
image_sources: List of URLs or local file paths of images
with_header: Whether to include MIME type header
Returns:
List[str]: Base64 encoded image strings, with MIME type prefix if with_header is True
Raises:
ValueError: When image source is invalid or image format is not supported
"""
if not image_sources:
raise ValueError("Image sources cannot be empty")
images = []
for image_source in image_sources:
try:
# Get file with validation (only image files allowed)
file_path, mime_type, content = get_file_from_source(
image_source,
allowed_mime_prefixes=["image/"],
max_size_mb=10.0, # 10MB limit for images
type="image",
)
# Optimize image
optimized_content = optimize_image(content)
# Encode to base64
image_base64 = base64.b64encode(optimized_content).decode()
# Format with header if requested
final_image = (
f"data:{mime_type};base64,{image_base64}"
if with_header
else image_base64
)
images.append(final_image)
# Clean up temporary file if it was created for a URL
if file_path != os.path.abspath(image_source) and os.path.exists(file_path):
os.unlink(file_path)
except Exception as e:
logger.error(f"Error encoding image from {image_source}: {str(e)}")
raise
return images
def image_to_base64(image_path):
try:
# todo 解析pdf或其他文件的图片
with Image.open(image_path) as image:
buffered = BytesIO()
image_format = image.format if image.format else "JPEG"
image.save(buffered, format=image_format)
image_bytes = buffered.getvalue()
base64_encoded = base64.b64encode(image_bytes).decode('utf-8')
return base64_encoded
except Exception as e:
print(f"Base64 error: {e}")
return None
def create_image_contents(prompt: str, image_base64: List[str]) -> List[Dict[str, Any]]:
"""Create uniform image format for querying llm."""
content = [
{"type": "text", "text": prompt},
]
content.extend(
[{"type": "image_url", "image_url": {"url": url}} for url in image_base64]
)
return content
@mcp.tool(description="solve the question by careful reasoning given the image(s) in given local filepath or url, including reasoning, ocr, etc.")
def mcp_image_recognition(
image_urls: List[str] = Field(
description="The input image(s) in given a list of local filepaths or urls."
),
question: str = Field(description="The question to ask."),
) -> str:
"""solve the question by careful reasoning given the image(s) in given filepath or url."""
try:
image_base64 = image_to_base64(image_urls[0])
logger.info(f"image_url: {image_urls[0]}")
reasoning_prompt = question
messages=[
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content":
[
{"type": "text", "text": reasoning_prompt},
{
"type": "image_url",
"image_url": {
"url": f"data:image/jpeg;base64,{image_base64}"
}
},
],
},
]
client = OpenAI(
api_key=os.getenv("LLM_API_KEY"),
base_url=os.getenv("LLM_BASE_URL")
)
response = client.chat.completions.create(
model=os.getenv("LLM_MODEL_NAME"),
messages=messages,
)
logger.info(f"response: {response}")
image_reasoning_result = response.choices[0].message.content
except Exception as e:
image_reasoning_result = ""
import traceback
traceback.print_exc()
logger.error(f"image_reasoning_result-Execute error: {e}")
logger.info(
f"---get_reasoning_by_image-image_reasoning_result:{image_reasoning_result}"
)
return image_reasoning_result
def main():
from dotenv import load_dotenv
load_dotenv()
print("Starting Image MCP Server...", file=sys.stderr)
mcp.run(transport='stdio')
# Make the module callable
def __call__():
"""
Make the module callable for uvx.
This function is called when the module is executed directly.
"""
main()
# Add this for compatibility with uvx
import sys
sys.modules[__name__].__call__ = __call__
# Run the server when the script is executed directly
if __name__ == "__main__":
main()
@@ -0,0 +1,180 @@
import asyncio
import json
import logging
import os
import sys
import aiohttp
from typing import List, Dict, Any, Optional
from dotenv import load_dotenv
from mcp.server import FastMCP
from pydantic import Field
from aworld.logs.util import logger
mcp = FastMCP("picsearch-server")
async def search_single(query: str, num: int = 5) -> Optional[Dict[str, Any]]:
"""Execute a single search query, returns None on error"""
try:
url = os.getenv('PIC_SEARCH_URL')
searchMode = os.getenv('PIC_SEARCH_SEARCHMODE')
source = os.getenv('PIC_SEARCH_SOURCE')
domain = os.getenv('PIC_SEARCH_DOMAIN')
uid = os.getenv('PIC_SEARCH_UID')
if not url or not searchMode or not source or not domain:
logger.warning(f"Query failed: url, searchMode, source, domain parameters incomplete")
return None
headers = {
'Content-Type': 'application/json'
}
data = {
"domain": domain,
"extParams": {
"contentType": "llmWholeImage"
},
"page": 0,
"pageSize": num,
"query": query,
"searchMode": searchMode,
"source": source,
"userId": uid
}
async with aiohttp.ClientSession() as session:
try:
async with session.post(url, headers=headers, json=data) as response:
if response.status != 200:
logger.warning(f"Query failed: {query}, status code: {response.status}")
return None
result = await response.json()
return result
except aiohttp.ClientError:
logger.warning(f"Request error: {query}")
return None
except Exception:
logger.warning(f"Query exception: {query}")
return None
def filter_valid_docs(result: Optional[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""Filter valid document results, returns empty list if input is None"""
if result is None:
return []
try:
valid_docs = []
# Check success field
if not result.get("success"):
return valid_docs
# Check searchDocs field
search_docs = result.get("searchImages", [])
if not search_docs:
return valid_docs
# Extract required fields
required_fields = ["title", "picUrl"]
for doc in search_docs:
# Check if all required fields exist and are not empty
is_valid = True
for field in required_fields:
if field not in doc or not doc[field]:
is_valid = False
break
if is_valid:
# Keep only required fields
filtered_doc = {field: doc[field] for field in required_fields}
valid_docs.append(filtered_doc)
return valid_docs
except Exception:
return []
@mcp.tool(description="Search Picture based on the user's input query")
async def search(
query: str = Field(
description="The query to search for picture"
),
num: int = Field(
5,
description="Maximum number of results to return, default is 5"
)
) -> Any:
"""Execute search function for a single query"""
try:
# Get configuration from environment variables
env_total_num = os.getenv('PIC_SEARCH_TOTAL_NUM')
if env_total_num and env_total_num.isdigit():
# Force override input num parameter with environment variable
num = int(env_total_num)
# If no query provided, return empty list
if not query:
return json.dumps([])
# Get actual number of results to return
slice_num = os.getenv('PIC_SEARCH_SLICE_NUM')
if slice_num and slice_num.isdigit():
actual_num = int(slice_num)
else:
actual_num = num
# Execute the query
result = await search_single(query, actual_num)
# Filter results
valid_docs = filter_valid_docs(result)
# Return results
result_json = json.dumps(valid_docs, ensure_ascii=False)
logger.info(f"Completed query: '{query}', found {len(valid_docs)} valid documents")
logger.info(result_json)
return result_json
except Exception as e:
# Return empty list on exception
logger.error(f"Error processing query: {str(e)}")
return json.dumps([])
def main():
from dotenv import load_dotenv
load_dotenv()
print("Starting Audio MCP picsearch-server...", file=sys.stderr)
mcp.run(transport="stdio")
# Make the module callable
def __call__():
"""
Make the module callable for uvx.
This function is called when the module is executed directly.
"""
main()
sys.modules[__name__].__call__ = __call__
if __name__ == "__main__":
main()
# if __name__ == "__main__":
# # Configure logging
# logging.basicConfig(
# level=logging.INFO,
# format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
# )
#
#
# # Test single query
# asyncio.run(search(query="Image search test"))
#
# # Test multiple queries no longer applies
@@ -0,0 +1,102 @@
import os
import sys
import traceback
from dotenv import load_dotenv
from mcp.server.fastmcp import FastMCP
from pydantic import Field
from aworld.config.conf import AgentConfig
from aworld.logs.util import logger
from aworld.models.llm import call_llm_model, get_llm_model
# Initialize MCP server
mcp = FastMCP("reasoning-server")
@mcp.tool(
description="Perform complex problem reasoning using powerful reasoning model."
)
def complex_problem_reasoning(
question: str = Field(
description="The input question for complex problem reasoning,"
+ " such as math and code contest problem",
),
original_task: str = Field(
default="",
description="The original task description."
+ " This argument could be fetched from the <task>TASK</task> tag",
),
) -> str:
"""
Perform complex problem reasoning using Powerful Reasoning model,
such as riddle, game or competition-level STEM(including code) problems.
Args:
question: The input question for complex problem reasoning
original_task: The original task description (optional)
Returns:
str: The reasoning result from the model
"""
try:
# Prepare the prompt with both the question and original task if provided
prompt = question
if original_task:
prompt = f"Original Task: {original_task}\n\nQuestion: {question}"
# Call the LLM model for reasoning
response = call_llm_model(
llm_model=get_llm_model(
conf=AgentConfig(
llm_provider="openai",
llm_model_name=os.getenv("LLM_MODEL_NAME", "your_openai_api_key"),
llm_api_key=os.getenv("LLM_API_KEY", "your_openai_api_key"),
llm_base_url=os.getenv("LLM_BASE_URL", "your_openai_base_url"),
)
),
messages=[
{
"role": "system",
"content": (
"You are an expert at solving complex problems including math,"
" code contests, riddles, and puzzles."
" Provide detailed step-by-step reasoning and a clear final answer."
),
},
{"role": "user", "content": prompt},
],
temperature=float(os.getenv("LLM_TEMPERATURE", "0.3")),
)
# Extract the reasoning result
reasoning_result = response.content
logger.info("Complex reasoning completed successfully")
return reasoning_result
except Exception as e:
logger.error(f"Error in complex problem reasoning: {traceback.format_exc()}")
return f"Error performing reasoning: {str(e)}"
def main():
load_dotenv()
print("Starting Reasoning MCP Server...", file=sys.stderr)
mcp.run(transport="stdio")
# Make the module callable
def __call__():
"""
Make the module callable for uvx.
This function is called when the module is executed directly.
"""
main()
sys.modules[__name__].__call__ = __call__
# Run the server when the script is executed directly
if __name__ == "__main__":
main()
@@ -0,0 +1,165 @@
"""
Search MCP Server
This module provides MCP server functionality for performing web searches using various search engines.
It supports structured queries and returns formatted search results.
Key features:
- Perform web searches using Exa, Google, and DuckDuckGo
- Filter and format search results
- Validate and process search queries
Main functions:
- mcpsearchexa: Searches the web using Exa
- mcpsearchgoogle: Searches the web using Google
- mcpsearchduckduckgo: Searches the web using DuckDuckGo
"""
import os
import sys
import traceback
from typing import List, Optional
import requests
from dotenv import load_dotenv
from mcp.server.fastmcp import FastMCP
from pydantic import BaseModel, Field
from aworld.logs.util import logger
# Initialize MCP server
mcp = FastMCP("search-server")
# Base search result model that all providers will use
class SearchResult(BaseModel):
"""Base search result model with common fields"""
id: str
title: str
url: str
snippet: str
source: str # Which search engine provided this result
class GoogleSearchResult(SearchResult):
"""Google-specific search result model"""
displayLink: str = ""
formattedUrl: str = ""
htmlSnippet: str = ""
htmlTitle: str = ""
kind: str = ""
link: str = ""
class SearchResponse(BaseModel):
"""Unified search response model"""
query: str
results: List[SearchResult]
count: int
source: str
error: Optional[str] = None
@mcp.tool(description="Search the web using Google Custom Search API.")
def mcpsearchgoogle(
query: str = Field(..., description="The search query string."),
num_results: int = Field(
10, description="Number of search results to return (default 10)."
),
safe_search: bool = Field(
True, description="Whether to enable safe search filtering."
),
language: str = Field("en", description="Language code for search results."),
country: str = Field("us", description="Country code for search results."),
) -> str:
"""
Search the web using Google Custom Search API.
Requires GOOGLE_API_KEY and GOOGLE_CSE_ID environment variables to be set.
"""
try:
api_key = os.environ.get("GOOGLE_API_KEY")
cse_id = os.environ.get("GOOGLE_CSE_ID")
if not api_key:
raise ValueError("GOOGLE_API_KEY environment variable not set")
if not cse_id:
raise ValueError("GOOGLE_CSE_ID environment variable not set")
# Ensure num_results is within valid range
num_results = max(1, num_results)
# Build the Google Custom Search API URL
url = "https://www.googleapis.com/customsearch/v1"
params = {
"key": api_key,
"cx": cse_id,
"q": query,
"num": num_results,
"safe": "active" if safe_search else "off",
"hl": language,
"gl": country,
}
logger.info(f"Google search starts for query: {query}")
response = requests.get(url, params=params, timeout=10)
response.raise_for_status()
data = response.json()
search_results = []
if "items" in data:
for i, item in enumerate(data["items"]):
result = GoogleSearchResult(
id=f"google-{i}",
title=item.get("title", ""),
url=item.get("link", ""),
snippet=item.get("snippet", ""),
source="google",
displayLink=item.get("displayLink", ""),
formattedUrl=item.get("formattedUrl", ""),
htmlSnippet=item.get("htmlSnippet", ""),
htmlTitle=item.get("htmlTitle", ""),
kind=item.get("kind", ""),
link=item.get("link", ""),
)
search_results.append(result)
return SearchResponse(
query=query,
results=search_results,
count=len(search_results),
source="google",
).model_dump_json()
except Exception as e:
logger.error(f"Google search error: {traceback.format_exc()}")
return SearchResponse(
query=query, results=[], count=0, source="google", error=str(e)
).model_dump_json()
def main():
load_dotenv()
print("Starting Search MCP Server...", file=sys.stderr)
mcp.run(transport="stdio")
# Make the module callable
def __call__():
"""
Make the module callable for uvx.
This function is called when the module is executed directly.
"""
main()
sys.modules[__name__].__call__ = __call__
# Run the server when the script is executed directly
if __name__ == "__main__":
main()
@@ -0,0 +1,193 @@
import asyncio
import json
import os
import tempfile
from typing import List, Optional, Tuple
from urllib.parse import urlparse
import requests
from mcp.server import FastMCP
from aworld.logs.util import logger
def get_mime_type(file_path: str, default_mime: Optional[str] = None) -> str:
"""
Detect MIME type of a file using python-magic if available,
otherwise fallback to extension-based detection.
Args:
file_path: Path to the file
default_mime: Default MIME type to return if detection fails
Returns:
str: Detected MIME type
"""
# Try using python-magic for accurate MIME type detection
try:
# mime = magic.Magic(mime=True)
# return mime.from_file(file_path)
return "audio/mpeg"
except (AttributeError, IOError):
# Fallback to extension-based detection
extension_mime_map = {
# Audio formats
".mp3": "audio/mpeg",
".wav": "audio/wav",
".ogg": "audio/ogg",
".m4a": "audio/mp4",
".flac": "audio/flac",
# Image formats
".jpg": "image/jpeg",
".jpeg": "image/jpeg",
".png": "image/png",
".gif": "image/gif",
".webp": "image/webp",
".bmp": "image/bmp",
".tiff": "image/tiff",
# Video formats
".mp4": "video/mp4",
".avi": "video/x-msvideo",
".mov": "video/quicktime",
".mkv": "video/x-matroska",
".webm": "video/webm",
}
ext = os.path.splitext(file_path)[1].lower()
return extension_mime_map.get(ext, default_mime or "application/octet-stream")
def is_url(path_or_url: str) -> bool:
"""
Check if the given string is a URL.
Args:
path_or_url: String to check
Returns:
bool: True if the string is a URL, False otherwise
"""
parsed = urlparse(path_or_url)
return bool(parsed.scheme and parsed.netloc)
def get_file_from_source(
source: str,
allowed_mime_prefixes: List[str] = None,
max_size_mb: float = 100.0,
timeout: int = 60,
type: str = "image",
) -> Tuple[str, str, bytes]:
"""
Unified function to get file content from a URL or local path with validation.
Args:
source: URL or local file path
allowed_mime_prefixes: List of allowed MIME type prefixes (e.g., ['audio/', 'video/'])
max_size_mb: Maximum allowed file size in MB
timeout: Timeout for URL requests in seconds
Returns:
Tuple[str, str, bytes]: (file_path, mime_type, file_content)
- For URLs, file_path will be a temporary file path
- For local files, file_path will be the original path
Raises:
ValueError: When file doesn't exist, exceeds size limit, or has invalid MIME type
IOError: When file cannot be read
requests.RequestException: When URL request fails
"""
max_size_bytes = max_size_mb * 1024 * 1024
temp_file = None
try:
if is_url(source):
# Handle URL
logger.info(f"Downloading file from URL: {source}")
response = requests.get(source, stream=True, timeout=timeout)
response.raise_for_status()
# Check Content-Length if available
content_length = response.headers.get("Content-Length")
if content_length and int(content_length) > max_size_bytes:
raise ValueError(f"File size exceeds limit of {max_size_mb}MB")
# Create a temporary file
temp_file = tempfile.NamedTemporaryFile(delete=False)
file_path = temp_file.name
# Download content in chunks to avoid memory issues
content = bytearray()
downloaded_size = 0
for chunk in response.iter_content(chunk_size=8192):
downloaded_size += len(chunk)
if downloaded_size > max_size_bytes:
raise ValueError(f"File size exceeds limit of {max_size_mb}MB")
temp_file.write(chunk)
content.extend(chunk)
temp_file.close()
# Get MIME type
if type == "audio":
mime_type = "audio/mpeg"
elif type == "image":
mime_type = "image/jpeg"
elif type == "video":
mime_type = "video/mp4"
# mime_type = get_mime_type(file_path)
# For URLs where magic fails, try to use Content-Type header
if mime_type == "application/octet-stream":
content_type = response.headers.get("Content-Type", "").split(";")[0]
if content_type:
mime_type = content_type
else:
# Handle local file
file_path = os.path.abspath(source)
# Check if file exists
if not os.path.exists(file_path):
raise ValueError(f"File not found: {file_path}")
# Check file size
file_size = os.path.getsize(file_path)
if file_size > max_size_bytes:
raise ValueError(f"File size exceeds limit of {max_size_mb}MB")
# Get MIME type
if type == "audio":
mime_type = "audio/mpeg"
elif type == "image":
mime_type = "image/jpeg"
elif type == "video":
mime_type = "video/mp4"
# mime_type = get_mime_type(file_path)
# Read file content
with open(file_path, "rb") as f:
content = f.read()
# Validate MIME type if allowed_mime_prefixes is provided
if allowed_mime_prefixes:
if not any(
mime_type.startswith(prefix) for prefix in allowed_mime_prefixes
):
allowed_types = ", ".join(allowed_mime_prefixes)
raise ValueError(
f"Invalid file type: {mime_type}. Allowed types: {allowed_types}"
)
return file_path, mime_type, content
except Exception as e:
# Clean up temporary file if an error occurs
if temp_file and os.path.exists(temp_file.name):
os.unlink(temp_file.name)
raise e
if __name__ == "__main__":
mcp_tools = []
logger.success(f"{json.dumps(mcp_tools, indent=4, ensure_ascii=False)}")
@@ -0,0 +1,484 @@
# pylint: disable=E1101
import base64
import os
import sys
import traceback
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple
import cv2
import numpy as np
from dotenv import load_dotenv
from mcp.server.fastmcp import FastMCP
from openai import OpenAI
from pydantic import Field
from aworld.logs.util import logger
from mcp_servers.utils import get_file_from_source
client = OpenAI(api_key=os.getenv("LLM_API_KEY"), base_url=os.getenv("LLM_BASE_URL"))
# Initialize MCP server
mcp = FastMCP("Video Server")
@dataclass
class KeyframeResult:
"""Result of keyframe extraction from a video.
Attributes:
frame_paths: List of file paths to the saved keyframes
frame_timestamps: List of timestamps (in seconds) corresponding to each frame
output_directory: Directory where frames were saved
frame_count: Number of frames extracted
success: Whether the extraction was successful
error_message: Error message if extraction failed, None otherwise
"""
frame_paths: List[str]
frame_timestamps: List[float]
output_directory: str
frame_count: int
success: bool
error_message: Optional[str] = None
VIDEO_ANALYZE = (
"Input is a sequence of video frames. Given user's task: {task}. "
"analyze the video content following these steps:\n"
"1. Temporal sequence understanding\n"
"2. Motion and action analysis\n"
"3. Scene context interpretation\n"
"4. Object and person tracking\n"
"Return a json string with the following format: "
'{{"video_analysis_result": "analysis result given task and video frames"}}'
)
VIDEO_EXTRACT_SUBTITLES = (
"Input is a sequence of video frames. "
"Extract all subtitles (if present) in the video. "
"Return a json string with the following format: "
'{"video_subtitles": "extracted subtitles from video"}'
)
VIDEO_SUMMARIZE = (
"Input is a sequence of video frames. "
"Summarize the main content of the video. "
"Include key points, main topics, and important visual elements. "
"Return a json string with the following format: "
'{"video_summary": "concise summary of the video content"}'
)
def get_video_frames(
video_source: str,
sample_rate: int = 2,
start_time: float = 0,
end_time: float = None,
) -> List[Dict[str, Any]]:
"""
Get frames from video with given sample rate using robust file handling
Args:
video_source: Path or URL to the video file
sample_rate: Number of frames to sample per second
start_time: Start time of the video segment in seconds (default: 0)
end_time: End time of the video segment in seconds (default: None, meaning the end of the video)
Returns:
List[Dict[str, Any]]: List of dictionaries containing frame data and timestamp
Raises:
ValueError: When video file cannot be opened or is not a valid video
"""
try:
# Get file with validation (only video files allowed)
file_path, _, _ = get_file_from_source(
video_source,
allowed_mime_prefixes=["video/"],
max_size_mb=2500.0, # 2500MB limit for videos
type="video", # Specify type as video to handle video files
)
# Open video file
video = cv2.VideoCapture(file_path)
if not video.isOpened():
raise ValueError(f"Could not open video file: {file_path}")
fps = video.get(cv2.CAP_PROP_FPS)
frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
video_duration = frame_count / fps # 30s
if end_time is None:
end_time = video_duration
if start_time > end_time:
raise ValueError("Start time cannot be greater than end time.")
if start_time < 0:
start_time = 0
if end_time > video_duration:
end_time = video_duration
start_frame = int(start_time * fps)
end_frame = int(end_time * fps)
all_frames = []
frames = []
# Calculate frame interval based on sample rate
frame_interval = max(1, int(fps / sample_rate))
# Set the video capture to the start frame
video.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
for i in range(start_frame, end_frame):
ret, frame = video.read()
if not ret:
break
# Convert frame to JPEG format
_, buffer = cv2.imencode(".jpg", frame)
frame_data = base64.b64encode(buffer).decode("utf-8")
# Add data URL prefix for JPEG image
frame_data = f"data:image/jpeg;base64,{frame_data}"
all_frames.append({"data": frame_data, "time": i / fps})
for i in range(0, len(all_frames), frame_interval):
frames.append(all_frames[i])
video.release()
# Clean up temporary file if it was created for a URL
if file_path != os.path.abspath(video_source) and os.path.exists(file_path):
os.unlink(file_path)
if not frames:
raise ValueError(f"Could not extract any frames from video: {video_source}")
return frames
except Exception as e:
logger.error(f"Error extracting frames from {video_source}: {str(e)}")
raise
def create_video_content(
prompt: str, video_frames: List[Dict[str, Any]]
) -> List[Dict[str, Any]]:
"""Create uniform video format for querying llm."""
content = [{"type": "text", "text": prompt}]
content.extend(
[
{"type": "image_url", "image_url": {"url": frame["data"]}}
for frame in video_frames
]
)
return content
@mcp.tool(description="Analyze the video content by the given question.")
def mcp_analyze_video(
video_url: str = Field(description="The input video in given filepath or url."),
question: str = Field(description="The question to analyze."),
sample_rate: int = Field(default=2, description="Sample n frames per second."),
start_time: float = Field(
default=0, description="Start time of the video segment in seconds."
),
end_time: float = Field(
default=None, description="End time of the video segment in seconds."
),
) -> str:
"""analyze the video content by the given question."""
try:
video_frames = get_video_frames(video_url, sample_rate, start_time, end_time)
logger.info(f"---len video_frames:{len(video_frames)}")
interval = 20
frame_nums = 30
all_res = []
for i in range(0, len(video_frames), interval):
inputs = []
cur_frames = video_frames[i : i + frame_nums]
content = create_video_content(
VIDEO_ANALYZE.format(task=question), cur_frames
)
inputs.append({"role": "user", "content": content})
try:
response = client.chat.completions.create(
model=os.getenv("LLM_MODEL_NAME"),
messages=inputs,
temperature=0,
)
cur_video_analysis_result = response.choices[0].message.content
except Exception:
cur_video_analysis_result = ""
all_res.append(
f"result of video part {int(i / interval + 1)}: {cur_video_analysis_result}"
)
if i + frame_nums >= len(video_frames):
break
video_analysis_result = "\n".join(all_res)
except (ValueError, IOError, RuntimeError):
video_analysis_result = ""
logger.error(f"video_analysis-Execute error: {traceback.format_exc()}")
logger.info(
f"---get_analysis_by_video-video_analysis_result:{video_analysis_result}"
)
return video_analysis_result
@mcp.tool(description="Extract subtitles from the video.")
def mcp_extract_video_subtitles(
video_url: str = Field(description="The input video in given filepath or url."),
sample_rate: int = Field(default=2, description="Sample n frames per second."),
start_time: float = Field(
default=0, description="Start time of the video segment in seconds."
),
end_time: float = Field(
default=None, description="End time of the video segment in seconds."
),
) -> str:
"""extract subtitles from the video."""
inputs = []
try:
video_frames = get_video_frames(video_url, sample_rate, start_time, end_time)
content = create_video_content(VIDEO_EXTRACT_SUBTITLES, video_frames)
inputs.append({"role": "user", "content": content})
response = client.chat.completions.create(
model=os.getenv("LLM_MODEL_NAME"),
messages=inputs,
temperature=0,
)
video_subtitles = response.choices[0].message.content
except (ValueError, IOError, RuntimeError):
video_subtitles = ""
logger.error(f"video_subtitles-Execute error: {traceback.format_exc()}")
logger.info(f"---get_subtitles_from_video-video_subtitles:{video_subtitles}")
return video_subtitles
@mcp.tool(description="Summarize the main content of the video.")
def mcp_summarize_video(
video_url: str = Field(description="The input video in given filepath or url."),
sample_rate: int = Field(default=2, description="Sample n frames per second."),
start_time: float = Field(
default=0, description="Start time of the video segment in seconds."
),
end_time: float = Field(
default=None, description="End time of the video segment in seconds."
),
) -> str:
"""summarize the main content of the video."""
try:
video_frames = get_video_frames(video_url, sample_rate, start_time, end_time)
logger.info(f"---len video_frames:{len(video_frames)}")
interval = 490
frame_nums = 500
all_res = []
for i in range(0, len(video_frames), interval):
inputs = []
cur_frames = video_frames[i : i + frame_nums]
content = create_video_content(VIDEO_SUMMARIZE, cur_frames)
inputs.append({"role": "user", "content": content})
try:
response = client.chat.completions.create(
model=os.getenv("LLM_MODEL_NAME"),
messages=inputs,
temperature=0,
)
logger.info(f"---response:{response}")
cur_video_summary = response.choices[0].message.content
except Exception:
cur_video_summary = ""
all_res.append(
f"summary of video part {int(i / interval + 1)}: {cur_video_summary}"
)
logger.info(
f"summary of video part {int(i / interval + 1)}: {cur_video_summary}"
)
video_summary = "\n".join(all_res)
except (ValueError, IOError, RuntimeError):
video_summary = ""
logger.error(f"video_summary-Execute error: {traceback.format_exc()}")
logger.info(f"---get_summary_from_video-video_summary:{video_summary}")
return video_summary
@mcp.tool(description="Extract key frames around the target time with scene detection")
def get_video_keyframes(
video_path: str = Field(description="The input video in given filepath or url."),
target_time: int = Field(
description=(
"The specific time point for extraction,"
" centered within the window_size argument,"
" the unit is of second."
)
),
window_size: int = Field(
default=5,
description="The window size for extraction, the unit is of second.",
),
cleanup: bool = Field(
default=False,
description="Whether to delete the original video file after processing.",
),
output_dir: str = Field(
default=os.getenv("FILESYSTEM_SERVER_WORKDIR", "./keyframes"),
description="Directory where extracted frames will be saved.",
),
) -> KeyframeResult:
"""Extract key frames around the target time with scene detection.
This function extracts frames from a video file around a specific time point,
using scene detection to identify significant changes between frames. Only frames
with substantial visual differences are saved, reducing redundancy.
Args:
video_path: Path or URL to the video file
target_time: Specific time point (in seconds) to extract frames around
window_size: Time window (in seconds) centered on target_time
cleanup: Whether to delete the original video file after processing
output_dir: Directory where extracted frames will be saved
Returns:
KeyframeResult: A dataclass containing paths to saved frames, timestamps,
and metadata about the extraction process
Raises:
Exception: Exceptions are caught internally and reported in the result
"""
def save_frames(frames, frame_times, output_dir) -> Tuple[List[str], List[float]]:
"""Save extracted frames to disk"""
os.makedirs(output_dir, exist_ok=True)
saved_paths = []
saved_timestamps = []
for _, (frame, timestamp) in enumerate(zip(frames, frame_times)):
filename = f"{output_dir}/frame_{timestamp:.2f}s.jpg"
os.makedirs(output_dir, exist_ok=True)
saved_paths = []
saved_timestamps = []
for _, (frame, timestamp) in enumerate(zip(frames, frame_times)):
filename = f"{output_dir}/frame_{timestamp:.2f}s.jpg"
cv2.imwrite(filename, frame)
saved_paths.append(filename)
saved_timestamps.append(timestamp)
return saved_paths, saved_timestamps
def extract_keyframes(
video_path, target_time, window_size
) -> Tuple[List[Any], List[float]]:
"""Extract key frames around the target time with scene detection"""
cap = cv2.VideoCapture(video_path)
fps = cap.get(cv2.CAP_PROP_FPS)
# Calculate frame numbers for the time window
start_frame = int((target_time - window_size / 2) * fps)
end_frame = int((target_time + window_size / 2) * fps)
frames = []
frame_times = []
# Set video position to start_frame
cap.set(cv2.CAP_PROP_POS_FRAMES, max(0, start_frame))
prev_frame = None
while cap.isOpened():
frame_pos = cap.get(cv2.CAP_PROP_POS_FRAMES)
if frame_pos >= end_frame:
break
ret, frame = cap.read()
if not ret:
break
# Convert frame to grayscale for scene detection
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
# If this is the first frame, save it
if prev_frame is None:
frames.append(frame)
frame_times.append(frame_pos / fps)
else:
# Calculate difference between current and previous frame
diff = cv2.absdiff(gray, prev_frame)
mean_diff = np.mean(diff)
# If significant change detected, save frame
if mean_diff > 20: # Threshold for scene change
frames.append(frame)
frame_times.append(frame_pos / fps)
prev_frame = gray
cap.release()
return frames, frame_times
try:
# Extract keyframes
frames, frame_times = extract_keyframes(video_path, target_time, window_size)
# Save frames
frame_paths, frame_timestamps = save_frames(frames, frame_times, output_dir)
# Cleanup
if cleanup and os.path.exists(video_path):
os.remove(video_path)
return KeyframeResult(
frame_paths=frame_paths,
frame_timestamps=frame_timestamps,
output_directory=output_dir,
frame_count=len(frame_paths),
success=True,
)
except Exception as e:
error_message = f"Error processing video: {str(e)}"
print(error_message)
return KeyframeResult(
frame_paths=[],
frame_timestamps=[],
output_directory=output_dir,
frame_count=0,
success=False,
error_message=error_message,
)
def main():
load_dotenv()
print("Starting Video MCP Server...", file=sys.stderr)
mcp.run(transport="stdio")
# Make the module callable
def __call__():
"""
Make the module callable for uvx.
This function is called when the module is executed directly.
"""
main()
# Add this for compatibility with uvx
sys.modules[__name__].__call__ = __call__
# Run the server when the script is executed directly
if __name__ == "__main__":
main()
@@ -0,0 +1,279 @@
"""
Youtube Download MCP Server
This module provides MCP server functionality for downloading files from Youtube URLs.
It handles various download scenarios with proper validation, error handling,
and progress tracking.
Key features:
- File downloading from Youtube HTTP/HTTPS URLs
- Download progress tracking
- File validation
- Safe file saving
Main functions:
- mcpyoutubedownload: Downloads files from URLs of Youtube to local filesystem
"""
import os
import sys
import time
import traceback
import urllib.parse
from datetime import datetime
from pathlib import Path
from typing import Optional
from dotenv import load_dotenv
from mcp.server.fastmcp import FastMCP
from pydantic import BaseModel, Field
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.by import By
from aworld.logs.util import logger
mcp = FastMCP("youtube-server")
_default_driver_path = os.environ.get(
"CHROME_DRIVER_PATH",
os.path.expanduser("~/Downloads/chromedriver-mac-arm64/chromedriver"),
)
class YoutubeDownloadResults(BaseModel):
"""Download result model with file information"""
file_path: str
file_name: str
file_size: int
content_type: Optional[str] = None
success: bool
error: Optional[str] = None
@mcp.tool(
description="Download the youtube file from the URL and save to the local filesystem."
)
def download_youtube_files(
url: str = Field(
description="The URL of youtube file to download. Must be a String."
),
output_dir: str = Field(
"/tmp/mcp_downloads",
description="Directory to save the downloaded files (default: /tmp/mcp_downloads).",
),
timeout: int = Field(
180, description="Download timeout in seconds (default: 180)."
),
) -> str:
"""Download the youtube file from the URL and save to the local filesystem.
Args:
url: The URL of youtube file to download, must be a String
output_dir: Directory to save the downloaded files
timeout: Download timeout in seconds
Returns:
JSON string with download results information
"""
# Handle Field objects if they're passed directly
if hasattr(url, "default") and not isinstance(url, str):
url = url.default
if hasattr(output_dir, "default") and not isinstance(output_dir, str):
output_dir = output_dir.default
if hasattr(timeout, "default") and not isinstance(timeout, int):
timeout = timeout.default
def _get_youtube_content(url: str, output_dir: str, timeout: int) -> None:
"""Use Selenium to download YouTube content via cobalt.tools"""
try:
options = webdriver.ChromeOptions()
options.add_argument("--disable-blink-features=AutomationControlled")
# Set download file default path
prefs = {
"download.default_directory": output_dir,
"download.prompt_for_download": False,
"download.directory_upgrade": True,
"safebrowsing.enabled": True,
}
options.add_experimental_option("prefs", prefs)
# Create WebDriver object and launch Chrome browser
service = Service(executable_path=_default_driver_path)
driver = webdriver.Chrome(service=service, options=options)
logger.info(f"Opening cobalt.tools to download from {url}")
# Open target webpage
driver.get("https://cobalt.tools/")
# Wait for page to load
time.sleep(5)
# Find input field and enter YouTube link
input_field = driver.find_element(By.ID, "link-area")
input_field.send_keys(url)
time.sleep(5)
# Find download button and click
download_button = driver.find_element(By.ID, "download-button")
download_button.click()
time.sleep(5)
try:
# Handle bot detection popup
driver.find_element(
By.CLASS_NAME,
"button.elevated.popup-button.undefined.svelte-nnawom.active",
).click()
except Exception as e:
logger.warning(f"Bot detection handling: {str(e)}")
# Wait for download to complete
cnt = 0
while (
len(os.listdir(output_dir)) == 0
or os.listdir(output_dir)[0].split(".")[-1] == "crdownload"
):
time.sleep(3)
cnt += 3
if cnt >= timeout:
logger.warning(f"Download timeout after {timeout} seconds")
break
logger.info("Download process completed")
except Exception as e:
logger.error(f"Error during YouTube content download: {str(e)}")
raise
finally:
# Close browser
if "driver" in locals():
driver.quit()
def _download_single_file(
url: str, output_dir: str, filename: str, timeout: int
) -> str:
"""Download a single file from URL and save it to the local filesystem."""
try:
# Validate URL
if not url.startswith(("http://", "https://")):
raise ValueError(
"Invalid URL format. URL must start with http:// or https://"
)
# Create output directory if it doesn't exist
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
# Determine filename if not provided
if not filename:
filename = os.path.basename(urllib.parse.urlparse(url).path)
if not filename:
filename = "downloaded_file"
filename += "_" + datetime.now().strftime("%Y%m%d_%H%M%S")
file_path = Path(os.path.join(output_path, filename))
file_path.mkdir(parents=True, exist_ok=True)
logger.info(f"Output path: {file_path}")
# check if video already exists with folder: /tmp/mcp_downloads
video_id = url.split("?v=")[-1].split("&")[0] if "?v=" in url else ""
base_path = os.getenv("FILESYSTEM_SERVER_WORKDIR")
# checker function
def find_existing_video(search_dir, video_id):
if not video_id:
return None
for item in os.listdir(search_dir):
item_path = os.path.join(search_dir, item)
if os.path.isfile(item_path) and video_id in item:
return item_path
elif os.path.isdir(item_path):
found = find_existing_video(item_path, video_id)
if found:
return found
return None
existing_file = find_existing_video(base_path, video_id)
if existing_file:
result = YoutubeDownloadResults(
file_path=existing_file,
file_name=os.path.basename(existing_file),
file_size=os.path.getsize(existing_file),
content_type="mp4",
success=True,
error=None,
)
logger.info(
f"Found {video_id} is already downloaded in: {existing_file}"
)
return result.model_dump_json()
logger.info(f"Downloading file from {url} to {file_path}")
_get_youtube_content(url, str(file_path), timeout)
# Check if download was successful
if len(os.listdir(file_path)) == 0:
raise FileNotFoundError("No files were downloaded")
download_file = os.path.join(file_path, os.listdir(file_path)[0])
# Get actual file size
actual_size = os.path.getsize(download_file)
logger.success(f"File downloaded successfully to {download_file}")
# Create result
result = YoutubeDownloadResults(
file_path=download_file,
file_name=os.listdir(file_path)[0],
file_size=actual_size,
content_type="mp4",
success=True,
error=None,
)
return result.model_dump_json()
except Exception as e:
error_msg = str(e)
logger.error(f"Download error: {traceback.format_exc()}")
result = YoutubeDownloadResults(
file_path="",
file_name="",
file_size=0,
content_type=None,
success=False,
error=error_msg,
)
return result.model_dump_json()
result_json = _download_single_file(url, output_dir, "", timeout)
result = YoutubeDownloadResults.model_validate_json(result_json)
return result.model_dump_json()
def main():
load_dotenv()
print("Starting YoutubeDownload MCP Server...", file=sys.stderr)
mcp.run(transport="stdio")
# Make the module callable
def __call__():
"""
Make the module callable for uvx.
This function is called when the module is executed directly.
"""
main()
sys.modules[__name__].__call__ = __call__
# Run the server when the script is executed directly
if __name__ == "__main__":
main()
@@ -0,0 +1,44 @@
--index-url https://mirrors.aliyun.com/pypi/simple/
fastapi==0.111.0
uvicorn[standard]==0.23.1
pydantic==2.9.2
python-multipart==0.0.9
python-socketio
grpcio
passlib==1.7.4
passlib[bcrypt]
PyJWT[crypto]
requests==2.32.4
aiohttp==3.9.5
httpx
datasets==3.3.2
executing
flask
openpyxl
selenium==4.32.0
fitz==0.0.1.dev2
tabulate==0.9.0
frontend
tools
PyPDF2
html2text
xmltodict
docx2markdown
python-pptx
browser_use
oss2
prometheus_client~=0.21.1
opentelemetry-sdk~=1.32.1
opentelemetry-api~=1.32.1
opentelemetry-exporter-otlp~=1.32.1
opentelemetry-instrumentation-system-metrics~=0.53b1
e2b_code_interpreter
sqlalchemy~=2.0.40
psycopg2-binary==2.9.9
bcrypt==4.3.0
+45
View File
@@ -0,0 +1,45 @@
#!/usr/bin/env bash
PORT="${PORT:-9099}"
HOST="${HOST:-0.0.0.0}"
# Default value for PIPELINES_DIR
PIPELINES_DIR=${PIPELINES_DIR:-./aworldspace/agents}
UVICORN_LOOP="${UVICORN_LOOP:-auto}"
# OSS mount configuration - read from environment variables
if [ -n "$OSS_BUCKET" ] && [ -n "$OSS_AK_ID" ] && [ -n "$OSS_AK_SECRET" ]; then
echo "Configuring OSS mount..."
# Create OSS credentials file
echo "${OSS_BUCKET}:${OSS_AK_ID}:${OSS_AK_SECRET}" >> /etc/passwd-ossfs
chmod 640 /etc/passwd-ossfs
# Create mount point directories if they don't exist
mkdir -p /app/logs
mkdir -p /app/trace_data
mkdir -p /app/aworldspace/datasets
# Mount OSS directories
echo "Mounting OSS directories..."
if [ -n "$OSS_REGION_URL" ] && [ -n "$OSS_BUCKET_URL" ]; then
# Use custom region and URL
ossfs ${OSS_BUCKET}:/aworld/logs /app/logs -odirect_read -ononempty -oregion=${OSS_REGION_URL} -ourl=${OSS_BUCKET_URL} &
ossfs ${OSS_BUCKET}:/aworld/trace_data /app/trace_data -odirect_read -ononempty -oregion=${OSS_REGION_URL} -ourl=${OSS_BUCKET_URL} &
ossfs ${OSS_BUCKET}:/aworld/datasets /app/aworldspace/datasets -odirect_read -ononempty -oregion=${OSS_REGION_URL} -ourl=${OSS_BUCKET_URL} &
else
# Use default configuration
ossfs ${OSS_BUCKET}:/aworld/logs /app/logs -odirect_read -ononempty &
ossfs ${OSS_BUCKET}:/aworld/trace_data /app/trace_data -odirect_read -ononempty &
ossfs ${OSS_BUCKET}:/aworld/datasets /app/aworldspace/datasets -odirect_read -ononempty &
fi
# Wait for mount to complete
sleep 2
echo "OSS mount configuration completed"
else
echo "OSS configuration incomplete, skipping OSS mount"
fi
uvicorn main:app --host "$HOST" --port "$PORT" --forwarded-allow-ips '*' --loop "$UVICORN_LOOP"