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
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:
@@ -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.
|
||||
|
||||

|
||||
|
||||
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 can’t 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()
|
||||
+794
@@ -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))
|
||||
+28
@@ -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
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user