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,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}")