ai-agent-book 精选快照(<2MB 代码与文档,来自 github.com/bojieli/ai-agent-book)
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
This commit is contained in:
@@ -0,0 +1,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}")
|
||||
Reference in New Issue
Block a user