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

This commit is contained in:
2026-08-20 13:12:50 +00:00
commit b119135836
10275 changed files with 3284984 additions and 0 deletions
@@ -0,0 +1,197 @@
# GAIA Agent Setup Guide
## 1. Overview
This guide will help you set up and run the GAIA agent for the AWorld framework. GAIA is a benchmark dataset for evaluating AI agents' capabilities.
## 2. Prerequisites
### 1. System Requirements
- **Operating System**: macOS or Linux (Windows not fully tested)
- **Node.js**: Version 22 LTS with npm
- **Conda**: For environment management
### 2. Required Software
- `libmagic1` - File type detection
- `libreoffice` - Document processing
- `ffmpeg` - Media processing
## 3. Installation Steps
### 1. Clone the Repository
```bash
git clone https://github.com/inclusionAI/AWorld.git
cd AWorld
```
### 2. Set Up Conda Environment
Create and activate a dedicated Conda environment for GAIA:
```bash
conda env create -f examples/gaia/aworld-gaia.yml
conda activate aworld-gaia
```
> **Note**: If you don't have Conda installed, download Miniconda from [here](https://www.anaconda.com/docs/getting-started/miniconda/install).
### 3. Install AWorld Framework
Install the AWorld framework and build the web UI:
```bash
# Install PDF processing dependencies
pip install "marker-pdf[full]" --no-deps
# Build web UI
sh -c "cd aworld/cmd/web/webui && npm install && npm run build"
# Install AWorld
python setup.py install
```
### 4. Install MCP Tool Dependencies
#### Install Playwright
```bash
playwright install chromium --with-deps --no-shell
```
#### Install System Dependencies
**For macOS:**
```bash
brew install libmagic
brew install ffmpeg
brew install --cask libreoffice
```
> **Note**: Install Homebrew from [brew.sh](https://brew.sh/) if not already installed.
**For Linux:**
```bash
apt-get install -y --no-install-recommends libmagic1 libreoffice ffmpeg
```
### 5. Prepare GAIA Dataset
Download the GAIA dataset from Hugging Face:
```bash
git clone git@hf.co:datasets/gaia-benchmark/GAIA examples/gaia/GAIA
```
> **⚠️ Important**:
> - You need to configure Hugging Face SSH keys to access the GAIA repository
> - The dataset path will be used as the `GAIA_DATASET_PATH` variable in your `.env` file
### 6. Configure Environment Variables
Create the environment configuration file:
```bash
cp examples/gaia/cmd/agent_deploy/gaia_agent/.env.template examples/gaia/cmd/agent_deploy/gaia_agent/.env
```
Edit the `.env` file and replace all `{YOUR_CONFIG}` placeholders with your actual configuration values.
## 3. Running the GAIA Agent
### 1. Web UI Interface
Start the GAIA agent web interface:
```bash
cd examples/gaia/cmd && aworld web
```
### 2. Command Line Interface
Run GAIA tasks using the command line interface:
```bash
python examples/gaia/run.py --split validation --q c61d22de-5f6c-4958-a7f6-5e9707bd3466
```
## 4. Command Line Arguments
### 1. Required Arguments
| Argument | Type | Description | Default | Example |
|----------|------|-------------|---------|---------|
| `--split` | str | Dataset split to use | `validation` | `--split test` |
| `--start` | int | Start index of the dataset | `0` | `--start 10` |
| `--end` | int | End index of the dataset | `165` | `--end 100` |
### 2. Optional Arguments
| Argument | Type | Description | Example |
|----------|------|-------------|---------|
| `--q` | str | Specific question index (overrides start/end) | `--q 0-0-0-0-0` |
| `--skip` | flag | Skip previously processed questions | `--skip` |
| `--blacklist_file_path` | str | Path to blacklist file | `--blacklist_file_path blacklist.txt` |
## 5. Usage Examples
### 1. Basic Usage
Process a range of questions in the validation split:
```bash
python examples/gaia/run.py --start 10 --end 50 --split validation
```
### 2. Process Specific Question
Run a single question by its index:
```bash
python examples/gaia/run.py --q 0-0-0-0-0
```
### 3. Skip Processed Questions
Process questions while skipping previously completed ones:
```bash
python examples/gaia/run.py --start 10 --end 50 --skip
```
### 4. Use Blacklist
Skip questions listed in a blacklist file:
```bash
python examples/gaia/run.py --start 10 --end 50 --blacklist_file_path blacklist.txt
```
## 6. Expected Output
When running successfully, you should see logs similar to:
```
YYYY-MM-DD HH:MM:SS - root - INFO - Agent answer: egalitarian
YYYY-MM-DD HH:MM:SS - root - INFO - Correct answer: egalitarian
YYYY-MM-DD HH:MM:SS - examples.gaia.utils - INFO - Evaluating egalitarian as a string.
YYYY-MM-DD HH:MM:SS - root - INFO - Question 0 Correct!
```
## 7. Troubleshooting
### 1. Common Issues
1. **Dataset Access Problems**
- Verify your Hugging Face SSH keys are correctly configured
- Ensure you have access to the GAIA repository
2. **Installation Issues**
- Set up a pip mirror if necessary for faster downloads
- Ensure all system dependencies are properly installed
3. **Environment Issues**
- Make sure you're using the correct Conda environment (`aworld-gaia`)
- Verify Node.js version is 22 LTS
### 2. Getting Help
If you encounter issues not covered in this guide:
- Check the project's main documentation
- Review the error logs for specific error messages
- Ensure all prerequisites are properly installed
@@ -0,0 +1,82 @@
# GAIA Agent Guard Functionality Setup
## 1. Overview
This guide covers the additional setup required for the enhanced guard agent in the GAIA agent. For basic installation and setup, please refer to [README.md](README.md).
## 2. Setting Up Guard Functionality
### 2.1 MCP Configuration
To enable the guard functionality, ensure the following MCP server is registered in your `mcp.json` file:
```json
{
"mcpServers": {
"maneuvering": {
"command": "python",
"args": [
"-m",
"examples.gaia.mcp_collections.intelligence.guard"
],
"env": {},
"client_session_timeout_seconds": 9999.0
}
}
}
```
### 2.2 Configure Environment Variables
Add the guard_llm API key to your `.env` file:
```bash
# Add this line to examples/gaia/cmd/agent_deploy/gaia_agent/.env
GUARD_LLM_API_KEY=your_guard_llm_api_key_here
```
### 2.3 Update Prompt Configuration
Replace the content of `prompt.py` with the enhanced version from `prompt_w_guard.py`:
```bash
cp examples/gaia/prompt_w_guard.py examples/gaia/prompt.py
```
### 2.4 Configure Task Subset Processing
To run specific subsets of GAIA tasks, add the following code to `run.py`:
```python
# load task subset from subset.txt
subset_file_path = Path(__file__).parent / "subset.txt"
if subset_file_path.exists():
with open(subset_file_path, "r", encoding="utf-8") as f:
task_subset = set(line.strip() for line in f if line.strip())
logging.info(f"Loaded {len(task_subset)} task IDs from subset.txt")
else:
task_subset = set() # Empty set if file doesn't exist
logging.warning("subset.txt file not found, using empty task subset")
```
And add the filtering logic in the main processing loop:
```python
# only process tasks that are in the subset
if dataset_i["task_id"] not in task_subset:
continue
```
## 3. Configuration Summary
To enable guard functionality, ensure you have:
1.**MCP Configuration**: `maneuvering` server registered in `mcp.json`
2.**Environment Variables**: `GUARD_LLM_API_KEY` added to `.env` file
3.**Prompt Update**: `prompt_w_guard.py` content copied to `prompt.py`
4.**Subset Processing**: Batch processing code added to `run.py`
This setup provides enhanced reasoning capabilities and efficient batch processing for GAIA task evaluation.
@@ -0,0 +1,8 @@
# coding: utf-8
# Copyright (c) 2025 inclusionAI.
from aworld.utils import import_packages
import_packages(["dotenv"])
from dotenv import load_dotenv
load_dotenv()
@@ -0,0 +1,248 @@
name: aworld-gaia
channels:
- defaults
dependencies:
- bzip2=1.0.8
- ca-certificates=2025.2.25
- expat=2.7.1
- libcxx=17.0.6
- libffi=3.4.4
- ncurses=6.4
- openssl=3.0.16
- pip=25.1
- python=3.12.11
- readline=8.2
- setuptools=78.1.1
- sqlite=3.45.3
- tk=8.6.14
- wheel=0.45.1
- xz=5.6.4
- zlib=1.2.13
- pip:
- aiofiles==24.1.0
- altair==5.5.0
- annotated-types==0.7.0
- anthropic
- anyio==4.9.0
- arxiv==2.2.0
- asttokens==3.0.0
- attrs==25.3.0
- Authlib==1.6.0
- backoff==2.2.1
- beautifulsoup4==4.13.4
- blinker==1.9.0
- Brotli==1.1.0
- browser-use==0.7.1
- bubus==1.2.1
- cachetools==5.5.2
- certifi==2025.6.15
- cffi==1.17.1
- cfgv==3.4.0
- chardet==5.2.0
- charset-normalizer==3.4.2
- chess==1.11.2
- click==8.2.1
- cobble==0.1.4
- cryptography==45.0.4
- cssselect2==0.8.0
- Cython==3.1.2
- decorator==5.2.1
- defusedxml==0.7.1
- distlib==0.3.9
- distro==1.9.0
- dotenv==0.9.9
- durationpy==0.10
- EbookLib==0.18
- einops==0.8.1
- et_xmlfile==2.0.0
- executing==2.2.0
- faiss-cpu==1.11.0
- feedparser==6.0.11
- filelock==3.18.0
- filetype==1.2.0
- Flask==3.1.1
- fonttools==4.58.4
- fsspec==2025.5.1
- ftfy==6.3.1
- gitdb==4.0.12
- GitPython==3.1.44
- google-ai-generativelanguage==0.6.18
- google-api-core==2.25.1
- google-auth==2.40.3
- google-genai==1.21.1
- googleapis-common-protos==1.70.0
- greenlet==3.2.3
- grpcio==1.73.0
- grpcio-status==1.73.0
- h11==0.16.0
- h2==4.2.0
- hf-xet==1.1.5
- hpack==4.1.0
- httpcore==1.0.9
- httpx==0.28.1
- httpx-sse==0.4.1
- huggingface-hub==0.33.0
- hyperframe==6.1.0
- identify==2.6.12
- idna==3.10
- ipython==9.3.0
- ipython_pygments_lexers==1.1.1
- itsdangerous==2.2.0
- jedi==0.19.2
- Jinja2==3.1.6
- jiter==0.10.0
- joblib==1.5.1
- jsonpatch==1.33
- jsonpickle==4.1.1
- jsonpointer==3.0.0
- jsonschema==4.24.0
- jsonschema-specifications==2025.4.1
- kubernetes==32.0.1
- langchain==0.3.25
- langchain-anthropic==0.3.15
- langchain-core==0.3.64
- langchain-deepseek==0.1.3
- langchain-google-genai==2.1.5
- langchain-ollama==0.3.3
- langchain-openai==0.3.21
- langchain-text-splitters==0.3.8
- langsmith==0.3.45
- lxml==5.4.0
- mammoth==1.9.1
- markdown
# - marker-pdf
- markdown-it-py==3.0.0
- markdown2==2.5.3
- markdownify==1.1.0
- MarkupSafe==3.0.2
- matplotlib-inline==0.1.7
- mcp==1.6.0
- mdurl==0.1.2
- mem0ai==0.1.111
- mpmath==1.3.0
- narwhals==1.44.0
- networkx==3.5
- nodeenv==1.9.1
- numpy==2.3.1
- oauthlib==3.3.1
- ollama==0.5.1
- openai
- opencv-python==4.11.0.86
- opencv-python-headless==4.11.0.86
- openpyxl==3.1.5
- orjson==3.10.18
- outcome==1.3.0.post0
- packaging==24.2
- pandas==2.3.0
- parso==0.8.4
- patchright==1.52.5
- pdftext==0.6.3
- pexpect==4.9.0
- pillow==10.4.0
- platformdirs==4.3.8
- playwright==1.52.0
- portalocker==2.10.1
- posthog==5.4.0
- prompt_toolkit==3.0.51
- proto-plus==1.26.1
- protobuf==6.31.1
- psutil==7.0.0
- ptyprocess==0.7.0
- pure_eval==0.2.3
- pyarrow==20.0.0
- pyasn1==0.6.1
- pyasn1_modules==0.4.2
- pycparser==2.22
- pydantic==2.11.5
- pydantic-settings
- pydantic_core
- pyautogui
- rubicon-objc
- pydeck==0.9.1
- pydyf==0.11.0
- pyee==13.0.0
- Pygments==2.19.2
- pypdfium2==4.30.0
- pyperclip==1.9.0
- pyphen==0.17.2
- PySocks==1.7.1
- pytesseract==0.3.13
- python-dateutil==2.9.0.post0
- python-docx==1.2.0
- python-dotenv==1.0.1
- python-pptx==1.0.2
- python-magic==0.4.27
- pytz==2025.2
- pyvis==0.3.2
- PyYAML==6.0.2
- qdrant-client==1.14.3
- RapidFuzz==3.13.0
- referencing==0.36.2
- regex==2024.11.6
- requests==2.32.4
- requests-oauthlib==2.0.0
- requests-toolbelt==1.0.0
- rich==14.0.0
- rpds-py==0.25.1
- rsa==4.9.1
- safetensors==0.5.3
- scikit-learn==1.7.0
- scipy==1.16.0
- screeninfo==0.8.1
- selenium==4.33.0
- sentence-transformers==4.1.0
- setuptools==78.1.1
- sgmllib3k==1.0.0
- shellingham==1.5.4
- six==1.17.0
- smmap==5.0.2
- sniffio==1.3.1
- sortedcontainers==2.4.0
- soupsieve==2.7
- SQLAlchemy==2.0.41
- sse-starlette==2.3.6
- stack-data==0.6.3
- starlette==0.47.1
- streamlit==1.46.0
- surya-ocr==0.14.6
- sympy==1.14.0
- tabulate==0.9.0
- tenacity==8.5.0
- threadpoolctl==3.6.0
- tiktoken==0.9.0
- tinycss2==1.4.0
- tinyhtml5==2.0.0
- tokenizers==0.21.2
- toml==0.10.2
- torch==2.7.1
- tornado==6.5.1
- tqdm==4.67.1
- traitlets==5.14.3
- transformers==4.52.4
- trio==0.30.0
- trio-websocket==0.12.2
- typer==0.16.0
- typing-inspection==0.4.1
- typing_extensions==4.13.2
- tzdata==2025.2
- urllib3==2.4.0
- uuid7==0.1.0
- uvicorn==0.34.3
- virtualenv==20.31.2
- waybackpy==3.0.6
- wcwidth==0.2.13
- weasyprint==63.1
- webencodings==0.5.1
- websocket-client==1.8.0
- websockets==15.0.1
- Werkzeug==3.1.3
- wheel==0.45.1
- wikipedia==1.4.0
- wsproto==1.2.0
- xlsxwriter==3.2.5
- youtube-transcript-api==1.1.0
- zopfli==0.2.3.post1
- zstandard==0.23.0
- httpx[socks]
- ffmpeg-python==0.2.0
prefix: /opt/homebrew/Caskroom/miniconda/base/envs/aworld-gaia
@@ -0,0 +1,42 @@
import logging
import os
from aworld.cmd.data_model import BaseAWorldAgent, ChatCompletionRequest
from examples.gaia.gaia_agent_runner import GaiaAgentRunner
logger = logging.getLogger(__name__)
class AWorldAgent(BaseAWorldAgent):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
os.makedirs(os.path.join(os.getcwd(), "static"), exist_ok=True)
async def run(self, prompt: str = None, request: ChatCompletionRequest = None):
llm_provider = os.getenv("LLM_PROVIDER", "openai")
llm_model_name = os.getenv("LLM_MODEL_NAME")
llm_api_key = os.getenv("LLM_API_KEY")
llm_base_url = os.getenv("LLM_BASE_URL")
llm_temperature = float(os.getenv("LLM_TEMPERATURE", 0.0))
if not llm_model_name or not llm_api_key or not llm_base_url:
raise ValueError(
"LLM_MODEL_NAME, LLM_API_KEY, LLM_BASE_URL must be set in your envrionment variables"
)
runner = GaiaAgentRunner(
llm_provider=llm_provider,
llm_model_name=llm_model_name,
llm_base_url=llm_base_url,
llm_api_key=llm_api_key,
llm_temperature=llm_temperature,
session_id=request.session_id,
)
if prompt is None and request is not None:
prompt = request.messages[-1].content
logger.info(f">>> Gaia Agent: prompt={prompt}, runner={runner}")
async for line in runner.run(prompt):
logger.info(f">>> Gaia Agent Line: {line}")
yield line
@@ -0,0 +1,346 @@
import json
import logging
import os
import re
import subprocess
import sys
import traceback
from typing import AsyncGenerator
import uuid
from aworld.cmd.utils.agent_ui_parser import (
AWorldWebAgentUI,
BaseToolResultParser,
ToolCard,
ToolResultParserFactory,
)
from aworld.config.conf import AgentConfig, TaskConfig
from aworld.agents.llm_agent import Agent
from aworld.core.task import Task
from aworld.output.artifact import ArtifactType
from aworld.output.workspace import WorkSpace
from aworld.runner import Runners
from aworld.output.ui.base import AworldUI
from aworld.output.base import Output, ToolResultOutput
from .utils import (
add_file_path,
load_dataset_meta_dict,
question_scorer,
)
from .prompt import system_prompt
logger = logging.getLogger(__name__)
class GaiaSearchToolResultParser(BaseToolResultParser):
async def parse(self, output: ToolResultOutput, workspace: WorkSpace):
tool_card = ToolCard.from_tool_result(output)
query = ""
try:
args = json.loads(tool_card.arguments)
query = args.get("query")
# aworld search server
if not query:
query = args.get("query_list")
except Exception:
pass
result_items = []
try:
results = json.loads(tool_card.results)
result_items = results.get("message", {}).get("results", [])
# aworld search server return url, not link
if result_items and isinstance(result_items, list):
for item in result_items:
if not item.get("link", None) and item.get("url", None):
item["link"] = item.get("url")
except Exception:
pass
if len(result_items) > 0:
tool_card.results = ""
tool_card.card_type = "tool_call_card_link_list"
tool_card.card_data = {
"title": "🔎 Gaia Search",
"query": query,
"search_items": result_items,
}
artifact_id = str(uuid.uuid4())
await workspace.create_artifact(
artifact_type=ArtifactType.WEB_PAGES,
artifact_id=artifact_id,
content=result_items,
metadata={
"query": query,
},
)
tool_card.artifacts.append(
{
"artifact_type": ArtifactType.WEB_PAGES.value,
"artifact_id": artifact_id,
}
)
return f"""
\n\n**🔎 Gaia Search**\n\n
```tool_card
{json.dumps(tool_card.model_dump(), ensure_ascii=False, indent=2)}
```\n
"""
class CustomToolResultParserFactory(ToolResultParserFactory):
def get_parser(self, tool_type: str, tool_name: str):
if tool_name in ("search_server", "search"):
return GaiaSearchToolResultParser()
return super().get_parser(tool_type, tool_name)
# Module-level flag to ensure dependencies are installed only once per program run
_install_dependencies_flag = False
class GaiaAgentRunner:
"""
Gaia Agent Runner
"""
def _install_dependencies(self):
try:
current_dir = os.path.dirname(os.path.abspath(__file__))
requirements_file = os.path.join(current_dir, "requirements.txt")
if os.path.exists(requirements_file):
logger.info(f"Installing dependencies from {requirements_file}")
subprocess.check_call(
[
sys.executable,
"-m",
"pip",
"install",
"-U",
"-r",
requirements_file,
]
)
subprocess.check_call(
[
sys.executable,
"-m",
"pip",
"install",
"--no-deps",
"-U",
"marker-pdf",
"anthropic==0.46.0",
]
)
logger.info("Dependencies installed successfully")
else:
logger.warning(f"Requirements file not found at {requirements_file}")
except Exception as e:
logger.error(f"Failed to install dependencies: {e}")
def __init__(
self,
llm_provider: str,
llm_model_name: str,
llm_base_url: str,
llm_api_key: str,
llm_temperature: float = 0.0,
mcp_config: dict = None,
session_id: str = None,
):
global _install_dependencies_flag
if not _install_dependencies_flag:
self._install_dependencies()
_install_dependencies_flag = True
self.session_id = session_id or str(uuid.uuid4())
self.agent_config = AgentConfig(
llm_provider=llm_provider,
llm_model_name=llm_model_name,
llm_api_key=llm_api_key,
llm_base_url=llm_base_url,
llm_temperature=llm_temperature,
)
if mcp_config is None:
mcp_path = os.path.join(
os.path.dirname(os.path.abspath(__file__)), "mcp.json"
)
with open(mcp_path, "r") as f:
mcp_config = json.load(f)
logger.info(f"Gaia Agent Runner mcp_config: {mcp_config}")
self.super_agent = Agent(
conf=self.agent_config,
name="gaia_super_agent",
system_prompt=system_prompt,
mcp_config=mcp_config,
mcp_servers=(
os.getenv("GAIA_MCP_SERVERS", "").split(",") if os.getenv("GAIA_MCP_SERVERS", "") else ""
or mcp_config.get("mcpServers", {}).keys()
),
)
self.gaia_dataset_path = os.path.abspath(
os.getenv(
"GAIA_DATASET_PATH",
os.path.join(
os.path.dirname(os.path.abspath(__file__)), "GAIA", "2023"
),
)
)
self.full_dataset = load_dataset_meta_dict(self.gaia_dataset_path)
logger.info(
f"Gaia Agent Runner initialized: super_agent={self.super_agent}, agent_config={self.agent_config}, gaia_dataset_path={self.gaia_dataset_path}, full_dataset={len(self.full_dataset)}"
)
async def run(self, prompt: str):
yield (f"\n### GAIA Agent Start!")
mcp_servers = "\n- ✅ ".join(self.super_agent.mcp_servers)
yield (f"\n```gaia_agent_status\n- ✅ {mcp_servers}\n```\n")
question = None
data_item = None
task_id = None
try:
json_data = json.loads(prompt)
task_id = json_data["task_id"]
data_item = self.full_dataset[task_id]
question = add_file_path(data_item, file_path=self.gaia_dataset_path)[
"Question"
]
yield (
f"\n### Gaia Question\n```gaia_question\n{json.dumps(data_item, indent=2)}\n```\n"
)
except Exception as e:
pass
if not question:
logger.warning(
"Could not find GAIA question for prompt, chat using prompt directly!"
)
yield (f"\n{prompt}\n")
question = prompt
try:
task = Task(
id=task_id + "." + uuid.uuid1().hex if task_id else uuid.uuid1().hex,
input=question,
agent=self.super_agent,
conf=TaskConfig(max_steps=20),
session_id=self.session_id,
endless_threshold=50,
)
last_output: Output = None
rich_ui = AWorldWebAgentUI(
session_id=self.session_id,
workspace=WorkSpace.from_local_storages(workspace_id=self.session_id),
tool_result_parser_factory=CustomToolResultParserFactory(),
)
async for output in Runners.streamed_run_task(task).stream_events():
logger.info(f"Gaia Agent Ouput: {output}")
res = await AworldUI.parse_output(output, rich_ui)
for item in res if isinstance(res, list) else [res]:
if isinstance(item, AsyncGenerator):
async for sub_item in item:
yield sub_item
if sub_item and str(sub_item).strip():
last_output = sub_item
else:
yield item
if item and str(item).strip():
last_output = item
logger.info(f"Gaia Agent Last Output: {last_output}")
if data_item and last_output:
final_response = self._judge_answer(data_item, last_output)
yield final_response
except Exception as e:
logger.error(f"Error processing {prompt}, error: {traceback.format_exc()}")
def _judge_answer(self, data_item: dict, result: Output):
answer = result
match = re.search(r"<answer>(.*?)</answer>", answer)
if match:
answer = match.group(1)
logger.info(f"Agent answer: {answer}")
logger.info(f"Correct answer: {data_item['Final answer']}")
if question_scorer(answer, data_item["Final answer"]):
logger.info(f"Question {data_item['task_id']} Correct!")
else:
logger.info(f"Question {data_item['task_id']} Incorrect!")
# Create the new result record
correct = question_scorer(answer, data_item["Final answer"])
new_result = {
"task_id": data_item["task_id"],
"level": data_item["Level"],
"question": data_item["Question"],
"answer": data_item["Final answer"],
"response": answer,
"is_correct": correct,
}
return f"\n## Final Result: {'' if correct else ''}\n \n```gaia_result\n{json.dumps(new_result, indent=2)}\n```"
else:
new_result = answer
return f"\n## Final Result:\n \n```gaia_result\n{json.dumps(new_result, indent=2)}\n```"
if __name__ == "__main__":
import asyncio
import argparse
from datetime import datetime
logger = logging.getLogger(__name__)
output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "output")
if not os.path.exists(output_dir):
os.makedirs(output_dir)
output_file = os.path.join(
output_dir, f"output_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md"
)
async def main():
parser = argparse.ArgumentParser()
parser.add_argument("--prompt", type=str, default="")
args = parser.parse_args()
try:
prompt = args.prompt
llm_provider = os.getenv("LLM_PROVIDER", "openai")
llm_model_name = os.getenv("LLM_MODEL_NAME")
llm_api_key = os.getenv("LLM_API_KEY")
llm_base_url = os.getenv("LLM_BASE_URL")
llm_temperature = os.getenv("LLM_TEMPERATURE", 0.0)
def send_output(output):
with open(output_file, "a") as f:
f.write(f"{output}\n")
async for i in GaiaAgentRunner(
llm_provider=llm_provider,
llm_model_name=llm_model_name,
llm_base_url=llm_base_url,
llm_api_key=llm_api_key,
llm_temperature=llm_temperature,
).run(prompt):
send_output(i)
except Exception as e:
logger.error(
f"Error processing {args.prompt}, error: {traceback.format_exc()}"
)
asyncio.run(main())
@@ -0,0 +1,136 @@
from typing import AsyncGenerator
import traceback
import logging
import os
import json
import sys
import uuid
import time
logger = logging.getLogger(__name__)
class GaiaAgentServer:
def __init__(self):
pass
def _get_model_config(self):
try:
llm_provider = os.getenv("LLM_PROVIDER", "openai")
llm_model_name = os.getenv("LLM_MODEL_NAME")
llm_api_key = os.getenv("LLM_API_KEY")
llm_base_url = os.getenv("LLM_BASE_URL")
llm_temperature = float(os.getenv("LLM_TEMPERATURE", 0.0))
return {
"provider": llm_provider,
"model": llm_model_name,
"api_key": llm_api_key,
"base_url": llm_base_url,
"temperature": llm_temperature,
}
except Exception as e:
logger.warning(
f">>> Gaia Agent: GAIA_MODEL_CONFIG is not configured, using LLM"
)
raise e
def models(self):
model = self._get_model_config()
return [
{
"id": f"{model['provider']}/{model['model']}",
"name": f"gaia_agent@{model['provider']}/{model['model']}",
}
]
async def chat_completions(self, body: dict) -> AsyncGenerator[str, None]:
def response_line(line: str, model: str):
return {
"object": "chat.completion.chunk",
"id": str(uuid.uuid4()).replace("-", ""),
"choices": [
{"index": 0, "delta": {"content": line, "role": "assistant"}}
],
"created": int(time.time()),
"model": model,
}
try:
logger.info(f">>> Gaia Agent: body={body}")
prompt = body["messages"][-1]["content"]
model = body["model"].replace("gaia_agent.", "")
logger.info(f">>> Gaia Agent: prompt={prompt}, model={model}")
selected_model = self._get_model_config()
logger.info(f">>> Gaia Agent: Using model configuration: {selected_model}")
logger.info(f">>> Gaia Agent Python Path: sys.path={sys.path}")
llm_provider = selected_model.get("provider")
llm_model_name = selected_model.get("model")
llm_api_key = selected_model.get("api_key")
llm_base_url = selected_model.get("base_url")
llm_temperature = selected_model.get("temperature", 0.0)
from examples.gaia.gaia_agent_runner import GaiaAgentRunner
runner = GaiaAgentRunner(
llm_provider=llm_provider,
llm_model_name=llm_model_name,
llm_base_url=llm_base_url,
llm_api_key=llm_api_key,
llm_temperature=llm_temperature,
)
logger.info(f">>> Gaia Agent: prompt={prompt}, runner={runner}")
async for i in runner.run(prompt):
line = response_line(i, model)
logger.info(f">>> Gaia Agent Line: {line}")
yield line
except Exception as e:
emsg = traceback.format_exc()
logger.error(f">>> Gaia Agent Error: exception {emsg}")
yield response_line(f"Gaia Agent Error: {emsg}", model)
finally:
logger.info(f">>> Gaia Agent Done")
import fastapi
from fastapi.responses import StreamingResponse
app = fastapi.FastAPI()
from examples.gaia.gaia_agent_server import GaiaAgentServer
agent_server = GaiaAgentServer()
@app.get("/v1/models")
async def models():
return agent_server.models()
@app.post("/v1/chat/completions")
async def chat_completions(request: fastapi.Request):
form_data = await request.json()
logger.info(f">>> Gaia Agent Server: form_data={form_data}")
async def event_generator():
async for chunk in agent_server.chat_completions(form_data):
# Format as SSE: each line needs to start with "data: " and end with two newlines
yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n"
return StreamingResponse(event_generator(), media_type="text/event-stream")
if __name__ == "__main__":
import uvicorn
uvicorn.run("gaia_agent_server:app", host="0.0.0.0", port=8888)
@@ -0,0 +1,206 @@
{
"mcpServers": {
"audio": {
"command": "python",
"args": [
"-m",
"examples.gaia.mcp_collections.media.audio"
],
"env": {
},
"client_session_timeout_seconds": 9999.0
},
"browser": {
"command": "python",
"args": [
"-m",
"examples.gaia.mcp_collections.tools.browser"
],
"env": {
"LLM_MODEL_NAME": "${LLM_MODEL_NAME}",
"LLM_API_KEY": "${LLM_API_KEY}",
"LLM_BASE_URL": "${LLM_BASE_URL}"
},
"client_session_timeout_seconds": 9999.0
},
"chess": {
"command": "python",
"args": [
"-m",
"examples.gaia.mcp_collections.tools.playchess"
],
"env": {
},
"client_session_timeout_seconds": 9999.0
},
"code": {
"command": "npx",
"args": [
"-y",
"@e2b/mcp-server"
],
"env": {
"E2B_API_KEY": "${E2B_API_KEY}"
}
},
"csv": {
"command": "python",
"args": [
"-m",
"examples.gaia.mcp_collections.documents.mscsv"
],
"env": {
},
"client_session_timeout_seconds": 9999.0
},
"docx": {
"command": "python",
"args": [
"-m",
"examples.gaia.mcp_collections.documents.msdocx"
],
"env": {
},
"client_session_timeout_seconds": 9999.0
},
"download": {
"command": "python",
"args": [
"-m",
"examples.gaia.mcp_collections.tools.download"
],
"env": {
},
"client_session_timeout_seconds": 9999.0
},
"xlsx": {
"command": "python",
"args": [
"-m",
"examples.gaia.mcp_collections.documents.msxlsx"
],
"env": {
},
"client_session_timeout_seconds": 9999.0
},
"image": {
"command": "python",
"args": [
"-m",
"examples.gaia.mcp_collections.media.image"
],
"env": {
},
"client_session_timeout_seconds": 9999.0
},
"pdf": {
"command": "python",
"args": [
"-m",
"examples.gaia.mcp_collections.documents.pdf"
],
"env": {
},
"client_session_timeout_seconds": 9999.0
},
"pptx": {
"command": "python",
"args": [
"-m",
"examples.gaia.mcp_collections.documents.mspptx"
],
"env": {
},
"client_session_timeout_seconds": 9999.0
},
"pubchem": {
"command": "python",
"args": [
"-m",
"examples.gaia.mcp_collections.tools.pubchem"
],
"env": {
},
"client_session_timeout_seconds": 9999.0
},
"reasoning": {
"command": "python",
"args": [
"-m",
"examples.gaia.mcp_collections.intelligence.think"
],
"env": {
},
"client_session_timeout_seconds": 9999.0
},
"search": {
"command": "python",
"args": [
"-m",
"examples.gaia.mcp_collections.tools.search"
],
"env": {
"GOOGLE_API_KEY": "${GOOGLE_API_KEY}",
"GOOGLE_CSE_ID": "${GOOGLE_CSE_ID}"
},
"client_session_timeout_seconds": 9999.0
},
"terminal": {
"command": "python",
"args": [
"-m",
"examples.gaia.mcp_collections.tools.terminal"
]
},
"video": {
"command": "python",
"args": [
"-m",
"examples.gaia.mcp_collections.media.video"
],
"env": {
},
"client_session_timeout_seconds": 9999.0
},
"wayback": {
"command": "python",
"args": [
"-m",
"examples.gaia.mcp_collections.tools.wayback"
],
"env": {
},
"client_session_timeout_seconds": 9999.0
},
"wikipedia": {
"command": "python",
"args": [
"-m",
"examples.gaia.mcp_collections.tools.wiki"
],
"env": {
},
"client_session_timeout_seconds": 9999.0
},
"youtube": {
"command": "python",
"args": [
"-m",
"examples.gaia.mcp_collections.tools.youtube"
],
"env": {
},
"client_session_timeout_seconds": 9999.0
},
"txt": {
"command": "python",
"args": [
"-m",
"examples.gaia.mcp_collections.documents.txt"
],
"env": {
},
"client_session_timeout_seconds": 9999.0
}
}
}
@@ -0,0 +1,102 @@
import logging
import os
from pathlib import Path
from typing import Any, Literal
from mcp.server import FastMCP
from pydantic import BaseModel, Field
from aworld.logs.util import Color
from examples.gaia.utils import color_log, setup_logger
class ActionArguments(BaseModel):
r"""Protocol: MCP Action Arguments"""
name: str = Field(description="The name of the action")
transport: Literal["stdio", "sse"] = Field(default="stdio", description="The transport of the action")
workspace: str | None = Field(
default=None,
description="The workspace of the action."
" If not specified or invalid, the workspace will be read from the environment variable AWORLD_WORKSPACE.",
)
unittest: bool = Field(default=False, description="Whether to run in unittest mode")
class ActionResponse(BaseModel):
r"""Protocol: MCP Action Response"""
success: bool = Field(default=False, description="Whether the action is successfully executed")
message: Any = Field(default=None, description="The execution result of the action")
metadata: dict[str, Any] = Field(default={}, description="The metadata of the action")
class ActionCollection:
r"""Base class for all ActionCollection."""
server: FastMCP
logger: logging.Logger
def __init__(self, arguments: ActionArguments) -> None:
self.unittest = arguments.unittest
self.transport = arguments.transport
self.supported_extensions = set()
self.workspace: Path = self._obtain_valid_workspace(arguments.workspace)
self.logger: logging.Logger = setup_logger(self.__class__.__name__, self.workspace)
self.server = FastMCP(arguments.name)
for tool_name in self.__class__.__dict__:
if tool_name.startswith("mcp_") and callable(getattr(self.__class__, tool_name)):
tool = getattr(self, tool_name)
self.server.add_tool(tool, description=tool.__doc__)
def run(self) -> None:
if not self.unittest:
self.server.run(transport=self.transport)
def _color_log(self, value: str, color: Color = None, level: str = "info"):
return color_log(self.logger, value, color, level=level)
def _obtain_valid_workspace(self, workspace: str | None = None) -> Path:
r"""
Obtain a valid workspace path.
Priority:
1. user defined workspace
2. environment variable AWORLD_WORKSPACE
3. home directory
"""
path = Path(workspace) if workspace else os.getenv("AWORLD_WORKSPACE", "~")
if path and path.expanduser().is_dir():
return path.expanduser().resolve()
# self._color_log("Invalid workspace path, using home directory instead.", Color.yellow)
return Path.home().expanduser().resolve()
def _validate_file_path(self, file_path: str) -> Path:
"""Validate and resolve file path. Rely on the predefined supported_extensions class variable.
Args:
file_path: Path to the document or media file
Returns:
Resolved Path object
Raises:
FileNotFoundError: If file doesn't exist
ValueError: If file type is not supported
"""
path = Path(file_path).expanduser()
if not path.is_absolute():
path = self.workspace / path
if not path.exists():
raise FileNotFoundError(f"File not found: {path}")
if path.suffix.lower() not in self.supported_extensions:
raise ValueError(
f"Unsupported file type: {path.suffix}. Supported types: {', '.join(self.supported_extensions)}"
)
return path
@@ -0,0 +1,24 @@
from pydantic import BaseModel, Field
class DocumentMetadata(BaseModel):
"""Metadata extracted from document processing."""
file_name: str = Field(description="Original file name")
file_size: int = Field(description="File size in bytes")
file_type: str = Field(description="Document file type/extension")
absolute_path: str = Field(description="Absolute path to the document file")
page_count: int | None = Field(default=None, description="Number of pages in document")
processing_time: float = Field(
description="Time taken to process the document in seconds", deprecated=True, exclude=True
)
extracted_images: list[str] = Field(default_factory=list, description="Paths to extracted image files")
extracted_media: list[dict[str, str]] = Field(
default_factory=list, description="list of extracted media files with type and path"
)
output_format: str = Field(description="Format of the extracted content")
llm_enhanced: bool = Field(default=False, description="Whether LLM enhancement was used", exclude=True)
ocr_applied: bool = Field(default=False, description="Whether OCR was applied", exclude=True)
extracted_text_file_path: str | None = Field(
default=None, description="Absolute path to the extracted text file (if applicable)"
)
@@ -0,0 +1,376 @@
import json
import os
import time
import traceback
from pathlib import Path
from typing import Any, Literal
import chardet
import pandas as pd
from dotenv import load_dotenv
from pydantic import Field
from pydantic.fields import FieldInfo
from aworld.logs.util import Color
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
from examples.gaia.mcp_collections.documents.models import DocumentMetadata
class CSVExtractionCollection(ActionCollection):
"""MCP service for CSV document content extraction using pandas.
Supports extraction from CSV files with various encodings and delimiters.
Provides LLM-friendly text output with structured metadata and data analysis.
"""
def __init__(self, arguments: ActionArguments) -> None:
super().__init__(arguments)
self._media_output_dir = self.workspace / "extracted_media"
self._media_output_dir.mkdir(exist_ok=True)
self.supported_extensions: set = {".csv", ".tsv", ".txt"}
self._color_log("CSV Extraction Service initialized", Color.green, "debug")
self._color_log(f"Media output directory: {self._media_output_dir}", Color.blue, "debug")
def _detect_encoding(self, file_path: Path) -> str:
"""Detect file encoding using chardet.
Args:
file_path: Path to the CSV file
Returns:
Detected encoding string
"""
try:
with open(file_path, "rb") as f:
raw_data = f.read(10000) # Read first 10KB for detection
result = chardet.detect(raw_data)
encoding = result.get("encoding", "utf-8")
confidence = result.get("confidence", 0)
self._color_log(f"Detected encoding: {encoding} (confidence: {confidence:.2f})", Color.blue)
return encoding if confidence > 0.7 else "utf-8"
except Exception as e:
self.logger.warning(f"Encoding detection failed: {e}, using utf-8")
return "utf-8"
def _detect_delimiter(self, file_path: Path, encoding: str) -> str:
"""Detect CSV delimiter by analyzing the first few lines.
Args:
file_path: Path to the CSV file
encoding: File encoding
Returns:
Detected delimiter character
"""
try:
with open(file_path, "r", encoding=encoding) as f:
sample = f.read(1024) # Read first 1KB
# Common delimiters to test
delimiters = [",", ";", "\t", "|", ":"]
delimiter_counts = {}
for delimiter in delimiters:
count = sample.count(delimiter)
if count > 0:
delimiter_counts[delimiter] = count
if delimiter_counts:
detected_delimiter = max(delimiter_counts, key=delimiter_counts.get)
self._color_log(f"Detected delimiter: '{detected_delimiter}'", Color.blue)
return detected_delimiter
else:
return ","
except Exception as e:
self.logger.warning(f"Delimiter detection failed: {e}, using comma")
return ","
def _extract_csv_content(
self, file_path: Path, max_rows: int | None = None, encoding: str | None = None, delimiter: str | None = None
) -> dict[str, Any]:
"""Extract content from CSV file using pandas.
Args:
file_path: Path to the CSV file
max_rows: Maximum number of rows to read
encoding: File encoding (auto-detected if None)
delimiter: CSV delimiter (auto-detected if None)
Returns:
Dictionary containing extracted content and metadata
"""
start_time = time.time()
# Auto-detect encoding and delimiter if not provided
if encoding is None:
encoding = self._detect_encoding(file_path)
if delimiter is None:
delimiter = self._detect_delimiter(file_path, encoding)
try:
# Read CSV with pandas
df = pd.read_csv(file_path, encoding=encoding, delimiter=delimiter, nrows=max_rows, low_memory=False)
# Get full file info for metadata
full_df_info = pd.read_csv(
file_path,
encoding=encoding,
delimiter=delimiter,
nrows=0, # Just get headers and shape info
)
# Count total rows efficiently
total_rows = sum(1 for _ in open(file_path, "r", encoding=encoding)) - 1 # Subtract header
processing_time = time.time() - start_time
return {
"dataframe": df,
"total_rows": total_rows,
"total_columns": len(full_df_info.columns),
"columns": list(df.columns),
"encoding": encoding,
"delimiter": delimiter,
"processing_time": processing_time,
"data_types": df.dtypes.to_dict(),
"memory_usage": df.memory_usage(deep=True).sum(),
}
except Exception as e:
self.logger.error(f"Failed to read CSV file: {e}")
raise
def _format_content_for_llm(self, df: pd.DataFrame, output_format: str, include_stats: bool = True) -> str:
"""Format extracted CSV content to be LLM-friendly.
Args:
df: Pandas DataFrame with CSV data
output_format: Desired output format
include_stats: Whether to include statistical summary
Returns:
Formatted content string
"""
if output_format.lower() == "markdown":
# Convert to markdown table
content = df.to_markdown(index=False, tablefmt="github")
if include_stats:
# Add statistical summary
stats_content = "\n\n## Data Summary\n\n"
stats_content += f"- **Rows**: {len(df)}\n"
stats_content += f"- **Columns**: {len(df.columns)}\n"
stats_content += f"- **Column Names**: {', '.join(df.columns)}\n\n"
# Add data types info
stats_content += "### Column Data Types\n\n"
for col, dtype in df.dtypes.items():
stats_content += f"- **{col}**: {dtype}\n"
# Add basic statistics for numeric columns
numeric_cols = df.select_dtypes(include=["number"]).columns
if len(numeric_cols) > 0:
stats_content += "\n### Numeric Column Statistics\n\n"
stats_df = df[numeric_cols].describe()
stats_content += stats_df.to_markdown(tablefmt="github")
content += stats_content
elif output_format.lower() == "json":
# Convert to JSON with metadata
data_dict = {
"data": df.to_dict(orient="records"),
"metadata": {
"rows": len(df),
"columns": len(df.columns),
"column_names": list(df.columns),
"data_types": {col: str(dtype) for col, dtype in df.dtypes.items()},
},
}
if include_stats:
numeric_cols = df.select_dtypes(include=["number"]).columns
if len(numeric_cols) > 0:
data_dict["statistics"] = df[numeric_cols].describe().to_dict()
content = json.dumps(data_dict, indent=2, default=str)
elif output_format.lower() == "html":
# Convert to HTML table
content = df.to_html(index=False, classes="table table-striped")
else:
# Plain text format
content = df.to_string(index=False)
return content
def mcp_extract_csv_content(
self,
file_path: str = Field(description="Path to the CSV document file to extract content from"),
output_format: Literal["markdown", "json", "html", "text"] = Field(
default="markdown", description="Output format: 'markdown', 'json', 'html', or 'text'"
),
max_rows: int | None = Field(default=None, description="Maximum number of rows to read (None for all rows)"),
include_statistics: bool = Field(default=True, description="Whether to include statistical summary in output"),
generate_visualizations: bool = Field(
default=False, description="Whether to generate and save data visualizations"
),
encoding: str | None = Field(default=None, description="File encoding (auto-detected if None)"),
delimiter: str | None = Field(default=None, description="CSV delimiter (auto-detected if None)"),
) -> ActionResponse:
"""Extract content from CSV documents using pandas.
This tool provides comprehensive CSV document content extraction with support for:
- CSV, TSV, and delimited text files
- Automatic encoding and delimiter detection
- Statistical analysis and data profiling
- Multiple output formats (Markdown, JSON, HTML, Text)
- Optional data visualizations
- Memory-efficient processing for large files
Args:
file_path: Path to the CSV file
output_format: Desired output format
max_rows: Maximum rows to process (None for all)
include_statistics: Include statistical summary
generate_visualizations: Generate data visualizations
encoding: File encoding (auto-detected if None)
delimiter: CSV delimiter (auto-detected if None)
Returns:
ActionResponse with extracted content, metadata, and optional visualizations
"""
try:
# Handle FieldInfo objects from pydantic
if isinstance(file_path, FieldInfo):
file_path = file_path.default
if isinstance(output_format, FieldInfo):
output_format = output_format.default
if isinstance(max_rows, FieldInfo):
max_rows = max_rows.default
if isinstance(include_statistics, FieldInfo):
include_statistics = include_statistics.default
if isinstance(generate_visualizations, FieldInfo):
generate_visualizations = generate_visualizations.default
if isinstance(encoding, FieldInfo):
encoding = encoding.default
if isinstance(delimiter, FieldInfo):
delimiter = delimiter.default
# Validate input file
file_path: Path = self._validate_file_path(file_path)
self._color_log(f"Processing CSV file: {file_path.name}", Color.cyan)
# Extract CSV content
extraction_result = self._extract_csv_content(
file_path, max_rows=max_rows, encoding=encoding, delimiter=delimiter
)
df: pd.DataFrame = extraction_result["dataframe"]
# Format content for LLM consumption
formatted_content = self._format_content_for_llm(df, output_format, include_stats=include_statistics)
# Prepare metadata
file_stats = file_path.stat()
document_metadata = DocumentMetadata(
file_name=file_path.name,
file_size=file_stats.st_size,
file_type=file_path.suffix.lower(),
absolute_path=str(file_path.absolute()),
page_count=None, # Not applicable for CSV
processing_time=extraction_result["processing_time"],
extracted_images=[], # CSV files don't contain images
extracted_media=None,
output_format=output_format,
llm_enhanced=False,
ocr_applied=False,
)
# Add CSV-specific metadata
csv_metadata = {
"total_rows": extraction_result["total_rows"],
"total_columns": extraction_result["total_columns"],
"rows_processed": len(df),
"columns_processed": len(df.columns),
"column_names": extraction_result["columns"],
"data_types": {k: str(v) for k, v in extraction_result["data_types"].items()},
"encoding": extraction_result["encoding"],
"delimiter": extraction_result["delimiter"],
"memory_usage_bytes": int(extraction_result["memory_usage"]),
}
# Merge metadata
final_metadata = {**document_metadata.model_dump(), **csv_metadata}
self._color_log(
f"Successfully extracted CSV content from {file_path.name} "
f"({extraction_result['total_rows']} rows, {extraction_result['total_columns']} columns",
Color.green,
)
return ActionResponse(success=True, message=formatted_content, metadata=final_metadata)
except FileNotFoundError as e:
self.logger.error(f"File not found: {str(e)}")
return ActionResponse(
success=False, message=f"File not found: {str(e)}", metadata={"error_type": "file_not_found"}
)
except ValueError as e:
self.logger.error(f"Invalid input: {str(e)}")
return ActionResponse(
success=False,
message=f"Invalid input: {str(e)}",
metadata={"error_type": "invalid_input"},
)
except Exception as e:
self.logger.error(f"CSV extraction failed: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=f"CSV extraction failed: {str(e)}",
metadata={"error_type": "extraction_error"},
)
def mcp_list_supported_formats(self) -> ActionResponse:
"""List all supported CSV formats for extraction.
Returns:
ActionResponse with list of supported file formats and their descriptions
"""
supported_formats = {
"CSV": "Comma-Separated Values files (.csv)",
"TSV": "Tab-Separated Values files (.tsv)",
"TXT": "Delimited text files (.txt)",
}
format_list = "\n".join(
[f"**{format_name}**: {description}" for format_name, description in supported_formats.items()]
)
return ActionResponse(
success=True,
message=f"Supported CSV formats:\n\n{format_list}",
metadata={"supported_formats": list(supported_formats.keys()), "total_formats": len(supported_formats)},
)
# Example usage and entry point
if __name__ == "__main__":
load_dotenv()
# Default arguments for testing
args = ActionArguments(
name="csv_extraction_service",
transport="stdio",
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
)
# Initialize and run the CSV extraction service
try:
service = CSVExtractionCollection(args)
service.run()
except Exception as e:
print(f"An error occurred: {e}: {traceback.format_exc()}")
@@ -0,0 +1,617 @@
import json
import os
import time
import traceback
import zipfile
from pathlib import Path
from typing import Any, Literal
from docx import Document
from docx.document import Document as DocumentType
from dotenv import load_dotenv
from pydantic import Field
from pydantic.fields import FieldInfo
from aworld.logs.util import Color
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
from examples.gaia.mcp_collections.documents.models import DocumentMetadata
class DOCXExtractionCollection(ActionCollection):
"""MCP service for DOCX/DOC document content extraction using python-docx.
Supports extraction from DOCX and DOC files with comprehensive content parsing.
Provides LLM-friendly text output with structured metadata and media file handling.
"""
def __init__(self, arguments: ActionArguments) -> None:
super().__init__(arguments)
self._media_output_dir = self.workspace / "extracted_media"
self._media_output_dir.mkdir(exist_ok=True)
self.supported_extensions = {".docx", ".doc"}
self._color_log("DOCX Extraction Service initialized", Color.green, "debug")
self._color_log(f"Media output directory: {self._media_output_dir}", Color.blue, "debug")
def _extract_images_from_docx(self, file_path: Path, file_stem: str) -> list[dict[str, str]]:
"""Extract embedded images from DOCX file.
Args:
file_path: Path to the DOCX file
file_stem: Base name for saving files
Returns:
List of dictionaries containing image file paths and metadata
"""
saved_media = []
try:
# DOCX files are ZIP archives
with zipfile.ZipFile(file_path, "r") as docx_zip:
# Look for media files in the word/media/ directory
media_files = [f for f in docx_zip.namelist() if f.startswith("word/media/")]
for idx, media_file in enumerate(media_files):
try:
# Extract file extension and create appropriate filename
original_name = Path(media_file).name
file_extension = Path(media_file).suffix
# Generate unique filename
media_filename = f"{file_stem}_media_{idx}_{original_name}"
media_path = self._media_output_dir / media_filename
# Extract and save the media file
with docx_zip.open(media_file) as source:
with open(media_path, "wb") as target:
target.write(source.read())
# Determine media type based on extension
media_type = "image"
if file_extension.lower() in [".mp3", ".wav", ".m4a", ".ogg"]:
media_type = "audio"
elif file_extension.lower() in [".mp4", ".avi", ".mov", ".wmv"]:
media_type = "video"
elif file_extension.lower() in [".png", ".jpg", ".jpeg", ".gif", ".bmp", ".tiff"]:
media_type = "image"
else:
media_type = "other"
saved_media.append(
{
"type": media_type,
"path": str(media_path),
"filename": media_filename,
"original_name": original_name,
"size_bytes": media_path.stat().st_size,
}
)
self._color_log(f"Extracted {media_type}: {media_filename}", Color.blue)
except Exception as e:
self.logger.error(f"Failed to extract media file {media_file}: {e}")
except Exception as e:
self.logger.warning(f"Could not extract media from DOCX: {e}")
return saved_media
def _extract_document_structure(self, doc: DocumentType) -> dict[str, Any]:
"""Extract document structure and metadata.
Args:
doc: python-docx Document object
Returns:
Dictionary containing document structure information
"""
structure = {
"paragraphs_count": len(doc.paragraphs),
"tables_count": len(doc.tables),
"sections_count": len(doc.sections),
"styles_used": [],
"has_headers_footers": False,
"page_count": None, # Not directly available in python-docx
}
# Collect unique styles used in the document
styles_used = set()
for paragraph in doc.paragraphs:
if paragraph.style and paragraph.style.name:
styles_used.add(paragraph.style.name)
structure["styles_used"] = list(styles_used)
# Check for headers and footers
for section in doc.sections:
if (section.header.paragraphs and any(p.text.strip() for p in section.header.paragraphs)) or (
section.footer.paragraphs and any(p.text.strip() for p in section.footer.paragraphs)
):
structure["has_headers_footers"] = True
break
return structure
def _extract_tables_content(self, doc: DocumentType) -> list[dict[str, Any]]:
"""Extract content from all tables in the document.
Args:
doc: python-docx Document object
Returns:
List of dictionaries containing table data
"""
tables_data = []
for table_idx, table in enumerate(doc.tables):
table_data = {
"table_index": table_idx,
"rows_count": len(table.rows),
"columns_count": len(table.columns) if table.rows else 0,
"data": [],
}
# Extract table content
for _, row in enumerate(table.rows):
row_data = []
for cell in row.cells:
cell_text = cell.text.strip()
row_data.append(cell_text)
table_data["data"].append(row_data)
tables_data.append(table_data)
return tables_data
def _convert_doc_to_docx(self, doc_path: Path) -> Path:
"""Convert .doc file to .docx using LibreOffice.
Args:
doc_path: Path to the .doc file
Returns:
Path to the converted .docx file
"""
import subprocess
output_dir = self._media_output_dir
docx_path = output_dir / f"{doc_path.stem}.docx"
# Check if already converted
if docx_path.exists():
self._color_log(f"Using existing converted file: {docx_path.name}", Color.blue)
return docx_path
self._color_log(f"Converting .doc to .docx: {doc_path.name}", Color.yellow)
cmd = ["libreoffice", "--headless", "--convert-to", "docx", "--outdir", str(output_dir), str(doc_path)]
try:
result = subprocess.run(cmd, capture_output=True, text=True, check=True, timeout=60) # pylint: disable=W0612
if docx_path.exists():
self._color_log(f"Conversion successful: {docx_path.name}", Color.green)
return docx_path
else:
raise RuntimeError("Conversion completed but output file not found")
except subprocess.CalledProcessError as e:
self.logger.error(f"LibreOffice conversion failed: {e.stderr}")
raise RuntimeError(f"Failed to convert .doc to .docx: {e.stderr}") from e
except subprocess.TimeoutExpired as e:
self.logger.error("LibreOffice conversion timed out")
raise RuntimeError("Conversion timed out after 60 seconds") from e
# Modify the _extract_content_from_docx method to handle .doc files
def _extract_content_from_docx(
self, file_path: Path, extract_tables: bool = True, extract_headers_footers: bool = True
) -> dict[str, Any]:
"""Extract content from DOCX file using python-docx.
Args:
file_path: Path to the DOCX/DOC file
extract_tables: Whether to extract table content
extract_headers_footers: Whether to extract headers and footers
Returns:
Dictionary containing extracted content and metadata
"""
start_time = time.time()
# Convert .doc to .docx if needed
if file_path.suffix.lower() == ".doc":
file_path = self._convert_doc_to_docx(file_path)
try:
# Load the document
doc = Document(str(file_path))
# Extract main text content
paragraphs = []
for paragraph in doc.paragraphs:
if paragraph.text.strip(): # Skip empty paragraphs
para_data = {"text": paragraph.text, "style": paragraph.style.name if paragraph.style else "Normal"}
paragraphs.append(para_data)
# Extract document structure
structure = self._extract_document_structure(doc)
# Extract tables if requested
tables = []
if extract_tables:
tables = self._extract_tables_content(doc)
# Extract headers and footers if requested
headers_footers = []
if extract_headers_footers:
for section_idx, section in enumerate(doc.sections):
# Extract header content
header_text = []
for para in section.header.paragraphs:
if para.text.strip():
header_text.append(para.text)
# Extract footer content
footer_text = []
for para in section.footer.paragraphs:
if para.text.strip():
footer_text.append(para.text)
if header_text or footer_text:
headers_footers.append(
{"section_index": section_idx, "header": header_text, "footer": footer_text}
)
processing_time = time.time() - start_time
return {
"paragraphs": paragraphs,
"tables": tables,
"headers_footers": headers_footers,
"structure": structure,
"processing_time": processing_time,
"word_count": sum(len(p["text"].split()) for p in paragraphs),
"character_count": sum(len(p["text"]) for p in paragraphs),
}
except Exception as e:
self.logger.error(f"Failed to extract content from DOCX: {e}")
raise
def _format_content_for_llm(
self, extraction_result: dict[str, Any], output_format: str, include_structure: bool = True
) -> str:
"""Format extracted DOCX content to be LLM-friendly.
Args:
extraction_result: Dictionary containing extracted content
output_format: Desired output format
include_structure: Whether to include document structure information
Returns:
Formatted content string
"""
if output_format.lower() == "markdown":
content_parts = []
# Add document structure info if requested
if include_structure:
structure = extraction_result["structure"]
content_parts.append("# Document Structure\n")
content_parts.append(f"- **Paragraphs**: {structure['paragraphs_count']}\n")
content_parts.append(f"- **Tables**: {structure['tables_count']}\n")
content_parts.append(f"- **Sections**: {structure['sections_count']}\n")
content_parts.append(f"- **Word Count**: {extraction_result['word_count']}\n")
content_parts.append(f"- **Character Count**: {extraction_result['character_count']}\n")
if structure["styles_used"]:
content_parts.append(f"- **Styles Used**: {', '.join(structure['styles_used'])}\n")
content_parts.append("\n---\n\n")
# Add main content
content_parts.append("# Document Content\n\n")
# Add paragraphs
for para in extraction_result["paragraphs"]:
# Format based on style
text = para["text"]
style = para["style"]
if "Heading" in style:
# Convert heading styles to markdown headers
if "Heading 1" in style:
content_parts.append(f"# {text}\n\n")
elif "Heading 2" in style:
content_parts.append(f"## {text}\n\n")
elif "Heading 3" in style:
content_parts.append(f"### {text}\n\n")
else:
content_parts.append(f"#### {text}\n\n")
else:
content_parts.append(f"{text}\n\n")
# Add tables
if extraction_result["tables"]:
content_parts.append("\n## Tables\n\n")
for table_idx, table in enumerate(extraction_result["tables"]):
content_parts.append(f"### Table {table_idx + 1}\n\n")
if table["data"]:
# Create markdown table
headers = table["data"][0] if table["data"] else []
if headers:
content_parts.append("| " + " | ".join(headers) + " |\n")
content_parts.append("|" + "---|" * len(headers) + "\n")
for row in table["data"][1:]:
content_parts.append("| " + " | ".join(row) + " |\n")
content_parts.append("\n")
# Add headers and footers
if extraction_result["headers_footers"]:
content_parts.append("\n## Headers and Footers\n\n")
for hf in extraction_result["headers_footers"]:
if hf["header"]:
content_parts.append(f"**Header (Section {hf['section_index'] + 1}):**\n")
for header_line in hf["header"]:
content_parts.append(f"{header_line}\n")
content_parts.append("\n")
if hf["footer"]:
content_parts.append(f"**Footer (Section {hf['section_index'] + 1}):**\n")
for footer_line in hf["footer"]:
content_parts.append(f"{footer_line}\n")
content_parts.append("\n")
return "".join(content_parts)
elif output_format.lower() == "json":
# Return structured JSON
return json.dumps(extraction_result, indent=2, ensure_ascii=False)
elif output_format.lower() == "html":
# Convert to HTML
html_parts = ["<html><body>"]
if include_structure:
html_parts.append("<h1>Document Structure</h1>")
structure = extraction_result["structure"]
html_parts.append(f"<p><strong>Paragraphs:</strong> {structure['paragraphs_count']}</p>")
html_parts.append(f"<p><strong>Tables:</strong> {structure['tables_count']}</p>")
html_parts.append(f"<p><strong>Word Count:</strong> {extraction_result['word_count']}</p>")
html_parts.append("<hr>")
html_parts.append("<h1>Document Content</h1>")
# Add paragraphs
for para in extraction_result["paragraphs"]:
text = para["text"].replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
style = para["style"]
if "Heading" in style:
if "Heading 1" in style:
html_parts.append(f"<h1>{text}</h1>")
elif "Heading 2" in style:
html_parts.append(f"<h2>{text}</h2>")
elif "Heading 3" in style:
html_parts.append(f"<h3>{text}</h3>")
else:
html_parts.append(f"<h4>{text}</h4>")
else:
html_parts.append(f"<p>{text}</p>")
# Add tables
if extraction_result["tables"]:
html_parts.append("<h2>Tables</h2>")
for table_idx, table in enumerate(extraction_result["tables"]):
html_parts.append(f"<h3>Table {table_idx + 1}</h3>")
html_parts.append("<table border='1'>")
for row_idx, row in enumerate(table["data"]):
html_parts.append("<tr>")
tag = "th" if row_idx == 0 else "td"
for cell in row:
cell_text = cell.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
html_parts.append(f"<{tag}>{cell_text}</{tag}>")
html_parts.append("</tr>")
html_parts.append("</table>")
html_parts.append("</body></html>")
return "".join(html_parts)
else:
# Plain text format
text_parts = []
# Add main content
for para in extraction_result["paragraphs"]:
text_parts.append(para["text"])
# Add tables as plain text
if extraction_result["tables"]:
text_parts.append("\n\n=== TABLES ===\n")
for table_idx, table in enumerate(extraction_result["tables"]):
text_parts.append(f"\nTable {table_idx + 1}:\n")
for row in table["data"]:
text_parts.append("\t".join(row) + "\n")
return "\n\n".join(text_parts)
def mcp_extract_docx_content(
self,
file_path: str = Field(description="Path to the DOCX/DOC document file to extract content from"),
output_format: Literal["markdown", "json", "html", "text"] = Field(
default="markdown", description="Output format: 'markdown', 'json', 'html', or 'text'"
),
extract_images: bool = Field(default=True, description="Whether to extract and save embedded images and media"),
extract_tables: bool = Field(default=True, description="Whether to extract table content"),
extract_headers_footers: bool = Field(default=True, description="Whether to extract headers and footers"),
include_structure: bool = Field(
default=True, description="Whether to include document structure information in output"
),
) -> ActionResponse:
"""Extract content from DOCX/DOC documents using python-docx.
This tool provides comprehensive DOCX/DOC document content extraction with support for:
- DOCX and DOC files
- Text extraction with style preservation
- Table extraction and formatting
- Headers and footers extraction
- Embedded media extraction (images, audio, etc.)
- Multiple output formats (Markdown, JSON, HTML, Text)
- Document structure analysis
Args:
file_path: Path to the DOCX/DOC file
output_format: Desired output format
extract_images: Extract embedded media files
extract_tables: Extract table content
extract_headers_footers: Extract headers and footers
include_structure: Include document structure info
Returns:
ActionResponse with extracted content, metadata, and media file paths
"""
try:
# Handle FieldInfo objects from pydantic
if isinstance(file_path, FieldInfo):
file_path = file_path.default
if isinstance(output_format, FieldInfo):
output_format = output_format.default
if isinstance(extract_images, FieldInfo):
extract_images = extract_images.default
if isinstance(extract_tables, FieldInfo):
extract_tables = extract_tables.default
if isinstance(extract_headers_footers, FieldInfo):
extract_headers_footers = extract_headers_footers.default
if isinstance(include_structure, FieldInfo):
include_structure = include_structure.default
# Validate input file
file_path: Path = self._validate_file_path(file_path)
self._color_log(f"Processing DOCX document: {file_path.name}", Color.cyan)
# Extract embedded media if requested
saved_media = []
if extract_images and file_path.suffix.lower() == ".docx":
saved_media = self._extract_images_from_docx(file_path, file_path.stem)
# Extract document content
extraction_result = self._extract_content_from_docx(
file_path, extract_tables=extract_tables, extract_headers_footers=extract_headers_footers
)
# Format content for LLM consumption
formatted_content = self._format_content_for_llm(
extraction_result, output_format, include_structure=include_structure
)
# Prepare metadata
file_stats = file_path.stat()
document_metadata = DocumentMetadata(
file_name=file_path.name,
file_size=file_stats.st_size,
file_type=file_path.suffix.lower(),
absolute_path=str(file_path.absolute()),
page_count=None, # Not directly available for DOCX
processing_time=extraction_result["processing_time"],
extracted_images=[media["path"] for media in saved_media if media["type"] == "image"],
extracted_media=saved_media,
output_format=output_format,
llm_enhanced=False,
ocr_applied=False,
)
# Add DOCX-specific metadata
docx_metadata = {
"paragraphs_count": extraction_result["structure"]["paragraphs_count"],
"tables_count": extraction_result["structure"]["tables_count"],
"sections_count": extraction_result["structure"]["sections_count"],
"word_count": extraction_result["word_count"],
"character_count": extraction_result["character_count"],
"styles_used": extraction_result["structure"]["styles_used"],
"has_headers_footers": extraction_result["structure"]["has_headers_footers"],
"has_embedded_media": len(saved_media) > 0,
"media_files_count": len(saved_media),
}
# Merge metadata
final_metadata = {**document_metadata.model_dump(), **docx_metadata}
self._color_log(
f"Successfully extracted DOCX content from {file_path.name} "
f"({extraction_result['word_count']} words, {extraction_result['structure']['tables_count']} tables, "
f"{len(saved_media)} media files)",
Color.green,
)
return ActionResponse(success=True, message=formatted_content, metadata=final_metadata)
except FileNotFoundError as e:
self.logger.error(f"File not found: {str(e)}")
return ActionResponse(
success=False, message=f"File not found: {str(e)}", metadata={"error_type": "file_not_found"}
)
except ValueError as e:
self.logger.error(f"Invalid input: {str(e)}")
return ActionResponse(
success=False,
message=f"Invalid input: {str(e)}",
metadata={"error_type": "invalid_input"},
)
except ImportError as e:
self.logger.error(f"Missing dependency: {str(e)}")
return ActionResponse(
success=False,
message=f"Missing dependency: {str(e)}. Please install python-docx: pip install python-docx",
metadata={"error_type": "missing_dependency"},
)
except Exception as e:
self.logger.error(f"DOCX extraction failed: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=f"DOCX extraction failed: {str(e)}",
metadata={"error_type": "extraction_error"},
)
def mcp_list_supported_formats(self) -> ActionResponse:
"""List all supported document formats for extraction.
Returns:
ActionResponse with list of supported file formats and their descriptions
"""
supported_formats = {
"DOCX": "Microsoft Word Open XML Document (.docx)",
"DOC": "Microsoft Word Document (.doc) - limited support",
}
format_list = "\n".join(
[f"**{format_name}**: {description}" for format_name, description in supported_formats.items()]
)
return ActionResponse(
success=True,
message=f"Supported document formats:\n\n{format_list}\n\n"
"**Note**: DOC format support is limited. For best results, convert to DOCX format.",
metadata={"supported_formats": list(supported_formats.keys()), "total_formats": len(supported_formats)},
)
# Example usage and entry point
if __name__ == "__main__":
load_dotenv()
# Default arguments for testing
args = ActionArguments(
name="docx_extraction_service",
transport="stdio",
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
)
# Initialize and run the DOCX extraction service
try:
service = DOCXExtractionCollection(args)
service.run()
except Exception as e:
print(f"An error occurred: {e}: {traceback.format_exc()}")
@@ -0,0 +1,545 @@
import json
import os
import time
import traceback
import zipfile
from pathlib import Path
from typing import Any, Literal
from dotenv import load_dotenv
from pptx import Presentation
from pptx.presentation import Presentation as PresentationType
from pydantic import Field
from pydantic.fields import FieldInfo
from aworld.logs.util import Color
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
from examples.gaia.mcp_collections.documents.models import DocumentMetadata
class PPTXExtractionCollection(ActionCollection):
"""MCP service for PPTX/PPT presentation content extraction using python-pptx.
Supports extraction from PPTX and PPT files with comprehensive content parsing.
Provides LLM-friendly text output with structured metadata and media file handling.
"""
def __init__(self, arguments: ActionArguments) -> None:
super().__init__(arguments)
self._media_output_dir = self.workspace / "extracted_media"
self._media_output_dir.mkdir(exist_ok=True)
self.supported_extensions = {".pptx", ".ppt"}
self._color_log("PPTX Extraction Service initialized", Color.green, "debug")
self._color_log(f"Media output directory: {self._media_output_dir}", Color.blue, "debug")
def _extract_images_from_pptx(self, file_path: Path, file_stem: str) -> list[dict[str, str]]:
"""Extract embedded images from PPTX file.
Args:
file_path: Path to the PPTX file
file_stem: Base name for saving files
Returns:
List of dictionaries containing image file paths and metadata
"""
saved_media = []
try:
# PPTX files are ZIP archives
with zipfile.ZipFile(file_path, "r") as zip_file:
# Find media files in the archive
media_files = [f for f in zip_file.namelist() if f.startswith("ppt/media/")]
for idx, media_file in enumerate(media_files):
try:
# Extract file extension
original_ext = Path(media_file).suffix
if not original_ext:
original_ext = ".png" # Default extension
# Generate unique filename
media_filename = f"{file_stem}_media_{idx}{original_ext}"
media_path = self._media_output_dir / media_filename
# Extract and save media file
with zip_file.open(media_file) as source:
with open(media_path, "wb") as target:
target.write(source.read())
# Determine media type
media_type = (
"image" if original_ext.lower() in {".png", ".jpg", ".jpeg", ".gif", ".bmp"} else "media"
)
saved_media.append(
{
"type": media_type,
"path": str(media_path),
"filename": media_filename,
"original_path": media_file,
}
)
self._color_log(f"Saved media: {media_filename}", Color.blue)
except Exception as e:
self.logger.error(f"Failed to extract media {media_file}: {e}")
except Exception as e:
self.logger.error(f"Failed to extract media from PPTX: {e}")
return saved_media
def _extract_slide_structure(self, presentation: PresentationType) -> dict[str, Any]:
"""Extract presentation structure information.
Args:
presentation: python-pptx Presentation object
Returns:
Dictionary containing structure metadata
"""
structure = {
"slide_count": len(presentation.slides),
"slide_layouts": [],
"has_notes": False,
"has_comments": False,
"slide_sizes": None,
}
# Get slide size information
if hasattr(presentation.slide_width, "inches") and hasattr(presentation.slide_height, "inches"):
structure["slide_sizes"] = {
"width_inches": presentation.slide_width.inches,
"height_inches": presentation.slide_height.inches,
}
# Analyze slide layouts and content
for slide_idx, slide in enumerate(presentation.slides):
layout_info = {
"slide_index": slide_idx,
"layout_name": slide.slide_layout.name if hasattr(slide.slide_layout, "name") else "Unknown",
"shape_count": len(slide.shapes),
"has_title": False,
"has_content": False,
}
# Check for title and content
for shape in slide.shapes:
if hasattr(shape, "text") and shape.text.strip():
if (
"title" in str(shape.placeholder_format.type).lower()
if hasattr(shape, "placeholder_format")
else False
):
layout_info["has_title"] = True
else:
layout_info["has_content"] = True
# Check for notes
if hasattr(slide, "notes_slide") and slide.notes_slide:
if hasattr(slide.notes_slide, "notes_text_frame") and slide.notes_slide.notes_text_frame.text.strip():
structure["has_notes"] = True
structure["slide_layouts"].append(layout_info)
return structure
def _extract_content_from_pptx(self, file_path: Path, extract_notes: bool = True) -> dict[str, Any]:
"""Extract content from PPTX file using python-pptx.
Args:
file_path: Path to the PPTX file
extract_notes: Whether to extract speaker notes
Returns:
Dictionary containing extracted content and metadata
"""
start_time = time.time()
try:
# Load the presentation
presentation = Presentation(str(file_path))
# Extract slide content
slides_content = []
total_text_length = 0
for slide_idx, slide in enumerate(presentation.slides):
slide_data = {
"slide_number": slide_idx + 1,
"title": "",
"content": [],
"notes": "",
"shapes_count": len(slide.shapes),
}
# Extract text from all shapes
for shape in slide.shapes:
if hasattr(shape, "text") and shape.text.strip():
text_content = shape.text.strip()
# Try to identify if this is a title
if (
hasattr(shape, "placeholder_format")
and "title" in str(shape.placeholder_format.type).lower()
):
slide_data["title"] = text_content
else:
slide_data["content"].append(text_content)
total_text_length += len(text_content)
# Extract text from tables if present
if hasattr(shape, "table"):
table_text = []
for row in shape.table.rows:
row_text = []
for cell in row.cells:
if cell.text.strip():
row_text.append(cell.text.strip())
if row_text:
table_text.append(" | ".join(row_text))
if table_text:
slide_data["content"].append("\n".join(table_text))
# Extract speaker notes if requested
if extract_notes and hasattr(slide, "notes_slide"):
try:
if (
hasattr(slide.notes_slide, "notes_text_frame")
and slide.notes_slide.notes_text_frame.text.strip()
):
slide_data["notes"] = slide.notes_slide.notes_text_frame.text.strip()
total_text_length += len(slide_data["notes"])
except Exception as e:
self.logger.warning(f"Failed to extract notes from slide {slide_idx + 1}: {e}")
slides_content.append(slide_data)
# Extract presentation structure
structure = self._extract_slide_structure(presentation)
processing_time = time.time() - start_time
return {
"slides": slides_content,
"structure": structure,
"processing_time": processing_time,
"total_text_length": total_text_length,
"slide_count": len(slides_content),
}
except Exception as e:
self.logger.error(f"Failed to extract content from PPTX: {e}")
raise
def _format_content_for_llm(
self, extraction_result: dict[str, Any], output_format: str, include_structure: bool = True
) -> str:
"""Format extracted PPTX content to be LLM-friendly.
Args:
extraction_result: Dictionary containing extracted content
output_format: Desired output format
include_structure: Whether to include presentation structure information
Returns:
Formatted content string
"""
slides = extraction_result["slides"]
structure = extraction_result["structure"]
if output_format.lower() == "markdown":
content_parts = []
if include_structure:
content_parts.append("# Presentation Overview")
content_parts.append(f"- **Total Slides**: {structure['slide_count']}")
if structure.get("slide_sizes"):
sizes = structure["slide_sizes"]
content_parts.append(
f'- **Slide Size**: {sizes["width_inches"]:.1f}" × {sizes["height_inches"]:.1f}"'
)
content_parts.append(f"- **Has Speaker Notes**: {'Yes' if structure['has_notes'] else 'No'}")
content_parts.append("")
# Format each slide
for slide in slides:
content_parts.append(f"## Slide {slide['slide_number']}")
if slide["title"]:
content_parts.append(f"**Title**: {slide['title']}")
content_parts.append("")
if slide["content"]:
content_parts.append("**Content**:")
for content_item in slide["content"]:
# Format multi-line content properly
for line in content_item.split("\n"):
if line.strip():
content_parts.append(f"- {line.strip()}")
content_parts.append("")
if slide["notes"]:
content_parts.append("**Speaker Notes**:")
content_parts.append(slide["notes"])
content_parts.append("")
content_parts.append("---")
content_parts.append("")
return "\n".join(content_parts)
elif output_format.lower() == "json":
return json.dumps(
{
"presentation_structure": structure if include_structure else None,
"slides": slides,
"metadata": {
"total_slides": len(slides),
"total_text_length": extraction_result["total_text_length"],
"processing_time": extraction_result["processing_time"],
},
},
indent=2,
)
elif output_format.lower() == "html":
html_parts = ["<div class='presentation'>"]
if include_structure:
html_parts.append("<div class='presentation-overview'>")
html_parts.append("<h1>Presentation Overview</h1>")
html_parts.append(f"<p><strong>Total Slides:</strong> {structure['slide_count']}</p>")
if structure.get("slide_sizes"):
sizes = structure["slide_sizes"]
html_parts.append(
"<p>"
"<strong>Slide Size:</strong> "
f"{sizes['width_inches']:.1f} × {sizes['height_inches']:.1f}"
"</p>"
)
html_parts.append(
f"<p><strong>Has Speaker Notes:</strong> {'Yes' if structure['has_notes'] else 'No'}</p>"
)
html_parts.append("</div>")
for slide in slides:
html_parts.append(f"<div class='slide' data-slide='{slide['slide_number']}'>")
html_parts.append(f"<h2>Slide {slide['slide_number']}</h2>")
if slide["title"]:
html_parts.append(f"<h3>{slide['title']}</h3>")
if slide["content"]:
html_parts.append("<div class='slide-content'>")
for content_item in slide["content"]:
html_parts.append(f"<p>{content_item.replace(chr(10), '<br>')}</p>")
html_parts.append("</div>")
if slide["notes"]:
html_parts.append(
"<div class='speaker-notes'>"
"<strong>Speaker Notes:</strong>"
f"<br>{slide['notes'].replace(chr(10), '<br>')}"
"</div>"
)
html_parts.append("</div>")
html_parts.append("</div>")
return "\n".join(html_parts)
else: # Plain text
text_parts = []
if include_structure:
text_parts.append("PRESENTATION OVERVIEW")
text_parts.append(f"Total Slides: {structure['slide_count']}")
if structure.get("slide_sizes"):
sizes = structure["slide_sizes"]
text_parts.append(f'Slide Size: {sizes["width_inches"]:.1f}" × {sizes["height_inches"]:.1f}"')
text_parts.append(f"Has Speaker Notes: {'Yes' if structure['has_notes'] else 'No'}")
text_parts.append("\n" + "=" * 50 + "\n")
for slide in slides:
text_parts.append(f"SLIDE {slide['slide_number']}")
if slide["title"]:
text_parts.append(f"Title: {slide['title']}")
if slide["content"]:
text_parts.append("Content:")
for content_item in slide["content"]:
text_parts.append(f" {content_item}")
if slide["notes"]:
text_parts.append(f"Speaker Notes: {slide['notes']}")
text_parts.append("\n" + "-" * 30 + "\n")
return "\n".join(text_parts)
def mcp_extract_pptx_content(
self,
file_path: str = Field(description="Path to the PPTX/PPT presentation file to extract content from"),
output_format: Literal["markdown", "json", "html", "text"] = Field(
default="markdown", description="Output format: 'markdown', 'json', 'html', or 'text'"
),
extract_images: bool = Field(
default=True, description="Whether to extract and save images from the presentation"
),
extract_notes: bool = Field(default=True, description="Whether to extract speaker notes"),
include_structure: bool = Field(
default=True, description="Whether to include presentation structure information"
),
) -> ActionResponse:
"""Extract content from PPTX/PPT presentations using python-pptx.
This tool provides comprehensive PowerPoint presentation content extraction with support for:
- PPTX and PPT files
- Text extraction from slides, titles, and content
- Speaker notes extraction
- Image and media extraction
- Presentation structure analysis
- Multiple output formats (Markdown, JSON, HTML, Text)
- LLM-optimized formatting
Args:
file_path: Path to the PPTX/PPT presentation file
output_format: Desired output format
extract_images: Extract embedded media files
extract_notes: Extract speaker notes
include_structure: Include presentation structure info
Returns:
ActionResponse with extracted content, metadata, and media file paths
"""
try:
# Handle FieldInfo objects from pydantic
if isinstance(file_path, FieldInfo):
file_path = file_path.default
if isinstance(output_format, FieldInfo):
output_format = output_format.default
if isinstance(extract_images, FieldInfo):
extract_images = extract_images.default
if isinstance(extract_notes, FieldInfo):
extract_notes = extract_notes.default
if isinstance(include_structure, FieldInfo):
include_structure = include_structure.default
# Validate input file
file_path: Path = self._validate_file_path(file_path)
self._color_log(f"Processing PPTX presentation: {file_path.name}", Color.cyan)
# Extract embedded media if requested
saved_media = []
if extract_images and file_path.suffix.lower() == ".pptx":
saved_media = self._extract_images_from_pptx(file_path, file_path.stem)
# Extract presentation content
extraction_result = self._extract_content_from_pptx(file_path, extract_notes=extract_notes)
# Format content for LLM consumption
formatted_content = self._format_content_for_llm(
extraction_result, output_format, include_structure=include_structure
)
# Prepare metadata
file_stats = file_path.stat()
document_metadata = DocumentMetadata(
file_name=file_path.name,
file_size=file_stats.st_size,
file_type=file_path.suffix.lower(),
absolute_path=str(file_path.absolute()),
page_count=extraction_result["slide_count"], # Use slide count as page count
processing_time=extraction_result["processing_time"],
extracted_images=[media["path"] for media in saved_media if media["type"] == "image"],
extracted_media=saved_media,
output_format=output_format,
llm_enhanced=False,
ocr_applied=False,
)
# Add PPTX-specific metadata
pptx_metadata = {
"slide_count": extraction_result["slide_count"],
"total_text_length": extraction_result["total_text_length"],
"has_speaker_notes": extraction_result["structure"]["has_notes"],
"slide_layouts": extraction_result["structure"]["slide_layouts"],
"presentation_size": extraction_result["structure"].get("slide_sizes"),
"media_files_count": len(saved_media),
}
# Merge metadata
final_metadata = {**document_metadata.model_dump(), **pptx_metadata}
self._color_log(
f"Successfully extracted content from {file_path.name} "
f"({extraction_result['slide_count']} slides, "
f"{len(formatted_content)} characters, {len(saved_media)} media files)",
Color.green,
)
return ActionResponse(success=True, message=formatted_content, metadata=final_metadata)
except FileNotFoundError as e:
self.logger.error(f"File not found: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False, message=f"File not found: {str(e)}", metadata={"error_type": "file_not_found"}
)
except ValueError as e:
self.logger.error(f"Invalid input: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False, message=f"Invalid input: {str(e)}", metadata={"error_type": "invalid_input"}
)
except Exception as e:
self.logger.error(f"PPTX extraction failed: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=f"PPTX extraction failed: {str(e)}",
metadata={"error_type": "extraction_error"},
)
def mcp_list_supported_formats(self) -> ActionResponse:
"""List all supported presentation formats for extraction.
Returns:
ActionResponse with list of supported file formats and their descriptions
"""
supported_formats = {
"PPTX": "Microsoft PowerPoint Presentation (.pptx) - full support",
"PPT": "Microsoft PowerPoint Presentation (.ppt) - limited support",
}
format_list = "\n".join(
[f"**{format_name}**: {description}" for format_name, description in supported_formats.items()]
)
return ActionResponse(
success=True,
message=f"Supported presentation formats:\n\n{format_list}",
metadata={"supported_formats": list(supported_formats.keys()), "total_formats": len(supported_formats)},
)
# Example usage and entry point
if __name__ == "__main__":
load_dotenv()
# Default arguments for testing
args = ActionArguments(
name="pptx_extraction_service",
transport="stdio",
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
)
# Initialize and run the PPTX extraction service
try:
service = PPTXExtractionCollection(args)
service.run()
except Exception as e:
print(f"An error occurred: {e}: {traceback.format_exc()}")
@@ -0,0 +1,680 @@
import json
import os
import subprocess
import sys
import time
import traceback
import zipfile
from pathlib import Path
from typing import Any, Literal
import pandas as pd
from dotenv import load_dotenv
from openpyxl import load_workbook
from pydantic import Field
from pydantic.fields import FieldInfo
from aworld.logs.util import Color
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
from examples.gaia.mcp_collections.documents.models import DocumentMetadata
class XLSXExtractionCollection(ActionCollection):
"""MCP service for Excel document content extraction using xlrd and pandas.
Supports extraction from XLSX and XLS files.
Provides LLM-friendly text output with structured metadata and media file handling.
Extracts worksheets, formulas, charts, and embedded images.
Includes screenshot functionality for visual representation of Excel data.
"""
def __init__(self, arguments: ActionArguments) -> None:
super().__init__(arguments)
self._media_output_dir = self.workspace / "extracted_media"
self._media_output_dir.mkdir(exist_ok=True)
# Create screenshots directory
self._screenshots_dir = self.workspace / "excel_screenshots"
self._screenshots_dir.mkdir(exist_ok=True)
self.supported_extensions: set = {
".xlsx",
".xls",
}
self._color_log("Excel Extraction Service initialized", Color.green, "debug")
self._color_log(f"Media output directory: {self._media_output_dir}", Color.blue, "debug")
self._color_log(f"Screenshots directory: {self._screenshots_dir}", Color.blue, "debug")
def _create_excel_screenshot(self, file_path: Path, sheet_name: str = None) -> str:
"""Create a JPEG screenshot of the valid Excel area using pyautogui.
Args:
file_path: Path to the Excel file
sheet_name: Specific sheet to screenshot (None for first sheet)
Returns:
Path to the generated JPEG screenshot
"""
try:
import pyautogui
# Generate unique filename
timestamp = int(time.time())
screenshot_filename = f"{file_path.stem}_{sheet_name or 'sheet'}_{timestamp}.jpg"
screenshot_path = self._screenshots_dir / screenshot_filename
# Open Excel file with default application
if sys.platform == "darwin": # macOS
subprocess.run(["open", str(file_path)], check=True)
elif sys.platform == "win32": # Windows
subprocess.run(["start", str(file_path)], shell=True, check=True)
else: # Linux
subprocess.run(["xdg-open", str(file_path)], check=True)
# Wait for Excel to open
time.sleep(3)
# Take screenshot of the entire screen
screenshot = pyautogui.screenshot()
# Convert RGBA to RGB before saving as JPEG
if screenshot.mode == "RGBA":
screenshot = screenshot.convert("RGB")
screenshot.save(screenshot_path, "JPEG", quality=95)
self._color_log(f"Created Excel screenshot: {screenshot_filename}", Color.green)
return str(screenshot_path)
except Exception as e:
self.logger.error(f"Failed to create Excel screenshot with pyautogui: {str(e)}")
raise
def _extract_embedded_media_xlsx(self, file_path: Path) -> list[dict[str, str]]:
"""Extract embedded media from XLSX files.
Args:
file_path: Path to the XLSX file
Returns:
List of dictionaries containing media information
"""
saved_media = []
try:
# Load workbook to extract images
workbook = load_workbook(file_path, data_only=False)
for sheet_name in workbook.sheetnames:
worksheet = workbook[sheet_name]
# Extract images from worksheet
if hasattr(worksheet, "_images"):
for idx, image in enumerate(worksheet._images):
try:
# Generate unique filename
image_filename = f"{file_path.stem}_{sheet_name}_img_{idx}.png"
image_path = self._media_output_dir / image_filename
# Save image
if hasattr(image, "ref"):
# Extract image data
img_data = image._data()
if img_data:
with open(image_path, "wb") as f:
f.write(img_data)
saved_media.append(
{
"type": "image",
"path": str(image_path),
"sheet": sheet_name,
"filename": image_filename,
}
)
self._color_log(f"Saved image: {image_filename}", Color.blue)
except Exception as e:
self.logger.warning(f"Failed to extract image {idx} from sheet {sheet_name}: {str(e)}")
# Also try to extract from ZIP structure for additional media
with zipfile.ZipFile(file_path, "r") as zip_file:
media_files = [f for f in zip_file.namelist() if f.startswith("xl/media/")]
for media_file in media_files:
try:
media_data = zip_file.read(media_file)
media_filename = f"{file_path.stem}_{Path(media_file).name}"
media_path = self._media_output_dir / media_filename
with open(media_path, "wb") as f:
f.write(media_data)
# Determine media type based on extension
media_ext = Path(media_file).suffix.lower()
if media_ext in [".png", ".jpg", ".jpeg", ".gif", ".bmp"]:
media_type = "image"
elif media_ext in [".mp3", ".wav", ".m4a"]:
media_type = "audio"
elif media_ext in [".mp4", ".avi", ".mov"]:
media_type = "video"
else:
media_type = "other"
saved_media.append(
{
"type": media_type,
"path": str(media_path),
"filename": media_filename,
"original_path": media_file,
}
)
self._color_log(f"Saved media: {media_filename}", Color.blue)
except Exception as e:
self.logger.warning(f"Failed to extract media {media_file}: {str(e)}")
except Exception as e:
self.logger.warning(f"Failed to extract media from XLSX: {str(e)}")
return saved_media
def _extract_excel_content(self, file_path: Path, sheet_names: list[str] | None = None) -> dict[str, Any]:
"""Extract content from Excel files using pandas and xlrd.
Args:
file_path: Path to the Excel file
sheet_names: Specific sheets to process (None for all sheets)
Returns:
Dictionary containing extracted content and metadata
"""
start_time = time.time()
try:
# Determine file type and read accordingly
if file_path.suffix.lower() == ".xlsx":
# Use openpyxl engine for XLSX files
excel_file = pd.ExcelFile(file_path, engine="openpyxl")
else:
# Use xlrd engine for XLS files
excel_file = pd.ExcelFile(file_path, engine="xlrd")
# Get all sheet names if not specified
if sheet_names is None:
sheet_names = excel_file.sheet_names
sheets_data = {}
total_rows = 0
total_cols = 0
# Extract data from each sheet
for sheet_name in sheet_names:
if sheet_name in excel_file.sheet_names:
try:
# Read sheet data
df = pd.read_excel(excel_file, sheet_name=sheet_name, header=None)
# Remove completely empty rows and columns
df = df.dropna(how="all").dropna(axis=1, how="all")
if not df.empty:
sheets_data[sheet_name] = {
"data": df,
"shape": df.shape,
"columns": df.columns.tolist(),
"non_empty_cells": df.count().sum(),
}
total_rows += df.shape[0]
total_cols = max(total_cols, df.shape[1])
else:
sheets_data[sheet_name] = {"data": df, "shape": (0, 0), "columns": [], "non_empty_cells": 0}
except Exception as e:
self.logger.warning(f"Failed to read sheet '{sheet_name}': {str(e)}")
sheets_data[sheet_name] = {
"error": str(e),
"shape": (0, 0),
"columns": [],
"non_empty_cells": 0,
}
processing_time = time.time() - start_time
return {
"sheets_data": sheets_data,
"sheet_names": list(sheets_data.keys()),
"total_sheets": len(sheets_data),
"total_rows": total_rows,
"total_columns": total_cols,
"processing_time": processing_time,
"file_engine": "openpyxl" if file_path.suffix.lower() == ".xlsx" else "xlrd",
}
except Exception as e:
self.logger.error(f"Failed to extract Excel content: {str(e)}")
raise
def _format_content_for_llm(
self, extraction_result: dict[str, Any], output_format: str, include_empty_cells: bool = False
) -> str:
"""Format extracted Excel content to be LLM-friendly.
Args:
extraction_result: Result from _extract_excel_content
output_format: Desired output format
include_empty_cells: Whether to include empty cells in output
Returns:
Formatted content string
"""
sheets_data = extraction_result["sheets_data"]
if output_format.lower() == "markdown":
content_parts = []
content_parts.append("# Excel Document Content\n")
content_parts.append(f"**Total Sheets:** {extraction_result['total_sheets']}\n")
content_parts.append(f"**Processing Engine:** {extraction_result['file_engine']}\n\n")
for sheet_name, sheet_info in sheets_data.items():
content_parts.append(f"## Sheet: {sheet_name}\n")
if "error" in sheet_info:
content_parts.append(f"**Error:** {sheet_info['error']}\n\n")
continue
df: pd.DataFrame = sheet_info["data"]
shape = sheet_info["shape"]
content_parts.append(f"**Dimensions:** {shape[0]} rows × {shape[1]} columns\n")
content_parts.append(f"**Non-empty cells:** {sheet_info['non_empty_cells']}\n\n")
if not df.empty:
# Convert DataFrame to markdown table
if include_empty_cells:
# Fill NaN values with empty string for display
df_display = df.fillna("")
else:
# Keep NaN values as they are
df_display = df
# Convert to markdown table
try:
markdown_table = df_display.to_markdown(index=False, tablefmt="pipe")
content_parts.append(f"### Data:\n{markdown_table}\n\n")
except Exception:
# Fallback to string representation
content_parts.append(f"### Data (text format):\n```\n{df_display.to_string()}\n```\n\n")
else:
content_parts.append("*Sheet is empty*\n\n")
return "".join(content_parts)
elif output_format.lower() == "json":
json_data = {
"document_info": {
"total_sheets": extraction_result["total_sheets"],
"total_rows": extraction_result["total_rows"],
"total_columns": extraction_result["total_columns"],
"processing_engine": extraction_result["file_engine"],
},
"sheets": {},
}
for sheet_name, sheet_info in sheets_data.items():
if "error" in sheet_info:
json_data["sheets"][sheet_name] = {"error": sheet_info["error"], "shape": sheet_info["shape"]}
continue
df = sheet_info["data"]
if not df.empty:
# Convert DataFrame to records
if include_empty_cells:
df_records = df.fillna("").to_dict("records")
else:
df_records = df.to_dict("records")
json_data["sheets"][sheet_name] = {
"shape": sheet_info["shape"],
"non_empty_cells": sheet_info["non_empty_cells"],
"data": df_records,
}
else:
json_data["sheets"][sheet_name] = {"shape": sheet_info["shape"], "non_empty_cells": 0, "data": []}
return json.dumps(json_data, indent=2, default=str)
elif output_format.lower() == "html":
html_parts = []
html_parts.append("<html><body>")
html_parts.append("<h1>Excel Document Content</h1>")
html_parts.append(f"<p><strong>Total Sheets:</strong> {extraction_result['total_sheets']}</p>")
html_parts.append(f"<p><strong>Processing Engine:</strong> {extraction_result['file_engine']}</p>")
for sheet_name, sheet_info in sheets_data.items():
html_parts.append(f"<h2>Sheet: {sheet_name}</h2>")
if "error" in sheet_info:
html_parts.append(f"<p><strong>Error:</strong> {sheet_info['error']}</p>")
continue
df = sheet_info["data"]
shape = sheet_info["shape"]
html_parts.append(f"<p><strong>Dimensions:</strong> {shape[0]} rows × {shape[1]} columns</p>")
html_parts.append(f"<p><strong>Non-empty cells:</strong> {sheet_info['non_empty_cells']}</p>")
if not df.empty:
# Convert DataFrame to HTML table
if include_empty_cells:
df_display = df.fillna("")
else:
df_display = df
html_table = df_display.to_html(index=False, escape=False, table_id=f"sheet_{sheet_name}")
html_parts.append(html_table)
else:
html_parts.append("<p><em>Sheet is empty</em></p>")
html_parts.append("</body></html>")
return "".join(html_parts)
else: # text format
content_parts = []
content_parts.append(f"Excel Document Content\n{'=' * 50}\n")
content_parts.append(f"Total Sheets: {extraction_result['total_sheets']}\n")
content_parts.append(f"Processing Engine: {extraction_result['file_engine']}\n\n")
for sheet_name, sheet_info in sheets_data.items():
content_parts.append(f"Sheet: {sheet_name}\n{'-' * 30}\n")
if "error" in sheet_info:
content_parts.append(f"Error: {sheet_info['error']}\n\n")
continue
df = sheet_info["data"]
shape = sheet_info["shape"]
content_parts.append(f"Dimensions: {shape[0]} rows × {shape[1]} columns\n")
content_parts.append(f"Non-empty cells: {sheet_info['non_empty_cells']}\n\n")
if not df.empty:
if include_empty_cells:
df_display = df.fillna("")
else:
df_display = df
content_parts.append(f"Data:\n{df_display.to_string()}\n\n")
else:
content_parts.append("Sheet is empty\n\n")
return "".join(content_parts)
def mcp_extract_excel_content(
self,
file_path: str = Field(description="Path to the Excel document file to extract content from"),
output_format: Literal["markdown", "json", "html", "text"] = Field(
default="markdown", description="Output format: 'markdown', 'json', 'html', or 'text'"
),
extract_images: bool = Field(default=True, description="Whether to extract and save images from the document"),
create_screenshot: bool = Field(
default=False, description="Whether to create a JPEG screenshot of the Excel data"
),
sheet_names: str | None = Field(
default=None, description="Comma-separated list of specific sheet names to process (None for all sheets)"
),
include_empty_cells: bool = Field(default=False, description="Whether to include empty cells in the output"),
screenshot_max_rows: int = Field(default=50, description="Maximum rows to include in screenshot"),
screenshot_max_cols: int = Field(default=20, description="Maximum columns to include in screenshot"),
) -> ActionResponse:
"""Extract content from Excel documents using pandas and xlrd.
This tool provides comprehensive Excel document content extraction with support for:
- XLSX and XLS files
- Multiple worksheets
- Text and numeric data extraction
- Image and media extraction (XLSX only)
- JPEG screenshot generation of Excel data
- Metadata collection
- LLM-optimized output formatting
Args:
file_path: Path to the Excel file
output_format: Desired output format
extract_images: Whether to extract embedded images
create_screenshot: Whether to create a JPEG screenshot
sheet_names: Specific sheets to process
include_empty_cells: Whether to include empty cells
screenshot_max_rows: Maximum rows in screenshot
screenshot_max_cols: Maximum columns in screenshot
Returns:
ActionResponse with extracted content, metadata, media file paths, and screenshot path
"""
try:
# Handle FieldInfo objects
if isinstance(file_path, FieldInfo):
file_path = file_path.default
if isinstance(output_format, FieldInfo):
output_format = output_format.default
if isinstance(extract_images, FieldInfo):
extract_images = extract_images.default
if isinstance(create_screenshot, FieldInfo):
create_screenshot = create_screenshot.default
if isinstance(sheet_names, FieldInfo):
sheet_names = sheet_names.default
if isinstance(include_empty_cells, FieldInfo):
include_empty_cells = include_empty_cells.default
if isinstance(screenshot_max_rows, FieldInfo):
screenshot_max_rows = screenshot_max_rows.default
if isinstance(screenshot_max_cols, FieldInfo):
screenshot_max_cols = screenshot_max_cols.default
# Validate input file
file_path: Path = self._validate_file_path(file_path)
self._color_log(f"Processing Excel document: {file_path.name}", Color.cyan)
# Parse sheet names if provided
target_sheets = None
if sheet_names:
target_sheets = [name.strip() for name in sheet_names.split(",")]
# Extract content from Excel file
extraction_result = self._extract_excel_content(file_path, target_sheets)
# Extract embedded media if requested (XLSX only)
saved_media = []
if extract_images and file_path.suffix.lower() == ".xlsx":
saved_media = self._extract_embedded_media_xlsx(file_path)
elif extract_images and file_path.suffix.lower() == ".xls":
self._color_log("Image extraction not supported for XLS files", Color.yellow)
# Create screenshot if requested
screenshot_path = None
if create_screenshot:
target_sheet = target_sheets[0] if target_sheets else None
screenshot_path = self._create_excel_screenshot(file_path, target_sheet)
# Format content for LLM consumption
formatted_content = self._format_content_for_llm(extraction_result, output_format, include_empty_cells)
# Prepare metadata
file_stats = file_path.stat()
# Create Excel-specific metadata
excel_metadata = {
"sheet_count": extraction_result["total_sheets"],
"sheet_names": extraction_result["sheet_names"],
"total_rows": extraction_result["total_rows"],
"total_columns": extraction_result["total_columns"],
"processing_engine": extraction_result["file_engine"],
"extracted_images": [media["path"] for media in saved_media if media["type"] == "image"],
"extracted_media": saved_media,
"screenshot_path": screenshot_path,
"include_empty_cells": include_empty_cells,
"processed_sheets": target_sheets or extraction_result["sheet_names"],
}
document_metadata = DocumentMetadata(
file_name=file_path.name,
file_size=file_stats.st_size,
file_type=file_path.suffix.lower(),
absolute_path=str(file_path.absolute()),
page_count=extraction_result["total_sheets"], # Use sheet count as "page" count
processing_time=extraction_result["processing_time"],
extracted_images=[media["path"] for media in saved_media if media["type"] == "image"],
extracted_media=saved_media,
output_format=output_format,
llm_enhanced=False,
ocr_applied=False,
)
# Combine standard and Excel-specific metadata
combined_metadata = document_metadata.model_dump()
combined_metadata.update(excel_metadata)
success_message = (
f"Successfully extracted content from {file_path.name} "
f"({len(formatted_content)} characters, {extraction_result['total_sheets']} sheets, "
f"{len(saved_media)} media files"
)
if screenshot_path:
success_message += f", screenshot saved to: {screenshot_path}"
self._color_log(success_message, Color.green)
return ActionResponse(success=True, message=formatted_content, metadata=combined_metadata)
except FileNotFoundError as e:
self.logger.error(f"File not found: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False, message=f"File not found: {str(e)}", metadata={"error_type": "file_not_found"}
)
except ValueError as e:
self.logger.error(f"Invalid input: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=f"Invalid input: {str(e)}",
metadata={"error_type": "invalid_input"},
)
except Exception as e:
self.logger.error(f"Excel extraction failed: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=f"Excel extraction failed: {str(e)}",
metadata={"error_type": "extraction_error"},
)
def mcp_create_excel_screenshot(
self,
file_path: str = Field(description="Path to the Excel document file"),
sheet_name: str | None = Field(
default=None, description="Specific sheet name to screenshot (None for first sheet)"
),
max_rows: int = Field(default=50, description="Maximum number of rows to include"),
max_cols: int = Field(default=20, description="Maximum number of columns to include"),
) -> ActionResponse:
"""Create a JPEG screenshot of the valid Excel area.
This tool creates a visual representation of Excel data as a JPEG image,
useful for further image processing or visual analysis.
Args:
file_path: Path to the Excel file
sheet_name: Specific sheet to screenshot
max_rows: Maximum rows to include in screenshot
max_cols: Maximum columns to include in screenshot
Returns:
ActionResponse with screenshot file path and metadata
"""
try:
# Handle FieldInfo objects
if isinstance(file_path, FieldInfo):
file_path = file_path.default
if isinstance(sheet_name, FieldInfo):
sheet_name = sheet_name.default
if isinstance(max_rows, FieldInfo):
max_rows = max_rows.default
if isinstance(max_cols, FieldInfo):
max_cols = max_cols.default
# Validate input file
file_path: Path = self._validate_file_path(file_path)
self._color_log(f"Creating screenshot for Excel document: {file_path.name}", Color.cyan)
# Create screenshot
screenshot_path = self._create_excel_screenshot(file_path, sheet_name)
# Prepare metadata
file_stats = file_path.stat()
screenshot_stats = Path(screenshot_path).stat()
metadata = {
"source_file": str(file_path.absolute()),
"source_file_size": file_stats.st_size,
"screenshot_path": screenshot_path,
"screenshot_size": screenshot_stats.st_size,
"sheet_name": sheet_name,
"max_rows_displayed": max_rows,
"max_cols_displayed": max_cols,
"format": "JPEG",
}
return ActionResponse(
success=True,
message=f"Excel screenshot created successfully. File saved to: {screenshot_path}",
metadata=metadata,
)
except Exception as e:
self.logger.error(f"Screenshot creation failed: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=f"Screenshot creation failed: {str(e)}",
metadata={"error_type": "screenshot_error"},
)
def mcp_list_supported_formats(self) -> ActionResponse:
"""List all supported Excel formats for extraction.
Returns:
ActionResponse with list of supported file formats and their descriptions
"""
supported_formats = {
"XLSX": "Excel 2007+ format files (.xlsx) - Full support including images",
"XLS": "Excel 97-2003 format files (.xls) - Text and data only",
}
format_list = "\n".join(
[f"**{format_name}**: {description}" for format_name, description in supported_formats.items()]
)
return ActionResponse(
success=True,
message=f"Supported Excel formats:\n\n{format_list}",
metadata={"supported_formats": list(supported_formats.keys()), "total_formats": len(supported_formats)},
)
# Example usage and entry point
if __name__ == "__main__":
load_dotenv()
# Default arguments for testing
args = ActionArguments(
name="excel_extraction_service",
transport="stdio",
workspace=os.getenv("AWORLD_WORKSPACE", "/tmp"),
)
# Initialize and run the Excel extraction service
try:
service = XLSXExtractionCollection(args)
service.run()
except Exception as e:
print(f"An error occurred: {e}: {traceback.format_exc()}")
@@ -0,0 +1,341 @@
import json
import os
import time
import traceback
from collections import defaultdict
from pathlib import Path
from typing import Any, Literal, Optional # Added Optional
import markdown
from dotenv import load_dotenv
from marker.converters.pdf import PdfConverter
from marker.models import create_model_dict
from marker.output import text_from_rendered
from marker.settings import settings
from pydantic import Field
from pydantic.fields import FieldInfo
from aworld.logs.util import Color
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
from examples.gaia.mcp_collections.documents.models import DocumentMetadata
class DocumentExtractionCollection(ActionCollection):
"""MCP service for PDF document content extraction using marker package.
Supports extraction from PDF files only.
Provides LLM-friendly text output with structured metadata and media file handling.
"""
def __init__(self, arguments: ActionArguments) -> None:
super().__init__(arguments)
self._models_loaded = False
self._marker_models = None
self._media_output_dir = self.workspace / "extracted_media"
self._media_output_dir.mkdir(exist_ok=True)
self._extracted_texts_dir = self.workspace / "extracted_texts" # New directory for text files
self._extracted_texts_dir.mkdir(exist_ok=True)
self.supported_extensions = {".pdf"}
self._color_log("PDF Extraction Service initialized", Color.green, "debug")
self._color_log(f"Media output directory: {self._media_output_dir}", Color.blue, "debug")
def _load_marker_models(self) -> None:
"""Load marker models for document processing.
Lazy loading to avoid unnecessary resource consumption.
"""
if not self._models_loaded:
try:
self._color_log("Loading marker models...", Color.yellow)
self._marker_models = create_model_dict()
self._models_loaded = True
self._color_log("Marker models loaded successfully", Color.green)
except Exception as e:
self.logger.error(f"Failed to load marker models: {str(e)}")
raise
def _extract_content_with_marker(
self, file_path: Path, page_range: str | None, force_ocr: bool = False
) -> dict[str, Any]:
"""Extract content using marker package.
Args:
file_path: Path to the document file
page_range: Specific pages to process (e.g., '0,5-10,20')
force_ocr: Use OCR to extract text from images if available
Returns:
Dictionary containing extracted content and metadata
"""
start_time = time.time()
# Prepare marker arguments
marker_args = {
"fname": str(file_path),
"model_lst": self._marker_models,
"max_pages": None,
"langs": None,
"batch_multiplier": 1,
"force_ocr": force_ocr,
}
# Handle page range
if page_range:
# Parse page range string (e.g., "0,5-10,20")
pages = []
for part in page_range.split(","):
if "-" in part:
start, end = map(int, part.split("-"))
pages.extend(range(start, end + 1))
else:
pages.append(int(part))
marker_args["page_range"] = pages
converter: PdfConverter = PdfConverter(artifact_dict=self._marker_models)
rendered = converter(str(file_path))
text, _, images = text_from_rendered(rendered)
text = text.encode(settings.OUTPUT_ENCODING, errors="replace").decode(settings.OUTPUT_ENCODING)
processing_time = time.time() - start_time
return {
"content": text,
"images": images or {},
"metadata": defaultdict(),
"processing_time": processing_time,
}
def _save_extracted_media(self, images: dict[str, Any], file_stem: str) -> list[dict[str, str]]:
"""Save extracted images and return their paths.
Args:
images: Dictionary of extracted images from marker
file_stem: Base name for saving files
Returns:
list of dictionaries containing media type and file paths
"""
saved_media = []
for idx, (page_num, image_data) in enumerate(images.items()):
try:
# Generate unique filename
image_filename = f"{file_stem}_page_{page_num}_img_{idx}.png"
image_path = self._media_output_dir / image_filename
# Save image data
if hasattr(image_data, "save"):
# PIL Image object
image_data.save(image_path)
elif isinstance(image_data, bytes):
# Raw image bytes
with open(image_path, "wb") as f:
f.write(image_data)
else:
# Handle other formats
self.logger.warning(f"Unknown image data type for page {page_num}: {type(image_data)}")
continue
saved_media.append(
{"type": "image", "path": str(image_path), "page": str(page_num), "filename": image_filename}
)
self._color_log(f"Saved image: {image_filename}", Color.blue)
except Exception as e:
self.logger.error(f"Failed to save image from page {page_num}: {str(e)}")
return saved_media
def _format_content_for_llm(self, content: str, output_format: str) -> str:
"""Format extracted content to be LLM-friendly.
Args:
content: Raw extracted content
output_format: Desired output format
Returns:
Formatted content string
"""
if output_format.lower() == "markdown":
# Content is already in markdown format from marker
return content
elif output_format.lower() == "json":
# Structure content as JSON
return json.dumps({"content": content, "format": "structured_text"}, indent=2)
elif output_format.lower() == "html":
# Convert markdown to HTML if needed
try:
return markdown.markdown(content)
except ImportError:
self.logger.warning("markdown package not available, returning raw content")
return content
else:
return content
def mcp_extract_document_content(
self,
file_path: str = Field(description="Path to the PDF document file to extract content from"),
output_format: Literal["markdown", "json", "html"] = Field(
default="markdown", description="Output format: 'markdown', 'json', or 'html'"
),
extract_images: bool = Field(default=True, description="Whether to extract and save images from the document"),
save_extracted_text_to_file: bool = Field(
default=False, description="Save extracted text to a local file"
), # New parameter
use_llm: bool = Field(default=False, description="Use LLM for enhanced accuracy (requires additional setup)"),
page_range: str | None = Field(default=None, description="Specific pages to process (e.g., '0,5-10,20')"),
force_ocr: bool = Field(default=False, description="Force OCR processing on the entire document"),
format_lines: bool = Field(
default=False, description="Reformat lines using local OCR model for better quality"
),
) -> ActionResponse:
"""Extract content from PDF documents using marker package.
This tool provides comprehensive PDF document content extraction with support for:
- PDF files
- Text extraction with proper formatting
- Image and media extraction
- Metadata collection
- LLM-optimized output formatting
Args:
args: Document extraction arguments including file path and options
Returns:
ActionResponse with extracted content, metadata, and media file paths
"""
try:
if isinstance(file_path, FieldInfo):
file_path = file_path.default
if isinstance(output_format, FieldInfo):
output_format = output_format.default
if isinstance(extract_images, FieldInfo):
extract_images = extract_images.default
if isinstance(save_extracted_text_to_file, FieldInfo): # Handle new parameter
save_extracted_text_to_file = save_extracted_text_to_file.default
if isinstance(page_range, FieldInfo):
page_range = page_range.default
if isinstance(use_llm, FieldInfo):
use_llm = use_llm.default
if isinstance(force_ocr, FieldInfo):
force_ocr = force_ocr.default
if isinstance(format_lines, FieldInfo):
format_lines = format_lines.default
# Validate input file
file_path: Path = self._validate_file_path(file_path)
self._color_log(f"Processing document: {file_path.name}", Color.cyan)
# Load marker models if needed
self._load_marker_models()
# Extract content using marker
extraction_result = self._extract_content_with_marker(file_path, page_range, force_ocr)
# Save extracted media if requested
saved_media = []
if extract_images and extraction_result["images"]:
saved_media = self._save_extracted_media(extraction_result["images"], file_path.stem)
# Format content for LLM consumption
formatted_content = self._format_content_for_llm(extraction_result["content"], output_format)
# Save extracted text to file if requested
saved_text_path_str: Optional[str] = None
if save_extracted_text_to_file:
text_file_name = f"{file_path.stem}_extracted_text.txt"
saved_text_path = self._extracted_texts_dir / text_file_name
try:
with open(saved_text_path, "w", encoding="utf-8") as f:
f.write(formatted_content)
saved_text_path_str = str(saved_text_path.absolute())
self._color_log(f"Saved extracted text to: {saved_text_path_str}", Color.blue)
except Exception as e:
self.logger.error(f"Failed to save extracted text to {saved_text_path}: {str(e)}")
# Optionally, you might want to reflect this failure in the response
# Prepare metadata
file_stats = file_path.stat()
document_metadata = DocumentMetadata(
file_name=file_path.name,
file_size=file_stats.st_size,
file_type=file_path.suffix.lower(),
absolute_path=str(file_path.absolute()),
page_count=extraction_result["metadata"].get("page_count"),
processing_time=extraction_result["processing_time"],
extracted_images=[media["path"] for media in saved_media if media["type"] == "image"],
extracted_media=saved_media,
output_format=output_format,
llm_enhanced=use_llm,
ocr_applied=force_ocr or format_lines,
extracted_text_file_path=saved_text_path_str,
)
self._color_log(
f"Successfully extracted content from {file_path.name} "
f"({len(formatted_content)} characters, {len(saved_media)} media files)",
Color.green,
)
return ActionResponse(success=True, message=formatted_content, metadata=document_metadata.model_dump())
except FileNotFoundError as e:
self.logger.error(f"File not found: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False, message=f"File not found: {str(e)}", metadata={"error_type": "file_not_found"}
)
except ValueError as e:
self.logger.error(f"Invalid input: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=f"Invalid input: {str(e)}: {traceback.format_exc()}",
metadata={"error_type": "invalid_input"},
)
except Exception as e:
self.logger.error(f"Document extraction failed: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=f"Document extraction failed: {str(e)}",
metadata={"error_type": "extraction_error"},
)
def mcp_list_supported_formats(self) -> ActionResponse:
"""list all supported document formats for extraction.
Returns:
ActionResponse with list of supported file formats and their descriptions
"""
supported_formats = {
"PDF": "Portable Document Format files (.pdf)",
}
format_list = "\n".join(
[f"**{format_name}**: {description}" for format_name, description in supported_formats.items()]
)
return ActionResponse(
success=True,
message=f"Supported document formats:\n\n{format_list}",
metadata={"supported_formats": list(supported_formats.keys()), "total_formats": len(supported_formats)},
)
# Example usage and entry point
if __name__ == "__main__":
load_dotenv()
# Default arguments for testing
args = ActionArguments(
name="document_extraction_service",
transport="stdio",
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
)
# Initialize and run the document extraction service
try:
service = DocumentExtractionCollection(args)
service.run()
except Exception as e:
print(f"An error occurred: {e}: {traceback.format_exc()}")
@@ -0,0 +1,648 @@
import json
import os
import time
import traceback
from pathlib import Path
from typing import Any, Literal
import chardet
from dotenv import load_dotenv
from pydantic import Field
from pydantic.fields import FieldInfo
from aworld.logs.util import Color
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
from examples.gaia.mcp_collections.documents.models import DocumentMetadata
from examples.gaia.mcp_collections.utils import get_mime_type
class TextExtractionCollection(ActionCollection):
"""MCP service for text document content extraction.
Supports extraction from TXT and other raw text format files.
Provides LLM-friendly text output with structured metadata and encoding detection.
Handles various text encodings and provides comprehensive file analysis.
"""
def __init__(self, arguments: ActionArguments) -> None:
super().__init__(arguments)
self._media_output_dir = self.workspace / "extracted_media"
self._media_output_dir.mkdir(exist_ok=True)
self.supported_extensions: set = {
".txt",
".text",
".log",
".md",
".markdown",
".rst",
".rtf",
".csv",
".tsv",
".json",
".xml",
".yaml",
".yml",
".ini",
".cfg",
".conf",
".properties",
".sql",
".py",
".js",
".html",
".htm",
".css",
".java",
".cpp",
".c",
".h",
".php",
".rb",
".go",
".rs",
".sh",
".bat",
".ps1",
".r",
".m",
".swift",
".kt",
".scala",
".pl",
".lua",
".vim",
".tex",
".bib",
}
self._color_log("Text Extraction Service initialized", Color.green, "debug")
self._color_log(f"Media output directory: {self._media_output_dir}", Color.blue, "debug")
def _validate_file_path(self, file_path: str) -> Path:
"""Validate and resolve file path.
Args:
file_path: Path to the text document file
Returns:
Resolved Path object
Raises:
FileNotFoundError: If file doesn't exist
ValueError: If file type is not supported
"""
path = Path(file_path)
if not path.is_absolute():
path = self.workspace / path
if not path.exists():
raise FileNotFoundError(f"File not found: {path}")
# Also check MIME type for files without extensions or unknown extensions
mime_type = get_mime_type(str(path), default_mime="text/plain")
is_text_mime = mime_type and mime_type.startswith("text/")
if path.suffix.lower() not in self.supported_extensions and not is_text_mime:
# Try to detect if it's a text file by reading a small sample
try:
with open(path, "rb") as f:
sample = f.read(1024)
# Check if the sample contains mostly printable characters
if self._is_likely_text(sample):
self._color_log(f"Detected text file without standard extension: {path.suffix}", Color.yellow)
else:
raise ValueError(
f"Unsupported file type: {path.suffix}. "
f"Supported types: {', '.join(sorted(self.supported_extensions))} or text MIME types"
)
except Exception as e:
raise ValueError(
f"Cannot determine if file is text: {str(e)}. "
f"Supported types: {', '.join(sorted(self.supported_extensions))}"
) from e
return path
def _is_likely_text(self, data: bytes) -> bool:
"""Check if binary data is likely to be text.
Args:
data: Binary data sample
Returns:
True if data appears to be text
"""
if not data:
return True
# Check for null bytes (common in binary files)
if b"\x00" in data:
return False
# Try to decode as text
try:
data.decode("utf-8")
return True
except UnicodeDecodeError:
pass
# Check if most bytes are printable ASCII
printable_count = sum(1 for byte in data if 32 <= byte <= 126 or byte in [9, 10, 13])
return printable_count / len(data) > 0.7
def _detect_encoding(self, file_path: Path) -> dict[str, Any]:
"""Detect file encoding and other characteristics.
Args:
file_path: Path to the text file
Returns:
Dictionary containing encoding information
"""
encoding_info = {
"detected_encoding": None,
"confidence": 0.0,
"bom_detected": False,
"line_endings": None,
"is_binary": False,
}
try:
# Read file in binary mode for encoding detection
with open(file_path, "rb") as f:
raw_data = f.read()
if not raw_data:
encoding_info["detected_encoding"] = "utf-8"
encoding_info["confidence"] = 1.0
return encoding_info
# Check for BOM (Byte Order Mark)
if raw_data.startswith(b"\xef\xbb\xbf"):
encoding_info["bom_detected"] = True
encoding_info["detected_encoding"] = "utf-8-sig"
encoding_info["confidence"] = 1.0
elif raw_data.startswith(b"\xff\xfe"):
encoding_info["bom_detected"] = True
encoding_info["detected_encoding"] = "utf-16-le"
encoding_info["confidence"] = 1.0
elif raw_data.startswith(b"\xfe\xff"):
encoding_info["bom_detected"] = True
encoding_info["detected_encoding"] = "utf-16-be"
encoding_info["confidence"] = 1.0
else:
# Use chardet for encoding detection
detection_result = chardet.detect(raw_data)
encoding_info["detected_encoding"] = detection_result.get("encoding", "utf-8")
encoding_info["confidence"] = detection_result.get("confidence", 0.0)
# Detect line endings
if b"\r\n" in raw_data:
encoding_info["line_endings"] = "CRLF (Windows)"
elif b"\n" in raw_data:
encoding_info["line_endings"] = "LF (Unix/Linux/Mac)"
elif b"\r" in raw_data:
encoding_info["line_endings"] = "CR (Classic Mac)"
else:
encoding_info["line_endings"] = "None detected"
# Check if file appears to be binary
encoding_info["is_binary"] = not self._is_likely_text(raw_data[:1024])
except Exception as e:
self.logger.warning(f"Failed to detect encoding: {str(e)}")
encoding_info["detected_encoding"] = "utf-8"
encoding_info["confidence"] = 0.0
return encoding_info
def _extract_text_content(self, file_path: Path, encoding: str | None = None) -> dict[str, Any]:
"""Extract content from text files.
Args:
file_path: Path to the text file
encoding: Specific encoding to use (None for auto-detection)
Returns:
Dictionary containing extracted content and metadata
"""
start_time = time.time()
# Detect encoding if not specified
encoding_info = self._detect_encoding(file_path)
if encoding:
# Use specified encoding
target_encoding = encoding
self._color_log(f"Using specified encoding: {encoding}", Color.blue)
else:
# Use detected encoding
target_encoding = encoding_info["detected_encoding"]
self._color_log(
f"Detected encoding: {target_encoding} (confidence: {encoding_info['confidence']:.2f})", Color.blue
)
try:
# Read file content
with open(file_path, "r", encoding=target_encoding, errors="replace") as f:
content = f.read()
# Analyze content
lines = content.splitlines()
# Calculate statistics
char_count = len(content)
line_count = len(lines)
word_count = len(content.split()) if content.strip() else 0
# Find longest and shortest lines
line_lengths = [len(line) for line in lines]
max_line_length = max(line_lengths) if line_lengths else 0
min_line_length = min(line_lengths) if line_lengths else 0
avg_line_length = sum(line_lengths) / len(line_lengths) if line_lengths else 0
# Count empty lines
empty_lines = sum(1 for line in lines if not line.strip())
# Detect file type based on content patterns
content_type = self._detect_content_type(content, file_path)
processing_time = time.time() - start_time
return {
"content": content,
"encoding_info": encoding_info,
"statistics": {
"character_count": char_count,
"line_count": line_count,
"word_count": word_count,
"empty_lines": empty_lines,
"max_line_length": max_line_length,
"min_line_length": min_line_length,
"avg_line_length": round(avg_line_length, 2),
},
"content_type": content_type,
"processing_time": processing_time,
"used_encoding": target_encoding,
}
except UnicodeDecodeError as e:
self.logger.error(f"Failed to decode file with encoding {target_encoding}: {str(e)}")
# Try with fallback encodings
fallback_encodings = ["utf-8", "latin-1", "cp1252", "iso-8859-1"]
for fallback_encoding in fallback_encodings:
if fallback_encoding != target_encoding:
try:
with open(file_path, "r", encoding=fallback_encoding, errors="replace") as f:
content = f.read()
self._color_log(f"Successfully read with fallback encoding: {fallback_encoding}", Color.yellow)
# Recalculate with fallback encoding
lines = content.splitlines()
char_count = len(content)
line_count = len(lines)
word_count = len(content.split()) if content.strip() else 0
processing_time = time.time() - start_time
return {
"content": content,
"encoding_info": encoding_info,
"statistics": {
"character_count": char_count,
"line_count": line_count,
"word_count": word_count,
"empty_lines": sum(1 for line in lines if not line.strip()),
"max_line_length": max(len(line) for line in lines) if lines else 0,
"min_line_length": min(len(line) for line in lines) if lines else 0,
"avg_line_length": round(sum(len(line) for line in lines) / len(lines), 2)
if lines
else 0,
},
"content_type": self._detect_content_type(content, file_path),
"processing_time": processing_time,
"used_encoding": fallback_encoding,
"encoding_fallback": True,
}
except Exception:
continue
raise ValueError("Unable to decode file with any supported encoding") from e
def _detect_content_type(self, content: str, file_path: Path) -> str:
"""Detect the type of content based on file extension and content patterns.
Args:
content: File content
file_path: Path to the file
Returns:
Detected content type
"""
extension = file_path.suffix.lower()
# Map extensions to content types
extension_map = {
".py": "Python source code",
".js": "JavaScript source code",
".html": "HTML document",
".htm": "HTML document",
".css": "CSS stylesheet",
".json": "JSON data",
".xml": "XML document",
".yaml": "YAML configuration",
".yml": "YAML configuration",
".md": "Markdown document",
".markdown": "Markdown document",
".rst": "reStructuredText document",
".csv": "CSV data",
".tsv": "TSV data",
".sql": "SQL script",
".log": "Log file",
".ini": "Configuration file",
".cfg": "Configuration file",
".conf": "Configuration file",
}
if extension in extension_map:
return extension_map[extension]
# Content-based detection
content_lower = content.lower().strip()
if content_lower.startswith("<?xml"):
return "XML document"
elif content_lower.startswith("{") and content_lower.endswith("}"):
return "JSON-like data"
elif content_lower.startswith("[") and content_lower.endswith("]"):
return "JSON array or configuration"
elif "#!/" in content[:50]:
return "Script file"
elif content.count(",") > content.count("\n") * 2:
return "CSV-like data"
else:
return "Plain text"
def _format_content_for_llm(
self, extraction_result: dict[str, Any], output_format: str, max_length: int | None = None
) -> str:
"""Format extracted text content to be LLM-friendly.
Args:
extraction_result: Result from _extract_text_content
output_format: Desired output format
max_length: Maximum length of content to include (None for no limit)
Returns:
Formatted content string
"""
content = extraction_result["content"]
stats = extraction_result["statistics"]
content_type = extraction_result["content_type"]
# Truncate content if needed
if max_length and len(content) > max_length:
content = (
content[:max_length]
+ f"\n\n[Content truncated - showing first {max_length} characters of {stats['character_count']} total]"
)
if output_format.lower() == "markdown":
formatted_parts = []
formatted_parts.append("# Text Document Content\n")
formatted_parts.append(f"**File Type:** {content_type}\n")
formatted_parts.append(f"**Encoding:** {extraction_result['used_encoding']}\n")
formatted_parts.append("**Statistics:**\n")
formatted_parts.append(f"- Characters: {stats['character_count']:,}\n")
formatted_parts.append(f"- Lines: {stats['line_count']:,}\n")
formatted_parts.append(f"- Words: {stats['word_count']:,}\n")
formatted_parts.append(f"- Empty lines: {stats['empty_lines']:,}\n")
formatted_parts.append(f"- Average line length: {stats['avg_line_length']} characters\n\n")
formatted_parts.append(f"## Content\n\n```\n{content}\n```")
return "".join(formatted_parts)
elif output_format.lower() == "json":
json_data = {
"document_info": {
"content_type": content_type,
"encoding": extraction_result["used_encoding"],
"statistics": stats,
},
"content": content,
}
return json.dumps(json_data, indent=2, ensure_ascii=False)
elif output_format.lower() == "html":
html_parts = []
html_parts.append("<html><head><meta charset='utf-8'></head><body>")
html_parts.append("<h1>Text Document Content</h1>")
html_parts.append(f"<p><strong>File Type:</strong> {content_type}</p>")
html_parts.append(f"<p><strong>Encoding:</strong> {extraction_result['used_encoding']}</p>")
html_parts.append("<h2>Statistics</h2>")
html_parts.append("<ul>")
html_parts.append(f"<li>Characters: {stats['character_count']:,}</li>")
html_parts.append(f"<li>Lines: {stats['line_count']:,}</li>")
html_parts.append(f"<li>Words: {stats['word_count']:,}</li>")
html_parts.append(f"<li>Empty lines: {stats['empty_lines']:,}</li>")
html_parts.append(f"<li>Average line length: {stats['avg_line_length']} characters</li>")
html_parts.append("</ul>")
html_parts.append("<h2>Content</h2>")
html_parts.append(f"<pre><code>{content}</code></pre>")
html_parts.append("</body></html>")
return "".join(html_parts)
else: # text format
text_parts = []
text_parts.append(f"Text Document Content\n{'=' * 50}\n")
text_parts.append(f"File Type: {content_type}\n")
text_parts.append(f"Encoding: {extraction_result['used_encoding']}\n")
text_parts.append("\nStatistics:\n")
text_parts.append(f" Characters: {stats['character_count']:,}\n")
text_parts.append(f" Lines: {stats['line_count']:,}\n")
text_parts.append(f" Words: {stats['word_count']:,}\n")
text_parts.append(f" Empty lines: {stats['empty_lines']:,}\n")
text_parts.append(f" Average line length: {stats['avg_line_length']} characters\n")
text_parts.append(f"\nContent:\n{'-' * 30}\n{content}")
return "".join(text_parts)
def mcp_extract_text_content(
self,
file_path: str = Field(description="Path to the text document file to extract content from"),
output_format: Literal["markdown", "json", "html", "text"] = Field(
default="markdown", description="Output format: 'markdown', 'json', 'html', or 'text'"
),
encoding: str | None = Field(default=None, description="Specific encoding to use (None for auto-detection)"),
max_content_length: int | None = Field(
default=None, description="Maximum length of content to include in output (None for no limit)"
),
) -> ActionResponse:
"""Extract content from text documents with encoding detection and analysis.
This tool provides comprehensive text document content extraction with support for:
- Various text file formats (TXT, MD, CSV, JSON, XML, source code, etc.)
- Automatic encoding detection with fallback options
- Content type detection and analysis
- Comprehensive text statistics
- LLM-optimized output formatting
- Binary file detection and handling
Args:
file_path: Path to the text file
output_format: Desired output format
encoding: Specific encoding to use
max_content_length: Maximum content length to include
Returns:
ActionResponse with extracted content, metadata, and file analysis
"""
try:
# Handle FieldInfo objects
if isinstance(file_path, FieldInfo):
file_path = file_path.default
if isinstance(output_format, FieldInfo):
output_format = output_format.default
if isinstance(encoding, FieldInfo):
encoding = encoding.default
if isinstance(max_content_length, FieldInfo):
max_content_length = max_content_length.default
# Validate input file
file_path: Path = self._validate_file_path(file_path)
self._color_log(f"Processing text document: {file_path.name}", Color.cyan)
# Extract content from text file
extraction_result = self._extract_text_content(file_path, encoding)
# Check if file appears to be binary
if extraction_result["encoding_info"]["is_binary"]:
self._color_log("Warning: File appears to contain binary data", Color.yellow)
# Format content for LLM consumption
formatted_content = self._format_content_for_llm(extraction_result, output_format, max_content_length)
# Prepare metadata
file_stats = file_path.stat()
# Create text-specific metadata
text_metadata = {
"content_type": extraction_result["content_type"],
"encoding_info": extraction_result["encoding_info"],
"text_statistics": extraction_result["statistics"],
"used_encoding": extraction_result["used_encoding"],
"encoding_fallback": extraction_result.get("encoding_fallback", False),
"content_truncated": max_content_length and len(extraction_result["content"]) > max_content_length,
"original_content_length": extraction_result["statistics"]["character_count"],
}
document_metadata = DocumentMetadata(
file_name=file_path.name,
file_size=file_stats.st_size,
file_type=file_path.suffix.lower() or ".txt",
absolute_path=str(file_path.absolute()),
page_count=extraction_result["statistics"]["line_count"], # Use line count as "page" count
processing_time=extraction_result["processing_time"],
extracted_images=[], # Text files don't contain images
extracted_media=[], # Text files don't contain media
output_format=output_format,
llm_enhanced=False,
ocr_applied=False,
)
# Combine standard and text-specific metadata
combined_metadata = document_metadata.model_dump()
combined_metadata.update(text_metadata)
self._color_log(
f"Successfully extracted content from {file_path.name} "
f"({extraction_result['statistics']['character_count']:,} characters, "
f"{extraction_result['statistics']['line_count']:,} lines, "
f"encoding: {extraction_result['used_encoding']})",
Color.green,
)
return ActionResponse(success=True, message=formatted_content, metadata=combined_metadata)
except FileNotFoundError as e:
self.logger.error(f"File not found: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False, message=f"File not found: {str(e)}", metadata={"error_type": "file_not_found"}
)
except ValueError as e:
self.logger.error(f"Invalid input: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=f"Invalid input: {str(e)}",
metadata={"error_type": "invalid_input"},
)
except Exception as e:
self.logger.error(f"Text extraction failed: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=f"Text extraction failed: {str(e)}",
metadata={"error_type": "extraction_error"},
)
def mcp_list_supported_formats(self) -> ActionResponse:
"""List all supported text formats for extraction.
Returns:
ActionResponse with list of supported file formats and their descriptions
"""
supported_formats = {
"TXT": "Plain text files (.txt, .text)",
"Markdown": "Markdown documents (.md, .markdown)",
"CSV/TSV": "Comma/Tab separated values (.csv, .tsv)",
"JSON": "JSON data files (.json)",
"XML": "XML documents (.xml)",
"YAML": "YAML configuration files (.yaml, .yml)",
"Source Code": "Programming language files (.py, .js, .html, .css, etc.)",
"Configuration": "Config files (.ini, .cfg, .conf, .properties)",
"Logs": "Log files (.log)",
"Documentation": "Documentation files (.rst, .rtf)",
"Scripts": "Script files (.sh, .bat, .ps1)",
"Other Text": "Any file with text MIME type or detectable text content",
}
format_list = "\n".join(
[f"**{format_name}**: {description}" for format_name, description in supported_formats.items()]
)
return ActionResponse(
success=True,
message=f"Supported text formats:\n\n{format_list}\n\n"
"**Note:** The service automatically detects encoding and "
"can handle files without standard extensions if they contain readable text.",
metadata={
"supported_formats": list(supported_formats.keys()),
"total_formats": len(supported_formats),
"encoding_detection": True,
"binary_detection": True,
},
)
# Example usage and entry point
if __name__ == "__main__":
load_dotenv()
# Default arguments for testing
args = ActionArguments(
name="text_extraction_service",
transport="stdio",
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
)
# Initialize and run the text extraction service
try:
service = TextExtractionCollection(args)
service.run()
except Exception as e:
print(f"An error occurred: {e}: {traceback.format_exc()}")
@@ -0,0 +1,349 @@
import os
import time
import traceback
from pathlib import Path
from typing import Literal
from dotenv import load_dotenv
from pydantic import BaseModel, Field
from pydantic.fields import FieldInfo
from aworld.config.conf import AgentConfig
from aworld.logs.util import Color
from aworld.models.llm import call_llm_model, get_llm_model
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
class CodeGenerationMetadata(BaseModel):
"""Metadata for code generation results."""
model_name: str | None = None
code_style: str | None = None
code_length: int | None = None
line_count: int | None = None
processing_time_seconds: float | None = None
temperature: float | None = None
has_requirements: bool | None = None
has_context: bool | None = None
saved_file_path: str | None = None
file_save_error: str | None = None
error_type: str | None = None
error_message: str | None = None
class CodeCollection(ActionCollection):
"""MCP service for generating executable Python code snippets using LLM.
Supports code generation for:
- Data processing and analysis tasks
- Algorithm implementations
- Utility functions and scripts
- Problem-solving code snippets
- Educational programming examples
"""
def __init__(self, arguments: ActionArguments) -> None:
super().__init__(arguments)
# Initialize code generation model configuration
self._llm_config = AgentConfig(
llm_provider="openai",
llm_model_name=os.getenv("CODE_LLM_MODEL_NAME", "anthropic/claude-sonnet-4"),
llm_api_key=os.getenv("CODE_LLM_API_KEY"),
llm_base_url=os.getenv("CODE_LLM_BASE_URL"),
)
self._color_log("Code Generation Service initialized", Color.green, "debug")
self._color_log(f"Using model: {self._llm_config.llm_model_name}", Color.blue, "debug")
def _prepare_code_prompt(self, task_description: str, requirements: str = "", context: str = "") -> str:
"""Prepare the code generation prompt with task description and optional requirements.
Args:
task_description: The main task for code generation
requirements: Optional specific requirements or constraints
context: Optional additional context or background information
Returns:
Formatted prompt string
"""
prompt_parts = [f"Task: {task_description}"]
if requirements:
prompt_parts.append(f"Requirements: {requirements}")
if context:
prompt_parts.append(f"Context: {context}")
return "\n\n".join(prompt_parts)
def _call_code_model(self, prompt: str, temperature: float = 0.1) -> str:
"""Call the code generation model with the prepared prompt.
Args:
prompt: The formatted prompt for code generation
temperature: Model temperature for response variability
Returns:
Generated code from the model
Raises:
Exception: If model call fails
"""
response = call_llm_model(
llm_model=get_llm_model(conf=self._llm_config),
messages=[
{
"role": "system",
"content": (
"You are an expert Python programmer. Generate clean, efficient, and "
"well-documented Python code that solves the given task. "
"Include proper error handling and follow Python best practices. "
"Return only executable Python code with minimal explanatory comments."
),
},
{"role": "user", "content": prompt},
],
temperature=temperature,
)
return response.content
def _extract_python_code(self, response: str) -> str:
"""Extract Python code from the model response.
Args:
response: Raw response from the model
Returns:
Extracted Python code
"""
# Remove markdown code blocks if present
lines = response.strip().split("\n")
# Find code block boundaries
start_idx = 0
end_idx = len(lines)
for i, line in enumerate(lines):
if line.strip().startswith("```python") or line.strip().startswith("```"):
start_idx = i + 1
break
for i in range(len(lines) - 1, -1, -1):
if lines[i].strip() == "```":
end_idx = i
break
# Extract the code
code_lines = lines[start_idx:end_idx]
return "\n".join(code_lines).strip()
def mcp_generate_python_code(
self,
task_description: str = Field(description="Description of the programming task or problem to solve"),
requirements: str = Field(
default="", description="Specific requirements, constraints, or specifications for the code"
),
context: str = Field(default="", description="Additional context or background information"),
temperature: float = Field(
default=0.1,
description="Model temperature for code generation (0.0-1.0, lower = more deterministic)",
ge=0.0,
le=1.0,
),
code_style: Literal["minimal", "documented", "verbose"] = Field(
default="documented",
description="Style of generated code: minimal (concise), documented (with comments), verbose (detailed)",
),
save_to_file_path: str | None = Field(
default=None,
description="Optional. Path to save the generated Python snippet. e.g., 'output/generated_script.py'",
),
) -> ActionResponse:
"""Generate executable Python code snippets based on task description.
This tool provides comprehensive code generation capabilities for:
- Solve simple math tasks and validations
- Data processing and analysis tasks
- Algorithm implementations and optimizations
- Utility functions and helper scripts
- Problem-solving code snippets
- Educational programming examples
- API integrations and automation scripts
Strengths:
- Generates clean, executable Python code
- Follows modern Python best practices (>=3.11)
- Includes proper error handling
- Supports various coding styles and complexity levels
Limitations:
- Cannot execute or test the generated code
- May require manual adjustments for specific environments
- Limited to Python programming language
Args:
task_description: Clear description of the programming task
requirements: Specific requirements or constraints
context: Additional context or background information
temperature: Model temperature controlling randomness
code_style: Style preference for the generated code
save_to_file_path: Optional. If provided, saves the generated code to this path within the workspace.
Returns:
ActionResponse with generated Python code and metadata
"""
try:
# Handle FieldInfo objects
if isinstance(task_description, FieldInfo):
task_description = task_description.default
if isinstance(requirements, FieldInfo):
requirements = requirements.default
if isinstance(context, FieldInfo):
context = context.default
if isinstance(temperature, FieldInfo):
temperature = temperature.default
if isinstance(code_style, FieldInfo):
code_style = code_style.default
if isinstance(save_to_file_path, FieldInfo):
save_to_file_path = save_to_file_path.default
# Validate input
if not task_description or not task_description.strip():
raise ValueError("Task description is required for code generation")
self._color_log(f"Generating code for: {task_description[:100]}...", Color.cyan)
start_time = time.time()
# Prepare the code generation prompt
prompt = self._prepare_code_prompt(task_description, requirements, context)
# Enhance prompt based on code style
if code_style == "minimal":
prompt += "\n\nGenerate concise, minimal code without extensive comments."
elif code_style == "verbose":
prompt += "\n\nGenerate detailed code with comprehensive comments and explanations."
elif code_style == "documented":
prompt += "\n\nGenerate well-documented code with clear comments and docstrings."
# Call the code generation model
raw_response = self._call_code_model(prompt, temperature)
# Extract clean Python code
generated_code = self._extract_python_code(raw_response)
processing_time = time.time() - start_time
# Populate metadata fields
metadata = CodeGenerationMetadata(
model_name=self._llm_config.llm_model_name,
code_style=code_style,
code_length=len(generated_code),
line_count=len(generated_code.split("\n")),
processing_time_seconds=round(processing_time, 2),
temperature=temperature,
has_requirements=bool(requirements.strip()),
has_context=bool(context.strip()),
)
# Save the generated code to a file if path is provided
if save_to_file_path:
try:
# Use _validate_file_path to ensure path is within workspace and get absolute path
# The check_existence=False allows creating a new file.
output_file_path_obj = Path(self._validate_file_path(save_to_file_path))
# Ensure parent directories exist
output_file_path_obj.parent.mkdir(parents=True, exist_ok=True)
with open(output_file_path_obj, "w", encoding="utf-8") as f:
f.write(generated_code)
metadata.saved_file_path = str(output_file_path_obj)
self._color_log(f"Generated code also saved to: {output_file_path_obj}", Color.blue)
except Exception as e:
self.logger.error(f"Failed to save code to file '{save_to_file_path}': {str(e)}")
metadata.file_save_error = str(e)
self._color_log(
f"Successfully generated code ({metadata.code_length} characters, "
f"{metadata.processing_time_seconds:.2f}s)",
Color.green,
)
return ActionResponse(success=True, message=generated_code, metadata=metadata.model_dump(exclude_none=True))
except ValueError as e:
self.logger.error(f"Invalid input: {str(e)}")
metadata.error_type = "invalid_input"
metadata.error_message = str(e)
return ActionResponse(
success=False,
message=f"Invalid input: {str(e)}",
metadata=metadata.model_dump(exclude_none=True),
)
except Exception as e:
self.logger.error(f"Code generation failed: {str(e)}: {traceback.format_exc()}")
metadata.error_type = "generation_error"
metadata.error_message = str(e)
return ActionResponse(
success=False,
message=f"Code generation failed: {str(e)}",
metadata=metadata.model_dump(exclude_none=True),
)
def mcp_get_code_capabilities(self) -> ActionResponse:
"""Get information about the code generation service capabilities.
Returns:
ActionResponse with service capabilities and configuration
"""
capabilities = {
"Data Processing": "Generate code for data manipulation, analysis, and visualization",
"Algorithm Implementation": "Create efficient algorithms and data structures",
"Utility Functions": "Build helper functions and reusable code components",
"Problem Solving": "Generate solutions for programming challenges and tasks",
"API Integration": "Create code for working with APIs and web services",
"Automation Scripts": "Build scripts for task automation and workflow optimization",
}
capability_list = "\n".join(
[f"**{capability}**: {description}" for capability, description in capabilities.items()]
)
metadata = {
"model_name": self._llm_config.llm_model_name,
"provider": self._llm_config.llm_provider,
"supported_capabilities": list(capabilities.keys()),
"total_capabilities": len(capabilities),
"code_styles": ["minimal", "documented", "verbose"],
"python_version": ">=3.11",
"supported_language": "Python",
}
return ActionResponse(
success=True,
message=f"Code Generation Service Capabilities:\n\n{capability_list}",
metadata=metadata,
)
# Example usage and entry point
if __name__ == "__main__":
load_dotenv()
# Default arguments for testing
args = ActionArguments(
name="code_generation_service",
transport="stdio",
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
)
# Initialize and run the code generation service
try:
service = CodeCollection(args)
service.run()
except Exception as e:
print(f"An error occurred: {e}: {traceback.format_exc()}")
@@ -0,0 +1,260 @@
import os
import time
import traceback
from typing import Literal
from dotenv import load_dotenv
from pydantic import Field
from pydantic.fields import FieldInfo
from aworld.config.conf import AgentConfig
from aworld.logs.util import Color
from aworld.models.llm import call_llm_model, get_llm_model
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
class GuardCollection(ActionCollection):
"""MCP service for diagnosing and correcting (if necessary) the reasoning/thinking process already existed in the currect context, or avoid the potential loopholes in the future, towards solving the complex problem correctly, through powerful guarding model with sophisticated experience.
The MUST Choice for the Thinking Process Reviewing phase, good at diagnosing the reasoning process in the context or giving valuable suggestions in advance.
Supports advanced guarding for reasoning process:
- Identify potential loopholes or oversights in the reasoning process already existed in the currect context, while solving the complex problem.
- If necessary, provide the corresponding supplements or guidance to the reasoning process in advance, to maneuver the reasoning/thinking process towards solving the complex problem in a proper direction.
"""
def __init__(self, arguments: ActionArguments) -> None:
super().__init__(arguments)
env_path = "/Users/zhitianxie/PycharmProjects/AWorld_gaia_July/AWorld/examples/gaia/cmd/agent_deploy/gaia_agent/.env"
load_dotenv(env_path, override=True, verbose=True)
# Initialize guarding model configuration
self._llm_config = AgentConfig(
llm_provider="openai",
# llm_model_name="google/gemini-2.5-flash-preview-05-20:thinking",
llm_model_name=os.getenv("GUARD_LLM_MODEL_NAME", "deepseek/deepseek-r1-0528:free"),
llm_api_key=os.getenv("GUARD_LLM_API_KEY"),
llm_base_url=os.getenv("GUARD_LLM_BASE_URL"),
)
self._color_log("Intelligence Guard Service initialized", Color.green, "debug")
self._color_log(f"Using model: {self._llm_config.llm_model_name}", Color.blue, "debug")
def _prepare_guarding_prompt(self, question: str, original_task: str = "") -> str:
"""Prepare the guarding prompt with question and optional context.
Args:
question: The main question for guarding the reasoning process, such as 'is there any potential loopholes or oversights in the reasoning process?'
original_task: Optional original task description for context
Returns:
Formatted prompt string
"""
if original_task:
return f"Original Task: {original_task}\n\nQuestion: {question}"
return f"Question: {question}"
def _call_guarding_model(self, prompt: str, temperature: float = 0.1) -> str:
"""Call the guarding model with the prepared prompt.
Args:
prompt: The formatted prompt for guarding the reasoning process
temperature: Model temperature for response variability
Returns:
guarding result from the model
Raises:
Exception: If model call fails
"""
response = call_llm_model(
llm_model=get_llm_model(conf=self._llm_config),
messages=[
{
"role": "system",
"content": (
"## Your Role\n"
"You are an expert at identifying the potential loopholes or oversights"
"of the current reasoning process while solving the complex problem.\n\n "
"## Your Task: \n"
"Based on the gathered information retrieved from the internet, and the reasoning process already"
"generated towards solving a complex task, you need to do the following 1 or 2 things, to guarntee the quality of the reasoning process, and a clear final answer: \n"
" 1. Provide your diagnosing result on the generated reasoning process and the corresponding the correction if necessary;\n"
" 2. Provide your insight and supplements in advance to avoid the potential loopholes or oversights in the future;\n\n"
"## Requirements: \n"
" 1. If the reasoning process already generated is complete and correct in your opinion, just say 'No loopholes or oversights found'. \n"
" 2. If the reasoning process already generated contains the materials that may lead to the potential logic mistake or lack of some important guardrails in your opinion, you may give a hint to the current reasoning process, with the necessary supplements.\n"
" 3. If the reasoning process already generated is seriously incorrect in your opinion, you may give the turn signal to the reasoning process, to maneuver the reasoning process towards solving the complex problem correctly. \n\n"
"## Restriction: \n"
" 1. Please do not make judgments about the authenticity of externally sourced information obtained through searches, as this is not part of your job responsibilities;\n"
" 2. Do not make additional inferences or assumptions about the content of such information itself.\n"
" 3. If the question lacks necessary details/data/clues in your opinion, you may ask for more details.\n\n"
"## Example 1:\n"
" Question: Is my reasoning process correct?\n"
" Reasoning Process: (nothing specified)\n"
" Your Identification Result: Your question lacks some information, please provide me more details so I can help you.\n\n"
),
},
{"role": "user", "content": prompt},
],
temperature=temperature,
)
return response.content
def mcp_guarding_reasoning_process(
self,
question: str = Field(
description="The input question for diagnosing the completeness/correctness of the reasoning process.\n"
"For example: based on the staged/phased information/data concluded as 1.xxxx 2. xxxx 3. xxxx...., is there any faults in the current reasoning process aaaaaa"
"that should be corrected? Or is there any loopholes or oversights that should be emphized in advance, towards solving the bbbbb problem?\n"
"Requirement: This input question should include a clear question with the necessary details/data/clues from the previous context"
"(such as the key information/data retrieved from the internet), to present more clues to help the diagnosing process."
"The more exact details/data contained in this input question, the better the diagnosing result will be."
),
original_task: str = Field(default="", description="The original task. This field is required and cannot be simplied, has to be true to the original task."),
temperature: float = Field(
default=0.1,
description="Model temperature for response variability (0.0-1.0)",
ge=0.0,
le=1.0,
),
guarding_style: Literal["detailed", "concise", "step-by-step"] = Field(
default="detailed",
description="Style of guarding output: detailed analysis, concise summary, or step-by-step breakdown",
),
) -> ActionResponse:
"""This tool provides advanced logic diagonsing and correcting ability, to improve the quality of the reasoning process that already exists in the current context, while solving the complex question:
- Identify potential loopholes or oversights in the current reasoning process while solving the complex problem.
- Providing the guidance, suggestions to the reasoning process, to correct the loopholes or oversights if identified in this shot.
Invoke Timing: During Thinking Process Reviewing, while diagnosing the reasoning process in the context or give valuable suggestions in advance, this tool is a reliable selection.
Strengths:
- Be relatively sensitive to common logical traps in some mathematics or logic problems.
Weakness:
- Inability to process media types: image, audio, or video.
- Inability to check the correctness of the retrieved information from the internet.
- Require precise description of problem context and settings, including the reasoning process, retrieved data and the complex task itself.
Args:
question: The input question that invokes this tool to diagnose and/or correct the suspected reasoning process in the context
original_task: Optional original task description for additional context
temperature: Model temperature controlling response variability
guarding_style: Style of guarding output format
Returns:
ActionResponse with guarding result and processing metadata
"""
try:
# Handle FieldInfo objects
if isinstance(question, FieldInfo):
question = question.default
if isinstance(original_task, FieldInfo):
original_task = original_task.default
if isinstance(temperature, FieldInfo):
temperature = temperature.default
if isinstance(guarding_style, FieldInfo):
guarding_style = guarding_style.default
# Validate input
if not question or not question.strip():
raise ValueError("Question is required for guarding the complex problem reasoning process")
self._color_log(f"Processing guarding request: {question[:100]}...", Color.cyan)
start_time = time.time()
# Prepare the guarding prompt
prompt = self._prepare_guarding_prompt(question, original_task) ## 简单的原始问题+分配给mcp server的问题
# Enhance prompt based on guarding style
if guarding_style == "step-by-step":
prompt += "\n\nPlease provide a clear step-by-step breakdown of your reviewing process of the reasoning process."
elif guarding_style == "concise":
prompt += "\n\nPlease provide a concise and final guarding answer."
elif guarding_style == "detailed":
prompt += "\n\nPlease provide detailed reviewing analysis with comprehensive guarding."
# Call the guarding model
guarding_result = self._call_guarding_model(prompt, temperature)
processing_time = time.time() - start_time
# Prepare metadata
metadata = {
"model_name": self._llm_config.llm_model_name,
"guarding_style": guarding_style,
"response_length": len(guarding_result),
}
self._color_log(
f"Successfully completed guarding ({len(guarding_result)} characters, {processing_time:.2f}s)",
Color.green,
)
return ActionResponse(success=True, message=guarding_result, metadata=metadata)
except ValueError as e:
self.logger.error(f"Invalid input: {str(e)}")
return ActionResponse(
success=False,
message=f"Invalid input: {str(e)}",
metadata={"error_type": "invalid_input", "error_message": str(e)},
)
except Exception as e:
self.logger.error(f"Guarding failed: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=f"Guarding failed: {str(e)}",
metadata={"error_type": "guarding_error", "error_message": str(e)},
)
def mcp_get_guarding_capabilities(self) -> ActionResponse:
"""Get information about the guarding reasoning process service capabilities.
Returns:
ActionResponse with service capabilities and configuration
"""
capabilities = {
"Logic Loopholes Detecting": "Identifying the logic loopholes in the reasoning process already generated previsouly",
"Detected Loopholes Correcting": "Correcting the logic loopholes identified in the reasoning process already generated previously",
"Oversights Prevention": "Providing necessary supplements as hints to the currect reasoning process, to prevent the possible oversights in the future",
}
capability_list = "\n".join(
[f"**{capability}**: {description}" for capability, description in capabilities.items()]
)
metadata = {
"model_name": self._llm_config.llm_model_name,
"provider": self._llm_config.llm_provider,
"supported_capabilities": list(capabilities.keys()),
"total_capabilities": len(capabilities),
"guarding_styles": ["detailed", "concise", "step-by-step"],
}
return ActionResponse(
success=True,
message=f"Intelligence guarding Service Capabilities:\n\n{capability_list}",
metadata=metadata,
)
# Example usage and entry point
if __name__ == "__main__":
load_dotenv()
# Default arguments for testing
args = ActionArguments(
name="intelligence_guarding_service",
transport="stdio",
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
)
# Initialize and run the intelligence guarding service
try:
service = GuardCollection(args)
service.run()
except Exception as e:
print(f"An error occurred: {e}: {traceback.format_exc()}")
@@ -0,0 +1,237 @@
import os
import time
import traceback
from typing import Literal
from dotenv import load_dotenv
from pydantic import Field
from pydantic.fields import FieldInfo
from aworld.config.conf import AgentConfig
from aworld.logs.util import Color
from aworld.models.llm import call_llm_model, get_llm_model
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
class ThinkCollection(ActionCollection):
"""MCP service for complex problem reasoning using powerful reasoning models.
Supports advanced reasoning for:
- Mathematical problems and proofs
- Code contest and programming challenges
- Logic puzzles and riddles
- Competition-level STEM problems
- Multi-step analytical reasoning
"""
def __init__(self, arguments: ActionArguments) -> None:
super().__init__(arguments)
# Initialize reasoning model configuration
self._llm_config = AgentConfig(
llm_provider="openai",
# llm_model_name="google/gemini-2.5-flash-preview-05-20:thinking",
llm_model_name=os.getenv("THINK_LLM_MODEL_NAME", "deepseek/deepseek-r1-0528:free"),
llm_api_key=os.getenv("THINK_LLM_API_KEY"),
llm_base_url=os.getenv("THINK_LLM_BASE_URL"),
)
self._color_log("Intelligence Reasoning Service initialized", Color.green, "debug")
self._color_log(f"Using model: {self._llm_config.llm_model_name}", Color.blue, "debug")
def _prepare_reasoning_prompt(self, question: str, original_task: str = "") -> str:
"""Prepare the reasoning prompt with question and optional context.
Args:
question: The main question for reasoning
original_task: Optional original task description for context
Returns:
Formatted prompt string
"""
if original_task:
return f"Original Task: {original_task}\n\nQuestion: {question}"
return f"Question: {question}"
def _call_reasoning_model(self, prompt: str, temperature: float = 0.3) -> str:
"""Call the reasoning model with the prepared prompt.
Args:
prompt: The formatted prompt for reasoning
temperature: Model temperature for response variability
Returns:
Reasoning result from the model
Raises:
Exception: If model call fails
"""
response = call_llm_model(
llm_model=get_llm_model(conf=self._llm_config),
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=temperature,
)
return response.content
def mcp_complex_problem_reasoning(
self,
question: str = Field(
description="The input question for complex problem reasoning, such as math and code contest problems"
),
original_task: str = Field(default="", description="The original task description."),
temperature: float = Field(
default=0.3,
description="Model temperature for response variability (0.0-1.0)",
ge=0.0,
le=1.0,
),
reasoning_style: Literal["detailed", "concise", "step-by-step"] = Field(
default="detailed",
description="Style of reasoning output: detailed analysis, concise summary, or step-by-step breakdown",
),
) -> ActionResponse:
"""This tool provides comprehensive reasoning capabilities for:
- Mathematical problems and proofs
- Programming and algorithm challenges
- Logic puzzles, brain teasers, and fun riddles
- Competition-level STEM problems
- Multi-step analytical reasoning tasks
Weakness:
- Inability to process media types: image, audio, or video
- Require precise description of problem context and settings
Args:
question: The input question requiring complex reasoning
original_task: Optional original task description for additional context
temperature: Model temperature controlling response variability
reasoning_style: Style of reasoning output format
Returns:
ActionResponse with reasoning result and processing metadata
"""
try:
# Handle FieldInfo objects
if isinstance(question, FieldInfo):
question = question.default
if isinstance(original_task, FieldInfo):
original_task = original_task.default
if isinstance(temperature, FieldInfo):
temperature = temperature.default
if isinstance(reasoning_style, FieldInfo):
reasoning_style = reasoning_style.default
# Validate input
if not question or not question.strip():
raise ValueError("Question is required for complex problem reasoning")
self._color_log(f"Processing reasoning request: {question[:100]}...", Color.cyan)
start_time = time.time()
# Prepare the reasoning prompt
prompt = self._prepare_reasoning_prompt(question, original_task)
# Enhance prompt based on reasoning style
if reasoning_style == "step-by-step":
prompt += "\n\nPlease provide a clear step-by-step breakdown of your reasoning process."
elif reasoning_style == "concise":
prompt += "\n\nPlease provide a concise but complete reasoning and final answer."
elif reasoning_style == "detailed":
prompt += "\n\nPlease provide detailed analysis with comprehensive reasoning."
# Call the reasoning model
reasoning_result = self._call_reasoning_model(prompt, temperature)
processing_time = time.time() - start_time
# Prepare metadata
metadata = {
"model_name": self._llm_config.llm_model_name,
"reasoning_style": reasoning_style,
"response_length": len(reasoning_result),
}
self._color_log(
f"Successfully completed reasoning ({len(reasoning_result)} characters, {processing_time:.2f}s)",
Color.green,
)
return ActionResponse(success=True, message=reasoning_result, metadata=metadata)
except ValueError as e:
self.logger.error(f"Invalid input: {str(e)}")
return ActionResponse(
success=False,
message=f"Invalid input: {str(e)}",
metadata={"error_type": "invalid_input", "error_message": str(e)},
)
except Exception as e:
self.logger.error(f"Reasoning failed: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=f"Reasoning failed: {str(e)}",
metadata={"error_type": "reasoning_error", "error_message": str(e)},
)
def mcp_get_reasoning_capabilities(self) -> ActionResponse:
"""Get information about the reasoning service capabilities.
Returns:
ActionResponse with service capabilities and configuration
"""
capabilities = {
"Mathematical Problems": "Advanced mathematical reasoning, proofs, and calculations",
"Code Contests": "Programming challenges, algorithm design, and optimization",
"Logic Puzzles": "Brain teasers, riddles, and logical reasoning problems",
"STEM Problems": "Competition-level science, technology, engineering, and math",
"Multi-step Analysis": "Complex analytical reasoning with multiple interconnected steps",
}
capability_list = "\n".join(
[f"**{capability}**: {description}" for capability, description in capabilities.items()]
)
metadata = {
"model_name": self._llm_config.llm_model_name,
"provider": self._llm_config.llm_provider,
"supported_capabilities": list(capabilities.keys()),
"total_capabilities": len(capabilities),
"reasoning_styles": ["detailed", "concise", "step-by-step"],
}
return ActionResponse(
success=True,
message=f"Intelligence Reasoning Service Capabilities:\n\n{capability_list}",
metadata=metadata,
)
# Example usage and entry point
if __name__ == "__main__":
load_dotenv()
# Default arguments for testing
args = ActionArguments(
name="intelligence_reasoning_service",
transport="stdio",
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
)
# Initialize and run the intelligence reasoning service
try:
service = ThinkCollection(args)
service.run()
except Exception as e:
print(f"An error occurred: {e}: {traceback.format_exc()}")
@@ -0,0 +1,555 @@
import json
import os
import subprocess
import time
import traceback
from pathlib import Path
from typing import Any, Literal
from dotenv import load_dotenv
from openai import OpenAI
from pydantic import BaseModel, Field
from pydantic.fields import FieldInfo
from aworld.logs.util import Color
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
class AudioMetadata(BaseModel):
"""Metadata extracted from audio processing."""
file_name: str = Field(description="Original audio file name")
file_size: int = Field(description="File size in bytes")
file_type: str = Field(description="Audio file type/extension")
absolute_path: str = Field(description="Absolute path to the audio file")
duration: float | None = Field(default=None, description="Duration of audio in seconds")
sample_rate: int | None = Field(default=None, description="Audio sample rate in Hz")
channels: int | None = Field(default=None, description="Number of audio channels")
bitrate: int | None = Field(default=None, description="Audio bitrate in kbps")
codec: str | None = Field(default=None, description="Audio codec used")
processing_time: float = Field(description="Time taken to process the audio in seconds")
output_files: list[str] = Field(default_factory=list, description="Paths to generated output files")
transcription: str | None = Field(default=None, description="Transcribed text from audio")
word_count: int | None = Field(default=None, description="Number of words in transcription")
output_format: str = Field(description="Format of the processed output")
class AudioCollection(ActionCollection):
"""MCP service for comprehensive audio processing using ffmpeg.
Supports various audio operations including:
- Audio format conversion
- Audio transcription (speech-to-text)
- Audio metadata extraction
- Audio quality enhancement
- Audio trimming and editing
- Audio analysis and feature extraction
"""
def __init__(self, arguments: ActionArguments) -> None:
super().__init__(arguments)
self._audio_output_dir = self.workspace / "processed_audio"
self._audio_output_dir.mkdir(exist_ok=True)
# Supported audio formats
self.supported_extensions = {
".mp3",
".wav",
".flac",
".aac",
".ogg",
".m4a",
".wma",
".opus",
".aiff",
".au",
".ra",
".amr",
}
self._color_log("Audio Processing Service initialized", Color.green, "debug")
self._color_log(f"Audio output directory: {self._audio_output_dir}", Color.blue, "debug")
# Check ffmpeg availability
self._check_ffmpeg_availability()
def _check_ffmpeg_availability(self) -> bool:
"""Check if ffmpeg is available in the system.
Returns:
bool: True if ffmpeg is available, False otherwise
"""
try:
result = subprocess.run(["ffmpeg", "-version"], capture_output=True, text=True, timeout=10, check=False)
if result.returncode == 0:
self._color_log("FFmpeg is available", Color.green, "debug")
else:
self._color_log("FFmpeg not found in system PATH", Color.red)
return result.returncode == 0
except (subprocess.TimeoutExpired, FileNotFoundError):
self._color_log(
"FFmpeg not available or timeout, Please refer to https://github.com/inclusionAI/AWorld/blob/main/examples/gaia/README.md#system-tools-setup for more details",
Color.red,
)
return False
def _prepare_audio_for_transcription(self, file_path: Path) -> Path:
"""Prepare audio file for transcription by converting to optimal format.
Args:
file_path: Path to the original audio file
Returns:
Path to the prepared audio file (WAV format, 16kHz)
"""
output_path = self._audio_output_dir / f"{file_path.stem}_for_transcription.wav"
# Convert to WAV format with 16kHz sample rate for optimal transcription
cmd = [
"ffmpeg",
"-i",
str(file_path),
"-ar",
"16000", # 16kHz sample rate
"-ac",
"1", # Mono channel
"-c:a",
"pcm_s16le", # 16-bit PCM
"-y",
str(output_path),
]
result = subprocess.run(cmd, capture_output=True, text=True, timeout=300, check=False)
if result.returncode != 0:
raise RuntimeError(f"Audio preparation for transcription failed: {result.stderr}")
return output_path
def _transcribe_with_whisper(self, audio_path: Path) -> dict[str, Any]:
"""Transcribe audio using OpenAI Whisper.
Args:
audio_path: Path to the audio file
Returns:
Dictionary containing transcription results
"""
try:
client: OpenAI = OpenAI(api_key=os.getenv("AUDIO_LLM_API_KEY"), base_url=os.getenv("AUDIO_LLM_BASE_URL"))
# Use the file for transcription
with open(audio_path, "rb") as audio_file:
transcription: str = client.audio.transcriptions.create(
file=audio_file,
model=os.getenv("AUDIO_LLM_MODEL_NAME"),
response_format="text",
)
return {"text": transcription.strip() if transcription else ""}
except Exception as e:
raise RuntimeError(f"Audio transcription failed: {e}: {traceback.format_exc()}") from e
def _get_audio_metadata(self, file_path: Path) -> dict[str, Any]:
"""Extract audio metadata using ffprobe.
Args:
file_path: Path to the audio file
Returns:
Dictionary containing audio metadata
"""
try:
cmd = ["ffprobe", "-v", "quiet", "-print_format", "json", "-show_format", "-show_streams", str(file_path)]
result = subprocess.run(cmd, capture_output=True, text=True, timeout=30, check=False)
if result.returncode == 0:
metadata = json.loads(result.stdout)
# Extract relevant audio information
format_info = metadata.get("format", {})
streams = metadata.get("streams", [])
audio_stream = next((s for s in streams if s.get("codec_type") == "audio"), {})
return {
"duration": float(format_info.get("duration", 0)),
"sample_rate": int(audio_stream.get("sample_rate", 0)) if audio_stream.get("sample_rate") else None,
"channels": int(audio_stream.get("channels", 0)) if audio_stream.get("channels") else None,
"bitrate": int(format_info.get("bit_rate", 0)) // 1000 if format_info.get("bit_rate") else None,
"codec": audio_stream.get("codec_name"),
}
else:
self.logger.warning(f"Failed to extract metadata: {result.stderr}")
return {}
except Exception as e:
self.logger.error(f"Error extracting audio metadata: {str(e)}")
return {}
def _trim_audio(self, input_path: Path, start_time: float, duration: float | None = None) -> Path:
"""Trim audio file to specified time range.
Args:
input_path: Path to input audio file
start_time: Start time in seconds
duration: Duration in seconds (if None, trim to end)
Returns:
Path to trimmed audio file
"""
output_path = self._audio_output_dir / f"{input_path.stem}_trimmed{input_path.suffix}"
cmd = ["ffmpeg", "-i", str(input_path), "-ss", str(start_time), "-y"]
if duration is not None:
cmd.extend(["-t", str(duration)])
cmd.extend(["-c", "copy", str(output_path)])
result = subprocess.run(cmd, capture_output=True, text=True, timeout=300, check=False)
if result.returncode != 0:
raise RuntimeError(f"Audio trimming failed: {result.stderr}")
return output_path
def mcp_transcribe_audio(
self,
file_path: str = Field(description="Path to the audio file to transcribe"),
model_size: Literal["tiny", "base", "small", "medium", "large"] = Field(
default="base",
description="Whisper model size: tiny (fastest), base (balanced), small, medium, large (most accurate)",
),
output_format: Literal["text", "detailed", "segments"] = Field(
default="text",
description="Output format: 'text' (plain text), 'detailed' (with metadata), 'segments' (timestamped)",
),
) -> ActionResponse:
"""Transcribe audio file to text using OpenAI Whisper.
This tool converts speech in audio files to text with high accuracy.
Supports multiple languages and provides various output formats including
timestamped segments for detailed analysis.
Args:
file_path: Path to the audio file to transcribe
model_size: Whisper model size affecting speed vs accuracy trade-off
output_format: Format of transcription output
Returns:
ActionResponse with transcribed text and detailed metadata
"""
try:
# Handle FieldInfo objects
if isinstance(file_path, FieldInfo):
file_path = file_path.default
if isinstance(model_size, FieldInfo):
model_size = model_size.default
if isinstance(output_format, FieldInfo):
output_format = output_format.default
start_time = time.time()
# Validate input file
file_path: Path = self._validate_file_path(file_path)
self._color_log(f"Transcribing audio: {file_path.name}", Color.cyan)
# Get original metadata
original_metadata = self._get_audio_metadata(file_path)
# Prepare audio for transcription
prepared_audio = self._prepare_audio_for_transcription(file_path)
# Perform transcription
transcription_result = self._transcribe_with_whisper(prepared_audio)
processing_time = time.time() - start_time
# Prepare file statistics
file_stats = file_path.stat()
# Count words in transcription
word_count = len(transcription_result["text"].split()) if transcription_result["text"] else 0
# Create metadata object
audio_metadata = AudioMetadata(
file_name=file_path.name,
file_size=file_stats.st_size,
file_type=file_path.suffix.lower(),
absolute_path=str(file_path.absolute()),
duration=original_metadata.get("duration"),
sample_rate=original_metadata.get("sample_rate"),
channels=original_metadata.get("channels"),
bitrate=original_metadata.get("bitrate"),
codec=original_metadata.get("codec"),
processing_time=processing_time,
output_files=[str(prepared_audio)],
transcription=transcription_result["text"],
word_count=word_count,
output_format=f"transcription_{output_format}",
)
# Format output based on requested format
if output_format == "text":
result_message = transcription_result["text"]
elif output_format == "detailed":
confidence_str = (
f"{transcription_result['confidence']:.2f}" if transcription_result["confidence"] else "N/A"
)
result_message = (
f"Transcription Results for {file_path.name}:\n\n"
f"**Text:** {transcription_result['text']}\n\n"
f"**Confidence:** {confidence_str}\n"
f"**Word Count:** {word_count}\n"
f"**Duration:** {original_metadata.get('duration', 0):.2f} seconds\n"
f"**Model:** {model_size}\n"
f"**Processing Time:** {processing_time:.2f} seconds"
)
elif output_format == "segments":
segments_text = "\n".join(
[
f"[{seg.get('start', 0):.2f}s - {seg.get('end', 0):.2f}s]: {seg.get('text', '').strip()}"
for seg in transcription_result.get("segments", [])
]
)
result_message = (
f"Timestamped Transcription for {file_path.name}:\n\n"
f"{segments_text}\n\n"
f"**Full Text:** {transcription_result['text']}"
)
else:
result_message = transcription_result["text"]
# Clean up temporary file
try:
prepared_audio.unlink()
except Exception:
pass # Ignore cleanup errors
self._color_log(f"Transcription completed: {word_count} words, {processing_time:.2f}s", Color.green)
return ActionResponse(success=True, message=result_message, metadata=audio_metadata.model_dump())
except Exception as e:
self.logger.error(f"Audio transcription failed: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=f"Audio transcription failed: {str(e)}",
metadata={"error_type": "transcription_error"},
)
def mcp_extract_audio_metadata(
self,
file_path: str = Field(description="Path to the audio file to analyze"),
) -> ActionResponse:
"""Extract comprehensive metadata from audio files.
This tool analyzes audio files and extracts detailed metadata including
duration, sample rate, channels, bitrate, codec, and other technical information.
Args:
file_path: Path to the audio file to analyze
Returns:
ActionResponse with detailed audio metadata
"""
try:
if isinstance(file_path, FieldInfo):
file_path = file_path.default
start_time = time.time()
# Validate input file
file_path: Path = self._validate_file_path(file_path)
self._color_log(f"Extracting metadata from: {file_path.name}", Color.cyan)
# Extract metadata
metadata = self._get_audio_metadata(file_path)
processing_time = time.time() - start_time
# Prepare file statistics
file_stats = file_path.stat()
# Create metadata object
audio_metadata = AudioMetadata(
file_name=file_path.name,
file_size=file_stats.st_size,
file_type=file_path.suffix.lower(),
absolute_path=str(file_path.absolute()),
duration=metadata.get("duration"),
sample_rate=metadata.get("sample_rate"),
channels=metadata.get("channels"),
bitrate=metadata.get("bitrate"),
codec=metadata.get("codec"),
processing_time=processing_time,
output_files=[],
output_format="metadata",
)
# Format metadata for LLM consumption
result_message = (
f"Audio Metadata for {file_path.name}:\n"
f"Duration: {metadata.get('duration', 'Unknown'):.2f} seconds\n"
f"Sample Rate: {metadata.get('sample_rate', 'Unknown')} Hz\n"
f"Channels: {metadata.get('channels', 'Unknown')}\n"
f"Bitrate: {metadata.get('bitrate', 'Unknown')} kbps\n"
f"Codec: {metadata.get('codec', 'Unknown')}\n"
f"File Size: {file_stats.st_size / 1024 / 1024:.2f} MB\n"
f"Format: {file_path.suffix.upper()}"
)
self._color_log(f"Metadata extraction completed in {processing_time:.2f}s", Color.green)
return ActionResponse(success=True, message=result_message, metadata=audio_metadata.model_dump())
except Exception as e:
self.logger.error(f"Metadata extraction failed: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=f"Metadata extraction failed: {str(e)}",
metadata={"error_type": "metadata_error"},
)
def mcp_trim_audio(
self,
file_path: str = Field(description="Path to the audio file to trim"),
start_time: float = Field(description="Start time in seconds"),
duration: float | None = Field(default=None, description="Duration in seconds (if None, trim to end)"),
) -> ActionResponse:
"""Trim audio file to specified time range.
This tool cuts audio files to extract specific segments based on start time
and duration. Useful for creating clips or removing unwanted sections.
Args:
file_path: Path to the source audio file
start_time: Start time in seconds for trimming
duration: Duration of the trimmed segment (optional)
Returns:
ActionResponse with trimmed audio file path and metadata
"""
try:
if isinstance(file_path, FieldInfo):
file_path = file_path.default
if isinstance(start_time, FieldInfo):
start_time = start_time.default
if isinstance(duration, FieldInfo):
duration = duration.default
process_start = time.time()
# Validate input file
file_path: Path = self._validate_file_path(file_path)
self._color_log(f"Trimming audio: {file_path.name}", Color.cyan)
# Get original metadata
original_metadata = self._get_audio_metadata(file_path)
# Validate time parameters
if start_time < 0:
raise ValueError("Start time cannot be negative")
if duration is not None and duration <= 0:
raise ValueError("Duration must be positive")
if original_metadata.get("duration") and start_time >= original_metadata["duration"]:
raise ValueError("Start time exceeds audio duration")
# Trim audio
output_path = self._trim_audio(file_path, start_time, duration)
# Get trimmed file metadata
trimmed_metadata = self._get_audio_metadata(output_path)
processing_time = time.time() - process_start
# Prepare file statistics
file_stats = file_path.stat()
# Create metadata object
audio_metadata = AudioMetadata(
file_name=file_path.name,
file_size=file_stats.st_size,
file_type=file_path.suffix.lower(),
absolute_path=str(file_path.absolute()),
duration=trimmed_metadata.get("duration"),
sample_rate=trimmed_metadata.get("sample_rate"),
channels=trimmed_metadata.get("channels"),
bitrate=trimmed_metadata.get("bitrate"),
codec=trimmed_metadata.get("codec"),
processing_time=processing_time,
output_files=[str(output_path)],
output_format="trimmed_audio",
)
end_time = start_time + (duration or (original_metadata.get("duration", 0) - start_time))
result_message = (
f"Successfully trimmed {file_path.name}\n"
f"Original duration: {original_metadata.get('duration', 0):.2f} seconds\n"
f"Trimmed segment: {start_time:.2f}s - {end_time:.2f}s\n"
f"New duration: {trimmed_metadata.get('duration', 0):.2f} seconds\n"
f"Output file: {output_path.name}"
)
self._color_log(f"Audio trimming completed in {processing_time:.2f}s", Color.green)
return ActionResponse(success=True, message=result_message, metadata=audio_metadata.model_dump())
except Exception as e:
self.logger.error(f"Audio trimming failed: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False, message=f"Audio trimming failed: {str(e)}", metadata={"error_type": "trimming_error"}
)
def mcp_list_supported_formats(self) -> ActionResponse:
"""List all supported audio formats for processing.
Returns:
ActionResponse with list of supported audio formats and their descriptions
"""
supported_formats = {
"MP3": "MPEG Audio Layer III (.mp3) - Most common compressed format",
"WAV": "Waveform Audio File Format (.wav) - Uncompressed, high quality",
"FLAC": "Free Lossless Audio Codec (.flac) - Lossless compression",
"AAC": "Advanced Audio Coding (.aac) - Efficient compression",
"OGG": "Ogg Vorbis (.ogg) - Open source compressed format",
"M4A": "MPEG-4 Audio (.m4a) - Apple's preferred format",
"WMA": "Windows Media Audio (.wma) - Microsoft format",
"OPUS": "Opus Audio (.opus) - Modern, efficient codec",
"AIFF": "Audio Interchange File Format (.aiff) - Apple's uncompressed format",
"AU": "Sun Audio (.au) - Unix audio format",
"RA": "RealAudio (.ra) - Streaming audio format",
"AMR": "Adaptive Multi-Rate (.amr) - Mobile audio format",
}
format_list = "\n".join(
[f"**{format_name}**: {description}" for format_name, description in supported_formats.items()]
)
return ActionResponse(
success=True,
message=f"Supported audio formats:\n\n{format_list}",
metadata={
"supported_formats": list(supported_formats.keys()),
"total_formats": len(supported_formats),
"ffmpeg_available": self._check_ffmpeg_availability(),
},
)
# Example usage and entry point
if __name__ == "__main__":
load_dotenv()
# Default arguments for testing
args = ActionArguments(
name="audio_processing_service",
transport="stdio",
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
)
# Initialize and run the audio processing service
try:
service = AudioCollection(args)
service.run()
print("Audio processing service started")
except Exception as e:
print(f"An error occurred: {e}: {traceback.format_exc()}")
@@ -0,0 +1,484 @@
import base64
import os
import time
import traceback
from io import BytesIO
from pathlib import Path
from typing import Any
import pytesseract
from dotenv import load_dotenv
from PIL import Image, ImageEnhance, ImageFilter
from pydantic import BaseModel, Field
from pydantic.fields import FieldInfo
from aworld.config.conf import AgentConfig
from aworld.logs.util import Color
from aworld.models.llm import call_llm_model, get_llm_model
from aworld.models.model_response import ModelResponse
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
class ImageMetadata(BaseModel):
"""Metadata extracted from image processing."""
file_name: str = Field(description="Original image file name")
file_size: int = Field(description="File size in bytes")
file_type: str = Field(description="Image file type/extension")
absolute_path: str = Field(description="Absolute path to the image file")
width: int | None = Field(default=None, description="Image width in pixels")
height: int | None = Field(default=None, description="Image height in pixels")
mode: str | None = Field(default=None, description="Image color mode (RGB, RGBA, L, etc.)")
format: str | None = Field(default=None, description="Image format (JPEG, PNG, etc.)")
has_transparency: bool = Field(default=False, description="Whether image has transparency")
processing_time: float = Field(description="Time taken to process the image in seconds", exclude=True)
output_files: list[str] = Field(default_factory=list, description="Paths to generated output files")
extracted_text: str | None = Field(default=None, description="Text extracted via OCR")
analysis_result: str | None = Field(default=None, description="AI analysis result")
compression_ratio: float | None = Field(default=None, description="Compression ratio if optimized")
output_format: str = Field(description="Format of the processed output")
class ImageCollection(ActionCollection):
"""MCP service for comprehensive image processing and analysis.
Supports various image operations including:
- Metadata extraction
- OCR (Optical Character Recognition)
- AI-powered image analysis and reasoning
"""
def __init__(self, arguments: ActionArguments) -> None:
super().__init__(arguments)
self._image_output_dir = self.workspace / "processed_images"
self._image_output_dir.mkdir(exist_ok=True)
# Supported image formats
self.supported_extensions = {
".jpg",
".jpeg",
".png",
".gif",
".webp",
".bmp",
".tiff",
".tif",
".ico",
".svg",
}
self._llm_config = AgentConfig(
llm_provider="openai",
llm_model_name=os.getenv("IMAGE_LLM_MODEL_NAME", "gpt-4o"),
llm_api_key=os.getenv("IMAGE_LLM_API_KEY"),
llm_base_url=os.getenv("IMAGE_LLM_BASE_URL"),
)
self._color_log("Image Processing Service initialized", Color.green, "debug")
self._color_log(f"Image output directory: {self._image_output_dir}", Color.blue, "debug")
def _load_image(self, file_path: Path) -> Image.Image:
"""Load image from file path.
Args:
file_path: Path to the image file
Returns:
PIL Image object
"""
try:
return Image.open(file_path)
except Exception as e:
raise RuntimeError(f"Failed to load image {file_path}: {str(e)}") from e
def _get_image_metadata(self, image: Image.Image, file_path: Path) -> dict[str, Any]:
"""Extract metadata from PIL Image object.
Args:
image: PIL Image object
file_path: Path to the original file
Returns:
Dictionary containing image metadata
"""
return {
"width": image.width,
"height": image.height,
"mode": image.mode,
"format": image.format or file_path.suffix.upper().lstrip("."),
"has_transparency": image.mode in ("RGBA", "LA") or "transparency" in image.info,
}
def _optimize_image(self, image: Image.Image, max_size: tuple[int, int] | None = None) -> Image.Image:
"""Optimize image for size and quality.
Args:
image: PIL Image object
max_size: Maximum dimensions (width, height)
Returns:
Optimized PIL Image object
"""
optimized = image.copy()
# Resize if max_size specified
if max_size:
optimized.thumbnail(max_size, Image.Resampling.LANCZOS)
return optimized
def _perform_ocr(self, image: Image.Image) -> str:
"""Perform OCR on image to extract text.
Args:
image: PIL Image object
Returns:
Extracted text string
"""
try:
return pytesseract.image_to_string(image).strip()
except ImportError:
return "OCR not available - pytesseract not installed"
except Exception as e:
return f"OCR failed: {str(e)}"
def _analyze_with_ai(self, image_base64: str, task: str) -> str:
"""Analyze image using AI model.
Args:
image_base64: Base64 encoded image
task: Analysis task description
Returns:
AI analysis result
"""
try:
messages = [
{
"role": "user",
"content": [
{"role": "text", "content": task},
{"type": "image_url", "image_url": {"url": image_base64}},
],
},
]
response: ModelResponse = call_llm_model(
llm_model=get_llm_model(conf=self._llm_config),
messages=messages,
temperature=float(os.getenv("LLM_TEMPERATURE", "1.0")),
)
self._color_log(f"{response.content=}", Color.green)
return response.content
except Exception as e:
return f"AI analysis failed: {str(e)}"
def _image_to_base64(self, image: Image.Image, output_format: str = "JPEG") -> str:
"""Convert PIL Image to base64 string.
Args:
image: PIL Image object
output_format: Output format (JPEG, PNG, etc.)
Returns:
Base64 encoded image string
"""
buffer = BytesIO()
# Convert RGBA or P to RGB for JPEG
if output_format.upper() == "JPEG":
if image.mode in ("RGBA", "LA"):
background = Image.new("RGB", image.size, (255, 255, 255))
background.paste(image, mask=image.split()[-1] if image.mode == "RGBA" else None)
image = background
elif image.mode == "P":
image = image.convert("RGB")
image.save(buffer, format=output_format, quality=85 if output_format.upper() == "JPEG" else None)
mime_type = f"image/{output_format.lower()}"
img_base64 = base64.b64encode(buffer.getvalue()).decode()
return f"data:{mime_type};base64,{img_base64}"
def mcp_extract_text_ocr(
self,
file_path: str = Field(description="Path to the image file for OCR"),
language: str = Field(default="eng", description="OCR language code (e.g., 'eng', 'spa', 'fra')"),
preprocess: bool = Field(default=True, description="Whether to preprocess image for better OCR"),
) -> ActionResponse:
"""Extract text from images using Optical Character Recognition (OCR).
This tool uses Tesseract OCR to extract text content from images,
with optional preprocessing to improve recognition accuracy.
Args:
file_path: Path to the image file
language: OCR language for better recognition
preprocess: Whether to enhance image for OCR
Returns:
ActionResponse with extracted text and metadata
"""
try:
# Handle FieldInfo objects
if isinstance(file_path, FieldInfo):
file_path = file_path.default
if isinstance(language, FieldInfo):
language = language.default
if isinstance(preprocess, FieldInfo):
preprocess = preprocess.default
start_time = time.time()
# Validate input file
file_path: Path = self._validate_file_path(file_path)
self._color_log(f"Performing OCR on: {file_path.name}", Color.cyan)
# Load image
image = self._load_image(file_path)
original_metadata = self._get_image_metadata(image, file_path)
# Preprocess image for better OCR if requested
if preprocess:
# Convert to grayscale
if image.mode != "L":
image = image.convert("L")
# Enhance contrast
enhancer = ImageEnhance.Contrast(image)
image = enhancer.enhance(2.0)
# Apply slight sharpening
image = image.filter(ImageFilter.SHARPEN)
# Perform OCR
extracted_text = self._perform_ocr(image)
processing_time = time.time() - start_time
# Count words and characters
word_count = len(extracted_text.split()) if extracted_text else 0
char_count = len(extracted_text) if extracted_text else 0
# Create metadata object
image_metadata = ImageMetadata(
file_name=file_path.name,
file_size=file_path.stat().st_size,
file_type=file_path.suffix.lower(),
absolute_path=str(file_path.absolute()),
width=original_metadata["width"],
height=original_metadata["height"],
mode=original_metadata["mode"],
format=original_metadata["format"],
has_transparency=original_metadata["has_transparency"],
processing_time=processing_time,
output_files=[],
extracted_text=extracted_text,
output_format="ocr_text",
)
if extracted_text:
result_message = (
f"OCR Results for {file_path.name}:\n\n"
f"**Extracted Text:**\n{extracted_text}\n\n"
f"**Statistics:**\n"
f"- Words: {word_count}\n"
f"- Characters: {char_count}\n"
f"- Language: {language}\n"
f"- Processing time: {processing_time:.2f}s"
)
else:
result_message = (
f"No text detected in {file_path.name}."
" The image may not contain readable text or OCR preprocessing may be needed."
)
self._color_log(f"OCR completed: {word_count} words extracted in {processing_time:.2f}s", Color.green)
return ActionResponse(success=True, message=result_message, metadata=image_metadata.model_dump())
except Exception as e:
self.logger.error(f"OCR failed: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=f"OCR failed: {str(e)}",
metadata={"error_type": "ocr_error"},
)
def mcp_analyze_image_ai(
self,
file_path: str = Field(description="Path to the image file for AI analysis"),
task: str = Field(
default="Describe what you see in this image",
description="Specific analysis task or question about the image",
),
) -> ActionResponse:
"""Analyze image content using AI vision models.
This tool uses advanced AI models to analyze and describe image content,
answer questions about images, or perform specific visual reasoning tasks.
Args:
file_path: Path to the image file
task: Specific analysis task or question
Returns:
ActionResponse with AI analysis results and metadata
"""
try:
# Handle FieldInfo objects
if isinstance(file_path, FieldInfo):
file_path = file_path.default
if isinstance(task, FieldInfo):
task = task.default
start_time = time.time()
# Validate input file
file_path: Path = self._validate_file_path(file_path)
self._color_log(f"Analyzing image with AI: {file_path.name}", Color.cyan)
# Load image
image = self._load_image(file_path)
original_metadata = self._get_image_metadata(image, file_path)
# Convert to base64 for AI analysis
image_base64 = self._image_to_base64(image, "JPEG")
# Perform AI analysis
analysis_result = self._analyze_with_ai(image_base64, task)
processing_time = time.time() - start_time
# Create metadata object
metadata_dict = {
"file_name": file_path.name,
"file_size": file_path.stat().st_size,
"file_type": file_path.suffix.lower(),
"absolute_path": str(file_path.absolute()),
"width": original_metadata["width"],
"height": original_metadata["height"],
"mode": original_metadata["mode"],
"format": original_metadata["format"],
"has_transparency": original_metadata["has_transparency"],
"processing_time": processing_time,
"output_files": [],
"analysis_result": analysis_result,
"output_format": "ai_analysis",
}
image_metadata = ImageMetadata(**metadata_dict)
result_message = (
f"AI Analysis Results for {file_path.name}:\n\n"
f"**Task:** {task}\n\n"
f"**Analysis:**\n{analysis_result}\n\n"
f"**Image Info:**\n"
f"- Dimensions: {original_metadata['width']}x{original_metadata['height']}\n"
f"- Format: {original_metadata['format']}\n"
f"- Processing time: {processing_time:.2f}s"
)
self._color_log(f"AI analysis completed in {processing_time:.2f}s", Color.green)
return ActionResponse(success=True, message=result_message, metadata=image_metadata.model_dump())
except Exception as e:
self.logger.error(f"AI image analysis failed: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=f"AI image analysis failed: {str(e)}",
metadata={"error_type": "ai_analysis_error"},
)
def mcp_get_image_metadata(
self,
file_path: str = Field(description="Path to the image file to analyze"),
) -> ActionResponse:
"""Extract comprehensive metadata from image files.
This tool analyzes image files and extracts detailed metadata including
dimensions, format, color mode, file size, and other technical information.
Args:
file_path: Path to the image file to analyze
Returns:
ActionResponse with detailed image metadata
"""
try:
if isinstance(file_path, FieldInfo):
file_path = file_path.default
start_time = time.time()
# Validate input file
file_path: Path = self._validate_file_path(file_path)
self._color_log(f"Extracting metadata from: {file_path.name}", Color.cyan)
# Load image
image = self._load_image(file_path)
metadata = self._get_image_metadata(image, file_path)
processing_time = time.time() - start_time
# Get file statistics
file_stats = file_path.stat()
# Create metadata object
image_metadata = ImageMetadata(
file_name=file_path.name,
file_size=file_stats.st_size,
file_type=file_path.suffix.lower(),
absolute_path=str(file_path.absolute()),
width=metadata["width"],
height=metadata["height"],
mode=metadata["mode"],
format=metadata["format"],
has_transparency=metadata["has_transparency"],
processing_time=processing_time,
output_files=[],
output_format="metadata",
)
# Calculate additional info
aspect_ratio = metadata["width"] / metadata["height"] if metadata["height"] > 0 else 0
megapixels = (metadata["width"] * metadata["height"]) / 1_000_000
result_message = (
f"Image Metadata for {file_path.name}:\n"
f"Dimensions: {metadata['width']}x{metadata['height']} pixels\n"
f"Aspect Ratio: {aspect_ratio:.2f}:1\n"
f"Megapixels: {megapixels:.2f} MP\n"
f"Color Mode: {metadata['mode']}\n"
f"Format: {metadata['format']}\n"
f"Has Transparency: {metadata['has_transparency']}\n"
f"File Size: {file_stats.st_size / 1024 / 1024:.2f} MB\n"
f"File Type: {file_path.suffix.upper()}"
)
self._color_log(f"Metadata extraction completed in {processing_time:.2f}s", Color.green)
return ActionResponse(success=True, message=result_message, metadata=image_metadata.model_dump())
except Exception as e:
self.logger.error(f"Metadata extraction failed: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=f"Metadata extraction failed: {str(e)}",
metadata={"error_type": "metadata_error"},
)
if __name__ == "__main__":
load_dotenv()
# Default arguments for testing
args = ActionArguments(
name="image_analysis_service",
transport="stdio",
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
)
# Initialize and run the image analysis service
try:
service = ImageCollection(args)
service.run()
except Exception as e:
print(f"An error occurred: {str(e)}")
@@ -0,0 +1,992 @@
"""
Video MCP Service
This module provides MCP service functionality for video operations including:
- Video content analysis with AI-powered insights
- Video summarization and key point extraction
- Keyframe extraction with scene detection
- Subtitle extraction from video content
It handles various video formats with proper validation, error handling,
and progress tracking while providing LLM-friendly formatted results.
"""
import base64
import os
import time
import traceback
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
import cv2
import numpy as np
from dotenv import load_dotenv
from pydantic import BaseModel, Field
from aworld.config.conf import AgentConfig
from aworld.logs.util import Color
from aworld.models.llm import call_llm_model, get_llm_model
from aworld.models.model_response import ModelResponse
from examples.gaia.mcp_collections.base import (
ActionArguments,
ActionCollection,
ActionResponse,
)
from ..utils import get_file_from_source
class VideoAnalysisResult(BaseModel):
"""Video analysis result model with structured data"""
video_source: str
analysis_result: str
frame_count: int
duration_analyzed: float
success: bool
error: str | None = None
class VideoSummaryResult(BaseModel):
"""Video summary result model with structured data"""
video_source: str
summary: str
frame_count: int
duration_analyzed: float
success: bool
error: str | None = None
class KeyframeResult(BaseModel):
"""Keyframe extraction result model with file information"""
frame_paths: list[str]
frame_timestamps: list[float]
output_directory: str
frame_count: int
target_time: float
window_size: float
success: bool
error: str | None = None
class VideoMetadata(BaseModel):
"""Metadata for video operation results"""
operation: str
video_source: str | None = None
sample_rate: int | None = None
start_time: float | None = None
end_time: float | None = None
target_time: float | None = None
window_size: float | None = None
output_directory: str | None = None
frame_count: int | None = None
execution_time: float | None = None
error_type: str | None = None
class VideoCollection(ActionCollection):
"""MCP service for video operations with AI-powered analysis.
Provides video processing capabilities including:
- AI-powered video content analysis
- Video summarization with key insights
- Keyframe extraction with scene detection
- Subtitle extraction from video content
- LLM-friendly result formatting
- Error handling and logging
"""
def __init__(self, arguments: ActionArguments) -> None:
super().__init__(arguments)
# Initialize supported video extensions
self.supported_extensions = {".mp4", ".avi", ".mov", ".mkv", ".webm", ".flv"}
# Video analysis prompts
self.video_analyze_prompt = (
"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"
)
self.video_summarize_prompt = (
"Input is a sequence of video frames. "
"Summarize the main content of the video. "
"Include key points, main topics, and important visual elements. "
)
self._color_log("Video service initialized", Color.green, "debug")
def _get_video_frames(
self,
video_source: str,
sample_rate: int = 2,
start_time: float = 0,
end_time: float | None = None,
) -> list[dict[str, any]]:
"""Extract frames from video with given sample rate.
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
end_time: End time of the video segment in seconds
Returns:
List of dictionaries containing frame data and timestamp
Raises:
ValueError: When video file cannot be opened or is not valid
"""
try:
# Get file with validation (only video files allowed)
file_path, _, _ = get_file_from_source(
video_source,
max_size_mb=2500.0, # 2500MB limit for videos
)
# Open video file
video = cv2.VideoCapture(file_path) # pylint: disable=E1101
if not video.isOpened():
raise ValueError(f"Could not open video file: {file_path}")
fps = video.get(cv2.CAP_PROP_FPS) # pylint: disable=E1101
frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) # pylint: disable=E1101
video_duration = frame_count / fps
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) # pylint: disable=E1101
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) # pylint: disable=E1101
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 != str(Path(video_source).resolve())
and Path(file_path).exists()
):
Path(file_path).unlink()
if not frames:
raise ValueError(
f"Could not extract any frames from video: {video_source}"
)
return frames
except Exception as e:
self._color_log(
f"Error extracting frames from {video_source}: {str(e)}", Color.red
)
raise
def _create_video_content(
self, 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
def _format_analysis_output(
self, result: VideoAnalysisResult, format_type: str = "markdown"
) -> str:
"""Format video analysis results for LLM consumption.
Args:
result: Video analysis result
format_type: Output format ('markdown', 'json', 'text')
Returns:
Formatted string suitable for LLM consumption
"""
if not result.success:
return f"Failed to analyze video: {result.error}"
if format_type == "json":
return result.model_dump_json(indent=2)
elif format_type == "text":
output_parts = [
"Video Analysis Results",
f"Source: {result.video_source}",
f"Frames Analyzed: {result.frame_count}",
f"Duration: {result.duration_analyzed:.2f} seconds",
"",
"Analysis:",
result.analysis_result,
]
return "\n".join(output_parts)
else: # markdown (default)
output_parts = [
"# Video Analysis Results ✅",
"",
"## Video Information",
f"**Source:** `{result.video_source}`",
f"**Frames Analyzed:** {result.frame_count}",
f"**Duration:** {result.duration_analyzed:.2f} seconds",
"",
"## Analysis Results",
result.analysis_result,
]
return "\n".join(output_parts)
def _format_summary_output(
self, result: VideoSummaryResult, format_type: str = "markdown"
) -> str:
"""Format video summary results for LLM consumption.
Args:
result: Video summary result
format_type: Output format ('markdown', 'json', 'text')
Returns:
Formatted string suitable for LLM consumption
"""
if not result.success:
return f"Failed to summarize video: {result.error}"
if format_type == "json":
return result.model_dump_json(indent=2)
elif format_type == "text":
output_parts = [
"Video Summary",
f"Source: {result.video_source}",
f"Frames Analyzed: {result.frame_count}",
f"Duration: {result.duration_analyzed:.2f} seconds",
"",
"Summary:",
result.summary,
]
return "\n".join(output_parts)
else: # markdown (default)
output_parts = [
"# Video Summary ✅",
"",
"## Video Information",
f"**Source:** `{result.video_source}`",
f"**Frames Analyzed:** {result.frame_count}",
f"**Duration:** {result.duration_analyzed:.2f} seconds",
"",
"## Summary",
result.summary,
]
return "\n".join(output_parts)
def _format_keyframe_output(
self, result: KeyframeResult, format_type: str = "markdown"
) -> str:
"""Format keyframe extraction results for LLM consumption.
Args:
result: Keyframe extraction result
format_type: Output format ('markdown', 'json', 'text')
Returns:
Formatted string suitable for LLM consumption
"""
if not result.success:
return f"Failed to extract keyframes: {result.error}"
if format_type == "json":
return result.model_dump_json(indent=2)
elif format_type == "text":
output_parts = [
"Keyframe Extraction Results",
f"Target Time: {result.target_time}s",
f"Window Size: {result.window_size}s",
f"Frames Extracted: {result.frame_count}",
f"Output Directory: {result.output_directory}",
"",
"Frame Files:",
]
for i, (path, timestamp) in enumerate(
zip(result.frame_paths, result.frame_timestamps), 1
):
output_parts.append(f"{i}. {path} (at {timestamp:.2f}s)")
return "\n".join(output_parts)
else: # markdown (default)
output_parts = [
"# Keyframe Extraction Results ✅",
"",
"## Extraction Parameters",
f"**Target Time:** {result.target_time}s",
f"**Window Size:** {result.window_size}s",
f"**Frames Extracted:** {result.frame_count}",
f"**Output Directory:** `{result.output_directory}`",
"",
"## Extracted Frames",
]
for i, (path, timestamp) in enumerate(
zip(result.frame_paths, result.frame_timestamps), 1
):
output_parts.append(f"{i}. `{path}` (at {timestamp:.2f}s)")
return "\n".join(output_parts)
def _analyze_frame_chunk(
self, chunk_data: tuple[int, list, str]
) -> tuple[int, str]:
"""Analyze a chunk of video frames using LLM.
Args:
chunk_data: Tuple containing (chunk_index, frames, question)
Returns:
Tuple of (chunk_index, analysis_result)
"""
chunk_index, frames, question = chunk_data
try:
content = self._create_video_content(
self.video_analyze_prompt.format(task=question), frames
)
inputs = [{"role": "user", "content": content}]
response: ModelResponse = call_llm_model(
get_llm_model(
conf=AgentConfig(
llm_provider="openai",
llm_model_name=os.getenv("VIDEO_LLM_MODEL_NAME"),
llm_api_key=os.getenv("VIDEO_LLM_API_KEY"),
llm_base_url=os.getenv("VIDEO_LLM_BASE_URL"),
)
),
inputs,
temperature=float(os.getenv("VIDEO_LLM_TEMPERATURE", "1.0")),
)
analysis_result = response.content
self._color_log(
f"✅ Completed analysis for chunk {chunk_index + 1}", Color.green
)
except Exception as e:
self._color_log(
f"❌ LLM analysis error for chunk {chunk_index + 1}: {str(e)}",
Color.yellow,
)
analysis_result = (
f"Analysis failed for video segment {chunk_index + 1}: {str(e)}"
)
return chunk_index, analysis_result
async def mcp_analyze_video(
self,
video_url: str = Field(description="Path or URL to the video file to analyze"),
question: str = Field(description="Question or task for video analysis"),
sample_rate: float = Field(
default=1.0, description="Frame sampling rate (frames per second)"
),
start_time: float = Field(default=0.0, description="Start time in seconds"),
end_time: float | None = Field(
default=None, description="End time in seconds (None for full video)"
),
output_format: str = Field(
default="markdown",
description="Output format: 'markdown', 'json', or 'text'",
),
max_workers: int = Field(
default=4, description="Maximum number of parallel workers for analysis"
),
) -> ActionResponse:
"""Analyze video content using AI with parallel processing.
This tool provides comprehensive video analysis capabilities including:
- Content understanding and description
- Object and scene detection
- Action and movement analysis
- Temporal event tracking
- Question-answering about video content
- Parallel processing for faster analysis
Args:
video_url: Path or URL to the video file
question: Specific question or analysis task
sample_rate: Frame sampling rate for analysis
start_time: Start time of the video segment in seconds
end_time: End time of the video segment in seconds
output_format: Format for the response output
max_workers: Maximum number of parallel workers
Returns:
ActionResponse with video analysis results and metadata
"""
start_exec_time = time.time()
try:
# Validate video file
video_path = self._validate_file_path(video_url)
self._color_log(f"🎬 Analyzing video: {video_url}", Color.cyan)
self._color_log(f"📋 Question: {question}", Color.blue)
# Extract video frames
video_frames = self._get_video_frames(
str(video_path), sample_rate, start_time, end_time
)
self._color_log(f"📸 Extracted {len(video_frames)} frames", Color.blue)
# Process frames in chunks of 64 frames for parallel analysis
chunk_size = 128
chunks = []
# Create chunks of `chunk_size` continuous frames
for i in range(0, len(video_frames), chunk_size):
chunk_frames = video_frames[i : i + chunk_size]
chunks.append((i // chunk_size, chunk_frames, question))
self._color_log(
f"🔄 Processing {len(chunks)} chunks with {max_workers} parallel workers",
Color.blue,
)
# Process chunks in parallel
all_results = [None] * len(chunks) # Pre-allocate to maintain order
with ThreadPoolExecutor(max_workers=max_workers) as executor:
# Submit all chunk analysis tasks
future_to_chunk = {
executor.submit(self._analyze_frame_chunk, chunk_data): chunk_data[
0
]
for chunk_data in chunks
}
# Collect results as they complete
for future in as_completed(future_to_chunk):
try:
chunk_index, result = future.result()
all_results[chunk_index] = (
f"Result of video part {chunk_index + 1}: {result}"
)
except Exception as e:
chunk_index = future_to_chunk[future]
self._color_log(
f"❌ Error processing chunk {chunk_index + 1}: {str(e)}",
Color.red,
)
all_results[chunk_index] = (
f"Result of video part {chunk_index + 1}: Analysis failed - {str(e)}"
)
# Filter out None results and join
analysis_result = "\n".join(
[result for result in all_results if result is not None]
)
duration_analyzed = (
end_time - start_time if end_time else len(video_frames) / sample_rate
)
# Create result
result = VideoAnalysisResult(
video_source=video_url,
analysis_result=analysis_result,
frame_count=len(video_frames),
duration_analyzed=duration_analyzed,
success=True,
error=None,
)
# Format output for LLM
message = self._format_analysis_output(result, output_format)
execution_time = time.time() - start_exec_time
# Create metadata
metadata = {
"video_source": video_url,
"frame_count": len(video_frames),
"chunks_processed": len(chunks),
"chunk_size": chunk_size,
"parallel_workers": max_workers,
"duration_analyzed": duration_analyzed,
"sample_rate": sample_rate,
"start_time": start_time,
"end_time": end_time,
"execution_time": execution_time,
"output_format": output_format,
"success": True,
}
self._color_log(
f"✅ Video analysis completed in {execution_time:.2f}s "
f"({len(chunks)} chunks, {len(video_frames)} frames)",
Color.green,
)
return ActionResponse(success=True, message=message, metadata=metadata)
except Exception as e:
execution_time = time.time() - start_exec_time
error_msg = f"Video analysis failed: {str(e)}"
self._color_log(f"{error_msg}", Color.red)
self.logger.error(f"{error_msg}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=error_msg,
metadata={
"video_source": video_url,
"execution_time": execution_time,
"error": str(e),
"success": False,
},
)
async def mcp_summarize_video(
self,
video_url: str = Field(
description="The input video filepath or URL to summarize."
),
sample_rate: int = Field(
default=1, description="Sample n frames per second (default: 1)."
),
start_time: float = Field(
default=0,
description="Start time of the video segment in seconds (default: 0).",
),
end_time: float | None = Field(
default=None,
description="End time of the video segment in seconds (default: None).",
),
output_format: str = Field(
default="markdown",
description="Output format: 'markdown', 'json', or 'text' (default: markdown).",
),
) -> ActionResponse:
"""Summarize the main content of a video using AI analysis.
This tool provides AI-powered video summarization with:
- Key point extraction
- Main topic identification
- Important visual element recognition
- LLM-optimized result formatting
Args:
video_url: The input video filepath or URL to summarize
sample_rate: Sample n frames per second
start_time: Start time of the video segment in seconds
end_time: End time of the video segment in seconds
output_format: Format for the response output
Returns:
ActionResponse with video summary results and metadata
"""
start_exec_time = time.time()
try:
# Validate video file
video_path = self._validate_file_path(video_url)
self._color_log(f"🎬 Summarizing video: {video_url}", Color.cyan)
# Extract video frames
video_frames = self._get_video_frames(
str(video_path), sample_rate, start_time, end_time
)
self._color_log(f"📸 Extracted {len(video_frames)} frames", Color.blue)
# Process frames in larger chunks for summarization
interval = 490
frame_nums = 500
all_results = []
for i in range(0, len(video_frames), interval):
cur_frames = video_frames[i : i + frame_nums]
content = self._create_video_content(
self.video_summarize_prompt, cur_frames
)
inputs = [{"role": "user", "content": content}]
try:
response: ModelResponse = call_llm_model(
get_llm_model(
conf=AgentConfig(
llm_provider="openai",
llm_model_name=os.getenv(
"VIDEO_LLM_MODEL_NAME", "gpt-4o"
),
llm_api_key=os.getenv("VIDEO_LLM_API_KEY"),
llm_base_url=os.getenv("VIDEO_LLM_BASE_URL"),
)
),
inputs,
temperature=float(os.getenv("VIDEO_LLM_TEMPERATURE", "1.0")),
)
cur_summary = response.content
except Exception as e:
self._color_log(
f"LLM summary error for chunk {i // interval + 1}: {str(e)}",
Color.yellow,
)
cur_summary = (
f"Summary failed for video segment {i // interval + 1}"
)
all_results.append(
f"Summary of video part {i // interval + 1}: {cur_summary}"
)
if i + frame_nums >= len(video_frames):
break
summary_result = "\n".join(all_results)
duration_analyzed = (
end_time - start_time if end_time else len(video_frames) / sample_rate
)
# Create result
result = VideoSummaryResult(
video_source=video_url,
summary=summary_result,
frame_count=len(video_frames),
duration_analyzed=duration_analyzed,
success=True,
error=None,
)
# Format output for LLM
message = self._format_summary_output(result, output_format)
execution_time = time.time() - start_exec_time
# Create metadata
metadata = VideoMetadata(
operation="summarize",
video_source=video_url,
sample_rate=sample_rate,
start_time=start_time,
end_time=end_time,
frame_count=len(video_frames),
execution_time=execution_time,
).model_dump()
self._color_log(
"✅ Video summarization completed successfully", Color.green
)
return ActionResponse(success=True, message=message, metadata=metadata)
except Exception as e:
error_msg = str(e)
self._color_log(
f"❌ Video summarization error: {traceback.format_exc()}", Color.red
)
# Format error for LLM
message = f"Failed to summarize video: {error_msg}"
execution_time = time.time() - start_exec_time
# Create metadata
metadata = VideoMetadata(
operation="summarize",
video_source=video_url,
error_type="summarization_failure",
execution_time=execution_time,
).model_dump()
return ActionResponse(success=False, message=message, metadata=metadata)
async def mcp_extract_keyframes(
self,
video_path: str = Field(description="The input video filepath or URL."),
target_time: int = Field(
description="The specific time point for extraction (in seconds), centered within the window_size."
),
window_size: int = Field(
default=5,
description="The window size for extraction (in seconds, default: 5).",
),
output_dir: str = Field(
default=None,
description="Directory where extracted frames will be saved (default: workspace/keyframes).",
),
output_format: str = Field(
default="markdown",
description="Output format: 'markdown', 'json', or 'text' (default: markdown).",
),
) -> ActionResponse:
"""Extract key frames around a target time with scene detection.
This tool provides keyframe extraction with:
- Scene detection for significant frame changes
- Configurable time windows
- Automatic output directory management
- LLM-optimized result formatting
Args:
video_path: The input video filepath or URL
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
output_format: Format for the response output
Returns:
ActionResponse with keyframe extraction results and metadata
"""
start_exec_time = time.time()
try:
# Validate video file
validated_path = self._validate_file_path(video_path)
# Set default output directory
output_dir = (
str(self.workspace / "keyframes") if output_dir is None else output_dir
)
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
self._color_log(f"🎬 Extracting keyframes from: {video_path}", Color.cyan)
self._color_log(
f"🎯 Target time: {target_time}s, Window: {window_size}s", Color.blue
)
# Extract keyframes with scene detection
frames, frame_times = self._extract_keyframes_with_scene_detection(
str(validated_path), target_time, window_size
)
# Save frames to disk
frame_paths, frame_timestamps = self._save_keyframes(
frames, frame_times, str(output_path)
)
# Cleanup if requested
# if cleanup and validated_path.exists():
# validated_path.unlink()
# self._color_log("🗑️ Cleaned up original video file", Color.yellow)
# Create result
result = KeyframeResult(
frame_paths=frame_paths,
frame_timestamps=frame_timestamps,
output_directory=str(output_path),
frame_count=len(frame_paths),
target_time=float(target_time),
window_size=float(window_size),
success=True,
error=None,
)
# Format output for LLM
message = self._format_keyframe_output(result, output_format)
execution_time = time.time() - start_exec_time
# Create metadata
metadata = VideoMetadata(
operation="extract_keyframes",
video_source=video_path,
target_time=float(target_time),
window_size=float(window_size),
output_directory=str(output_path),
frame_count=len(frame_paths),
execution_time=execution_time,
).model_dump()
self._color_log(
f"✅ Extracted {len(frame_paths)} keyframes successfully", Color.green
)
return ActionResponse(success=True, message=message, metadata=metadata)
except Exception as e:
error_msg = str(e)
self._color_log(
f"❌ Keyframe extraction error: {traceback.format_exc()}", Color.red
)
# Format error for LLM
message = f"Failed to extract keyframes: {error_msg}"
execution_time = time.time() - start_exec_time
# Create metadata
metadata = VideoMetadata(
operation="extract_keyframes",
video_source=video_path,
target_time=float(target_time) if target_time else None,
window_size=float(window_size) if window_size else None,
error_type="keyframe_extraction_failure",
execution_time=execution_time,
).model_dump()
return ActionResponse(success=False, message=message, metadata=metadata)
def _extract_keyframes_with_scene_detection(
self, video_path: str, target_time: int, window_size: int
) -> tuple[list[any], list[float]]:
"""Extract key frames around the target time with scene detection.
Args:
video_path: Path to the video file
target_time: Target time in seconds
window_size: Window size in seconds
Returns:
Tuple of (frames, frame_times)
"""
cap = cv2.VideoCapture(video_path) # pylint: disable=E1101
fps = cap.get(cv2.CAP_PROP_FPS) # pylint: disable=E1101
# 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)
total_frames_in_window = end_frame - start_frame
max_frames = 384 # Maximum allowed frames to prevent memory issues
# Calculate sampling interval for even distribution
if total_frames_in_window <= max_frames:
# If total frames is within limit, use scene detection normally
frame_interval = 1
use_scene_detection = True
else:
# If exceeds limit, sample evenly across the window
frame_interval = total_frames_in_window // max_frames
use_scene_detection = False # Skip scene detection for even sampling
frames = []
frame_times = []
# Set video position to start_frame
cap.set(cv2.CAP_PROP_POS_FRAMES, max(0, start_frame)) # pylint: disable=E1101
prev_frame = None
frame_count = 0
while cap.isOpened() and len(frames) < max_frames:
frame_pos = cap.get(cv2.CAP_PROP_POS_FRAMES) # pylint: disable=E1101
if frame_pos >= end_frame:
break
ret, frame = cap.read()
if not ret:
break
if use_scene_detection:
# Use original scene detection logic
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) # pylint: disable=E1101
# 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) # pylint: disable=E1101
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
else:
# Use even sampling for large windows
if frame_count % frame_interval == 0:
frames.append(frame)
frame_times.append(frame_pos / fps)
frame_count += 1
cap.release()
return frames, frame_times
def _save_keyframes(
self, frames: list[any], frame_times: list[float], output_dir: str
) -> tuple[list[str], list[float]]:
"""Save extracted frames to disk.
Args:
frames: List of frame objects
frame_times: List of frame timestamps
output_dir: Output directory path
Returns:
Tuple of (saved_paths, saved_timestamps)
"""
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
saved_paths = []
saved_timestamps = []
for _, (frame, timestamp) in enumerate(zip(frames, frame_times)):
filename = output_path / f"frame_{timestamp:.2f}s.jpg"
cv2.imwrite(str(filename), frame) # pylint: disable=E1101
saved_paths.append(str(filename))
saved_timestamps.append(timestamp)
return saved_paths, saved_timestamps
# Default arguments for testing
if __name__ == "__main__":
load_dotenv()
arguments = ActionArguments(
name="video",
transport="stdio",
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
)
try:
service = VideoCollection(arguments)
service.run()
except Exception as e:
print(f"An error occurred: {str(e)}")
@@ -0,0 +1,374 @@
"""
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 using browser-use package.
Main functions:
- mcp_browser_use: Performs browser automation tasks with LLM-friendly output
"""
import json
import os
import re
import sys
import time
import traceback
try:
from browser_use import Agent, AgentHistoryList, BrowserProfile
from browser_use.llm import ChatOpenAI
from dotenv import load_dotenv
from pydantic import BaseModel, Field
from aworld.logs.util import Color
from ..base import ActionArguments, ActionCollection, ActionResponse
except Exception as e:
print(f"Failed to import browser tool: {traceback.format_exc()}")
raise e
print(f"Browser tool sys.path: {sys.path}")
class BrowserMetadata(BaseModel):
"""Metadata for browser automation results."""
task: str
execution_successful: bool
steps_taken: int | None = None
downloaded_files: list[str] = Field(default_factory=list)
visited_urls: list[str] = Field(default_factory=list)
execution_time: float | None = None
error_type: str | None = None
trace_log_path: str | None = None
class BrowserActionCollection(ActionCollection):
"""MCP service for browser automation using browser-use package.
Provides comprehensive web automation capabilities including:
- Web scraping and content extraction
- Form submission and interaction
- File downloads and media handling
- LLM-enhanced browsing with memory
- Robot detection and paywall handling
"""
def __init__(self, arguments: ActionArguments) -> None:
super().__init__(arguments)
# Load environment variables
load_dotenv()
# Extended system prompt for browser automation
self.extended_browser_system_prompt = """
10. URL ends with .pdf
- If the go_to_url function with `https://any_url/any_file_name.pdf` as the parameter, just report the url link and hint the user to download using `download` mcp tool or `curl`, then execute `done` action.
11. Robot Detection:
- If the page is a robot detection page, abort immediately. Then navigate to the most authoritative source for similar information instead
# Efficiency Guidelines
0. if download option is available, always **DOWNLOAD** as possible! Also, report the download url link in your result.
1. Use specific search queries with key terms from the task
2. Avoid getting distracted by tangential information
3. If blocked by paywalls, try archive.org or similar alternatives
4. Document each significant finding clearly and concisely
5. Precisely extract the necessary information with minimal browsing steps.
"""
# Initialize LLM configuration
self.llm_config = ChatOpenAI(
model=os.getenv("LLM_MODEL_NAME"),
api_key=os.getenv("LLM_API_KEY"),
base_url=os.getenv("LLM_BASE_URL"),
temperature=1.0,
)
self._color_log(f"Browser llm_config: {self.llm_config}", Color.green)
# Browser profile configuration
self.browser_profile = BrowserProfile(
cookies_file=os.getenv("COOKIES_FILE_PATH"),
downloads_dir=str(self.workspace),
downloads_path=str(self.workspace),
save_recording_path=str(self.workspace),
save_downloads_path=str(self.workspace),
chromium_sandbox=False,
headless=True,
)
self._color_log(f"Browser browser_profile: {self.browser_profile}", Color.green)
# Log configuration
self.trace_log_dir = str(self.workspace / "logs")
os.makedirs(f"{self.trace_log_dir}/browser_log", exist_ok=True)
self._color_log("Browser automation service initialized", Color.green)
self._color_log(
f"Downloads directory: {self.browser_profile.downloads_path}", Color.blue
)
self._color_log(
f"Trace logs directory: {self.trace_log_dir}/browser_log", Color.blue
)
def _create_browser_agent(self, task: str) -> Agent:
"""Create a browser agent instance with configured settings.
Args:
task: The task description for the browser agent
Returns:
Configured Agent instance
"""
return Agent(
task=task,
llm=self.llm_config,
extend_system_message=self.extended_browser_system_prompt,
use_vision=True,
enable_memory=False,
browser_profile=self.browser_profile,
save_conversation_path=f"{self.trace_log_dir}/browser_log/trace.log",
)
def _extract_visited_urls(self, extracted_content: list[str]) -> list[str]:
"""Inner method to extract URLs from content using regex.
Args:
content_list: List of content strings to search for URLs
Returns:
List of unique URLs found in the content
"""
url_pattern = r'https?://[^\s<>"\[\]{}|\\^`]+'
visited_urls = set()
for content in extracted_content:
if content and isinstance(content, str):
urls = re.findall(url_pattern, content)
visited_urls.update(urls)
return list(visited_urls)
def _format_extracted_content(self, extracted_content: list[str]) -> str:
"""Format extracted content to be LLM-friendly.
Args:
extracted_content: List of extracted content strings from browser execution
Returns:
Formatted string suitable for LLM consumption
"""
if not extracted_content:
return "No content extracted from browser execution."
# Handle list of strings
if len(extracted_content) == 1:
# Single item - return it directly with formatting
return f"**Extracted Content:**\n{extracted_content[0]}"
else:
# Multiple items - format as numbered list
formatted_parts = ["**Extracted Content:**"]
for i, content in enumerate(extracted_content, 1):
if content.strip(): # Only include non-empty content
formatted_parts.append(f"{i}. {content}")
return (
"\n".join(formatted_parts)
if len(formatted_parts) > 1
else "No meaningful content extracted from browser execution."
)
async def mcp_browser_use(
self,
task: str = Field(
description="The task to perform using the browser automation agent"
),
max_steps: int = Field(
default=50, description="Maximum number of steps for browser execution"
),
extract_format: str = Field(
default="markdown",
description="Format for extracted content: 'markdown', 'json', or 'text'",
),
) -> ActionResponse:
"""Perform browser automation tasks using the browser-use package.
This tool provides comprehensive browser automation capabilities including:
- Web scraping and content extraction
- Form submission and automated interactions
- File downloads and media handling
- LLM-enhanced browsing with memory and vision
- Automatic handling of robot detection and paywalls
Args:
task: Description of the browser automation task to perform
max_steps: Maximum number of execution steps (default: 50)
extract_format: Output format for extracted content
Returns:
ActionResponse with LLM-friendly extracted content and execution metadata
"""
try:
self._color_log(f"🎯 Starting browser task: {task}", Color.cyan)
# Create browser agent
agent = self._create_browser_agent(task)
start_time = time.time()
browser_execution: AgentHistoryList = await agent.run(max_steps=max_steps)
execution_time = time.time() - start_time
if (
browser_execution is not None
and browser_execution.is_done()
and browser_execution.is_successful()
):
# Extract and format content
extracted_content = browser_execution.extracted_content()
final_result = browser_execution.final_result()
# Format content based on requested format
if extract_format.lower() == "json":
formatted_content = json.dumps(
{"summary": final_result, "extracted_data": extracted_content},
indent=2,
)
elif extract_format.lower() == "text":
formatted_content = f"{final_result}\n\n{self._format_extracted_content(extracted_content)}"
else: # markdown (default)
formatted_content = (
f"## Browser Automation Result\n\n**Summary:** {final_result}\n\n"
f"{self._format_extracted_content(extracted_content)}"
)
# Prepare metadata
metadata = BrowserMetadata(
task=task,
execution_successful=True,
steps_taken=(
len(browser_execution.history)
if hasattr(browser_execution, "history")
else None
),
downloaded_files=[],
visited_urls=self._extract_visited_urls(extracted_content),
execution_time=execution_time,
trace_log_path=f"{self.trace_log_dir}/browser_log/trace.log",
)
self._color_log(f"🗒️ Detail: {extracted_content}", Color.lightgrey)
self._color_log(f"🌏 Result: {final_result}", Color.green)
return ActionResponse(
success=True,
message=formatted_content,
metadata=metadata.model_dump(),
)
else:
# Handle execution failure
error_msg = "Browser execution failed or was not completed successfully"
metadata = BrowserMetadata(
task=task,
execution_successful=False,
execution_time=execution_time,
error_type="execution_failure",
trace_log_path=f"{self.trace_log_dir}/browser_log/trace.log",
)
self._color_log(f"{error_msg}", Color.red)
return ActionResponse(
success=False, message=error_msg, metadata=metadata.model_dump()
)
except Exception as e:
error_msg = f"Browser automation failed: {str(e)}"
error_trace = traceback.format_exc()
self.logger.error(f"Browser execution error: {error_trace}")
metadata = BrowserMetadata(
task=task,
execution_successful=False,
error_type="exception",
trace_log_path=f"{self.trace_log_dir}/browser_log/trace.log",
)
self._color_log(f"{error_msg}", Color.red)
return ActionResponse(
success=False,
message=f"{error_msg}\n\nError details: {error_trace}",
metadata=metadata.model_dump(),
)
def mcp_get_browser_capabilities(self) -> ActionResponse:
"""Get information about browser automation capabilities and configuration.
Returns:
ActionResponse with browser service capabilities and current configuration
"""
capabilities = {
"automation_features": [
"Web scraping and content extraction",
"Form submission and interaction",
"File downloads and media handling",
"LLM-enhanced browsing with vision",
"Memory-enabled browsing sessions",
"Robot detection and paywall handling",
],
"supported_formats": ["markdown", "json", "text"],
"configuration": {
"llm_model": os.getenv("LLM_MODEL_NAME", "Not configured"),
"downloads_directory": self.browser_profile.downloads_path,
"cookies_enabled": bool(os.getenv("COOKIES_FILE_PATH")),
"trace_logging": True,
"vision_enabled": True,
"headless": True,
},
}
formatted_info = f"""# Browser Automation Service Capabilities
## Features
{chr(10).join(f"- {feature}" for feature in capabilities["automation_features"])}
## Supported Output Formats
{chr(10).join(f"- {fmt}" for fmt in capabilities["supported_formats"])}
## Current Configuration
- **LLM Model:** {capabilities["configuration"]["llm_model"]}
- **Downloads Directory:** {capabilities["configuration"]["downloads_directory"]}
- **Cookies Enabled:** {capabilities["configuration"]["cookies_enabled"]}
- **Vision Enabled:** {capabilities["configuration"]["vision_enabled"]}
- **Memory Enabled:** {capabilities["configuration"]["memory_enabled"]}
- **Trace Logging:** {capabilities["configuration"]["trace_logging"]}
"""
return ActionResponse(
success=True, message=formatted_info, metadata=capabilities
)
# Example usage and entry point
if __name__ == "__main__":
load_dotenv()
# Default arguments for testing
args = ActionArguments(
name="browser_automation_service",
transport="stdio",
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
)
# Initialize and run the browser automation service
try:
service = BrowserActionCollection(args)
service.run()
except Exception as e:
print(f"An error occurred: {e}: {traceback.format_exc()}")
@@ -0,0 +1,479 @@
"""
Download MCP Server
This module provides MCP server functionality for downloading files from URLs.
It supports HTTP/HTTPS downloads with configurable options and returns LLM-friendly formatted results.
Key features:
- Download files from HTTP/HTTPS URLs
- Configurable timeout and overwrite options
- Custom headers support for authentication
- LLM-optimized output formatting
- Comprehensive error handling and logging
- Path validation and directory creation
Main functions:
- mcp_download_file: Download files from URLs with comprehensive options
- mcp_get_download_capabilities: Get download service capabilities
"""
import json
import shutil
import time
import traceback
from datetime import datetime
from pathlib import Path
from urllib.parse import urlparse
import requests
from dotenv import load_dotenv
from pydantic import BaseModel, Field
from pydantic.fields import FieldInfo
from aworld.logs.util import Color
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
class DownloadResult(BaseModel):
"""Individual download operation result with structured data."""
url: str
file_path: str
success: bool
file_size: int | None = None
duration: str
timestamp: str
error_message: str | None = None
class DownloadMetadata(BaseModel):
"""Metadata for download operation results."""
url: str
output_path: str
timeout_seconds: int
overwrite_enabled: bool
execution_time: float | None = None
file_size_bytes: int | None = None
content_type: str | None = None
status_code: int | None = None
error_type: str | None = None
headers_used: bool = False
class DownloadCollection(ActionCollection):
"""MCP service for file download operations with comprehensive controls.
Provides secure file download capabilities including:
- HTTP/HTTPS URL support
- Configurable timeout controls
- Custom headers for authentication
- Path validation and directory creation
- LLM-friendly result formatting
- Error handling and logging
"""
def __init__(self, arguments: ActionArguments) -> None:
super().__init__(arguments)
# Configuration
self.default_timeout = 60 * 3 # 3 minutes timeout
self.max_file_size = 1024 * 1024 * 1024 # 1GB limit
self.supported_schemes = {"http", "https"}
self.headers = {
"User-Agent": (
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
"AppleWebKit/537.36 (KHTML, like Gecko) "
"Chrome/91.0.4472.124 Safari/537.36"
),
"Accept-Language": "en-US,en;q=0.9",
}
self._color_log("Download service initialized", Color.green, "debug")
self._color_log(f"Workspace: {self.workspace}", Color.blue, "debug")
def _validate_url(self, url: str) -> tuple[bool, str | None]:
"""Validate URL format and scheme.
Args:
url: URL to validate
Returns:
Tuple of (is_valid, error_message)
"""
try:
parsed = urlparse(url)
if not parsed.scheme:
return False, "URL must include a scheme (http:// or https://)"
if parsed.scheme.lower() not in self.supported_schemes:
return False, f"Unsupported URL scheme: {parsed.scheme}. Supported: {', '.join(self.supported_schemes)}"
if not parsed.netloc:
return False, "URL must include a valid domain"
return True, None
except Exception as e:
return False, f"Invalid URL format: {str(e)}"
def _resolve_output_path(self, output_path: str) -> Path:
"""Resolve and validate output file path.
Args:
output_path: Output file path (absolute or relative)
Returns:
Resolved Path object
"""
path = Path(output_path).expanduser()
if not path.is_absolute():
path = self.workspace / path
# Ensure parent directory exists
path.parent.mkdir(parents=True, exist_ok=True)
return path.resolve()
def _format_download_output(self, result: DownloadResult, output_format: str = "markdown") -> str:
"""Format download results for LLM consumption.
Args:
result: Download execution result
output_format: Format type ('markdown', 'json', 'text')
Returns:
Formatted string suitable for LLM consumption
"""
if output_format == "json":
return json.dumps(result.model_dump(), indent=2)
elif output_format == "text":
output_parts = [
f"URL: {result.url}",
f"File Path: {result.file_path}",
f"Status: {'SUCCESS' if result.success else 'FAILED'}",
f"Duration: {result.duration}",
f"Timestamp: {result.timestamp}",
]
if result.file_size is not None:
output_parts.append(f"File Size: {result.file_size:,} bytes")
if result.error_message:
output_parts.append(f"Error: {result.error_message}")
return "\n".join(output_parts)
else: # markdown (default)
status_emoji = "" if result.success else ""
output_parts = [
f"# File Download {status_emoji}",
f"**URL:** `{result.url}`",
f"**File Path:** `{result.file_path}`",
f"**Status:** {'SUCCESS' if result.success else 'FAILED'}",
f"**Duration:** {result.duration}",
f"**Timestamp:** {result.timestamp}",
]
if result.file_size is not None:
size_mb = result.file_size / (1024 * 1024)
output_parts.append(f"**File Size:** {result.file_size:,} bytes ({size_mb:.2f} MB)")
if result.error_message:
output_parts.extend(["\n## Error Details", f"```\n{result.error_message}\n```"])
return "\n".join(output_parts)
async def _download_file_async(
self, url: str, output_path: Path, timeout: int, headers: dict[str, str] | None
) -> DownloadResult:
"""Download file asynchronously with comprehensive error handling.
Args:
url: URL to download from
output_path: Local path to save file
timeout: Request timeout in seconds
headers: Optional custom headers
Returns:
DownloadResult with execution details
"""
start_time = datetime.now()
try:
self._color_log(f"📥 Starting download: {url}", Color.cyan)
with requests.get(url, stream=True, timeout=timeout, headers=headers) as response:
response.raise_for_status()
# Check content length if available
content_length = response.headers.get("content-length")
if content_length and int(content_length) > self.max_file_size:
raise ValueError(f"File too large: {content_length} bytes (max: {self.max_file_size})")
# Download file
with open(output_path, "wb") as f:
shutil.copyfileobj(response.raw, f)
file_size = output_path.stat().st_size
duration = str(datetime.now() - start_time)
self._color_log(f"✅ Download completed: {file_size:,} bytes", Color.green)
return DownloadResult(
url=url,
file_path=str(output_path),
success=True,
file_size=file_size,
duration=duration,
timestamp=start_time.isoformat(),
)
except requests.exceptions.Timeout:
duration = str(datetime.now() - start_time)
error_msg = f"Download timed out after {timeout} seconds"
self._color_log(f"{error_msg}", Color.red)
return DownloadResult(
url=url,
file_path=str(output_path),
success=False,
duration=duration,
timestamp=start_time.isoformat(),
error_message=error_msg,
)
except requests.exceptions.RequestException as e:
duration = str(datetime.now() - start_time)
error_msg = f"Request failed: {str(e)}"
self._color_log(f"{error_msg}", Color.red)
return DownloadResult(
url=url,
file_path=str(output_path),
success=False,
duration=duration,
timestamp=start_time.isoformat(),
error_message=error_msg,
)
except Exception as e:
duration = str(datetime.now() - start_time)
error_msg = f"Unexpected error: {str(e)}"
self._color_log(f"💥 {error_msg}", Color.red)
return DownloadResult(
url=url,
file_path=str(output_path),
success=False,
duration=duration,
timestamp=start_time.isoformat(),
error_message=error_msg,
)
async def mcp_download_file(
self,
url: str = Field(description="HTTP/HTTPS URL of the file to download"),
output_file_path: str = Field(
description="Local path where the file should be saved (absolute or relative to workspace)"
),
overwrite: bool = Field(default=False, description="Whether to overwrite existing files (default: False)"),
timeout: int = Field(default=60, description="Download timeout in seconds (default: 60)"),
output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'"),
) -> ActionResponse:
"""Download a file from a URL with comprehensive options and controls.
This tool provides secure file download capabilities with:
- HTTP/HTTPS URL support
- Configurable timeout controls
- Path validation and directory creation
- File size limits and safety checks
- LLM-optimized result formatting
Args:
url: The HTTP/HTTPS URL of the file to download
output_file_path: Local path to save the downloaded file
overwrite: Whether to overwrite existing files
timeout: Maximum download time in seconds
output_format: Format for the response output
Returns:
ActionResponse with download results and metadata
"""
# Handle FieldInfo objects
if isinstance(url, FieldInfo):
url = url.default
if isinstance(output_file_path, FieldInfo):
output_file_path = output_file_path.default
if isinstance(overwrite, FieldInfo):
overwrite = overwrite.default
if isinstance(timeout, FieldInfo):
timeout = timeout.default
if isinstance(output_format, FieldInfo):
output_format = output_format.default
try:
# Validate URL
url_valid, url_error = self._validate_url(url)
if not url_valid:
return ActionResponse(
success=False,
message=f"Invalid URL: {url_error}",
metadata=DownloadMetadata(
url=url,
output_path=output_file_path,
timeout_seconds=timeout,
overwrite_enabled=overwrite,
error_type="invalid_url",
).model_dump(),
)
# Resolve output path
output_path = self._resolve_output_path(output_file_path)
# Check if file exists and overwrite setting
if output_path.exists() and not overwrite:
existing_size = output_path.stat().st_size
return ActionResponse(
success=False,
message=f"File already exists at {output_path} ({existing_size:,} bytes) and overwrite is disabled",
metadata=DownloadMetadata(
url=url,
output_path=str(output_path),
timeout_seconds=timeout,
overwrite_enabled=overwrite,
file_size_bytes=existing_size,
error_type="file_exists",
).model_dump(),
)
# Perform download
start_time = time.time()
result = await self._download_file_async(url, output_path, timeout, self.headers)
execution_time = time.time() - start_time
# Format output
formatted_output = self._format_download_output(result, output_format)
# Create metadata
metadata = DownloadMetadata(
url=url,
output_path=str(output_path),
timeout_seconds=timeout,
overwrite_enabled=overwrite,
execution_time=execution_time,
file_size_bytes=result.file_size,
headers_used=self.headers is not None,
)
if not result.success:
metadata.error_type = "download_failure"
return ActionResponse(
success=result.success,
message=formatted_output,
metadata=metadata.model_dump(),
)
except Exception as e:
error_msg = f"Failed to download file: {str(e)}"
self.logger.error(f"Download error: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=error_msg,
metadata=DownloadMetadata(
url=url,
output_path=output_file_path,
timeout_seconds=timeout,
overwrite_enabled=overwrite,
error_type="internal_error",
).model_dump(),
)
def mcp_get_download_capabilities(self) -> ActionResponse:
"""Get information about download service capabilities and configuration.
Returns:
ActionResponse with download service capabilities and current configuration
"""
capabilities = {
"requests_available": requests is not None,
"supported_schemes": list(self.supported_schemes),
"supported_features": [
"HTTP/HTTPS URL downloads",
"Configurable timeout controls",
"Custom headers support",
"Path validation and directory creation",
"File size limits and safety checks",
"Multiple output formats (markdown, json, text)",
"LLM-optimized result formatting",
"Comprehensive error handling",
],
"supported_formats": ["markdown", "json", "text"],
"configuration": {
"default_timeout": self.default_timeout,
"max_file_size_bytes": self.max_file_size,
"workspace": str(self.workspace),
},
"safety_features": [
"URL validation",
"File size limits",
"Timeout controls",
"Path validation",
"Overwrite protection",
"Error handling and logging",
],
}
max_size_mb = self.max_file_size / (1024 * 1024)
formatted_info = f"""# Download Service Capabilities
## Status
- **Workspace:** `{self.workspace}`
## Supported Features
{chr(10).join(f"- {feature}" for feature in capabilities["supported_features"])}
## Supported URL Schemes
{chr(10).join(f"- {scheme}://" for scheme in capabilities["supported_schemes"])}
## Supported Output Formats
{chr(10).join(f"- {fmt}" for fmt in capabilities["supported_formats"])}
## Configuration
- **Default Timeout:** {capabilities["configuration"]["default_timeout"]} seconds
- **Max File Size:** {self.max_file_size:,} bytes ({max_size_mb:.1f} MB)
## Safety Features
{chr(10).join(f"- {feature}" for feature in capabilities["safety_features"])}
"""
return ActionResponse(
success=True,
message=formatted_info,
metadata=capabilities,
)
# Default arguments for testing
if __name__ == "__main__":
import os
load_dotenv()
arguments = ActionArguments(
name="download",
transport="stdio",
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
)
try:
service = DownloadCollection(arguments)
service.run()
except Exception as e:
print(f"An error occurred: {e}: {traceback.format_exc()}")
@@ -0,0 +1,723 @@
"""
ArXiv MCP Server
This module provides MCP server functionality for ArXiv academic paper operations.
It supports paper search, metadata extraction, and content retrieval with LLM-friendly formatting.
Key features:
- Search ArXiv papers by query, author, category, or ID
- Extract paper metadata (title, authors, abstract, etc.)
- Download and process paper PDFs
- LLM-optimized content formatting
- Comprehensive error handling and logging
Main functions:
- mcp_search_papers: Search ArXiv papers with flexible criteria
- mcp_get_paper_details: Get detailed information about specific papers
- mcp_download_paper: Download paper PDF and extract text content
- mcp_get_categories: Get available ArXiv subject categories
- mcp_get_arxiv_capabilities: Get service capabilities and configuration
"""
import json
import traceback
from datetime import datetime
import arxiv
from dotenv import load_dotenv
from pydantic import BaseModel, Field
from pydantic.fields import FieldInfo
from aworld.logs.util import Color
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
class PaperResult(BaseModel):
"""Individual paper search result with structured data."""
entry_id: str
title: str
authors: list[str]
summary: str
published: str
updated: str | None = None
categories: list[str]
primary_category: str
pdf_url: str | None = None
doi: str | None = None
journal_ref: str | None = None
comment: str | None = None
class ArxivMetadata(BaseModel):
"""Metadata for ArXiv operation results."""
operation: str
query: str | None = None
max_results: int | None = None
sort_by: str | None = None
sort_order: str | None = None
total_results: int | None = None
execution_time: float | None = None
error_type: str | None = None
paper_id: str | None = None
download_path: str | None = None
file_size: int | None = None
class ArxivActionCollection(ActionCollection):
"""MCP service for ArXiv academic paper operations.
Provides comprehensive ArXiv functionality including:
- Paper search with flexible criteria (query, author, category, ID)
- Detailed paper metadata extraction
- PDF download and text content extraction
- Subject category information
- LLM-optimized result formatting
- Error handling and logging
"""
def __init__(self, arguments: ActionArguments) -> None:
super().__init__(arguments)
# Initialize supported file extensions for PDF processing
self.supported_extensions = {".pdf"}
# ArXiv client configuration
self.client = arxiv.Client(
page_size=100,
delay_seconds=3.0, # Be respectful to ArXiv servers
num_retries=3,
)
# Create downloads directory
self._downloads_dir = self.workspace / "arxiv_downloads"
self._downloads_dir.mkdir(exist_ok=True)
# ArXiv subject categories mapping
self.subject_categories = {
"cs": "Computer Science",
"math": "Mathematics",
"physics": "Physics",
"astro-ph": "Astrophysics",
"cond-mat": "Condensed Matter",
"gr-qc": "General Relativity and Quantum Cosmology",
"hep-ex": "High Energy Physics - Experiment",
"hep-lat": "High Energy Physics - Lattice",
"hep-ph": "High Energy Physics - Phenomenology",
"hep-th": "High Energy Physics - Theory",
"math-ph": "Mathematical Physics",
"nlin": "Nonlinear Sciences",
"nucl-ex": "Nuclear Experiment",
"nucl-th": "Nuclear Theory",
"quant-ph": "Quantum Physics",
"q-bio": "Quantitative Biology",
"q-fin": "Quantitative Finance",
"stat": "Statistics",
"econ": "Economics",
"eess": "Electrical Engineering and Systems Science",
}
self._color_log("ArXiv service initialized", Color.green, "debug")
self._color_log(f"Downloads directory: {self._downloads_dir}", Color.blue, "debug")
def _format_paper_result(self, paper: arxiv.Result) -> PaperResult:
"""Convert arxiv.Result to structured PaperResult.
Args:
paper: ArXiv paper result object
Returns:
Structured PaperResult object
"""
return PaperResult(
entry_id=paper.entry_id,
title=paper.title.strip(),
authors=[author.name for author in paper.authors],
summary=paper.summary.strip(),
published=paper.published.isoformat(),
updated=paper.updated.isoformat() if paper.updated else None,
categories=paper.categories,
primary_category=paper.primary_category,
pdf_url=paper.pdf_url,
doi=paper.doi,
journal_ref=paper.journal_ref,
comment=paper.comment,
)
def _format_search_results(self, results: list[PaperResult], output_format: str = "markdown") -> str:
"""Format paper search results for LLM consumption.
Args:
results: List of paper results
output_format: Format type ('markdown', 'json', 'text')
Returns:
Formatted string suitable for LLM consumption
"""
if not results:
return "No papers found matching the search criteria."
if output_format == "json":
return json.dumps([result.model_dump() for result in results], indent=2)
elif output_format == "text":
output_parts = [f"Found {len(results)} papers:\n"]
for i, paper in enumerate(results, 1):
authors_str = ", ".join(paper.authors[:3])
if len(paper.authors) > 3:
authors_str += f" et al. ({len(paper.authors)} total)"
output_parts.extend(
[
f"{i}. {paper.title}",
f" Authors: {authors_str}",
f" Published: {paper.published[:10]}",
f" Categories: {', '.join(paper.categories)}",
f" ArXiv ID: {paper.entry_id.split('/')[-1]}",
f" Abstract: {paper.summary[:200]}...",
"",
]
)
return "\n".join(output_parts)
else: # markdown (default)
output_parts = [f"# ArXiv Search Results\n\nFound **{len(results)}** papers:\n"]
for i, paper in enumerate(results, 1):
authors_str = ", ".join(paper.authors[:3])
if len(paper.authors) > 3:
authors_str += f" *et al.* ({len(paper.authors)} total)"
arxiv_id = paper.entry_id.split("/")[-1]
output_parts.extend(
[
f"## {i}. {paper.title}",
f"**Authors:** {authors_str}",
f"**Published:** {paper.published[:10]}",
f"**Categories:** {', '.join(paper.categories)}",
f"**ArXiv ID:** `{arxiv_id}`",
f"**PDF:** [Download]({paper.pdf_url})" if paper.pdf_url else "",
"",
f"**Abstract:** {paper.summary}",
"",
"---",
"",
]
)
return "\n".join(output_parts)
def _format_paper_details(self, paper: PaperResult, output_format: str = "markdown") -> str:
"""Format detailed paper information for LLM consumption.
Args:
paper: Paper result object
output_format: Format type ('markdown', 'json', 'text')
Returns:
Formatted string with detailed paper information
"""
if output_format == "json":
return json.dumps(paper.model_dump(), indent=2)
elif output_format == "text":
output_parts = [
f"Title: {paper.title}",
f"Authors: {', '.join(paper.authors)}",
f"Published: {paper.published}",
f"Updated: {paper.updated or 'N/A'}",
f"Primary Category: {paper.primary_category}",
f"All Categories: {', '.join(paper.categories)}",
f"ArXiv ID: {paper.entry_id.split('/')[-1]}",
f"PDF URL: {paper.pdf_url or 'N/A'}",
f"DOI: {paper.doi or 'N/A'}",
f"Journal Reference: {paper.journal_ref or 'N/A'}",
f"Comment: {paper.comment or 'N/A'}",
"",
"Abstract:",
paper.summary,
]
return "\n".join(output_parts)
else: # markdown (default)
arxiv_id = paper.entry_id.split("/")[-1]
output_parts = [
f"# {paper.title}",
"",
f"**Authors:** {', '.join(paper.authors)}",
f"**Published:** {paper.published[:10]}",
f"**Updated:** {paper.updated[:10] if paper.updated else 'N/A'}",
f"**Primary Category:** {paper.primary_category}",
f"**All Categories:** {', '.join(paper.categories)}",
f"**ArXiv ID:** `{arxiv_id}`",
f"**PDF:** [Download]({paper.pdf_url})" if paper.pdf_url else "**PDF:** N/A",
f"**DOI:** {paper.doi}" if paper.doi else "**DOI:** N/A",
f"**Journal Reference:** {paper.journal_ref}" if paper.journal_ref else "**Journal Reference:** N/A",
f"**Comment:** {paper.comment}" if paper.comment else "**Comment:** N/A",
"",
"## Abstract",
"",
paper.summary,
]
return "\n".join(output_parts)
async def mcp_search_papers(
self,
query: str = Field(description="Search query (keywords, title, author, etc.)"),
sort_by: str = Field(
default="relevance", description="Sort by: 'relevance', 'lastUpdatedDate', 'submittedDate'"
),
sort_order: str = Field(default="descending", description="Sort order: 'ascending' or 'descending'"),
category: str | None = Field(default=None, description="Filter by ArXiv category (e.g., 'cs.AI', 'math.CO')"),
output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'"),
) -> ActionResponse:
"""Search ArXiv papers with flexible criteria.
This tool provides comprehensive ArXiv paper search with:
- Keyword, title, and author search capabilities
- Category filtering for specific subject areas
- Flexible sorting options (relevance, date)
- Configurable result limits
- LLM-optimized result formatting
Args:
query: Search terms (can include keywords, titles, author names)
sort_by: Sorting criteria for results
sort_order: Order of sorting (ascending/descending)
category: Optional category filter (e.g., 'cs.AI' for AI papers)
output_format: Format for the response output
Returns:
ActionResponse with search results and metadata
"""
# Handle FieldInfo objects
if isinstance(query, FieldInfo):
query = query.default
if isinstance(sort_by, FieldInfo):
sort_by = sort_by.default
if isinstance(sort_order, FieldInfo):
sort_order = sort_order.default
if isinstance(category, FieldInfo):
category = category.default
if isinstance(output_format, FieldInfo):
output_format = output_format.default
try:
self._color_log(f"🔍 Searching ArXiv for: {query}", Color.cyan)
start_time = datetime.now()
# Build search query
search_query = query
if category:
search_query = f"cat:{category} AND ({query})"
# Configure sort criteria
sort_criterion = arxiv.SortCriterion.Relevance
if sort_by == "lastUpdatedDate":
sort_criterion = arxiv.SortCriterion.LastUpdatedDate
elif sort_by == "submittedDate":
sort_criterion = arxiv.SortCriterion.SubmittedDate
sort_order_enum = arxiv.SortOrder.Descending
if sort_order == "ascending":
sort_order_enum = arxiv.SortOrder.Ascending
# Perform search
search = arxiv.Search(
query=search_query, max_results=300000, sort_by=sort_criterion, sort_order=sort_order_enum
)
# Execute search and collect results
results = []
for paper in self.client.results(search):
results.append(self._format_paper_result(paper))
execution_time = (datetime.now() - start_time).total_seconds()
# Format output
formatted_output = self._format_search_results(results, output_format)
# Create metadata
metadata = ArxivMetadata(
operation="search_papers",
query=query,
sort_by=sort_by,
sort_order=sort_order,
total_results=len(results),
execution_time=execution_time,
)
self._color_log(f"✅ Found {len(results)} papers in {execution_time:.2f}s", Color.green)
return ActionResponse(success=True, message=formatted_output, metadata=metadata.model_dump())
except Exception as e:
error_msg = f"Failed to search ArXiv papers: {str(e)}"
self.logger.error(f"ArXiv search error: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=error_msg,
metadata=ArxivMetadata(operation="search_papers", query=query, error_type="search_error").model_dump(),
)
async def mcp_get_paper_details(
self,
paper_id: str = Field(description="ArXiv paper ID (e.g., '2301.07041' or 'arxiv:2301.07041')"),
output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'"),
) -> ActionResponse:
"""Get detailed information about a specific ArXiv paper.
Args:
paper_id: ArXiv paper identifier
output_format: Format for the response output
Returns:
ActionResponse with detailed paper information and metadata
"""
# Handle FieldInfo objects
if isinstance(paper_id, FieldInfo):
paper_id = paper_id.default
if isinstance(output_format, FieldInfo):
output_format = output_format.default
try:
# Clean paper ID
clean_id = paper_id.replace("arxiv:", "").strip()
self._color_log(f"📄 Getting details for paper: {clean_id}", Color.cyan)
start_time = datetime.now()
# Search for the specific paper
search = arxiv.Search(id_list=[clean_id])
paper = next(self.client.results(search), None)
if not paper:
return ActionResponse(
success=False,
message=f"Paper not found: {clean_id}",
metadata=ArxivMetadata(
operation="get_paper_details", paper_id=clean_id, error_type="paper_not_found"
).model_dump(),
)
execution_time = (datetime.now() - start_time).total_seconds()
# Format paper details
paper_result = self._format_paper_result(paper)
formatted_output = self._format_paper_details(paper_result, output_format)
# Create metadata
metadata = ArxivMetadata(operation="get_paper_details", paper_id=clean_id, execution_time=execution_time)
self._color_log(f"✅ Retrieved paper details in {execution_time:.2f}s", Color.green)
return ActionResponse(success=True, message=formatted_output, metadata=metadata.model_dump())
except Exception as e:
error_msg = f"Failed to get paper details: {str(e)}"
self.logger.error(f"Paper details error: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=error_msg,
metadata=ArxivMetadata(
operation="get_paper_details", paper_id=paper_id, error_type="retrieval_error"
).model_dump(),
)
async def mcp_download_paper(
self,
paper_id: str = Field(description="ArXiv paper ID (e.g., '2301.07041' or 'arxiv:2301.07041')"),
extract_text: bool = Field(default=True, description="Whether to extract text content from PDF"),
output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'"),
) -> ActionResponse:
"""Download ArXiv paper PDF and optionally extract text content.
Args:
paper_id: ArXiv paper identifier
extract_text: Whether to extract and return text content
output_format: Format for the response output
Returns:
ActionResponse with download status and optional text content
"""
# Handle FieldInfo objects
if isinstance(paper_id, FieldInfo):
paper_id = paper_id.default
if isinstance(extract_text, FieldInfo):
extract_text = extract_text.default
if isinstance(output_format, FieldInfo):
output_format = output_format.default
try:
# Clean paper ID
clean_id = paper_id.replace("arxiv:", "").strip()
self._color_log(f"📥 Downloading paper: {clean_id}", Color.cyan)
start_time = datetime.now()
# Search for the paper
search = arxiv.Search(id_list=[clean_id])
paper = next(self.client.results(search), None)
if not paper:
return ActionResponse(
success=False,
message=f"Paper not found: {clean_id}",
metadata=ArxivMetadata(
operation="download_paper", paper_id=clean_id, error_type="paper_not_found"
).model_dump(),
)
# Download PDF
filename = f"{clean_id.replace('/', '_')}.pdf"
download_path = self._downloads_dir / filename
paper.download_pdf(dirpath=str(self._downloads_dir), filename=filename)
execution_time = (datetime.now() - start_time).total_seconds()
file_size = download_path.stat().st_size if download_path.exists() else 0
# Prepare response message
if output_format == "json":
response_data = {
"paper_id": clean_id,
"title": paper.title,
"download_path": str(download_path),
"file_size": file_size,
"download_time": execution_time,
}
if extract_text:
try:
# Basic text extraction (would need additional libraries like PyPDF2 or pdfplumber)
response_data["text_extraction"] = (
"Text extraction requires additional PDF processing libraries"
)
except Exception:
response_data["text_extraction"] = "Text extraction failed"
formatted_output = json.dumps(response_data, indent=2)
elif output_format == "text":
output_parts = [
"Paper Downloaded Successfully",
f"Paper ID: {clean_id}",
f"Title: {paper.title}",
f"Download Path: {download_path}",
f"File Size: {file_size:,} bytes",
f"Download Time: {execution_time:.2f} seconds",
]
if extract_text:
output_parts.append("\nNote: Text extraction requires additional PDF processing libraries")
formatted_output = "\n".join(output_parts)
else: # markdown (default)
output_parts = [
"# 📥 Paper Download Complete",
"",
f"**Paper ID:** `{clean_id}`",
f"**Title:** {paper.title}",
f"**Download Path:** `{download_path}`",
f"**File Size:** {file_size:,} bytes",
f"**Download Time:** {execution_time:.2f} seconds",
]
if extract_text:
output_parts.extend(
[
"",
"## 📄 Text Extraction",
(
"*Note: Text extraction requires additional "
"PDF processing libraries like PyPDF2 or pdfplumber*"
),
]
)
formatted_output = "\n".join(output_parts)
# Create metadata
metadata = ArxivMetadata(
operation="download_paper",
paper_id=clean_id,
download_path=str(download_path),
file_size=file_size,
execution_time=execution_time,
)
self._color_log(f"✅ Downloaded paper in {execution_time:.2f}s ({file_size:,} bytes)", Color.green)
return ActionResponse(success=True, message=formatted_output, metadata=metadata.model_dump())
except Exception as e:
error_msg = f"Failed to download paper: {str(e)}"
self.logger.error(f"Paper download error: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=error_msg,
metadata=ArxivMetadata(
operation="download_paper", paper_id=paper_id, error_type="download_error"
).model_dump(),
)
def mcp_get_categories(
self,
output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'"),
) -> ActionResponse:
"""Get available ArXiv subject categories.
Args:
output_format: Format for the response output
Returns:
ActionResponse with category information
"""
# Handle FieldInfo objects
if isinstance(output_format, FieldInfo):
output_format = output_format.default
try:
if output_format == "json":
formatted_output = json.dumps(self.subject_categories, indent=2)
elif output_format == "text":
output_parts = ["ArXiv Subject Categories:\n"]
for code, name in self.subject_categories.items():
output_parts.append(f"{code}: {name}")
formatted_output = "\n".join(output_parts)
else: # markdown (default)
output_parts = [
"# ArXiv Subject Categories",
"",
"Available categories for filtering search results:",
"",
]
for code, name in self.subject_categories.items():
output_parts.append(f"- **`{code}`**: {name}")
output_parts.extend(
[
"",
"## Usage Examples",
"- `cs.AI` - Artificial Intelligence",
"- `cs.LG` - Machine Learning",
"- `math.CO` - Combinatorics",
"- `physics.gen-ph` - General Physics",
]
)
formatted_output = "\n".join(output_parts)
metadata = ArxivMetadata(operation="get_categories", total_results=len(self.subject_categories))
return ActionResponse(success=True, message=formatted_output, metadata=metadata.model_dump())
except Exception as e:
error_msg = f"Failed to get categories: {str(e)}"
self.logger.error(f"Categories error: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=error_msg,
metadata=ArxivMetadata(operation="get_categories", error_type="internal_error").model_dump(),
)
def mcp_get_arxiv_capabilities(self) -> ActionResponse:
"""Get information about ArXiv service capabilities and configuration.
Returns:
ActionResponse with service capabilities and current configuration
"""
capabilities = {
"supported_operations": [
"Paper search with flexible criteria",
"Detailed paper metadata retrieval",
"PDF download and storage",
"Subject category filtering",
"Multiple output formats (markdown, json, text)",
"LLM-optimized result formatting",
],
"search_capabilities": [
"Keyword and phrase search",
"Author name search",
"Title search",
"Category filtering",
"Date-based sorting",
"Relevance-based sorting",
],
"supported_formats": ["markdown", "json", "text"],
"configuration": {
"downloads_directory": str(self._downloads_dir),
"client_page_size": 100,
"client_delay_seconds": 3.0,
"client_num_retries": 3,
"supported_categories_count": len(self.subject_categories),
},
"rate_limiting": {
"delay_between_requests": "3.0 seconds",
"retry_attempts": 3,
"respectful_usage": "Configured for ArXiv server guidelines",
},
}
formatted_info = f"""# ArXiv Service Capabilities
## Supported Operations
{chr(10).join(f"- {op}" for op in capabilities["supported_operations"])}
## Search Capabilities
{chr(10).join(f"- {cap}" for cap in capabilities["search_capabilities"])}
## Supported Output Formats
{chr(10).join(f"- {fmt}" for fmt in capabilities["supported_formats"])}
## Current Configuration
- **Downloads Directory:** {capabilities["configuration"]["downloads_directory"]}
- **Client Page Size:** {capabilities["configuration"]["client_page_size"]}
- **Request Delay:** {capabilities["configuration"]["client_delay_seconds"]} seconds
- **Retry Attempts:** {capabilities["configuration"]["client_num_retries"]}
- **Available Categories:** {capabilities["configuration"]["supported_categories_count"]}
## Rate Limiting & Ethics
- **Delay Between Requests:** {capabilities["rate_limiting"]["delay_between_requests"]}
- **Retry Policy:** {capabilities["rate_limiting"]["retry_attempts"]} attempts
- **Server Respect:** {capabilities["rate_limiting"]["respectful_usage"]}
"""
return ActionResponse(success=True, message=formatted_info, metadata=capabilities)
# Default arguments for testing
if __name__ == "__main__":
import os
load_dotenv()
arguments = ActionArguments(
name="arxiv",
transport="stdio",
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
)
try:
service = ArxivActionCollection(arguments)
service.run()
except Exception as e:
print(f"An error occurred: {e}: {traceback.format_exc()}")
@@ -0,0 +1,535 @@
"""
Chess MCP Server
This module provides MCP server functionality for chess game operations and analysis.
It utilizes the 'python-chess' library to support various chess-related tasks.
Key features:
- Manage chess game states (new game, load FEN, make moves)
- Validate and execute moves in UCI or SAN format
- Get legal moves for the current position
- Check game status (checkmate, stalemate, draw, etc.)
- Basic board visualization (ASCII)
- LLM-optimized output formatting for game states and analysis
Main functions:
- mcp_new_game: Start a new chess game
- mcp_load_fen: Load a game state from FEN notation
- mcp_make_move: Make a move on the current board
- mcp_get_legal_moves: List all legal moves in the current position
- mcp_get_board_state: Get the current board state (FEN, ASCII, status)
- mcp_get_game_status: Check the current game status (e.g., checkmate)
- mcp_get_chess_capabilities: Get service capabilities
"""
import json
import traceback
from datetime import datetime
from typing import Any
import chess
from dotenv import load_dotenv
from pydantic import BaseModel, Field
from pydantic.fields import FieldInfo
from aworld.logs.util import Color
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
class ChessBoardState(BaseModel):
"""Structured representation of the chess board state."""
fen: str
turn: str # 'white' or 'black'
castling_rights: str
ep_square: str | None = None # The target square if an en passant capture is possible *right now*
halfmove_clock: int
fullmove_number: int
is_check: bool
is_checkmate: bool
is_stalemate: bool
is_insufficient_material: bool
is_seventyfive_moves: bool
is_fivefold_repetition: bool
is_game_over: bool
ascii_board: str
legal_moves_uci: list[str]
legal_moves_san: list[str]
is_en_passant_possible: bool # True if there is a legal en passant capture
en_passant_capture_square: str | None = None # The square a pawn would move TO for en passant
class ChessMoveResult(BaseModel):
"""Result of making a chess move."""
move_uci: str
move_san: str
is_capture: bool
is_check: bool
is_kingside_castling: bool
is_queenside_castling: bool
board_after_move: ChessBoardState
class ChessMetadata(BaseModel):
"""Metadata for Chess operation results."""
operation: str
fen_before: str | None = None
fen_after: str | None = None
move_played: str | None = None
execution_time: float | None = None
error_type: str | None = None
engine_analysis_depth: int | None = None
class ChessCollection(ActionCollection):
"""MCP service for chess game operations and analysis.
Provides capabilities to manage chess games, make moves, analyze positions,
and get game status, all formatted for LLM interaction.
"""
def __init__(self, arguments: ActionArguments) -> None:
super().__init__(arguments)
self.board = chess.Board()
# For more advanced analysis, you might initialize a chess engine here
# Example: self.engine = chess.engine.SimpleEngine.popen_uci("/path/to/stockfish")
# Ensure Stockfish or another UCI engine is installed and path is correct.
self._color_log("Chess service initialized", Color.green, "debug")
self._color_log(f"Initial board FEN: {self.board.fen()}", Color.blue, "debug")
def _get_current_board_state(self) -> ChessBoardState:
"""Helper to get the current board state in a structured format."""
legal_moves_uci = [move.uci() for move in self.board.legal_moves]
legal_moves_san = []
# Generating SAN for all moves can be slow, do it carefully or on demand
# for move in self.board.legal_moves:
# try:
# legal_moves_san.append(self.board.san(move))
# except Exception:
# legal_moves_san.append(move.uci()) # Fallback to UCI if SAN fails
has_legal_ep = self.board.has_legal_en_passant()
ep_sq_name = chess.square_name(self.board.ep_square) if self.board.ep_square else None
return ChessBoardState(
fen=self.board.fen(),
turn="white" if self.board.turn == chess.WHITE else "black",
castling_rights=self.board.castling_xfen(),
ep_square=ep_sq_name, # This is the target square from FEN, might not be a legal capture
halfmove_clock=self.board.halfmove_clock,
fullmove_number=self.board.fullmove_number,
is_check=self.board.is_check(),
is_checkmate=self.board.is_checkmate(),
is_stalemate=self.board.is_stalemate(),
is_insufficient_material=self.board.is_insufficient_material(),
is_seventyfive_moves=self.board.is_seventyfive_moves(),
is_fivefold_repetition=self.board.is_fivefold_repetition(),
is_game_over=self.board.is_game_over(),
ascii_board=str(self.board),
legal_moves_uci=legal_moves_uci,
legal_moves_san=legal_moves_san, # Populate if SAN generation is enabled
is_en_passant_possible=has_legal_ep,
en_passant_capture_square=ep_sq_name if has_legal_ep else None,
)
def _format_board_state_output(self, state: ChessBoardState, output_format: str = "markdown") -> str:
"""Format board state for LLM consumption."""
if output_format == "json":
return json.dumps(state.model_dump(), indent=2)
status_parts = []
if state.is_checkmate:
status_parts.append("Checkmate!")
elif state.is_stalemate:
status_parts.append("Stalemate!")
elif state.is_insufficient_material:
status_parts.append("Draw by insufficient material.")
elif state.is_seventyfive_moves:
status_parts.append("Draw by 75-move rule.")
elif state.is_fivefold_repetition:
status_parts.append("Draw by fivefold repetition.")
elif state.is_check:
status_parts.append("Check!")
game_status = " ".join(status_parts) if status_parts else "Game in progress."
en_passant_info = "N/A"
if state.is_en_passant_possible and state.en_passant_capture_square:
en_passant_info = f"Yes, capture on {state.en_passant_capture_square}"
elif state.ep_square: # FEN might list an ep_square even if no legal ep move
en_passant_info = f"Target square {state.ep_square} (no legal en passant capture)"
if output_format == "text":
return (
f"Board FEN: {state.fen}\n"
f"Turn: {state.turn.capitalize()}\n"
f"Status: {game_status}\n"
f"Castling: {state.castling_rights}\n"
f"En Passant Possible: {en_passant_info}\n" # Updated line
f"Halfmove Clock: {state.halfmove_clock}\n"
f"Fullmove Number: {state.fullmove_number}\n"
f"Game Over: {'Yes' if state.is_game_over else 'No'}\n"
f"Legal Moves (UCI): {', '.join(state.legal_moves_uci[:10])}... ({len(state.legal_moves_uci)} total)\n"
f"Board:\n{state.ascii_board}"
)
else: # markdown (default)
return (
f"### Chess Board State\n"
f"**FEN:** `{state.fen}`\n"
f"**Turn:** {state.turn.capitalize()}\n"
f"**Status:** {game_status}\n"
f"**Castling Rights:** {state.castling_rights}\n"
f"**En Passant Possible:** {en_passant_info}\n" # Updated line
f"**Game Over:** {'Yes' if state.is_game_over else 'No'}\n"
f"**Legal Moves (UCI, sample):** `{', '.join(state.legal_moves_uci[:5])}`... ({len(state.legal_moves_uci)} total)\n"
f"```\n{state.ascii_board}\n```"
)
async def mcp_new_game(self) -> ActionResponse:
"""Starts a new standard chess game, resetting the board.
Returns:
ActionResponse with the initial board state.
"""
start_time = datetime.now()
self.board.reset()
self._color_log("🚀 New chess game started", Color.green)
current_state = self._get_current_board_state()
formatted_output = self._format_board_state_output(current_state)
execution_time = (datetime.now() - start_time).total_seconds()
metadata = ChessMetadata(
operation="new_game", fen_after=self.board.fen(), execution_time=execution_time
).model_dump()
return ActionResponse(success=True, message=formatted_output, metadata=metadata)
async def mcp_load_fen(
self, fen_string: str = Field(description="FEN string representing the board state.")
) -> ActionResponse:
"""Loads a chess game from a FEN (Forsyth-Edwards Notation) string.
Args:
fen_string: The FEN string to load.
Returns:
ActionResponse with the board state after loading the FEN.
"""
# Handle FieldInfo
if isinstance(fen_string, FieldInfo):
fen_string = fen_string.default
start_time = datetime.now()
try:
self.board.set_fen(fen_string)
self._color_log(f"🔄 Board loaded from FEN: {fen_string}", Color.blue)
current_state = self._get_current_board_state()
formatted_output = self._format_board_state_output(current_state)
execution_time = (datetime.now() - start_time).total_seconds()
metadata = ChessMetadata(
operation="load_fen", fen_after=self.board.fen(), execution_time=execution_time
).model_dump()
return ActionResponse(success=True, message=formatted_output, metadata=metadata)
except ValueError as e:
error_msg = f"Invalid FEN string: {str(e)}"
self.logger.error(f"FEN loading error: {traceback.format_exc()}")
execution_time = (datetime.now() - start_time).total_seconds()
metadata = ChessMetadata(
operation="load_fen", error_type="invalid_fen", execution_time=execution_time
).model_dump()
return ActionResponse(success=False, message=error_msg, metadata=metadata)
async def mcp_make_move(
self, move_str: str = Field(description="Move in UCI (e.g., 'e2e4') or SAN (e.g., 'Nf3') format.")
) -> ActionResponse:
"""Makes a move on the current chess board.
The move can be in UCI (Universal Chess Interface) format (e.g., 'g1f3')
or SAN (Standard Algebraic Notation) format (e.g., 'Nf3').
Args:
move_str: The move to make.
Returns:
ActionResponse with the result of the move and new board state.
"""
# Handle FieldInfo
if isinstance(move_str, FieldInfo):
move_str = move_str.default
start_time = datetime.now()
fen_before = self.board.fen()
try:
move = None
# Try parsing as UCI first, then SAN
try:
move = self.board.parse_uci(move_str)
except ValueError:
try:
move = self.board.parse_san(move_str)
except ValueError as e_san:
raise ValueError(f"Invalid move format. UCI error: N/A, SAN error: {e_san}") from e_san
if move not in self.board.legal_moves:
raise ValueError(f"Illegal move: {move_str}")
move_san = self.board.san(move)
is_capture = self.board.is_capture(move)
is_kingside_castling = self.board.is_kingside_castling(move)
is_queenside_castling = self.board.is_queenside_castling(move)
self.board.push(move)
is_check_after_move = self.board.is_check()
self._color_log(f"♟️ Move made: {move_str} (UCI: {move.uci()}, SAN: {move_san})", Color.cyan)
current_state = self._get_current_board_state()
move_result = ChessMoveResult(
move_uci=move.uci(),
move_san=move_san,
is_capture=is_capture,
is_check=is_check_after_move, # Check status *after* the move
is_kingside_castling=is_kingside_castling,
is_queenside_castling=is_queenside_castling,
board_after_move=current_state,
)
# Format output (can be customized)
formatted_output = f"Move {move_result.move_san} (UCI: {move_result.move_uci}) played.\n"
formatted_output += self._format_board_state_output(current_state)
execution_time = (datetime.now() - start_time).total_seconds()
metadata = ChessMetadata(
operation="make_move",
fen_before=fen_before,
fen_after=self.board.fen(),
move_played=move.uci(),
execution_time=execution_time,
).model_dump()
return ActionResponse(success=True, message=formatted_output, metadata=metadata)
except ValueError as e:
error_msg = f"Failed to make move '{move_str}': {str(e)}"
self.logger.error(f"Move error: {traceback.format_exc()}")
execution_time = (datetime.now() - start_time).total_seconds()
metadata = ChessMetadata(
operation="make_move",
fen_before=fen_before,
move_played=move_str,
error_type="invalid_or_illegal_move",
execution_time=execution_time,
).model_dump()
return ActionResponse(success=False, message=error_msg, metadata=metadata)
async def mcp_get_legal_moves(
self,
output_format: str = Field(
default="markdown", description="Output format: 'uci_list', 'san_list', 'markdown', 'json'"
),
) -> ActionResponse:
"""Gets all legal moves for the current board position.
Args:
output_format: 'uci_list' (simple list of UCI moves),
'san_list' (simple list of SAN moves),
'markdown' (formatted list),
'json' (structured list).
Returns:
ActionResponse with the list of legal moves.
"""
# Handle FieldInfo
if isinstance(output_format, FieldInfo):
output_format = output_format.default
start_time = datetime.now()
legal_moves_uci = [move.uci() for move in self.board.legal_moves]
message_content: Any
if output_format == "uci_list":
message_content = legal_moves_uci
elif output_format == "san_list":
try:
message_content = [self.board.san(move) for move in self.board.legal_moves]
except Exception as e:
self.logger.warning(
f"Could not generate all SAN moves: {e}. Falling back to UCI for problematic moves."
)
san_moves = []
for move in self.board.legal_moves:
try:
san_moves.append(self.board.san(move))
except Exception:
san_moves.append(move.uci() + " (SAN failed)")
message_content = san_moves
elif output_format == "json":
moves_data = []
for move in self.board.legal_moves:
try:
san = self.board.san(move)
except Exception:
san = move.uci() + " (SAN failed)"
moves_data.append({"uci": move.uci(), "san": san})
message_content = json.dumps(moves_data, indent=2)
else: # markdown
if not legal_moves_uci:
message_content = "No legal moves available (game might be over)."
else:
san_formatted_moves = []
for move_uci in legal_moves_uci[:20]: # Display sample for markdown
try:
move_obj = self.board.parse_uci(move_uci)
san_formatted_moves.append(f"`{self.board.san(move_obj)}` ({move_uci})")
except Exception:
san_formatted_moves.append(f"`{move_uci}` (SAN failed)")
header = f"### Legal Moves ({len(legal_moves_uci)} total)\n"
moves_list_md = "\n".join([f"- {m}" for m in san_formatted_moves])
if len(legal_moves_uci) > 20:
moves_list_md += "\n- ... (and more)"
message_content = header + moves_list_md
execution_time = (datetime.now() - start_time).total_seconds()
metadata = ChessMetadata(
operation="get_legal_moves", fen_before=self.board.fen(), execution_time=execution_time
).model_dump()
return ActionResponse(success=True, message=message_content, metadata=metadata)
async def mcp_get_board_state(
self, output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'")
) -> ActionResponse:
"""Gets the current state of the chess board.
Includes FEN, turn, game status, ASCII board, and legal moves.
Args:
output_format: Desired format for the board state.
Returns:
ActionResponse with the current board state.
"""
# Handle FieldInfo
if isinstance(output_format, FieldInfo):
output_format = output_format.default
start_time = datetime.now()
current_state = self._get_current_board_state()
formatted_output = self._format_board_state_output(current_state, output_format)
execution_time = (datetime.now() - start_time).total_seconds()
metadata = ChessMetadata(
operation="get_board_state",
fen_before=self.board.fen(), # FEN is part of the state, so 'before' and 'after' are same here
execution_time=execution_time,
).model_dump()
return ActionResponse(success=True, message=formatted_output, metadata=metadata)
async def mcp_get_game_status(self) -> ActionResponse:
"""Checks and returns the current game status (e.g., checkmate, stalemate).
Returns:
ActionResponse with a human-readable game status and structured data.
"""
start_time = datetime.now()
state = self._get_current_board_state()
status_message = "Game in progress."
if state.is_checkmate:
status_message = f"Checkmate! {state.turn.capitalize()} is mated."
elif state.is_stalemate:
status_message = "Stalemate! The game is a draw."
elif state.is_insufficient_material:
status_message = "Draw by insufficient material."
elif state.is_seventyfive_moves:
status_message = "Draw by 75-move rule."
elif state.is_fivefold_repetition:
status_message = "Draw by fivefold repetition."
elif state.is_check:
status_message = f"{state.turn.capitalize()} is in check."
status_data = {
"status_message": status_message,
"is_game_over": state.is_game_over,
"is_check": state.is_check,
"is_checkmate": state.is_checkmate,
"is_stalemate": state.is_stalemate,
"is_draw": state.is_stalemate
or state.is_insufficient_material
or state.is_seventyfive_moves
or state.is_fivefold_repetition,
"winner": None, # Could be determined if checkmate
}
if state.is_checkmate:
status_data["winner"] = "black" if self.board.turn == chess.WHITE else "white"
execution_time = (datetime.now() - start_time).total_seconds()
metadata = ChessMetadata(
operation="get_game_status", fen_before=self.board.fen(), execution_time=execution_time
).model_dump()
metadata.update(status_data) # Add specific status flags to metadata
return ActionResponse(success=True, message=status_message, metadata=metadata)
def mcp_get_chess_capabilities(self) -> ActionResponse:
"""Get information about the Chess service capabilities.
Returns:
ActionResponse with service capabilities.
"""
capabilities_info = {
"service_name": "Chess MCP Service",
"library_used": "python-chess",
"supported_operations": [
"new_game: Start a new chess game.",
"load_fen: Load game state from FEN string.",
"make_move: Make a move (UCI or SAN).",
"get_legal_moves: List legal moves.",
"get_board_state: Get current board FEN, ASCII, status, etc.",
"get_game_status: Check for checkmate, stalemate, draw conditions.",
],
"output_formats": ["markdown", "json", "text"],
"move_input_formats": ["UCI (e.g., e2e4)", "SAN (e.g., Nf3)"],
"fen_support": "Full FEN loading and generation.",
"engine_integration": "Basic structure for UCI engine integration (not fully implemented by default).",
}
formatted_message = "# Chess Service Capabilities\n\n"
formatted_message += f"**Service Name:** {capabilities_info['service_name']}\n"
formatted_message += f"**Core Library:** {capabilities_info['library_used']}\n\n"
formatted_message += "**Supported Operations:**\n"
for op in capabilities_info["supported_operations"]:
formatted_message += f"- {op}\n"
formatted_message += "\n**Supported Output Formats:** " + ", ".join(capabilities_info["output_formats"]) + "\n"
formatted_message += "**Move Input Formats:** " + ", ".join(capabilities_info["move_input_formats"]) + "\n"
return ActionResponse(success=True, message=formatted_message, metadata=capabilities_info)
# Optional: Method to close engine if it was initialized
# def __del__(self):
# if hasattr(self, 'engine') and self.engine:
# self.engine.quit()
# self._color_log("Chess engine quit.", Color.yellow)
# Default arguments for testing
if __name__ == "__main__":
import os
load_dotenv()
arguments = ActionArguments(
name="chess_service",
transport="stdio",
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
)
try:
service = ChessCollection(arguments)
service.run()
except Exception as e:
print(f"An error occurred: {e}: {traceback.format_exc()}")
@@ -0,0 +1,641 @@
"""
PubChem MCP Server
This module provides MCP server functionality for accessing PubChem database programmatically.
It supports compound searches, property retrieval, and structure-based queries using PubChem's REST API.
Key features:
- Compound search by name, CID, SMILES, or InChI
- Property retrieval (molecular weight, formula, etc.)
- Structure similarity searches
- Bioactivity data access
- 3D structure downloads
- Rate limiting compliance (max 5 requests/second)
Main functions:
- mcp_search_compounds: Search for compounds by various identifiers
- mcp_get_compound_properties: Retrieve compound properties
- mcp_get_compound_synonyms: Get compound names and synonyms
- mcp_search_similar_compounds: Find structurally similar compounds
- mcp_get_bioactivity_data: Retrieve bioactivity assay data
- mcp_download_structure: Download 2D/3D structure files
"""
import time
import traceback
from datetime import datetime
from typing import Literal
from urllib.parse import quote
import requests
from dotenv import load_dotenv
from pydantic import BaseModel, Field
from pydantic.fields import FieldInfo
from aworld.logs.util import Color
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
# pylint: disable=C0301
class CompoundData(BaseModel):
"""Structured compound data from PubChem."""
cid: int | None = None
name: str | None = None
molecular_formula: str | None = None
molecular_weight: float | None = None
smiles: str | None = None
inchi: str | None = None
synonyms: list[str] = []
class PubChemMetadata(BaseModel):
"""Metadata for PubChem operation results."""
query_type: str
query_value: str
api_endpoint: str
response_time: float
total_results: int | None = None
rate_limit_delay: float | None = None
error_type: str | None = None
timestamp: str
class PubChemCollection(ActionCollection):
"""MCP service for PubChem database access with comprehensive chemical data retrieval.
Provides access to PubChem's extensive chemical database including:
- Compound identification and search capabilities
- Chemical property and structure data
- Bioactivity and assay information
- Structure similarity searches
- 2D/3D molecular structure downloads
- Synonym and nomenclature data
Complies with PubChem usage policies:
- Maximum 5 requests per second
- Automatic rate limiting
- Proper error handling for timeouts
"""
def __init__(self, arguments: ActionArguments) -> None:
super().__init__(arguments)
# PubChem API configuration
self.base_url = "https://pubchem.ncbi.nlm.nih.gov/rest/pug"
self.base_url_view = "https://pubchem.ncbi.nlm.nih.gov/rest/pug_view"
self.request_delay = 0.2 # 200ms delay to stay under 5 req/sec limit
self.last_request_time = 0.0
# Request timeout settings
self.timeout = 30 # PubChem's 30-second limit
# Initialize request session with headers
self.session = requests.Session()
self.session.headers.update(
{"User-Agent": "AWorld-PubChem-MCP/1.0 (https://github.com/aworld-framework)", "Accept": "application/json"}
)
self._color_log("PubChem MCP Service initialized", Color.green, "debug")
self._color_log(f"Base URL: {self.base_url}", Color.blue, "debug")
def _rate_limit(self) -> float:
"""Enforce rate limiting to comply with PubChem usage policy.
Returns:
Actual delay time applied
"""
current_time = time.time()
time_since_last = current_time - self.last_request_time
if time_since_last < self.request_delay:
delay = self.request_delay - time_since_last
time.sleep(delay)
self.last_request_time = time.time()
return delay
self.last_request_time = current_time
return 0.0
def _make_request(self, url: str, params: dict = None) -> tuple[dict | None, float]:
"""Make a rate-limited request to PubChem API.
Args:
url: API endpoint URL
params: Query parameters
Returns:
Tuple of (response_data, response_time)
Raises:
requests.RequestException: For API request failures
"""
start_time = time.time()
try:
response = self.session.get(url, params=params, timeout=self.timeout)
response_time = time.time() - start_time
if response.status_code == 200:
return response.json(), response_time
elif response.status_code == 503:
raise requests.RequestException("PubChem service temporarily unavailable (503)")
else:
raise requests.RequestException(f"HTTP {response.status_code}: {response.text}")
except requests.Timeout as e:
response_time = time.time() - start_time
raise requests.RequestException(f"Request timeout after {self.timeout}s") from e
except requests.RequestException:
response_time = time.time() - start_time
raise
def mcp_search_compounds(
self,
query: str = Field(description="Search query (compound name, CID, SMILES, InChI, etc.)"),
search_type: Literal["name", "cid", "smiles", "inchi", "formula"] = Field(
default="name",
description="Type of search: name (compound name), cid (PubChem ID), smiles, inchi, or formula",
),
max_results: int = Field(default=10, description="Maximum number of results to return (1-100)", ge=1, le=100),
) -> ActionResponse:
"""Search for chemical compounds in PubChem database.
Supports multiple search types:
- Name: Search by common or IUPAC names
- CID: Search by PubChem Compound ID
- SMILES: Search by SMILES notation
- InChI: Search by InChI identifier
- Formula: Search by molecular formula
Args:
query: Search term or identifier
search_type: Type of search to perform
max_results: Maximum number of compounds to return
Returns:
ActionResponse with compound search results and metadata
"""
try:
# Handle FieldInfo objects
if isinstance(query, FieldInfo):
query = query.default
if isinstance(search_type, FieldInfo):
search_type = search_type.default
if isinstance(max_results, FieldInfo):
max_results = max_results.default
if not query or not query.strip():
raise ValueError("Search query is required")
self._color_log(f"Searching PubChem for: {query} (type: {search_type})", Color.cyan)
# Build API URL based on search type
if search_type == "cid":
url = f"{self.base_url}/compound/cid/{quote(str(query))}/property/Title,MolecularFormula,MolecularWeight,CanonicalSMILES,InChI/JSON"
elif search_type == "name":
url = f"{self.base_url}/compound/name/{quote(query)}/property/Title,MolecularFormula,MolecularWeight,CanonicalSMILES,InChI/JSON"
elif search_type == "smiles":
url = f"{self.base_url}/compound/smiles/{quote(query)}/property/Title,MolecularFormula,MolecularWeight,CanonicalSMILES,InChI/JSON"
elif search_type == "inchi":
url = f"{self.base_url}/compound/inchi/{quote(query)}/property/Title,MolecularFormula,MolecularWeight,CanonicalSMILES,InChI/JSON"
elif search_type == "formula":
url = f"{self.base_url}/compound/formula/{quote(query)}/property/Title,MolecularFormula,MolecularWeight,CanonicalSMILES,InChI/JSON"
else:
raise ValueError(f"Unsupported search type: {search_type}")
# Make API request
data, response_time = self._make_request(url)
# Parse results
compounds = []
if data and "PropertyTable" in data and "Properties" in data["PropertyTable"]:
properties_list = data["PropertyTable"]["Properties"][:max_results]
for prop in properties_list:
compound = CompoundData(
cid=prop.get("CID"),
name=prop.get("Title"),
molecular_formula=prop.get("MolecularFormula"),
molecular_weight=prop.get("MolecularWeight"),
smiles=prop.get("CanonicalSMILES"),
inchi=prop.get("InChI"),
)
compounds.append(compound)
# Format results for LLM
if compounds:
result_lines = [f"Found {len(compounds)} compound(s) for query '{query}':\n"]
for i, compound in enumerate(compounds, 1):
result_lines.append(f"{i}. **{compound.name}** (CID: {compound.cid})")
result_lines.append(f" - Formula: {compound.molecular_formula}")
result_lines.append(f" - Molecular Weight: {compound.molecular_weight} g/mol")
result_lines.append(f" - SMILES: {compound.smiles}")
if compound.inchi:
result_lines.append(
f" - InChI: {compound.inchi[:100]}..."
if len(compound.inchi) > 100
else f" - InChI: {compound.inchi}"
)
result_lines.append("")
message = "\n".join(result_lines)
else:
message = f"No compounds found for query '{query}' using search type '{search_type}'"
# Prepare metadata
metadata = PubChemMetadata(
query_type=search_type,
query_value=query,
api_endpoint=url,
response_time=response_time,
total_results=len(compounds),
timestamp=datetime.now().isoformat(),
)
self._color_log(f"Found {len(compounds)} compounds ({response_time:.2f}s)", Color.green)
return ActionResponse(success=True, message=message, metadata=metadata.model_dump())
except ValueError as e:
self.logger.error(f"Invalid input: {str(e)}")
return ActionResponse(
success=False,
message=f"Invalid input: {str(e)}",
metadata={"error_type": "invalid_input", "error_message": str(e)},
)
except requests.RequestException as e:
self.logger.error(f"PubChem API error: {str(e)}")
return ActionResponse(
success=False,
message=f"PubChem API error: {str(e)}",
metadata={"error_type": "api_error", "error_message": str(e)},
)
except Exception as e:
self.logger.error(f"Search failed: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=f"Search failed: {str(e)}",
metadata={"error_type": "general_error", "error_message": str(e)},
)
def mcp_get_compound_synonyms(
self,
cid: int = Field(description="PubChem Compound ID (CID)"),
max_synonyms: int = Field(default=20, description="Maximum number of synonyms to return (1-100)", ge=1, le=100),
) -> ActionResponse:
"""Retrieve synonyms and alternative names for a PubChem compound.
Args:
cid: PubChem Compound ID
max_synonyms: Maximum number of synonyms to return
Returns:
ActionResponse with compound synonyms and metadata
"""
try:
# Handle FieldInfo objects
if isinstance(cid, FieldInfo):
cid = cid.default
if isinstance(max_synonyms, FieldInfo):
max_synonyms = max_synonyms.default
if not cid or cid <= 0:
raise ValueError("Valid PubChem CID is required")
self._color_log(f"Retrieving synonyms for CID: {cid}", Color.cyan)
# Build API URL for synonyms
url = f"{self.base_url}/compound/cid/{cid}/synonyms/JSON"
# Make API request
data, response_time = self._make_request(url)
# Parse synonyms
synonyms = []
if data and "InformationList" in data and "Information" in data["InformationList"]:
info_list = data["InformationList"]["Information"]
if info_list and "Synonym" in info_list[0]:
synonyms = info_list[0]["Synonym"][:max_synonyms]
# Format results for LLM
if synonyms:
result_lines = [f"Found {len(synonyms)} synonym(s) for CID {cid}:\n"]
for i, synonym in enumerate(synonyms, 1):
result_lines.append(f"{i}. {synonym}")
message = "\n".join(result_lines)
else:
message = f"No synonyms found for CID {cid}"
# Prepare metadata
metadata = PubChemMetadata(
query_type="synonyms",
query_value=str(cid),
api_endpoint=url,
response_time=response_time,
total_results=len(synonyms),
timestamp=datetime.now().isoformat(),
)
self._color_log(f"Retrieved {len(synonyms)} synonyms ({response_time:.2f}s)", Color.green)
return ActionResponse(success=True, message=message, metadata=metadata.model_dump())
except ValueError as e:
self.logger.error(f"Invalid input: {str(e)}")
return ActionResponse(
success=False,
message=f"Invalid input: {str(e)}",
metadata={"error_type": "invalid_input", "error_message": str(e)},
)
except requests.RequestException as e:
self.logger.error(f"PubChem API error: {str(e)}")
return ActionResponse(
success=False,
message=f"PubChem API error: {str(e)}",
metadata={"error_type": "api_error", "error_message": str(e)},
)
except Exception as e:
self.logger.error(f"Synonym retrieval failed: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=f"Synonym retrieval failed: {str(e)}",
metadata={"error_type": "general_error", "error_message": str(e)},
)
def mcp_get_compound_properties(
self,
cid: int = Field(description="PubChem Compound ID (CID)"),
properties: list[str] = Field(
default=[
"MolecularWeight",
"MolecularFormula",
"CanonicalSMILES",
"InChI",
"XLogP",
"TPSA",
"HBondDonorCount",
"HBondAcceptorCount",
],
description="List of properties to retrieve (e.g., MolecularWeight, XLogP, TPSA)",
),
) -> ActionResponse:
"""Retrieve detailed chemical properties for a PubChem compound.
Common properties include:
- MolecularWeight: Molecular weight in g/mol
- MolecularFormula: Chemical formula
- CanonicalSMILES: SMILES notation
- InChI: InChI identifier
- XLogP: Partition coefficient
- TPSA: Topological polar surface area
- HBondDonorCount: Hydrogen bond donor count
- HBondAcceptorCount: Hydrogen bond acceptor count
Args:
cid: PubChem Compound ID
properties: List of property names to retrieve
Returns:
ActionResponse with compound properties and metadata
"""
try:
# Handle FieldInfo objects
if isinstance(cid, FieldInfo):
cid = cid.default
if isinstance(properties, FieldInfo):
properties = properties.default
if not cid or cid <= 0:
raise ValueError("Valid PubChem CID is required")
if not properties:
properties = ["MolecularWeight", "MolecularFormula", "CanonicalSMILES"]
self._color_log(f"Retrieving properties for CID: {cid}", Color.cyan)
# Build API URL for properties
props_str = ",".join(properties)
url = f"{self.base_url}/compound/cid/{cid}/property/{props_str}/JSON"
# Make API request
data, response_time = self._make_request(url)
# Parse properties
compound_props = {}
if data and "PropertyTable" in data and "Properties" in data["PropertyTable"]:
props_data = data["PropertyTable"]["Properties"][0]
compound_props = {k: v for k, v in props_data.items() if k != "CID"}
# Format results for LLM
if compound_props:
result_lines = [f"Properties for PubChem CID {cid}:\n"]
for prop_name, prop_value in compound_props.items():
if prop_name == "InChI" and isinstance(prop_value, str) and len(prop_value) > 100:
result_lines.append(f"**{prop_name}**: {prop_value[:100]}...")
else:
result_lines.append(f"**{prop_name}**: {prop_value}")
message = "\n".join(result_lines)
else:
message = f"No properties found for CID {cid}"
# Prepare metadata
metadata = PubChemMetadata(
query_type="properties",
query_value=str(cid),
api_endpoint=url,
response_time=response_time,
total_results=len(compound_props),
timestamp=datetime.now().isoformat(),
)
self._color_log(f"Retrieved {len(compound_props)} properties ({response_time:.2f}s)", Color.green)
return ActionResponse(success=True, message=message, metadata=metadata.model_dump())
except ValueError as e:
self.logger.error(f"Invalid input: {str(e)}")
return ActionResponse(
success=False,
message=f"Invalid input: {str(e)}",
metadata={"error_type": "invalid_input", "error_message": str(e)},
)
except requests.RequestException as e:
self.logger.error(f"PubChem API error: {str(e)}")
return ActionResponse(
success=False,
message=f"PubChem API error: {str(e)}",
metadata={"error_type": "api_error", "error_message": str(e)},
)
except Exception as e:
self.logger.error(f"Property retrieval failed: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=f"Property retrieval failed: {str(e)}",
metadata={"error_type": "general_error", "error_message": str(e)},
)
def mcp_search_similar_compounds(
self,
cid: int = Field(description="PubChem Compound ID to find similar compounds for"),
similarity_threshold: float = Field(
default=0.9, description="Similarity threshold (0.0-1.0, higher = more similar)", ge=0.0, le=1.0
),
max_results: int = Field(
default=10, description="Maximum number of similar compounds to return (1-50)", ge=1, le=50
),
) -> ActionResponse:
"""Find structurally similar compounds using PubChem's similarity search.
Uses Tanimoto similarity coefficient for 2D structure comparison.
Args:
cid: Reference compound CID for similarity search
similarity_threshold: Minimum similarity score (0.0-1.0)
max_results: Maximum number of similar compounds to return
Returns:
ActionResponse with similar compounds and metadata
"""
try:
# Handle FieldInfo objects
if isinstance(cid, FieldInfo):
cid = cid.default
if isinstance(similarity_threshold, FieldInfo):
similarity_threshold = similarity_threshold.default
if isinstance(max_results, FieldInfo):
max_results = max_results.default
if not cid or cid <= 0:
raise ValueError("Valid PubChem CID is required")
self._color_log(f"Searching for compounds similar to CID: {cid}", Color.cyan)
# Build API URL for similarity search
threshold_percent = int(similarity_threshold * 100)
url = f"{self.base_url}/compound/fastsimilarity_2d/cid/{cid}/property/Title,MolecularFormula,MolecularWeight/JSON"
params = {"Threshold": threshold_percent, "MaxRecords": max_results}
# Make API request
data, response_time = self._make_request(url, params)
# Parse similar compounds
similar_compounds: list[CompoundData] = []
if data and "PropertyTable" in data and "Properties" in data["PropertyTable"]:
properties_list = data["PropertyTable"]["Properties"]
for prop in properties_list:
if prop.get("CID") != cid: # Exclude the query compound itself
compound = CompoundData(
cid=prop.get("CID"),
name=prop.get("Title"),
molecular_formula=prop.get("MolecularFormula"),
molecular_weight=prop.get("MolecularWeight"),
)
similar_compounds.append(compound)
# Format results for LLM
if similar_compounds:
result_lines = [
f"Found {len(similar_compounds)} compound(s) similar to CID {cid} (threshold: {similarity_threshold}):\n"
]
for i, compound in enumerate(similar_compounds, 1):
result_lines.append(f"{i}. **{compound.name}** (CID: {compound.cid})")
result_lines.append(f" - Formula: {compound.molecular_formula}")
result_lines.append(f" - Molecular Weight: {compound.molecular_weight} g/mol")
result_lines.append("")
message = "\n".join(result_lines)
else:
message = f"No similar compounds found for CID {cid} with similarity threshold {similarity_threshold}"
# Prepare metadata
metadata = PubChemMetadata(
query_type="similarity",
query_value=str(cid),
api_endpoint=url,
response_time=response_time,
total_results=len(similar_compounds),
timestamp=datetime.now().isoformat(),
)
self._color_log(f"Found {len(similar_compounds)} similar compounds ({response_time:.2f}s)", Color.green)
return ActionResponse(success=True, message=message, metadata=metadata.model_dump())
except ValueError as e:
self.logger.error(f"Invalid input: {str(e)}")
return ActionResponse(
success=False,
message=f"Invalid input: {str(e)}",
metadata={"error_type": "invalid_input", "error_message": str(e)},
)
except requests.RequestException as e:
self.logger.error(f"PubChem API error: {str(e)}")
return ActionResponse(
success=False,
message=f"PubChem API error: {str(e)}",
metadata={"error_type": "api_error", "error_message": str(e)},
)
except Exception as e:
self.logger.error(f"Similarity search failed: {str(e)}: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=f"Similarity search failed: {str(e)}",
metadata={"error_type": "general_error", "error_message": str(e)},
)
def mcp_get_pubchem_capabilities(self) -> ActionResponse:
"""Get information about the PubChem service capabilities.
Returns:
ActionResponse with service capabilities and configuration
"""
capabilities = {
"Compound Search": "Search by name, CID, SMILES, InChI, or molecular formula",
"Property Retrieval": "Get molecular weight, formula, SMILES, physicochemical properties",
"Synonym Lookup": "Retrieve alternative names and identifiers for compounds",
"Similarity Search": "Find structurally similar compounds using Tanimoto similarity",
"Rate Limiting": "Compliant with PubChem's 5 requests/second limit",
"Data Formats": "JSON responses with structured compound data",
}
capability_list = "\n".join(
[f"**{capability}**: {description}" for capability, description in capabilities.items()]
)
metadata = {
"base_url": self.base_url,
"rate_limit": "5 requests/second",
"timeout": f"{self.timeout} seconds",
"supported_capabilities": list(capabilities.keys()),
"total_capabilities": len(capabilities),
"search_types": ["name", "cid", "smiles", "inchi", "formula"],
"data_source": "PubChem (NCBI)",
}
return ActionResponse(
success=True, message=f"PubChem MCP Service Capabilities:\n\n{capability_list}", metadata=metadata
)
# Example usage and entry point
if __name__ == "__main__":
load_dotenv()
# Default arguments for testing
args = ActionArguments(name="pubchem_service", transport="stdio", workspace="~")
# Initialize and run the PubChem service
try:
service = PubChemCollection(args)
service.run()
except Exception as e:
print(f"An error occurred: {e}: {traceback.format_exc()}")
@@ -0,0 +1,387 @@
"""
Search MCP Server
This module provides MCP server functionality for performing web searches using various search engines.
It supports structured queries and returns LLM-friendly formatted search results.
Key features:
- Perform web searches using Google Custom Search API
- Filter and format search results for LLM consumption
- Validate and process search queries with metadata tracking
Main functions:
- mcp_search_google: Searches the web using Google Custom Search API
- mcp_get_search_capabilities: Returns information about search service capabilities
"""
import json
import os
import time
import traceback
import requests
from dotenv import load_dotenv
from pydantic import BaseModel, Field
from pydantic.fields import FieldInfo
from aworld.logs.util import Color
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
class SearchResult(BaseModel):
"""Individual search result with structured data."""
id: str
title: str
url: str
snippet: str
source: str
display_link: str | None = None
formatted_url: str | None = None
class SearchMetadata(BaseModel):
"""Metadata for search operation results."""
query: str
search_engine: str
total_results: int
search_time: float | None = None
language: str = "en"
country: str = "us"
safe_search: bool = True
error_type: str | None = None
api_quota_used: bool = False
class SearchCollection(ActionCollection):
"""MCP service for web search operations using various search engines.
Provides comprehensive web search capabilities including:
- Google Custom Search API integration
- LLM-friendly result formatting
- Search result filtering and validation
- Metadata tracking for search operations
- Error handling and quota management
"""
def __init__(self, arguments: ActionArguments) -> None:
super().__init__(arguments)
# Load environment variables
load_dotenv()
# Validate required API credentials
self.google_api_key = os.getenv("GOOGLE_API_KEY")
self.google_cse_id = os.getenv("GOOGLE_CSE_ID")
# Log initialization status
self._color_log("Search service initialized", Color.green, "debug")
if self.google_api_key and self.google_cse_id:
self._color_log("Google Search API credentials found", Color.blue, "debug")
else:
self._color_log("Google Search API credentials missing - some features may be unavailable", Color.yellow)
def _format_search_results_for_llm(self, results: list[SearchResult], query: str) -> str:
"""Format search results to be LLM-friendly.
Args:
results: List of search results
query: Original search query
Returns:
Formatted string suitable for LLM consumption
"""
if not results:
return f"No search results found for query: '{query}'"
formatted_parts = [f"# Search Results for: '{query}'", f"Found {len(results)} results:\n"]
for i, result in enumerate(results, 1):
result_section = [
f"## Result {i}: {result.title}",
f"**URL:** {result.url}",
f"**Source:** {result.source}",
]
if result.display_link:
result_section.append(f"**Domain:** {result.display_link}")
result_section.append(f"**Summary:** {result.snippet}")
result_section.append("") # Empty line for spacing
formatted_parts.append("\n".join(result_section))
return "\n".join(formatted_parts)
def _validate_search_parameters(self, query: str, num_results: int) -> tuple[str, int]:
"""Validate and normalize search parameters.
Args:
query: Search query string
num_results: Number of results requested
Returns:
Tuple of (validated_query, validated_num_results)
Raises:
ValueError: If parameters are invalid
"""
if not query or not query.strip():
raise ValueError("Search query cannot be empty")
# Normalize query
validated_query = query.strip()
# Validate and clamp num_results
validated_num_results = max(1, min(num_results, 10)) # Google CSE limit is 10
return validated_query, validated_num_results
def mcp_search_google(
self,
query: str = Field(description="The search query string to search for"),
num_results: int = Field(default=5, description="Number of search results to return (1-10, default: 5)"),
safe_search: bool = Field(default=True, description="Whether to enable safe search filtering"),
language: str = Field(default="en", description="Language code for search results (e.g., 'en', 'es', 'fr')"),
country: str = Field(default="us", description="Country code for search results (e.g., 'us', 'uk', 'ca')"),
output_format: str = Field(default="json", description="Output format: 'markdown', 'json', or 'text'"),
) -> ActionResponse:
"""Search the web using Google Custom Search API.
This tool provides comprehensive web search capabilities with:
- Google Custom Search API integration
- Configurable result count and filtering
- Safe search and localization options
- LLM-optimized result formatting
- Detailed metadata tracking
Args:
query: The search query string
num_results: Number of results to return (1-10)
safe_search: Enable safe search filtering
language: Language code for results
country: Country code for results
output_format: Format for the response
Returns:
ActionResponse with formatted search results and metadata
"""
if isinstance(query, FieldInfo):
query = query.default
if isinstance(num_results, FieldInfo):
num_results = num_results.default
if isinstance(safe_search, FieldInfo):
safe_search = safe_search.default
if isinstance(language, FieldInfo):
language = language.default
if isinstance(country, FieldInfo):
country = country.default
if isinstance(output_format, FieldInfo):
output_format = output_format.default
try:
# Validate API credentials
if not self.google_api_key or not self.google_cse_id:
return ActionResponse(
success=False,
message=(
"Google Search API credentials not configured. "
"Please set GOOGLE_API_KEY and GOOGLE_CSE_ID environment variables."
),
metadata={"error_type": "missing_credentials"},
)
# Validate parameters
validated_query, validated_num_results = self._validate_search_parameters(query, num_results)
self._color_log(f"🔍 Searching Google for: '{validated_query}'", Color.cyan)
# Prepare API request
start_time = time.time()
url = "https://www.googleapis.com/customsearch/v1"
params = {
"key": self.google_api_key,
"cx": self.google_cse_id,
"q": validated_query,
"num": validated_num_results,
"safe": "active" if safe_search else "off",
"hl": language,
"gl": country,
}
# Make API request
response = requests.get(url, params=params, timeout=10)
response.raise_for_status()
search_time = time.time() - start_time
data = response.json()
# Parse search results
search_results = []
if "items" in data:
for i, item in enumerate(data["items"]):
result = SearchResult(
id=f"google-{i}",
title=item.get("title", ""),
url=item.get("link", ""),
snippet=item.get("snippet", ""),
source="google",
display_link=item.get("displayLink", ""),
formatted_url=item.get("formattedUrl", ""),
)
search_results.append(result)
# Format results based on requested format
if "json" == "json":
formatted_content = {
"query": validated_query,
"results": [result.model_dump() for result in search_results],
"count": len(search_results),
}
message_content = formatted_content
elif output_format.lower() == "text":
if search_results:
result_lines = []
for i, result in enumerate(search_results, 1):
result_lines.append(f"{i}. {result.title}")
result_lines.append(f" URL: {result.url}")
result_lines.append(f" Summary: {result.snippet}")
result_lines.append("") # Empty line
message_content = "\n".join(result_lines)
else:
message_content = f"No results found for: {validated_query}"
else: # markdown (default)
message_content = self._format_search_results_for_llm(search_results, validated_query)
# Prepare metadata
metadata = SearchMetadata(
query=validated_query,
search_engine="google",
total_results=len(search_results),
search_time=search_time,
language=language,
country=country,
safe_search=safe_search,
api_quota_used=True,
)
self._color_log(f"✅ Found {len(search_results)} results in {search_time:.2f}s", Color.green)
return ActionResponse(success=True, message=message_content, metadata=metadata.model_dump())
except requests.exceptions.RequestException as e:
error_msg = f"Google Search API request failed: {str(e)}"
self.logger.error(f"Search API error: {traceback.format_exc()}")
metadata = SearchMetadata(
query=query, search_engine="google", total_results=0, error_type="api_request_failed"
)
self._color_log(f"{error_msg}", Color.red)
return ActionResponse(success=False, message=error_msg, metadata=metadata.model_dump())
except ValueError as e:
error_msg = f"Invalid search parameters: {str(e)}"
metadata = SearchMetadata(
query=query, search_engine="google", total_results=0, error_type="invalid_parameters"
)
self._color_log(f"{error_msg}", Color.red)
return ActionResponse(success=False, message=error_msg, metadata=metadata.model_dump())
except Exception as e:
error_msg = f"Search operation failed: {str(e)}"
error_trace = traceback.format_exc()
self.logger.error(f"Unexpected search error: {error_trace}")
metadata = SearchMetadata(
query=query, search_engine="google", total_results=0, error_type="unexpected_error"
)
self._color_log(f"{error_msg}", Color.red)
return ActionResponse(
success=False, message=f"{error_msg}\n\nError details: {error_trace}", metadata=metadata.model_dump()
)
def mcp_get_search_capabilities(self) -> ActionResponse:
"""Get information about search service capabilities and configuration.
Returns:
ActionResponse with search service capabilities and current configuration
"""
capabilities = {
"search_engines": ["Google Custom Search API"],
"supported_features": [
"Web search with customizable result count",
"Safe search filtering",
"Language and country localization",
"Multiple output formats (markdown, json, text)",
"LLM-optimized result formatting",
"Detailed metadata tracking",
],
"supported_formats": ["markdown", "json", "text"],
"configuration": {
"google_api_configured": bool(self.google_api_key and self.google_cse_id),
"max_results_per_query": 10,
"default_language": "en",
"default_country": "us",
"safe_search_default": True,
},
"limitations": [
"Google CSE has daily quota limits",
"Maximum 10 results per query",
"Requires valid API credentials",
],
}
formatted_info = f"""# Search Service Capabilities
## Available Search Engines
{chr(10).join(f"- {engine}" for engine in capabilities["search_engines"])}
## Features
{chr(10).join(f"- {feature}" for feature in capabilities["supported_features"])}
## Supported Output Formats
{chr(10).join(f"- {fmt}" for fmt in capabilities["supported_formats"])}
## Current Configuration
- **Google API Configured:** {capabilities["configuration"]["google_api_configured"]}
- **Max Results Per Query:** {capabilities["configuration"]["max_results_per_query"]}
- **Default Language:** {capabilities["configuration"]["default_language"]}
- **Default Country:** {capabilities["configuration"]["default_country"]}
- **Safe Search Default:** {capabilities["configuration"]["safe_search_default"]}
## Limitations
{chr(10).join(f"- {limitation}" for limitation in capabilities["limitations"])}
"""
return ActionResponse(success=True, message=formatted_info, metadata=capabilities)
# Example usage and entry point
if __name__ == "__main__":
load_dotenv()
# Default arguments for testing
args = ActionArguments(
name="search_service",
transport="stdio",
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
)
# Initialize and run the search service
try:
service = SearchCollection(args)
service.run()
except Exception as e:
print(f"An error occurred: {e}: {traceback.format_exc()}")
@@ -0,0 +1,512 @@
"""
Terminal MCP Server
This module provides MCP server functionality for executing terminal commands safely.
It supports command execution with timeout controls and returns LLM-friendly formatted results.
Key features:
- Execute terminal commands with configurable timeouts
- Cross-platform command execution support
- Command history tracking and retrieval
- Safety checks for dangerous commands
- LLM-optimized output formatting
Main functions:
- mcp_execute_command: Execute terminal commands with safety checks
- mcp_get_command_history: Retrieve recent command execution history
- mcp_get_terminal_capabilities: Get terminal service capabilities
"""
import asyncio
import json
import os
import platform
import time
import traceback
from datetime import datetime
from dotenv import load_dotenv
from pydantic import BaseModel, Field
from pydantic.fields import FieldInfo
from aworld.logs.util import Color
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
# pylint: disable=C0301
class CommandResult(BaseModel):
"""Individual command execution result with structured data."""
command: str
success: bool
stdout: str
stderr: str
return_code: int
duration: str
timestamp: str
class TerminalMetadata(BaseModel):
"""Metadata for terminal operation results."""
command: str
platform: str
working_directory: str
timeout_seconds: int
execution_time: float | None = None
return_code: int | None = None
safety_check_passed: bool = True
error_type: str | None = None
history_count: int | None = None
class TerminalActionCollection(ActionCollection):
"""MCP service for terminal command execution with safety controls.
Provides secure terminal command execution capabilities including:
- Cross-platform command execution
- Configurable timeout controls
- Command history tracking
- Safety checks for dangerous operations
- LLM-friendly result formatting
- Error handling and logging
"""
def __init__(self, arguments: ActionArguments) -> None:
super().__init__(arguments)
# Initialize command history
self.command_history: list[dict] = []
self.max_history_size = 50
# Define dangerous commands for safety
self.dangerous_commands = [
"rm -rf /",
"mkfs",
"dd if=",
":(){ :|:& };:", # Unix
"del /f /s /q",
"format",
"diskpart", # Windows
"sudo rm",
"sudo dd",
"sudo mkfs", # Sudo variants
]
# Get current platform info
self.platform_info = {
"system": platform.system(),
"platform": platform.platform(),
"architecture": platform.architecture()[0],
}
self._color_log("Terminal service initialized", Color.green, "debug")
self._color_log(f"Platform: {self.platform_info['system']}", Color.blue, "debug")
def _check_command_safety(self, command: str) -> tuple[bool, str | None]:
"""Check if command is safe to execute.
Args:
command: Command string to check
Returns:
Tuple of (is_safe, reason_if_unsafe)
"""
command_lower = command.lower().strip()
for dangerous_cmd in self.dangerous_commands:
if dangerous_cmd.lower() in command_lower:
return False, f"Command contains dangerous pattern: {dangerous_cmd}"
return True, None
def _format_command_output(self, result: CommandResult, output_format: str = "markdown") -> str:
"""Format command execution results for LLM consumption.
Args:
result: Command execution result
output_format: Format type ('markdown', 'json', 'text')
Returns:
Formatted string suitable for LLM consumption
"""
if output_format == "json":
return json.dumps(result.model_dump(), indent=2)
elif output_format == "text":
output_parts = [
f"Command: {result.command}",
f"Status: {'SUCCESS' if result.success else 'FAILED'}",
f"Duration: {result.duration}",
f"Return Code: {result.return_code}",
]
if result.stdout:
output_parts.extend(["\nOutput:", result.stdout])
if result.stderr:
output_parts.extend(["\nErrors/Warnings:", result.stderr])
return "\n".join(output_parts)
else: # markdown (default)
status_emoji = "" if result.success else ""
output_parts = [
f"# Terminal Command Execution {status_emoji}",
f"**Command:** `{result.command}`",
f"**Status:** {'SUCCESS' if result.success else 'FAILED'}",
f"**Duration:** {result.duration}",
f"**Return Code:** {result.return_code}",
f"**Timestamp:** {result.timestamp}",
]
if result.stdout:
output_parts.extend(["\n## Output", "```", result.stdout.strip(), "```"])
if result.stderr:
output_parts.extend(["\n## Errors/Warnings", "```", result.stderr.strip(), "```"])
return "\n".join(output_parts)
async def _execute_command_async(self, command: str, timeout: int) -> CommandResult:
"""Execute command asynchronously with timeout.
Args:
command: Command to execute
timeout: Timeout in seconds
Returns:
CommandResult with execution details
"""
start_time = datetime.now()
try:
# Create appropriate subprocess for platform
if self.platform_info["system"] == "Windows":
process = await asyncio.create_subprocess_shell(
command, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, shell=True
)
else:
process = await asyncio.create_subprocess_shell(
command,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
shell=True,
executable="/bin/bash",
)
try:
stdout, stderr = await asyncio.wait_for(process.communicate(), timeout)
stdout = stdout.decode("utf-8", errors="replace")
stderr = stderr.decode("utf-8", errors="replace")
return_code = process.returncode
except asyncio.TimeoutError:
try:
process.kill()
except Exception:
pass
duration = str(datetime.now() - start_time)
return CommandResult(
command=command,
success=False,
stdout="",
stderr=f"Command timed out after {timeout} seconds",
return_code=-1,
duration=duration,
timestamp=start_time.isoformat(),
)
duration = str(datetime.now() - start_time)
result = CommandResult(
command=command,
success=return_code == 0,
stdout=stdout,
stderr=stderr,
return_code=return_code,
duration=duration,
timestamp=start_time.isoformat(),
)
# Add to history
self.command_history.append(
{
"timestamp": start_time.isoformat(),
"command": command,
"success": return_code == 0,
"duration": duration,
}
)
# Maintain history size limit
if len(self.command_history) > self.max_history_size:
self.command_history.pop(0)
return result
except Exception as e:
duration = str(datetime.now() - start_time)
return CommandResult(
command=command,
success=False,
stdout="",
stderr=f"Error executing command: {str(e)}",
return_code=-1,
duration=duration,
timestamp=start_time.isoformat(),
)
async def mcp_execute_command(
self,
command: str = Field(description="Terminal command to execute"),
timeout: int = Field(default=30, description="Command timeout in seconds (default: 30)"),
output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'"),
) -> ActionResponse:
"""Execute a terminal command with safety checks and timeout controls.
This tool provides secure command execution with:
- Cross-platform compatibility (Windows, macOS, Linux)
- Configurable timeout controls
- Safety checks for dangerous commands
- LLM-optimized result formatting
- Command history tracking
Specialized Feature:
- Execute Python code and output the result to stdout
- Example (Directly execute simple Python code): `python -c "nums = [1, 2, 3, 4]\nsum_of_nums = sum(nums)\nprint(f'{sum_of_nums=}')"`
- Example (Execute code from a file): `python my_script.py`
Args:
command: The terminal command to execute
timeout: Maximum execution time in seconds
output_format: Format for the response output
Returns:
ActionResponse with command execution results and metadata
"""
# Handle FieldInfo objects
if isinstance(command, FieldInfo):
command = command.default
if isinstance(timeout, FieldInfo):
timeout = timeout.default
if isinstance(output_format, FieldInfo):
output_format = output_format.default
try:
# Safety check
is_safe, safety_reason = self._check_command_safety(command)
if not is_safe:
return ActionResponse(
success=False,
message=f"Command rejected for security reasons: {safety_reason}",
metadata=TerminalMetadata(
command=command,
platform=self.platform_info["system"],
working_directory=str(self.workspace),
timeout_seconds=timeout,
safety_check_passed=False,
error_type="security_violation",
).model_dump(),
)
self._color_log(f"🔧 Executing command: {command}", Color.cyan)
# Execute command
start_time = time.time()
result = await self._execute_command_async(command, timeout)
execution_time = time.time() - start_time
# Format output
formatted_output = self._format_command_output(result, output_format)
# Create metadata
metadata = TerminalMetadata(
command=command,
platform=self.platform_info["system"],
working_directory=str(self.workspace),
timeout_seconds=timeout,
execution_time=execution_time,
return_code=result.return_code,
safety_check_passed=True,
)
if result.success:
self._color_log("✅ Command completed successfully", Color.green)
else:
self._color_log(f"❌ Command failed with return code {result.return_code}", Color.red)
metadata.error_type = "execution_failure"
return ActionResponse(success=result.success, message=formatted_output, metadata=metadata.model_dump())
except Exception as e:
error_msg = f"Failed to execute command: {str(e)}"
self.logger.error(f"Command execution error: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=error_msg,
metadata=TerminalMetadata(
command=command,
platform=self.platform_info["system"],
working_directory=str(self.workspace),
timeout_seconds=timeout,
safety_check_passed=True,
error_type="internal_error",
).model_dump(),
)
def mcp_get_command_history(
self,
count: int = Field(default=10, description="Number of recent commands to return (default: 10)"),
output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'"),
) -> ActionResponse:
"""Retrieve recent command execution history.
Args:
count: Number of recent commands to return
output_format: Format for the response output
Returns:
ActionResponse with command history and metadata
"""
# Handle FieldInfo objects
if isinstance(count, FieldInfo):
count = count.default
if isinstance(output_format, FieldInfo):
output_format = output_format.default
try:
# Get recent history
recent_history = self.command_history[-count:] if self.command_history else []
if not recent_history:
message = "No command history available."
else:
if output_format == "json":
message = json.dumps(recent_history, indent=2)
elif output_format == "text":
history_lines = []
for i, entry in enumerate(recent_history, 1):
status = "SUCCESS" if entry["success"] else "FAILED"
history_lines.append(
f"{i}. [{entry['timestamp']}] {entry['command']} - {status}"
f" ({entry.get('duration', 'N/A')})"
)
message = "\n".join(history_lines)
else: # markdown
history_lines = ["# Command History", f"Showing {len(recent_history)} recent commands:\n"]
for i, entry in enumerate(recent_history, 1):
status_emoji = "" if entry["success"] else ""
history_lines.extend(
[
f"## {i}. {status_emoji} `{entry['command']}`",
f"- **Timestamp:** {entry['timestamp']}",
f"- **Duration:** {entry.get('duration', 'N/A')}",
"",
]
)
message = "\n".join(history_lines)
metadata = TerminalMetadata(
command="get_command_history",
platform=self.platform_info["system"],
working_directory=str(self.workspace),
timeout_seconds=0,
history_count=len(recent_history),
)
return ActionResponse(success=True, message=message, metadata=metadata.model_dump())
except Exception as e:
error_msg = f"Failed to retrieve command history: {str(e)}"
self.logger.error(f"History retrieval error: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=error_msg,
metadata=TerminalMetadata(
command="get_command_history",
platform=self.platform_info["system"],
working_directory=str(self.workspace),
timeout_seconds=0,
error_type="internal_error",
).model_dump(),
)
def mcp_get_terminal_capabilities(self) -> ActionResponse:
"""Get information about terminal service capabilities and configuration.
Returns:
ActionResponse with terminal service capabilities and current configuration
"""
capabilities = {
"platform_info": self.platform_info,
"supported_features": [
"Cross-platform command execution",
"Configurable timeout controls",
"Command history tracking",
"Safety checks for dangerous commands",
"Multiple output formats (markdown, json, text)",
"LLM-optimized result formatting",
"Async command execution",
],
"supported_formats": ["markdown", "json", "text"],
"configuration": {
"max_history_size": self.max_history_size,
"current_history_count": len(self.command_history),
"working_directory": str(self.workspace),
"dangerous_commands_count": len(self.dangerous_commands),
},
"safety_features": [
"Dangerous command detection",
"Timeout controls",
"Error handling and logging",
"Command validation",
],
}
formatted_info = f"""# Terminal Service Capabilities
## Platform Information
- **System:** {self.platform_info["system"]}
- **Platform:** {self.platform_info["platform"]}
- **Architecture:** {self.platform_info["architecture"]}
## Features
{chr(10).join(f"- {feature}" for feature in capabilities["supported_features"])}
## Supported Output Formats
{chr(10).join(f"- {fmt}" for fmt in capabilities["supported_formats"])}
## Current Configuration
- **Max History Size:** {capabilities["configuration"]["max_history_size"]}
- **Current History Count:** {capabilities["configuration"]["current_history_count"]}
- **Working Directory:** {capabilities["configuration"]["working_directory"]}
- **Dangerous Commands Monitored:** {capabilities["configuration"]["dangerous_commands_count"]}
## Safety Features
{chr(10).join(f"- {feature}" for feature in capabilities["safety_features"])}
"""
return ActionResponse(success=True, message=formatted_info, metadata=capabilities)
# Default arguments for testing
if __name__ == "__main__":
load_dotenv()
arguments = ActionArguments(
name="terminal",
transport="stdio",
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
)
try:
service = TerminalActionCollection(arguments)
service.run()
except Exception as e:
print(f"An error occurred: {e}: {traceback.format_exc()}")
@@ -0,0 +1,541 @@
"""
Wayback Machine MCP Server
This module provides MCP server functionality for interacting with the Wayback Machine.
It supports listing archived versions, fetching archived content, and saving pages to the archive.
Key features:
- List available archived versions of URLs with date filtering
- Fetch content from specific archived page versions
- Save current pages to the Wayback Machine
- LLM-optimized output formatting with text extraction
- Comprehensive error handling and logging
Main functions:
- mcp_list_archived_versions: List available snapshots for a URL
- mcp_get_archived_content: Fetch content from a specific archived version
- mcp_get_wayback_capabilities: Get service capabilities information
"""
import json
import os
import time
import traceback
from datetime import datetime
import requests
from bs4 import BeautifulSoup
from dotenv import load_dotenv
from pydantic import BaseModel, Field
from pydantic.fields import FieldInfo
from waybackpy import WaybackMachineCDXServerAPI
from aworld.logs.util import Color
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
class ArchivedVersion(BaseModel):
"""Individual archived version with structured data."""
timestamp: str
url: str
status_code: str
digest: str
length: str
mime_type: str
class WaybackMetadata(BaseModel):
"""Metadata for Wayback Machine operation results."""
url: str
operation: str # 'list_versions', 'get_content', 'save_page'
timestamp: str | None = None
total_versions: int | None = None
date_range: dict[str, str | None] | None = None
content_length: int | None = None
text_extracted: bool = False
truncated: bool = False
execution_time: float | None = None
error_type: str | None = None
user_agent: str = "AWorld/1.0 (https://github.com/inclusionAI/AWorld; qintong.wqt@antgroup.com)"
class WaybackActionCollection(ActionCollection):
"""MCP service for Wayback Machine operations.
Provides comprehensive Wayback Machine functionality including:
- Listing archived versions of URLs with flexible filtering
- Fetching content from specific archived snapshots
- LLM-friendly content formatting and text extraction
- Robust error handling and detailed logging
- Multiple output formats (markdown, JSON, text)
"""
def __init__(self, arguments: ActionArguments) -> None:
super().__init__(arguments)
# Configuration
self.user_agent = "AWorld/1.0 (https://github.com/inclusionAI/AWorld; qintong.wqt@antgroup.com)"
self.default_timeout = 30
self.max_content_length = 8192
self._color_log("Wayback Machine service initialized", Color.green, "debug")
self._color_log(f"User Agent: {self.user_agent}", Color.blue, "debug")
def _format_versions_for_llm(self, versions: list[ArchivedVersion], query_info: dict) -> str:
"""Format archived versions list for LLM consumption.
Args:
versions: List of archived versions
query_info: Query information including URL and filters
Returns:
Formatted string suitable for LLM consumption
"""
if not versions:
return f"No archived versions found for URL: {query_info.get('url', 'Unknown')}"
output_parts = [
f"# Wayback Machine Archives for {query_info.get('url', 'Unknown')}",
f"\nFound **{len(versions)}** archived versions",
]
if query_info.get("from_date") or query_info.get("to_date"):
date_filter = []
if query_info.get("from_date"):
date_filter.append(f"From: {query_info['from_date']}")
if query_info.get("to_date"):
date_filter.append(f"To: {query_info['to_date']}")
output_parts.append(f"\n**Date Filter:** {' | '.join(date_filter)}")
output_parts.append("\n## Available Versions:")
for i, version in enumerate(versions[:10], 1): # Show first 10
timestamp_formatted = self._format_timestamp(version.timestamp)
output_parts.append(
f"\n{i}. **{timestamp_formatted}**\n"
f" - Archive URL: {version.url}\n"
f" - Status: {version.status_code} | Size: {version.length} bytes\n"
f" - Type: {version.mime_type}"
)
if len(versions) > 10:
output_parts.append(f"\n... and {len(versions) - 10} more versions")
return "\n".join(output_parts)
def _format_content_for_llm(self, content_data: dict, output_format: str = "markdown") -> str:
"""Format archived content for LLM consumption.
Args:
content_data: Content data dictionary
output_format: Format type ('markdown', 'json', 'text')
Returns:
Formatted string suitable for LLM consumption
"""
if output_format == "json":
return json.dumps(content_data, indent=2)
elif output_format == "text":
return content_data.get("content", "")
else: # markdown (default)
output_parts = [
f"# Archived Content from {content_data.get('url', 'Unknown')}",
f"\n**Requested Timestamp:** {content_data.get('timestamp', 'Unknown')}",
f"**Actual Timestamp:** {self._format_timestamp(content_data.get('fetched_timestamp', ''))}",
f"**Content Length:** {content_data.get('original_content_length', 0):,} characters",
]
if content_data.get("truncated"):
output_parts.append(f"**Note:** Content truncated to {self.max_content_length:,} characters")
if content_data.get("extract_text_only"):
output_parts.append("**Note:** Text-only extraction applied")
output_parts.extend(["\n## Content:", "\n---\n", content_data.get("content", ""), "\n---"])
return "\n".join(output_parts)
def _format_timestamp(self, timestamp: str) -> str:
"""Format Wayback Machine timestamp to human-readable format.
Args:
timestamp: Wayback timestamp (YYYYMMDDhhmmss)
Returns:
Human-readable timestamp
"""
try:
if len(timestamp) >= 14:
dt = datetime.strptime(timestamp[:14], "%Y%m%d%H%M%S")
return dt.strftime("%Y-%m-%d %H:%M:%S UTC")
return timestamp
except (ValueError, TypeError):
return timestamp or "Unknown"
def _validate_wayback_parameters(self, url: str, timestamp: str = None) -> tuple[str, str | None]:
"""Validate and normalize Wayback Machine parameters.
Args:
url: URL to validate
timestamp: Optional timestamp to validate
Returns:
Tuple of (validated_url, validated_timestamp)
Raises:
ValueError: If parameters are invalid
"""
if not url or not url.strip():
raise ValueError("URL cannot be empty")
url = url.strip()
if not url.startswith(("http://", "https://")):
url = "https://" + url
validated_timestamp = None
if timestamp:
timestamp = timestamp.strip()
if len(timestamp) < 8:
raise ValueError("Timestamp must be at least 8 characters (YYYYMMDD)")
validated_timestamp = timestamp
return url, validated_timestamp
async def mcp_list_archived_versions(
self,
url: str = Field(description="The URL of the website to check for archived versions"),
limit: int = Field(default=10, description="Maximum number of versions to return (0 for all)"),
from_date: str | None = Field(default=None, description="Start date filter (YYYYMMDDhhmmss)"),
to_date: str | None = Field(default=None, description="End date filter (YYYYMMDDhhmmss)"),
output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'"),
) -> ActionResponse:
"""List available archived versions of a URL from the Wayback Machine.
This function queries the Wayback Machine CDX API to retrieve all available
archived snapshots for a given URL, with optional date range filtering.
Args:
url: The URL to search for archived versions
limit: Maximum number of versions to return (0 for all, default: 10)
from_date: Start date for filtering versions (YYYYMMDDhhmmss format)
to_date: End date for filtering versions (YYYYMMDDhhmmss format)
output_format: Format for the response ('markdown', 'json', or 'text')
Returns:
ActionResponse with archived versions list and metadata
"""
start_time = time.time()
# Handle FieldInfo objects
if isinstance(url, FieldInfo):
url = url.default
if isinstance(limit, FieldInfo):
limit = limit.default
if isinstance(from_date, FieldInfo):
from_date = from_date.default
if isinstance(to_date, FieldInfo):
to_date = to_date.default
if isinstance(output_format, FieldInfo):
output_format = output_format.default
try:
# Validate parameters
url, _ = self._validate_wayback_parameters(url)
self._color_log(f"Listing archived versions for: {url}", Color.blue)
# Query Wayback Machine CDX API
cdx_api = WaybackMachineCDXServerAPI(url, user_agent=self.user_agent)
all_snapshots = list(cdx_api.snapshots())
# Apply date filtering
if from_date or to_date:
snapshots = [
s
for s in all_snapshots
if (not from_date or s.timestamp >= from_date) and (not to_date or s.timestamp <= to_date)
]
else:
snapshots = all_snapshots
if not snapshots:
return ActionResponse(
success=False,
message="No archived versions found for the specified URL and date range.",
metadata=WaybackMetadata(
url=url,
operation="list_versions",
total_versions=0,
date_range={"from_date": from_date, "to_date": to_date},
execution_time=time.time() - start_time,
error_type="no_results",
).model_dump(),
)
# Convert to structured format
versions = [
ArchivedVersion(
timestamp=snapshot.timestamp,
url=snapshot.archive_url,
status_code=snapshot.statuscode,
digest=snapshot.digest,
length=snapshot.length,
mime_type=snapshot.mimetype,
)
for snapshot in snapshots
]
# Apply limit
if limit > 0 and len(versions) > limit:
versions = versions[:limit]
# Format output
query_info = {"url": url, "from_date": from_date, "to_date": to_date, "total_found": len(snapshots)}
if output_format == "json":
message = [version.model_dump() for version in versions]
else:
message = self._format_versions_for_llm(versions, query_info)
execution_time = time.time() - start_time
self._color_log(f"Found {len(versions)} archived versions in {execution_time:.2f}s", Color.green)
return ActionResponse(
success=True,
message=message,
metadata=WaybackMetadata(
url=url,
operation="list_versions",
total_versions=len(snapshots),
date_range={"from_date": from_date, "to_date": to_date},
execution_time=execution_time,
).model_dump(),
)
except Exception as e:
error_msg = f"Failed to list archived versions: {str(e)}"
self._color_log(error_msg, Color.red)
self.logger.error(f"Error in mcp_list_archived_versions: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=error_msg,
metadata=WaybackMetadata(
url=url or "unknown",
operation="list_versions",
execution_time=time.time() - start_time,
error_type=type(e).__name__,
).model_dump(),
)
async def mcp_get_archived_content(
self,
url: str = Field(description="The URL of the website to fetch archived content from"),
timestamp: str = Field(description="The timestamp of the desired version (YYYYMMDDhhmmss)"),
extract_text_only: bool = Field(default=True, description="Extract only text content, removing HTML tags"),
truncate_content: bool = Field(default=False, description="Truncate content to manageable length for LLMs"),
output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'"),
) -> ActionResponse:
"""Fetch content from a specific archived page version.
This function retrieves the content of a specific archived snapshot from the
Wayback Machine, with options for text extraction and content truncation.
Args:
url: The URL of the website to fetch
timestamp: The timestamp of the desired version (YYYYMMDDhhmmss)
extract_text_only: Whether to extract only text content (default: True)
truncate_content: Whether to truncate content for LLM consumption (default: False)
output_format: Format for the response ('markdown', 'json', or 'text')
Returns:
ActionResponse with archived content and metadata
"""
start_time = time.time()
# Handle FieldInfo objects
if isinstance(url, FieldInfo):
url = url.default
if isinstance(timestamp, FieldInfo):
timestamp = timestamp.default
if isinstance(extract_text_only, FieldInfo):
extract_text_only = extract_text_only.default
if isinstance(truncate_content, FieldInfo):
truncate_content = truncate_content.default
if isinstance(output_format, FieldInfo):
output_format = output_format.default
try:
# Validate parameters
url, timestamp = self._validate_wayback_parameters(url, timestamp)
self._color_log(f"Fetching archived content: {url} at {timestamp}", Color.blue)
# Query Wayback Machine for closest snapshot
cdx_api = WaybackMachineCDXServerAPI(url, user_agent=self.user_agent)
snapshot = cdx_api.near(wayback_machine_timestamp=timestamp)
if not snapshot or not snapshot.archive_url:
return ActionResponse(
success=False,
message=f"No archived version found for {url} at timestamp {timestamp}",
metadata=WaybackMetadata(
url=url,
operation="get_content",
timestamp=timestamp,
execution_time=time.time() - start_time,
error_type="no_snapshot",
).model_dump(),
)
# Fetch content
response = requests.get(snapshot.archive_url, timeout=self.default_timeout)
response.raise_for_status()
content = response.text
original_length = len(content)
# Extract text if requested
if extract_text_only:
soup = BeautifulSoup(content, "html.parser")
content = soup.get_text(separator=" ", strip=True)
# Truncate if requested
truncated = False
if truncate_content and len(content) > self.max_content_length:
content = content[: self.max_content_length] + "..."
truncated = True
# Prepare content data
content_data = {
"url": url,
"timestamp": timestamp,
"fetched_timestamp": snapshot.timestamp,
"content": content,
"original_content_length": original_length,
"truncated": truncated,
"extract_text_only": extract_text_only,
}
# Format output
if output_format == "json":
message = content_data
elif output_format == "text":
message = content
else: # markdown
message = self._format_content_for_llm(content_data, output_format)
execution_time = time.time() - start_time
self._color_log(f"Retrieved {len(content):,} characters in {execution_time:.2f}s", Color.green)
return ActionResponse(
success=True,
message=message,
metadata=WaybackMetadata(
url=url,
operation="get_content",
timestamp=snapshot.timestamp,
content_length=len(content),
text_extracted=extract_text_only,
truncated=truncated,
execution_time=execution_time,
).model_dump(),
)
except Exception as e:
error_msg = f"Failed to fetch archived content: {str(e)}"
self._color_log(error_msg, Color.red)
self.logger.error(f"Error in mcp_get_archived_content: {traceback.format_exc()}")
return ActionResponse(
success=False,
message=error_msg,
metadata=WaybackMetadata(
url=url or "unknown",
operation="get_content",
timestamp=timestamp or "unknown",
execution_time=time.time() - start_time,
error_type=type(e).__name__,
).model_dump(),
)
def mcp_get_wayback_capabilities(self) -> ActionResponse:
"""Get Wayback Machine service capabilities and configuration.
Returns:
ActionResponse with service capabilities information
"""
capabilities = {
"service": "Wayback Machine MCP Server",
"version": "1.0.0",
"description": "Interact with the Internet Archive's Wayback Machine",
"operations": {
"list_versions": "List archived versions of URLs with date filtering",
"get_content": "Fetch content from specific archived snapshots",
},
"features": {
"date_filtering": True,
"text_extraction": True,
"content_truncation": True,
"multiple_formats": ["markdown", "json", "text"],
"error_handling": True,
"logging": True,
},
"configuration": {
"user_agent": self.user_agent,
"default_timeout": self.default_timeout,
"max_content_length": self.max_content_length,
},
"limits": {"max_content_length": self.max_content_length, "request_timeout": self.default_timeout},
}
message = f"""# Wayback Machine Service Capabilities
## Service Information
- **Service:** {capabilities["service"]}
- **Version:** {capabilities["version"]}
- **Description:** {capabilities["description"]}
## Available Operations
- **List Versions:** {capabilities["operations"]["list_versions"]}
- **Get Content:** {capabilities["operations"]["get_content"]}
- **Save Page:** {capabilities["operations"]["save_page"]}
## Features
- **Date Filtering:** {capabilities["features"]["date_filtering"]}
- **Text Extraction:** {capabilities["features"]["text_extraction"]}
- **Content Truncation:** {capabilities["features"]["content_truncation"]}
- **Output Formats:** {", ".join(capabilities["features"]["multiple_formats"])}
- **Error Handling:** {capabilities["features"]["error_handling"]}
- **Logging:** {capabilities["features"]["logging"]}
## Configuration
- **User Agent:** {capabilities["configuration"]["user_agent"]}
- **Default Timeout:** {capabilities["configuration"]["default_timeout"]} seconds
- **Max Content Length:** {capabilities["configuration"]["max_content_length"]:,} characters
## Limits
- **Max Content Length:** {capabilities["limits"]["max_content_length"]:,} characters
- **Request Timeout:** {capabilities["limits"]["request_timeout"]} seconds
"""
return ActionResponse(success=True, message=message, metadata=capabilities)
# Default arguments for testing
if __name__ == "__main__":
load_dotenv()
arguments = ActionArguments(
name="wayback-machine-server",
transport="stdio",
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
)
try:
service = WaybackActionCollection(arguments)
service.run()
except Exception as e:
print(f"An error occurred: {e}")
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,621 @@
"""
Yahoo Finance MCP Action Collection
This module provides Yahoo Finance data access through the ActionCollection framework.
It supports stock quotes, historical data, company information, financial statements,
news search, and market summaries with LLM-optimized output formatting.
Key features:
- Real-time stock quotes and market data
- Historical price data with configurable intervals
- Company information and financial statements
- Financial news search
- Market indices summaries
- LLM-friendly data formatting
- Comprehensive error handling
Main functions:
- mcp_get_stock_quote: Get current stock quote information
- mcp_get_historical_data: Retrieve historical OHLCV data
- mcp_get_company_info: Fetch company details and business information
- mcp_get_financial_statements: Access income statements, balance sheets, cash flow
"""
import os
import time
import traceback
from datetime import datetime
from typing import Any
import yfinance as yf
from dotenv import load_dotenv
from pydantic import BaseModel, Field
from pydantic.fields import FieldInfo
from aworld.logs.util import Color
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
class YFinanceMetadata(BaseModel):
"""Metadata for Yahoo Finance operation results."""
symbol: str
operation: str
execution_time: float | None = None
data_points: int | None = None
error_type: str | None = None
timestamp: str | None = None
class YahooFinanceActionCollection(ActionCollection):
"""Yahoo Finance MCP service for financial data access.
Provides comprehensive financial data capabilities including:
- Real-time stock quotes and market data
- Historical price data with flexible time ranges
- Company information and business details
- Financial statements (income, balance sheet, cash flow)
- Financial news search and aggregation
- Market indices summaries and overviews
- LLM-optimized data formatting
- Error handling and validation
"""
def _format_financial_data(self, data: Any, data_type: str) -> str:
"""Format financial data for LLM consumption.
Args:
data: Raw financial data
data_type: Type of data for context
Returns:
LLM-friendly formatted string
"""
if isinstance(data, dict):
if data_type == "quote":
return self._format_quote_data(data)
elif data_type == "company":
return self._format_company_data(data)
elif isinstance(data, list):
if data_type == "historical":
return self._format_historical_data(data)
elif data_type == "market_summary":
return self._format_market_summary_data(data)
elif data_type == "news":
return self._format_news_list_data(data)
return str(data)
def _format_quote_data(self, quote: dict[str, Any]) -> str:
"""Format stock quote data for LLM."""
lines = [f"# Stock Quote: {quote.get('symbol', 'N/A')}"]
if quote.get("companyName"):
lines.append(f"**Company:** {quote['companyName']}")
if quote.get("currentPrice"):
lines.append(f"**Current Price:** ${quote['currentPrice']:.2f} {quote.get('currency', '')}")
if quote.get("previousClose"):
change = quote.get("currentPrice", 0) - quote.get("previousClose", 0)
change_pct = (change / quote["previousClose"]) * 100 if quote.get("previousClose") else 0
direction = "📈" if change >= 0 else "📉"
lines.append(f"**Change:** {direction} ${change:.2f} ({change_pct:.2f}%)")
if quote.get("dayHigh") and quote.get("dayLow"):
lines.append(f"**Day Range:** ${quote['dayLow']:.2f} - ${quote['dayHigh']:.2f}")
if quote.get("volume"):
lines.append(f"**Volume:** {quote['volume']:,}")
if quote.get("marketCap"):
lines.append(f"**Market Cap:** ${quote['marketCap']:,}")
return "\n".join(lines)
def _format_company_data(self, company: dict[str, Any]) -> str:
"""Format company information for LLM."""
lines = [f"# Company Information: {company.get('symbol', 'N/A')}"]
if company.get("longName"):
lines.append(f"**Company Name:** {company['longName']}")
if company.get("sector"):
lines.append(f"**Sector:** {company['sector']}")
if company.get("industry"):
lines.append(f"**Industry:** {company['industry']}")
if company.get("fullTimeEmployees"):
lines.append(f"**Employees:** {company['fullTimeEmployees']:,}")
if company.get("city") and company.get("country"):
location = f"{company['city']}, {company['country']}"
if company.get("state"):
location = f"{company['city']}, {company['state']}, {company['country']}"
lines.append(f"**Location:** {location}")
if company.get("website"):
lines.append(f"**Website:** {company['website']}")
if company.get("longBusinessSummary"):
summary = (
company["longBusinessSummary"][:500] + "..."
if len(company["longBusinessSummary"]) > 500
else company["longBusinessSummary"]
)
lines.extend(["\n**Business Summary:**", summary])
return "\n".join(lines)
def _format_historical_data(self, data: list[dict[str, Any]]) -> str:
"""Format historical data for LLM."""
if not data:
return "No historical data available."
lines = [f"# Historical Data ({len(data)} records)"]
# Show first few and last few records
preview_count = min(3, len(data))
lines.append("\n**Recent Data:**")
for record in data[-preview_count:]:
date = record.get("Date", record.get("Datetime", "N/A"))
close = record.get("Close", 0)
volume = record.get("Volume", 0)
lines.append(f"- {date}: Close ${close:.2f}, Volume {volume:,}")
if len(data) > preview_count * 2:
lines.append(f"\n... {len(data) - preview_count * 2} more records ...")
if len(data) > preview_count:
lines.append("\n**Earliest Data:**")
for record in data[:preview_count]:
date = record.get("Date", record.get("Datetime", "N/A"))
close = record.get("Close", 0)
volume = record.get("Volume", 0)
lines.append(f"- {date}: Close ${close:.2f}, Volume {volume:,}")
return "\n".join(lines)
def _format_news_list_data(self, news_list: list[dict[str, Any]]) -> str:
"""Format news list for LLM."""
if not news_list:
return "No news articles found."
lines = [f"# Financial News ({len(news_list)} articles)"]
for i, article in enumerate(news_list, 1):
lines.append(f"\n## {i}. {article.get('title', 'No Title')}")
if article.get("publisher"):
lines.append(f"**Publisher:** {article['publisher']}")
if article.get("providerPublishTime"):
lines.append(f"**Published:** {article['providerPublishTime']}")
if article.get("link"):
lines.append(f"**Link:** {article['link']}")
return "\n".join(lines)
def _format_market_summary_data(self, summaries: list[dict[str, Any]]) -> str:
"""Format market summary for LLM."""
if not summaries:
return "No market data available."
lines = ["# Market Summary"]
for summary in summaries:
symbol = summary.get("symbol", "N/A")
name = summary.get("name", symbol)
price = summary.get("currentPrice", 0)
change = summary.get("change", 0)
change_pct = summary.get("percentChange", 0)
direction = "📈" if change >= 0 else "📉"
lines.append(f"\n**{name} ({symbol})**")
lines.append(f"Price: ${price:.2f} {direction} {change:+.2f} ({change_pct:+.2f}%)")
return "\n".join(lines)
async def mcp_get_stock_quote(
self,
symbol: str = Field(description="Stock ticker symbol (e.g., AAPL, MSFT)"),
) -> ActionResponse:
"""Get current stock quote information.
Fetches real-time stock quote data including current price, daily changes,
volume, market cap, and other key metrics for the specified ticker symbol.
Args:
symbol: The stock ticker symbol to fetch quote for
Returns:
ActionResponse with formatted quote data and metadata
"""
# Handle FieldInfo objects
if isinstance(symbol, FieldInfo):
symbol = symbol.default
try:
start_time = time.time()
self._color_log(f"📊 Fetching stock quote for: {symbol}", Color.cyan)
ticker = yf.Ticker(symbol)
info = ticker.info
if not info or (info.get("regularMarketPrice") is None and info.get("currentPrice") is None):
# Try to get basic history to validate symbol
hist = ticker.history(period="1d")
if hist.empty:
raise ValueError(f"No data found for symbol: {symbol}. It might be invalid or delisted.")
raise ValueError(f"Could not retrieve detailed quote for symbol: {symbol}. Limited data available.")
# Extract key quote information
quote_data = {
"symbol": symbol.upper(),
"companyName": info.get("shortName", info.get("longName")),
"currentPrice": info.get("regularMarketPrice", info.get("currentPrice")),
"previousClose": info.get("previousClose"),
"open": info.get("regularMarketOpen", info.get("open")),
"dayHigh": info.get("regularMarketDayHigh", info.get("dayHigh")),
"dayLow": info.get("regularMarketDayLow", info.get("dayLow")),
"volume": info.get("regularMarketVolume", info.get("volume")),
"averageVolume": info.get("averageVolume"),
"marketCap": info.get("marketCap"),
"fiftyTwoWeekHigh": info.get("fiftyTwoWeekHigh"),
"fiftyTwoWeekLow": info.get("fiftyTwoWeekLow"),
"currency": info.get("currency"),
"exchange": info.get("exchange"),
}
# Filter out None values
quote_data = {k: v for k, v in quote_data.items() if v is not None}
execution_time = time.time() - start_time
formatted_message = self._format_financial_data(quote_data, "quote")
metadata = YFinanceMetadata(
symbol=symbol.upper(),
operation="get_stock_quote",
execution_time=execution_time,
data_points=len(quote_data),
yfinance_available=True,
timestamp=datetime.now().isoformat(),
)
self._color_log("✅ Stock quote retrieved successfully", Color.green)
return ActionResponse(success=True, message=formatted_message, metadata=metadata.model_dump())
except Exception as e:
error_msg = f"Failed to fetch stock quote for {symbol}: {str(e)}"
self.logger.error(f"Stock quote error: {traceback.format_exc()}")
metadata = YFinanceMetadata(
symbol=symbol.upper(),
operation="get_stock_quote",
error_type=type(e).__name__,
timestamp=datetime.now().isoformat(),
)
return ActionResponse(success=False, message=error_msg, metadata=metadata.model_dump())
async def mcp_get_historical_data(
self,
symbol: str = Field(description="Stock ticker symbol (e.g., AAPL, MSFT)"),
start: str = Field(description="Start date (YYYY-MM-DD)"),
end: str = Field(description="End date (YYYY-MM-DD)"),
interval: str = Field(default="1d", description="Data interval (1d, 1wk, 1mo, etc.)"),
max_rows_preview: int = Field(default=10, description="Max rows for preview (0 for all data)"),
) -> ActionResponse:
"""Retrieve historical stock data.
Fetches historical OHLCV (Open, High, Low, Close, Volume) data for the
specified ticker symbol within the given date range and interval.
Args:
symbol: Stock ticker symbol
start: Start date in YYYY-MM-DD format
end: End date in YYYY-MM-DD format
interval: Data interval (1d, 1wk, 1mo, etc.)
max_rows_preview: Maximum rows to show in preview
Returns:
ActionResponse with historical data and metadata
"""
# Handle FieldInfo objects
if isinstance(symbol, FieldInfo):
symbol = symbol.default
if isinstance(start, FieldInfo):
start = start.default
if isinstance(end, FieldInfo):
end = end.default
if isinstance(interval, FieldInfo):
interval = interval.default
if isinstance(max_rows_preview, FieldInfo):
max_rows_preview = max_rows_preview.default
try:
start_time = time.time()
self._color_log(f"📈 Fetching historical data for: {symbol} ({start} to {end})", Color.cyan)
ticker = yf.Ticker(symbol)
hist_df = ticker.history(start=start, end=end, interval=interval)
if hist_df.empty:
raise ValueError(
f"No historical data found for {symbol} with start={start}, end={end}, interval={interval}"
)
# Convert DataFrame to list of dictionaries
hist_df.reset_index(inplace=True)
# Ensure date columns are strings for JSON serialization
if "Date" in hist_df.columns:
hist_df["Date"] = hist_df["Date"].astype(str)
if "Datetime" in hist_df.columns:
hist_df["Datetime"] = hist_df["Datetime"].astype(str)
# Clean column names
hist_df.columns = hist_df.columns.str.replace(" ", "")
historical_data = hist_df.to_dict(orient="records")
execution_time = time.time() - start_time
# Format message based on data size
if max_rows_preview > 0 and len(historical_data) > max_rows_preview:
preview_count = max_rows_preview // 2
preview_count = max(1, preview_count)
preview_data = historical_data[:preview_count] + historical_data[-preview_count:]
formatted_message = self._format_financial_data(preview_data, "historical")
formatted_message += (
f"\n\n*Note: Showing preview of {len(preview_data)} out of {len(historical_data)} total records*"
)
else:
formatted_message = self._format_financial_data(historical_data, "historical")
metadata = YFinanceMetadata(
symbol=symbol.upper(),
operation="get_historical_data",
execution_time=execution_time,
data_points=len(historical_data),
yfinance_available=True,
timestamp=datetime.now().isoformat(),
)
self._color_log(f"✅ Historical data retrieved: {len(historical_data)} records", Color.green)
return ActionResponse(success=True, message=formatted_message, metadata=metadata.model_dump())
except Exception as e:
error_msg = f"Failed to fetch historical data for {symbol}: {str(e)}"
self.logger.error(f"Historical data error: {traceback.format_exc()}")
metadata = YFinanceMetadata(
symbol=symbol.upper(),
operation="get_historical_data",
error_type=type(e).__name__,
timestamp=datetime.now().isoformat(),
)
return ActionResponse(success=False, message=error_msg, metadata=metadata.model_dump())
async def mcp_get_company_info(
self,
symbol: str = Field(description="Stock ticker symbol (e.g., AAPL, MSFT)"),
) -> ActionResponse:
"""Get company information and business details.
Fetches comprehensive company information including sector, industry,
employee count, business summary, location, and other key details.
Args:
symbol: Stock ticker symbol
Returns:
ActionResponse with company information and metadata
"""
# Handle FieldInfo objects
if isinstance(symbol, FieldInfo):
symbol = symbol.default
try:
start_time = time.time()
self._color_log(f"🏢 Fetching company info for: {symbol}", Color.cyan)
ticker = yf.Ticker(symbol)
info = ticker.info
if not info or not info.get("symbol"):
raise ValueError(f"No company information found for symbol: {symbol}. It might be invalid.")
# Extract key company information
company_data = {
"symbol": info.get("symbol"),
"shortName": info.get("shortName"),
"longName": info.get("longName"),
"sector": info.get("sector"),
"industry": info.get("industry"),
"fullTimeEmployees": info.get("fullTimeEmployees"),
"longBusinessSummary": info.get("longBusinessSummary"),
"city": info.get("city"),
"state": info.get("state"),
"country": info.get("country"),
"website": info.get("website"),
"exchange": info.get("exchange"),
"currency": info.get("currency"),
"marketCap": info.get("marketCap"),
}
# Filter out None values
company_data = {k: v for k, v in company_data.items() if v is not None}
execution_time = time.time() - start_time
formatted_message = self._format_financial_data(company_data, "company")
metadata = YFinanceMetadata(
symbol=symbol.upper(),
operation="get_company_info",
execution_time=execution_time,
data_points=len(company_data),
yfinance_available=True,
timestamp=datetime.now().isoformat(),
)
self._color_log("✅ Company information retrieved successfully", Color.green)
return ActionResponse(success=True, message=formatted_message, metadata=metadata.model_dump())
except Exception as e:
error_msg = f"Failed to fetch company info for {symbol}: {str(e)}"
self.logger.error(f"Company info error: {traceback.format_exc()}")
metadata = YFinanceMetadata(
symbol=symbol.upper(),
operation="get_company_info",
error_type=type(e).__name__,
timestamp=datetime.now().isoformat(),
)
return ActionResponse(success=False, message=error_msg, metadata=metadata.model_dump())
async def mcp_get_financial_statements(
self,
symbol: str = Field(description="Stock ticker symbol (e.g., AAPL, MSFT)"),
statement_type: str = Field(description="Statement type: income_statement, balance_sheet, or cash_flow"),
period_type: str = Field(default="annual", description="Period type: annual or quarterly"),
max_columns_preview: int = Field(default=4, description="Max periods to show (0 for all)"),
) -> ActionResponse:
"""Get financial statements for a company.
Fetches financial statements including income statement, balance sheet,
or cash flow statement for the specified company and period.
Args:
symbol: Stock ticker symbol
statement_type: Type of statement (income_statement, balance_sheet, cash_flow)
period_type: Period type (annual or quarterly)
max_columns_preview: Maximum periods to show in preview
Returns:
ActionResponse with financial statement data and metadata
"""
# Handle FieldInfo objects
if isinstance(symbol, FieldInfo):
symbol = symbol.default
if isinstance(statement_type, FieldInfo):
statement_type = statement_type.default
if isinstance(period_type, FieldInfo):
period_type = period_type.default
if isinstance(max_columns_preview, FieldInfo):
max_columns_preview = max_columns_preview.default
try:
start_time = time.time()
self._color_log(f"📋 Fetching {statement_type} for: {symbol} ({period_type})", Color.cyan)
ticker = yf.Ticker(symbol)
statement_df = None
# Get appropriate statement
if statement_type == "income_statement":
statement_df = ticker.income_stmt if period_type == "annual" else ticker.quarterly_income_stmt
elif statement_type == "balance_sheet":
statement_df = ticker.balance_sheet if period_type == "annual" else ticker.quarterly_balance_sheet
elif statement_type == "cash_flow":
statement_df = ticker.cashflow if period_type == "annual" else ticker.quarterly_cashflow
else:
raise ValueError(
f"Invalid statement_type: {statement_type}. "
"Must be one of: income_statement, balance_sheet, cash_flow"
)
if statement_df is None or statement_df.empty:
raise ValueError(f"No {period_type} {statement_type} data found for symbol {symbol}")
# Process DataFrame
statement_df.reset_index(inplace=True)
statement_df.rename(columns={"index": "Item"}, inplace=True)
# Convert date columns to strings
for col in statement_df.columns:
if col != "Item":
try:
if hasattr(col, "strftime"):
statement_df.rename(columns={col: col.strftime("%Y-%m-%d")}, inplace=True)
except Exception:
pass
statement_data = statement_df.to_dict(orient="records")
execution_time = time.time() - start_time
# Format message
if max_columns_preview > 0 and len(statement_df.columns) > (max_columns_preview + 1):
columns_to_keep = ["Item"] + list(statement_df.columns[1 : max_columns_preview + 1])
preview_df = statement_df[columns_to_keep]
preview_data = preview_df.to_dict(orient="records")
formatted_message = f"# {statement_type.replace('_', ' ').title()} ({period_type.title()})\n\n"
formatted_message += (
f"Showing preview of most recent {max_columns_preview} periods "
f"out of {len(statement_df.columns) - 1} available.\n\n"
)
# Show key financial items
for item in preview_data[:10]: # Show first 10 items
item_name = item.get("Item", "N/A")
formatted_message += f"**{item_name}:**\n"
for col, value in item.items():
if col != "Item" and value is not None:
formatted_message += (
f" - {col}: {value:,}\n"
if isinstance(value, (int, float))
else f" - {col}: {value}\n"
)
formatted_message += "\n"
else:
formatted_message = f"# {statement_type.replace('_', ' ').title()} ({period_type.title()})\n\n"
formatted_message += f"Complete financial statement with {len(statement_data)} line items.\n"
metadata = YFinanceMetadata(
symbol=symbol.upper(),
operation="get_financial_statements",
execution_time=execution_time,
data_points=len(statement_data),
yfinance_available=True,
timestamp=datetime.now().isoformat(),
)
self._color_log(f"✅ Financial statements retrieved: {len(statement_data)} items", Color.green)
return ActionResponse(success=True, message=formatted_message, metadata=metadata.model_dump())
except Exception as e:
error_msg = f"Failed to fetch {statement_type} for {symbol}: {str(e)}"
self.logger.error(f"Financial statements error: {traceback.format_exc()}")
metadata = YFinanceMetadata(
symbol=symbol.upper(),
operation="get_financial_statements",
error_type=type(e).__name__,
timestamp=datetime.now().isoformat(),
)
return ActionResponse(success=False, message=error_msg, metadata=metadata.model_dump())
# Default arguments for testing
if __name__ == "__main__":
load_dotenv()
arguments = ActionArguments(
name="yahoo-finance",
transport="stdio",
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
)
try:
service = YahooFinanceActionCollection(arguments)
service.run()
except Exception as e:
print(f"Error: {e}")
@@ -0,0 +1,549 @@
"""
YouTube MCP Service
This module provides MCP service functionality for YouTube operations including:
- Downloading videos from YouTube URLs
- Extracting transcripts from YouTube videos
It handles various scenarios with proper validation, error handling,
and progress tracking while providing LLM-friendly formatted results.
"""
import os
import time
import traceback
from datetime import datetime
from pathlib import Path
from dotenv import load_dotenv
from pydantic import BaseModel, Field
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.by import By
from youtube_transcript_api import FetchedTranscript, YouTubeTranscriptApi
from aworld.logs.util import Color
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
# Default driver path for Chrome WebDriver
_DEFAULT_DRIVER_PATH = os.environ.get(
"CHROME_DRIVER_PATH", str(Path("~/Downloads/chromedriver-mac-arm64/chromedriver").expanduser())
)
class YoutubeDownloadResults(BaseModel):
"""Download result model with file information"""
file_path: str
file_name: str
file_size: int
content_type: str | None = None
success: bool
error: str | None = None
class TranscriptResult(BaseModel):
"""Transcript result model with transcript information"""
video_id: str
transcript: FetchedTranscript
success: bool
error: str | None = None
class YouTubeMetadata(BaseModel):
"""Metadata for YouTube operation results"""
operation: str
url: str | None = None
video_id: str | None = None
file_path: str | None = None
file_name: str | None = None
file_size: int | None = None
content_type: str | None = None
language_code: str | None = None
translate_to_language: str | None = None
execution_time: float | None = None
error_type: str | None = None
class YouTubeActionCollection(ActionCollection):
"""MCP service for YouTube operations.
Provides YouTube capabilities including:
- Video downloading with Selenium automation
- Transcript extraction and translation
- LLM-friendly result formatting
- Error handling and logging
"""
def __init__(self, arguments: ActionArguments) -> None:
super().__init__(arguments)
# Initialize supported file extensions
self.supported_extensions = {".mp4", ".webm", ".mkv"}
self._color_log("YouTube service initialized", Color.green, "debug")
def _format_transcript_output(self, result: TranscriptResult, format_type: str = "markdown") -> str:
"""Format transcript results for LLM consumption.
Args:
result: Transcript extraction result
format_type: Output format ('markdown', 'json', 'text')
Returns:
Formatted string suitable for LLM consumption
"""
if result is None or not result.success:
return f"Failed to extract transcript: {result.error}"
if format_type == "json":
return result.model_dump()
elif format_type == "text":
output = [f"Transcript for video ID: {result.video_id}\n"]
# Access snippets from FetchedTranscript
for entry in result.transcript.snippets:
start_time = entry["start"]
text = entry["text"]
minutes, seconds = divmod(int(start_time), 60)
timestamp = f"{minutes:02d}:{seconds:02d}"
output.append(f"[{timestamp}] {text}")
return "\n".join(output)
else: # markdown (default)
output = [f"# Transcript for YouTube Video: {result.video_id}\n"]
output.append("| Timestamp | Text |")
output.append("| --- | --- |")
# Access snippets from FetchedTranscript
for entry in result.transcript.snippets:
start_time = entry["start"]
text: str = entry["text"]
minutes, seconds = divmod(int(start_time), 60)
timestamp = f"{minutes:02d}:{seconds:02d}"
# Escape pipe characters in markdown table
safe_text = text.replace("|", "\\|")
output.append(f"| {timestamp} | {safe_text} |")
return "\n".join(output)
def _format_download_output(self, result: YoutubeDownloadResults, format_type: str = "markdown") -> str:
"""Format download results for LLM consumption.
Args:
result: Download result
format_type: Output format ('markdown', 'json', 'text')
Returns:
Formatted string suitable for LLM consumption
"""
if not result.success:
return f"Failed to download video: {result.error}"
if format_type == "json":
return result.model_dump()
elif format_type == "text":
output_parts = [
"Download completed successfully",
f"File: {result.file_name}",
f"Path: {result.file_path}",
f"Size: {result.file_size} bytes",
]
if result.content_type:
output_parts.append(f"Content Type: {result.content_type}")
return "\n".join(output_parts)
else: # markdown (default)
output_parts = [
"# YouTube Download Results ✅",
"",
"## File Information",
f"**Filename:** `{result.file_name}`",
f"**Path:** `{result.file_path}`",
f"**Size:** {result.file_size} bytes",
]
if result.content_type:
output_parts.append(f"**Content Type:** {result.content_type}")
return "\n".join(output_parts)
def _get_youtube_content(self, url: str, output_dir: str, timeout: int) -> None:
"""Use Selenium to download YouTube content via cobalt.tools
Args:
url: YouTube video URL
output_dir: Directory to save downloaded content
timeout: Maximum time to wait for download in seconds
"""
driver = None
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)
self._color_log(f"Opening cobalt.tools to download from {url}", Color.blue)
# 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:
self._color_log(f"Bot detection handling: {str(e)}", Color.yellow)
# try:
# t = 0
# while t < timeout:
# if (
# "downloading" not in driver.find_element(By.CLASS_NAME, "status-text.svelte-dmosdd").text
# and "starting" not in driver.find_element(By.CLASS_NAME, "status-text.svelte-dmosdd").text
# ):
# driver.find_element(By.CLASS_NAME, "button.action-button.svelte-dmosdd").click()
# break
# t += 3
# time.sleep(3)
# except Exception as e:
# self._color_log(f"Bot detection handling: {str(e)}", Color.yellow)
# 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:
self._color_log(f"Download timeout after {timeout} seconds", Color.yellow)
break
self._color_log("Download process completed", Color.green)
except Exception as e:
self._color_log(f"Error during YouTube content download: {str(e)}", Color.red)
raise
finally:
# Close browser
if driver:
driver.quit()
def _find_existing_video(self, search_dir: str, video_id: str) -> str | None:
"""Recursively search for an existing video file with the given ID.
Args:
search_dir: Directory to search in
video_id: YouTube video ID to look for
Returns:
Path to existing file if found, None otherwise
"""
if not video_id:
return None
search_path = Path(search_dir)
if not search_path.exists():
return None
for item in search_path.iterdir():
if item.is_file() and video_id in item.name:
return str(item)
elif item.is_dir():
found = self._find_existing_video(str(item), video_id)
if found:
return found
return None
async def mcp_download_youtube_video(
self,
url: str = Field(description="The URL of YouTube video to download."),
timeout: int = Field(180, description="Download timeout in seconds (default: 180)."),
output_format: str = Field(
"markdown", description="Output format: 'markdown', 'json', or 'text' (default: markdown)."
),
) -> ActionResponse:
"""Download a YouTube video from URL and save it to the local filesystem.
This tool provides YouTube video downloading with:
- Selenium-based automation via cobalt.tools
- Configurable timeout controls
- Existing file detection to avoid redundant downloads
- LLM-optimized result formatting
Args:
url: The URL of YouTube video to download
timeout: Maximum download time in seconds
output_format: Format for the response output
Returns:
ActionResponse with download results and metadata
"""
start_time = time.time()
try:
# Validate URL
if not url.startswith(("http://", "https://")):
raise ValueError("Invalid URL format. URL must start with http:// or https://")
if not ("youtube.com" in url or "youtu.be" in url):
raise ValueError("URL must be a valid YouTube URL")
# Create output directory if it doesn't exist
output_path = self.workspace / "youtube_downloads"
output_path.mkdir(parents=True, exist_ok=True)
# Generate filename based on timestamp
filename = f"youtube_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
file_path = output_path / filename
file_path.mkdir(parents=True, exist_ok=True)
self._color_log(f"Output path: {file_path}", Color.blue)
# Extract video ID for existing file check
video_id = url.split("?v=")[-1].split("&")[0] if "?v=" in url else ""
if "youtu.be/" in url and not video_id:
video_id = url.split("youtu.be/")[-1].split("?")[0]
# Check if video already exists
base_path = self.workspace
existing_file = self._find_existing_video(str(base_path), video_id)
if existing_file:
existing_path = Path(existing_file)
result = YoutubeDownloadResults(
file_path=str(existing_path),
file_name=existing_path.name,
file_size=existing_path.stat().st_size,
content_type="mp4",
success=True,
error=None,
)
self._color_log(f"Found {video_id} already downloaded in: {existing_file}", Color.green)
# Format output for LLM
message = self._format_download_output(result, output_format)
execution_time = time.time() - start_time
# Create metadata
metadata = YouTubeMetadata(
operation="download",
url=url,
video_id=video_id,
file_path=str(existing_path),
file_name=existing_path.name,
file_size=existing_path.stat().st_size,
content_type="mp4",
execution_time=execution_time,
).model_dump()
return ActionResponse(success=True, message=message, metadata=metadata)
# Download the video
self._color_log(f"Downloading video from {url} to {file_path}", Color.blue)
self._get_youtube_content(url, str(file_path), timeout)
# Check if download was successful
downloaded_files = list(file_path.iterdir())
if not downloaded_files:
raise FileNotFoundError("No files were downloaded")
download_file = downloaded_files[0]
file_size = download_file.stat().st_size
self._color_log(f"File downloaded successfully to {download_file}", Color.green)
# Create result
result = YoutubeDownloadResults(
file_path=str(download_file),
file_name=download_file.name,
file_size=file_size,
content_type="mp4",
success=True,
error=None,
)
# Format output for LLM
message = self._format_download_output(result, output_format)
execution_time = time.time() - start_time
# Create metadata
metadata = YouTubeMetadata(
operation="download",
url=url,
video_id=video_id,
file_path=str(download_file),
file_name=download_file.name,
file_size=file_size,
content_type="mp4",
execution_time=execution_time,
).model_dump()
return ActionResponse(success=True, message=message, metadata=metadata)
except Exception as e:
error_msg = str(e)
self._color_log(f"Download error: {traceback.format_exc()}", Color.red)
# Format error for LLM
message = f"Failed to download YouTube video: {error_msg}"
execution_time = time.time() - start_time
# Create metadata
metadata = YouTubeMetadata(
operation="download",
url=url,
error_type="download_failure",
execution_time=execution_time,
).model_dump()
return ActionResponse(success=False, message=message, metadata=metadata)
async def mcp_extract_youtube_transcript(
self,
video_id: str = Field(description="The YouTube video ID or URL to extract transcript from."),
language_code: str = Field("en", description="Language code for the transcript (default: en)."),
translate_to_language: str | None = Field(
None, description="Translate transcript to this language code if provided."
),
output_format: str = Field(
"markdown", description="Output format: 'markdown', 'json', or 'text' (default: markdown)."
),
) -> ActionResponse:
"""Extract transcript from a YouTube video given its video ID or URL.
This tool provides transcript extraction with:
- Support for multiple languages
- Translation capabilities
- URL or video ID input handling
- LLM-optimized result formatting
Args:
video_id: The YouTube video ID or URL to extract transcript from
language_code: Language code for the transcript
translate_to_language: Translate transcript to this language code if provided
output_format: Format for the response output
Returns:
ActionResponse with transcript data and metadata
"""
start_time = time.time()
try:
# Clean video_id if full URL was provided
if "youtube.com" in video_id or "youtu.be" in video_id:
if "?v=" in video_id:
video_id = video_id.split("?v=")[-1].split("&")[0]
elif "youtu.be/" in video_id:
video_id = video_id.split("youtu.be/")[-1].split("?")[0]
self._color_log(f"Extracting transcript for video ID: {video_id}", Color.blue)
# Get transcript in specified language
if translate_to_language:
# Get transcript and translate it
y_api = YouTubeTranscriptApi()
transcript_list = y_api.list(video_id)
transcript = None
try:
# Try to get transcript in specified language
transcript = transcript_list.find_transcript([language_code])
except Exception:
# If specified language not found, get any available transcript
transcript = transcript_list.find_generated_transcript(["en"])
# Translate to target language
transcript_data = transcript.translate(translate_to_language).fetch()
else:
try:
# Get transcript without translation
transcript_data: FetchedTranscript = (
YouTubeTranscriptApi()
.list(video_id)
.find_transcript((language_code,))
.fetch(preserve_formatting=False)
)
except Exception:
transcript_data = None
result = TranscriptResult(video_id=video_id, transcript=transcript_data, success=True, error=None)
self._color_log(f"Successfully extracted transcript for video ID: {video_id}", Color.green)
# Format output for LLM
message = self._format_transcript_output(result, output_format)
execution_time = time.time() - start_time
# Create metadata
metadata = YouTubeMetadata(
operation="transcript",
video_id=video_id,
language_code=language_code,
translate_to_language=translate_to_language,
execution_time=execution_time,
).model_dump()
return ActionResponse(success=True, message=message, metadata=metadata)
except Exception as e:
error_msg = str(e)
self._color_log(f"Transcript extraction error: {traceback.format_exc()}", Color.red)
# Format error for LLM
message = f"Failed to extract transcript: {error_msg}"
execution_time = time.time() - start_time
# Create metadata
metadata = YouTubeMetadata(
operation="transcript",
video_id=video_id,
language_code=language_code,
translate_to_language=translate_to_language,
error_type="transcript_failure",
execution_time=execution_time,
).model_dump()
return ActionResponse(success=False, message=message, metadata=metadata)
# Default arguments for testing
if __name__ == "__main__":
load_dotenv()
arguments = ActionArguments(
name="youtube_service",
transport="stdio",
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
)
try:
youtube_service = YouTubeActionCollection(arguments)
youtube_service.run()
except Exception as e:
print(f"Error: {e}")
@@ -0,0 +1,169 @@
import os
import tempfile
import traceback
from pathlib import Path
from urllib.parse import urlparse
import magic
import requests
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_mime_type(file_path: str, default_mime: str | None = 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)
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 = Path(file_path).suffix.lower()
return extension_mime_map.get(ext, default_mime or "application/octet-stream")
def get_file_from_source(
source: str,
max_size_mb: float = 100.0,
timeout: int = 60,
) -> 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
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
if is_url(source):
# Handle URL source
try:
# Make a HEAD request first to check content length
head_response = requests.head(source, timeout=timeout, allow_redirects=True)
head_response.raise_for_status()
# Check content length if available
content_length = head_response.headers.get("content-length")
if content_length and int(content_length) > max_size_bytes:
raise ValueError(
f"File size ({int(content_length) / (1024 * 1024):.2f} MB) "
f"exceeds maximum allowed size ({max_size_mb} MB)"
)
# Download the file
response = requests.get(source, timeout=timeout, stream=True)
response.raise_for_status()
# Read content with size checking
content = b""
for chunk in response.iter_content(chunk_size=8192):
if len(content) + len(chunk) > max_size_bytes:
raise ValueError(f"File size exceeds maximum allowed size ({max_size_mb} MB)")
content += chunk
# Create temporary file
parsed_url = urlparse(source)
filename = os.path.basename(parsed_url.path) or "downloaded_file"
# Create temporary file with proper extension
suffix = Path(filename).suffix or ".tmp"
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
temp_file.write(content)
temp_path = temp_file.name
# Get MIME type
mime_type = get_mime_type(temp_path)
return temp_path, mime_type, content
except requests.RequestException as e:
raise requests.RequestException(f"Failed to download file from URL: {e}: {traceback.format_exc()}")
except Exception as e:
raise IOError(f"Error processing URL: {e}: {traceback.format_exc()}") from e
else:
# Handle local file path
file_path = Path(source)
# Check if file exists
if not file_path.exists():
raise ValueError(f"File does not exist: {source}")
if not file_path.is_file():
raise ValueError(f"Path is not a file: {source}")
# Check file size
file_size = file_path.stat().st_size
if file_size > max_size_bytes:
raise ValueError(
f"File size ({file_size / (1024 * 1024):.2f} MB) exceeds maximum allowed size ({max_size_mb} MB)"
)
# Read file content
try:
with open(file_path, "rb") as f:
content = f.read()
except Exception as e:
raise IOError(f"Cannot read file {source}: {e}: {traceback.format_exc()}") from e
# Get MIME type
mime_type = get_mime_type(str(file_path))
return str(file_path), mime_type, content
@@ -0,0 +1,72 @@
"""Add config table
Revision ID: ca81bd47c050
Revises: 7e5b5dc7342b
Create Date: 2024-08-25 15:26:35.241684
"""
from typing import Sequence, Union
import sqlalchemy as sa
from alembic import op
from sqlalchemy import text
# revision identifiers, used by Alembic.
revision: str = "ca81bd47c050"
down_revision: Union[str, None] = "7e5b5dc7342b"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade():
op.create_table(
"config",
sa.Column("id", sa.Integer, primary_key=True),
sa.Column("data", sa.JSON(), nullable=False),
sa.Column("version", sa.Integer, nullable=False),
sa.Column(
"created_at", sa.DateTime(), nullable=False, server_default=sa.func.now()
),
sa.Column(
"updated_at",
sa.DateTime(),
nullable=True,
server_default=sa.func.now(),
onupdate=sa.func.now(),
),
)
init_gaia_config()
def downgrade():
op.drop_table("config")
def init_gaia_config():
import json
prompt_suggestions = None
with open(
"/app/aworld/examples/gaia/GAIA/2023/validation/metadata.jsonl",
"r",
encoding="utf-8",
) as f:
data_set = [json.loads(line) for line in f]
prompt_suggestions = [
{
"title": [
i["task_id"],
i["Question"][:100],
],
"content": json.dumps({"task_id": i["task_id"]}),
}
for i in data_set
]
conn = op.get_bind()
conn.execute(
text("INSERT INTO config (id, data, version) VALUES (1, :data, 0)"),
{"data": json.dumps({"ui": {"prompt_suggestions": prompt_suggestions}})},
)
print(f">>> patch gaia_agent: add prompt_suggestions success!")
@@ -0,0 +1,37 @@
system_prompt = """You are an all-capable AI assistant, aimed at solving any task presented by the user.
## Task Description:
Please note that the task can be very 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 to verify correctness.
## Workflow:
1. **Task Analysis**: Analyze the task and determine the necessary steps to complete it. Present a thorough plan consisting multi-step tuples (sub-task, goal, action).
2. **Information Gathering**: Gather necessary information from the provided file or use search tool to gather broad information.
3. **Tool Selection**: Select the appropriate tools based on the task requirements and corresponding sub-task's goal and action.
4. **Result Analysis**: Analyze the results obtained from sub-tasks and determine if the original task has been solved.
5. **Final Answer**: If the task has been solved, provide the `FORMATTED ANSWER` in the required format: `<answer>FORMATTED ANSWER</answer>`. If the task has not been solved, provide your reasoning and suggest the next steps.
## Guardrails:
1. Do not use any tools outside of the provided tools list.
2. Always use only one tool at a time in each step of your execution.
3. Even if the task is complex, there is always a solution.
4. If you can't find the answer using one method, try another approach or use different tools to find the solution.
## Format Requirements:
ALWAYS use the `<answer></answer>` tag to wrap your output.
Your `FORMATTED ANSWER` should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
- **Number**: 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.
- **String**: 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.
- **List**: 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.
- **Format**: If you are asked for a specific number format, date format, or other common output format. Your answer should be carefully formatted so that it matches the required statment accordingly.
- `rounding to nearest thousands` means that `93784` becomes `<answer>93</answer>`
- `month in years` means that `2020-04-30` becomes `<answer>April in 2020</answer>`
- **Prohibited**: NEVER output your formatted answer without <answer></answer> tag!
### Examples
1. <answer>apple tree</answer>
2. <answer>3, 4, 5</answer>
3. <answer>(.*?)</answer>
"""
@@ -0,0 +1,45 @@
system_prompt = """You are an all-capable AI assistant, aimed at solving any task presented by the user.
## Task Description:
Please note that the task can be very 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 (there are guarding/reasoning maneuvering tools that will help you analysis and improve the reasoning process). Always use available tools to verify correctness.
## Workflow:
1. **Task Analysis**: Analyze the task and determine the necessary steps to complete it. Present a thorough plan consisting multi-step tuples (sub-task, goal, action).
2. **Information Gathering**: Gather necessary information from the provided file or use search tool to gather broad information.
3. **Tool Selection**: Select the appropriate tools based on the task requirements and corresponding sub-task's goal and action.
4. **Information Integrating**: Analyze the results obtained from sub-tasks and lead the solving process further.
5. **Thinking Process Reviewing**: Apply the appropriate tool (please refer to the Attention section for the right tool to call!) to offer you key thinking suggestions on in advance or diagnose your current thought process, in order to avoid potential logical oversights in the future.
6. **Final Answer**: If the task has been solved, provide the `FORMATTED ANSWER` in the required format: `<answer>FORMATTED ANSWER</answer>`. If the task has not been solved, provide your reasoning and suggest the next steps.
## Guardrails:
1. Do not use any tools outside of the provided tools list.
2. Always use only one tool at a time in each step of your execution.
3. Even if the task is complex, there is always a solution.
4. If you can't find the answer using one method, try another approach or use different tools to find the solution.
5. In the phase of Thinking Process Reviewing, be patient! Don't rush to conclude the Final Answer directly! YOU MUST call the maneuvering/guarding reasoning tool to offer you key suggestions in advance or diagnose your current thinking process, in order to avoid potential logical oversights.
## Mandatory Requirement:
1. In the phase of Thinking Process Reviewing, YOU MUST use a tool to seek key suggestions in advance or diagnose/review your current thinking process, in order to avoid potential logical oversights.
2. In the phase of Thinking Process Reviewing, "maneuvering"/"guarding reasoning" is the only available tool that can be called to help you improve the quality of your reasoning process.
## Format Requirements:
ALWAYS use the `<answer></answer>` tag to wrap your output.
Your `FORMATTED ANSWER` should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
- **Number**: 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.
- **String**: 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.
- **List**: 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.
- **Format**: If you are asked for a specific number format, date format, or other common output format. Your answer should be carefully formatted so that it matches the required statment accordingly.
- `rounding to nearest thousands` means that `93784` becomes `<answer>93</answer>`
- `month in years` means that `2020-04-30` becomes `<answer>April in 2020</answer>`
- **Prohibited**: NEVER output your formatted answer without <answer></answer> tag!
### Formatted Answer Examples
1. <answer>apple tree</answer>
2. <answer>3, 4, 5</answer>
3. <answer>(.*?)</answer>
Now, please read the task in the following carefully, keep the Task Description, Workflow, Guardrails, Mandatory Requirement and Format Requirements in mind, start your execution.
"""
@@ -0,0 +1,8 @@
#!/bin/bash
BASE_DIR=$(dirname "$(readlink -f "$0")")
cd $BASE_DIR
sh -c "cd ../../ && python setup.py install"
sh -c "cd cmd && aworld web"
@@ -0,0 +1,228 @@
import argparse
import json
import logging
import os
import re
import traceback
from pathlib import Path
from typing import Any, Dict, List
from dotenv import load_dotenv
from aworld.agents.llm_agent import Agent
from aworld.config.conf import AgentConfig, TaskConfig
from aworld.core.task import Task
from aworld.runner import Runners
from aworld.core.task import Task
from examples.gaia.prompt import system_prompt
from examples.gaia.utils import (
add_file_path,
load_dataset_meta,
question_scorer,
report_results,
)
# Create log directory if it doesn't exist
if not os.path.exists(os.getenv("AWORLD_WORKSPACE", "~")):
os.makedirs(os.getenv("AWORLD_WORKSPACE", "~"))
parser = argparse.ArgumentParser()
parser.add_argument(
"--start",
type=int,
default=0,
help="Start index of the dataset",
)
parser.add_argument(
"--end",
type=int,
default=20,
help="End index of the dataset",
)
parser.add_argument(
"--q",
type=str,
help="Question Index, e.g., 0-0-0-0-0. Highest priority: override other arguments if provided.",
)
parser.add_argument(
"--skip",
action="store_true",
help="Skip the question if it has been processed before.",
)
parser.add_argument(
"--split",
type=str,
default="validation",
help="Split of the dataset, e.g., validation, test",
)
parser.add_argument(
"--blacklist_file_path",
type=str,
nargs="?",
help="Blacklist file path, e.g., blacklist.txt",
)
args = parser.parse_args()
def setup_logging():
logging_logger = logging.getLogger()
logging_logger.setLevel(logging.INFO)
log_file_name = f"/super_agent_{args.q}.log" if args.q else f"/super_agent_{args.start}_{args.end}.log"
file_handler = logging.FileHandler(
os.getenv("AWORLD_WORKSPACE", "~") + log_file_name,
mode="a",
encoding="utf-8",
)
file_handler.setLevel(logging.INFO)
formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
file_handler.setFormatter(formatter)
logging_logger.addHandler(file_handler)
if __name__ == "__main__":
load_dotenv()
setup_logging()
gaia_dataset_path = os.getenv("GAIA_DATASET_PATH", "./gaia_dataset")
full_dataset = load_dataset_meta(gaia_dataset_path, split=args.split)
logging.info(f"Total questions: {len(full_dataset)}")
try:
with open(Path(__file__).parent / "mcp.json", mode="r", encoding="utf-8") as f:
mcp_config: dict[dict[str, Any]] = json.loads(f.read())
available_servers: list[str] = list(server_name for server_name in mcp_config.get("mcpServers", {}).keys())
logging.info(f"🔧 MCP Available Servers: {available_servers}")
except json.JSONDecodeError as e:
logging.error(f"Error loading mcp_collections.json: {e}")
mcp_config = {}
agent_config = AgentConfig(
llm_provider=os.getenv("LLM_PROVIDER", "openai"),
llm_model_name=os.getenv("LLM_MODEL_NAME", "gpt-4o"),
llm_base_url=os.getenv("LLM_BASE_URL"),
llm_api_key=os.getenv("LLM_API_KEY"),
llm_temperature=os.getenv("LLM_TEMPERATURE", 0.0)
)
super_agent = Agent(
conf=agent_config,
name="gaia_super_agent",
system_prompt=system_prompt,
mcp_config=mcp_config,
mcp_servers=available_servers,
)
# load results from the checkpoint file
if os.path.exists(os.getenv("AWORLD_WORKSPACE", "~") + "/results.json"):
with open(os.getenv("AWORLD_WORKSPACE", "~") + "/results.json", "r", encoding="utf-8") as results_f:
results: List[Dict[str, Any]] = json.load(results_f)
else:
results: List[Dict[str, Any]] = []
# load blacklist `task_id`
if args.blacklist_file_path and os.path.exists(args.blacklist_file_path):
with open(args.blacklist_file_path, "r", encoding="utf-8") as f:
blacklist = set(f.read().splitlines())
else:
blacklist = set() # Empty set if file doesn't exist
try:
# slice dataset by args.start and args.end, overrided by args.q (single `task_id`)
dataset_slice = (
[dataset_record for idx, dataset_record in enumerate(full_dataset) if dataset_record["task_id"] in args.q]
if args.q is not None
else full_dataset[args.start : args.end]
)
# main loop to execute questions
for i, dataset_i in enumerate(dataset_slice):
# specify `task_id`
if args.q and args.q != dataset_i["task_id"]:
continue
# only valid for args.q==None
if not args.q:
# blacklist
if dataset_i["task_id"] in blacklist:
continue
# pass
if any(
# Question Done and Correct
(result["task_id"] == dataset_i["task_id"] and result["is_correct"])
for result in results
) or any(
# Question Done and Incorrect, but Level is 3
(result["task_id"] == dataset_i["task_id"] and not result["is_correct"] and dataset_i["Level"] == 3)
for result in results
):
continue
# skip
if args.skip and any(
# Question Done and Correct
(result["task_id"] == dataset_i["task_id"])
for result in results
):
continue
# run
try:
logging.info(f"Start to process: {dataset_i['task_id']}")
logging.info(f"Detail: {dataset_i}")
logging.info(f"Question: {dataset_i['Question']}")
logging.info(f"Level: {dataset_i['Level']}")
logging.info(f"Tools: {dataset_i['Annotator Metadata']['Tools']}")
question = add_file_path(dataset_i, file_path=gaia_dataset_path, split=args.split)["Question"]
task = Task(input=question, agent=super_agent, conf=TaskConfig())
result = Runners.sync_run_task(task=task)
match = re.search(r"<answer>(.*?)</answer>", result[task.id].answer)
if match:
answer = match.group(1)
logging.info(f"Agent answer: {answer}")
logging.info(f"Correct answer: {dataset_i['Final answer']}")
if question_scorer(answer, dataset_i["Final answer"]):
logging.info(f"Question {i} Correct!")
else:
logging.info("Incorrect!")
# Create the new result record
new_result = {
"task_id": dataset_i["task_id"],
"level": dataset_i["Level"],
"question": question,
"answer": dataset_i["Final answer"],
"response": answer,
"is_correct": question_scorer(answer, dataset_i["Final answer"]),
}
# Check if this task_id already exists in results
existing_index = next(
(i for i, result in enumerate(results) if result["task_id"] == dataset_i["task_id"]),
None,
)
if existing_index is not None:
# Update existing record
results[existing_index] = new_result
logging.info(f"Updated existing record for task_id: {dataset_i['task_id']}")
else:
# Append new record
results.append(new_result)
logging.info(f"Added new record for task_id: {dataset_i['task_id']}")
except Exception:
logging.error(f"Error processing {i}: {traceback.format_exc()}")
continue
except KeyboardInterrupt:
pass
finally:
# report
report_results(results)
with open(os.getenv("AWORLD_WORKSPACE", "~") + "/results.json", "w", encoding="utf-8") as f:
json.dump(results, f, indent=4, ensure_ascii=False)
@@ -0,0 +1,109 @@
6af95c8f-8cbf-4c12-b02c-f9a23cc1ecb9
198ffd8f-6041-458d-bacc-fe49872cfa43
6583799b-573a-4e95-8b28-4f0397bd45c2
6ace0798-0d0f-4122-8705-78acebf23191
7b10295b-5287-41d7-897c-663fe194c04d
4033181f-1988-476b-bc33-6da0f96d7bd0
7707f3dd-1aa6-42f5-847a-b66f3eaf2ee4
cf0682db-1e31-4a36-913f-aedd501e44d1
07ed8ebc-535a-4c2f-9677-3e434a08f7fd
cd886ddd-2d12-4347-9c7a-64774f66a3d3
d89733a3-7d86-4ed8-b5a3-bf4831b06e3c
cfd773c8-371f-425c-b081-f254f96c0530
c7d40b93-4ee9-4643-8a25-1bb1d5c00f06
fece4d7a-4093-47b8-b58c-68ec591ec4c8
fa7bc619-bf17-4ef7-be80-c0e32144140f
d6f433d5-4203-47d4-990a-28dbf99e6afd
cca708f8-277f-4d9c-920c-26e8065b82ad
3157eb2e-2b7c-48ae-8676-437c64d813d5
c1e3311d-5112-4690-bcdc-7da181657835
55aeb0ce-9170-4959-befb-59f6116887a4
56376d48-f456-4c24-a917-834be04c7608
543eadae-e80a-44ba-b98d-6e8c4040dae8
97f452b6-f144-4224-bfe4-b6a8e0bbce6a
fe8f4748-5d00-4a27-9070-090a0cfdeac4
e8b226bf-2ccd-4748-973d-08dad7aa8253
2f14e114-55de-464f-af9b-82443c6ad5ba
44cb62e5-b0c5-4f12-b96f-21e80ed2062d
e196e50c-09fb-4a3a-81a9-c81ef41b0708
4d8fc0e1-69c0-4fa9-9db1-05935ec5633e
bb505145-7a10-4d23-8af9-39d2b833be3f
2372640f-a119-4960-b624-5ffbe81d6f3f
8aa22ddc-fe80-464c-9c96-8d534fa5c7d5
ffad7d0a-e20d-461f-b5be-07804917235d
04893fc3-34fc-4117-8457-a717ad01a6a9
8770c47d-56ad-42cf-b0a6-7e48d0d19019
0a8063da-eae2-4f84-8603-e71e8714ec23
5bbf523f-b902-4d7d-8e8d-212d00018733
021f63bc-b8cf-484d-a373-ca468fa8ab37
70510d87-5004-4e4a-b078-21abf699dc12
405c2f42-5d01-46d1-97eb-2e96b7cee542
0a3d99c6-d130-49a8-a8bc-6cd5999cc3b4
55caccf7-8c3b-4d40-842a-825789f55dd7
4d317450-948c-4445-b15a-a170813a05e5
86ca62df-b518-48e7-9115-1b0b800e5453
7805912b-c8da-4134-9b54-b590f884352a
e8df5ee8-e804-4ab4-a5b5-f737ffb7c37a
82b89810-1217-4ad8-aa9f-26e7c74ba6e5
68d05bcc-e6cd-43f1-8e9e-e73d507cd02f
b1a37359-2e2e-4521-bb3d-cecdf3ec5419
91d929df-7666-4eed-894c-1bafe5b2e883
171dd6d2-d1d4-439b-8d4e-7507018a816b
e013fc18-abfc-4101-9acc-6273112517f6
220a2b08-ffdc-4665-af4e-025670f5408b
dbb02ff7-f947-491b-9ce2-41e3df16dbb8
e58ea654-994a-44fd-aaff-bf033c715c56
e47bd787-8024-425e-985c-7809bd992fa5
355b827f-fff0-4e0c-9ff0-65dea0609838
a15f8d90-115d-4376-98b3-c866be21ef8f
beeee382-3b29-4be2-b6c7-e54e3543c93f
7245af7c-404e-4d60-9ef4-94ed301e5315
8488761b-6d5a-4be0-86d6-5cf8d6c13efe
cadc593c-ebde-43ec-afa6-a8a4e95c0e38
cbdb17dc-62d9-4463-b648-2eaacfeba4e5
b78952b8-af19-4ffa-9275-cfa41a255b8a
07c3029f-7095-455d-a9e9-cd5e34001b38
01af5e27-e389-4c69-a133-412e16532b67
d543b75f-96e6-4b04-bf2c-60acd0d48246
68ccf11a-bcd3-41e5-a5ee-3e29253449e9
f1ba834a-3bcb-4e55-836c-06cc1e2ccb9f
dd024dd9-8da6-4d4e-aee1-ed0d999035a9
314b3c26-3e8d-4f8c-a9aa-c77275072e0b
27ee81e0-9f94-4b2d-8ab5-3c7033fcefca
8f515e77-1b54-4183-bd5d-3c5f71674931
afd1efe6-03dd-478c-9eb1-e562355ee94e
f30a836f-420f-4738-80af-d2491e072dab
3cc53dbf-1ab9-4d21-a56a-fc0151c10f89
b548db79-fabd-4bbc-81f3-43097159a3a5
bdb2f8ff-83a3-4288-a019-68c7968198f6
d50b8ecb-a8aa-4696-ad84-403ef15e2c8b
b01a3c5f-56ce-4fee-bc15-afd2047bc976
5f2b2e54-5047-4394-81be-198230c3b508
f08e2efd-1885-4437-8c88-90a20558ad1d
bb686773-a454-44d5-b193-cfccfdb90dee
2b2e4e18-a825-4285-aeb1-8d93d2521627
d62cbee6-47c7-4918-825d-3b73b1af7e85
91f2bf12-5280-4efc-b9a7-26e67ca850b4
063800f6-8832-4856-972b-17b877612533
462413b8-642c-4a4d-be27-1a9406100e31
dfa03d6c-402b-43fc-9222-5738f8bdfd0c
7cf5f072-4c37-4ccb-8c4e-a8f66fcdd9e1
52e8ce1c-09bd-4537-8e2d-67d1648779b9
e4b6ca51-c70d-45f7-90dd-cbd1790df1cb
b74b4ce7-4f03-42b5-b60e-62da7ffa282e
4dccff35-158b-4647-8ad0-b7432112077d
0a5a5fd6-97b7-4cce-ad9f-bdaaee042bec
d31b3ec0-bf91-42dd-82fb-43655964505d
ebfd308e-59c1-4389-8812-7ea4cda7f4bc
8c638aa7-81d9-4d06-9ae4-0c0a5540ea5b
aea1ea38-dfd0-41ab-ad79-badc3c69c784
0870ef1f-c712-4e67-8940-8581cc9455dd
4c386922-5c65-4f33-86d5-107a97ad6654
b74aa3b0-16a7-4a79-a8df-0cbdcf1d2eb8
b4395888-4e6d-4d02-b7ee-11dbccdb802d
4b52450e-0c0d-434e-8358-85e80cb2514c
d8434132-f196-4048-82c3-c06facff53c0
f6d29ef1-0e4d-41cb-ac25-e60023b3bd96
e395eb42-8b78-4995-9f8b-25e3a56ab359
70e0a9c6-24bf-48ed-afa1-f0d0eaaa0209
d0e1cd7d-0af2-4ca0-8e9d-7a3fa0687e7f
@@ -0,0 +1,250 @@
import json
import logging
import os
import re
import string
from pathlib import Path
from typing import Any, Dict, List, Optional
from tabulate import tabulate
from aworld.logs.util import Color
logger = logging.getLogger(__name__)
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.info(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.info(tabulate(level_table, headers=headers, tablefmt="grid"))
def setup_logger(logger_name, output_folder_path, file_name="main.log"):
"""
Set up a logger with the given name that writes to the specified file.
Returns a configured logger instance.
"""
if not os.path.exists(output_folder_path):
os.makedirs(output_folder_path)
formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
log_file = os.path.join(output_folder_path, file_name)
# Check if the logger already has handlers to avoid duplicates
logger = logging.getLogger(logger_name)
# Remove existing handlers if any
if logger.hasHandlers():
for handler in logger.handlers[:]:
logger.removeHandler(handler)
# Add file handler
handler = logging.FileHandler(log_file, mode="a", encoding="utf-8")
handler.setLevel(logging.INFO)
handler.setFormatter(formatter)
logger.setLevel(logging.INFO)
logger.addHandler(handler)
return logger
def color_log(logger: logging.Logger, value: str, color: Color | None, level: str | None = None):
# Default to 'info' level if none specified
if level is None:
level = "info"
# Format the message with color
if color is None:
message = f"{Color.black} {value} {Color.reset}"
else:
message = f"{color} {value} {Color.reset}"
# Log according to the specified level
level_lower = level.lower()
if level_lower == "debug":
logger.debug(message)
elif level_lower == "info":
logger.info(message)
elif level_lower == "warning" or level_lower == "warn":
logger.warning(message)
elif level_lower == "error":
logger.error(message)
elif level_lower == "critical":
logger.critical(message)
else:
# Default to info for unknown levels
logger.info(message)