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,102 @@
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from aworld.logs.util import Color
|
||||
from examples.gaia.utils import color_log, setup_logger
|
||||
|
||||
|
||||
class ActionArguments(BaseModel):
|
||||
r"""Protocol: MCP Action Arguments"""
|
||||
|
||||
name: str = Field(description="The name of the action")
|
||||
transport: Literal["stdio", "sse"] = Field(default="stdio", description="The transport of the action")
|
||||
workspace: str | None = Field(
|
||||
default=None,
|
||||
description="The workspace of the action."
|
||||
" If not specified or invalid, the workspace will be read from the environment variable AWORLD_WORKSPACE.",
|
||||
)
|
||||
unittest: bool = Field(default=False, description="Whether to run in unittest mode")
|
||||
|
||||
|
||||
class ActionResponse(BaseModel):
|
||||
r"""Protocol: MCP Action Response"""
|
||||
|
||||
success: bool = Field(default=False, description="Whether the action is successfully executed")
|
||||
message: Any = Field(default=None, description="The execution result of the action")
|
||||
metadata: dict[str, Any] = Field(default={}, description="The metadata of the action")
|
||||
|
||||
|
||||
class ActionCollection:
|
||||
r"""Base class for all ActionCollection."""
|
||||
|
||||
server: FastMCP
|
||||
logger: logging.Logger
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
self.unittest = arguments.unittest
|
||||
self.transport = arguments.transport
|
||||
self.supported_extensions = set()
|
||||
|
||||
self.workspace: Path = self._obtain_valid_workspace(arguments.workspace)
|
||||
|
||||
self.logger: logging.Logger = setup_logger(self.__class__.__name__, self.workspace)
|
||||
|
||||
self.server = FastMCP(arguments.name)
|
||||
for tool_name in self.__class__.__dict__:
|
||||
if tool_name.startswith("mcp_") and callable(getattr(self.__class__, tool_name)):
|
||||
tool = getattr(self, tool_name)
|
||||
self.server.add_tool(tool, description=tool.__doc__)
|
||||
|
||||
def run(self) -> None:
|
||||
if not self.unittest:
|
||||
self.server.run(transport=self.transport)
|
||||
|
||||
def _color_log(self, value: str, color: Color = None, level: str = "info"):
|
||||
return color_log(self.logger, value, color, level=level)
|
||||
|
||||
def _obtain_valid_workspace(self, workspace: str | None = None) -> Path:
|
||||
r"""
|
||||
Obtain a valid workspace path.
|
||||
Priority:
|
||||
1. user defined workspace
|
||||
2. environment variable AWORLD_WORKSPACE
|
||||
3. home directory
|
||||
"""
|
||||
path = Path(workspace) if workspace else os.getenv("AWORLD_WORKSPACE", "~")
|
||||
if path and path.expanduser().is_dir():
|
||||
return path.expanduser().resolve()
|
||||
|
||||
# self._color_log("Invalid workspace path, using home directory instead.", Color.yellow)
|
||||
return Path.home().expanduser().resolve()
|
||||
|
||||
def _validate_file_path(self, file_path: str) -> Path:
|
||||
"""Validate and resolve file path. Rely on the predefined supported_extensions class variable.
|
||||
|
||||
Args:
|
||||
file_path: Path to the document or media file
|
||||
|
||||
Returns:
|
||||
Resolved Path object
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If file doesn't exist
|
||||
ValueError: If file type is not supported
|
||||
"""
|
||||
path = Path(file_path).expanduser()
|
||||
if not path.is_absolute():
|
||||
path = self.workspace / path
|
||||
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"File not found: {path}")
|
||||
|
||||
if path.suffix.lower() not in self.supported_extensions:
|
||||
raise ValueError(
|
||||
f"Unsupported file type: {path.suffix}. Supported types: {', '.join(self.supported_extensions)}"
|
||||
)
|
||||
|
||||
return path
|
||||
@@ -0,0 +1,24 @@
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class DocumentMetadata(BaseModel):
|
||||
"""Metadata extracted from document processing."""
|
||||
|
||||
file_name: str = Field(description="Original file name")
|
||||
file_size: int = Field(description="File size in bytes")
|
||||
file_type: str = Field(description="Document file type/extension")
|
||||
absolute_path: str = Field(description="Absolute path to the document file")
|
||||
page_count: int | None = Field(default=None, description="Number of pages in document")
|
||||
processing_time: float = Field(
|
||||
description="Time taken to process the document in seconds", deprecated=True, exclude=True
|
||||
)
|
||||
extracted_images: list[str] = Field(default_factory=list, description="Paths to extracted image files")
|
||||
extracted_media: list[dict[str, str]] = Field(
|
||||
default_factory=list, description="list of extracted media files with type and path"
|
||||
)
|
||||
output_format: str = Field(description="Format of the extracted content")
|
||||
llm_enhanced: bool = Field(default=False, description="Whether LLM enhancement was used", exclude=True)
|
||||
ocr_applied: bool = Field(default=False, description="Whether OCR was applied", exclude=True)
|
||||
extracted_text_file_path: str | None = Field(
|
||||
default=None, description="Absolute path to the extracted text file (if applicable)"
|
||||
)
|
||||
@@ -0,0 +1,376 @@
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
import chardet
|
||||
import pandas as pd
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import Field
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from aworld.logs.util import Color
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
from examples.gaia.mcp_collections.documents.models import DocumentMetadata
|
||||
|
||||
|
||||
class CSVExtractionCollection(ActionCollection):
|
||||
"""MCP service for CSV document content extraction using pandas.
|
||||
|
||||
Supports extraction from CSV files with various encodings and delimiters.
|
||||
Provides LLM-friendly text output with structured metadata and data analysis.
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
self._media_output_dir = self.workspace / "extracted_media"
|
||||
self._media_output_dir.mkdir(exist_ok=True)
|
||||
|
||||
self.supported_extensions: set = {".csv", ".tsv", ".txt"}
|
||||
|
||||
self._color_log("CSV Extraction Service initialized", Color.green, "debug")
|
||||
self._color_log(f"Media output directory: {self._media_output_dir}", Color.blue, "debug")
|
||||
|
||||
def _detect_encoding(self, file_path: Path) -> str:
|
||||
"""Detect file encoding using chardet.
|
||||
|
||||
Args:
|
||||
file_path: Path to the CSV file
|
||||
|
||||
Returns:
|
||||
Detected encoding string
|
||||
"""
|
||||
try:
|
||||
with open(file_path, "rb") as f:
|
||||
raw_data = f.read(10000) # Read first 10KB for detection
|
||||
result = chardet.detect(raw_data)
|
||||
encoding = result.get("encoding", "utf-8")
|
||||
confidence = result.get("confidence", 0)
|
||||
|
||||
self._color_log(f"Detected encoding: {encoding} (confidence: {confidence:.2f})", Color.blue)
|
||||
return encoding if confidence > 0.7 else "utf-8"
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Encoding detection failed: {e}, using utf-8")
|
||||
return "utf-8"
|
||||
|
||||
def _detect_delimiter(self, file_path: Path, encoding: str) -> str:
|
||||
"""Detect CSV delimiter by analyzing the first few lines.
|
||||
|
||||
Args:
|
||||
file_path: Path to the CSV file
|
||||
encoding: File encoding
|
||||
|
||||
Returns:
|
||||
Detected delimiter character
|
||||
"""
|
||||
try:
|
||||
with open(file_path, "r", encoding=encoding) as f:
|
||||
sample = f.read(1024) # Read first 1KB
|
||||
|
||||
# Common delimiters to test
|
||||
delimiters = [",", ";", "\t", "|", ":"]
|
||||
delimiter_counts = {}
|
||||
|
||||
for delimiter in delimiters:
|
||||
count = sample.count(delimiter)
|
||||
if count > 0:
|
||||
delimiter_counts[delimiter] = count
|
||||
|
||||
if delimiter_counts:
|
||||
detected_delimiter = max(delimiter_counts, key=delimiter_counts.get)
|
||||
self._color_log(f"Detected delimiter: '{detected_delimiter}'", Color.blue)
|
||||
return detected_delimiter
|
||||
else:
|
||||
return ","
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Delimiter detection failed: {e}, using comma")
|
||||
return ","
|
||||
|
||||
def _extract_csv_content(
|
||||
self, file_path: Path, max_rows: int | None = None, encoding: str | None = None, delimiter: str | None = None
|
||||
) -> dict[str, Any]:
|
||||
"""Extract content from CSV file using pandas.
|
||||
|
||||
Args:
|
||||
file_path: Path to the CSV file
|
||||
max_rows: Maximum number of rows to read
|
||||
encoding: File encoding (auto-detected if None)
|
||||
delimiter: CSV delimiter (auto-detected if None)
|
||||
|
||||
Returns:
|
||||
Dictionary containing extracted content and metadata
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Auto-detect encoding and delimiter if not provided
|
||||
if encoding is None:
|
||||
encoding = self._detect_encoding(file_path)
|
||||
if delimiter is None:
|
||||
delimiter = self._detect_delimiter(file_path, encoding)
|
||||
|
||||
try:
|
||||
# Read CSV with pandas
|
||||
df = pd.read_csv(file_path, encoding=encoding, delimiter=delimiter, nrows=max_rows, low_memory=False)
|
||||
|
||||
# Get full file info for metadata
|
||||
full_df_info = pd.read_csv(
|
||||
file_path,
|
||||
encoding=encoding,
|
||||
delimiter=delimiter,
|
||||
nrows=0, # Just get headers and shape info
|
||||
)
|
||||
|
||||
# Count total rows efficiently
|
||||
total_rows = sum(1 for _ in open(file_path, "r", encoding=encoding)) - 1 # Subtract header
|
||||
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
return {
|
||||
"dataframe": df,
|
||||
"total_rows": total_rows,
|
||||
"total_columns": len(full_df_info.columns),
|
||||
"columns": list(df.columns),
|
||||
"encoding": encoding,
|
||||
"delimiter": delimiter,
|
||||
"processing_time": processing_time,
|
||||
"data_types": df.dtypes.to_dict(),
|
||||
"memory_usage": df.memory_usage(deep=True).sum(),
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to read CSV file: {e}")
|
||||
raise
|
||||
|
||||
def _format_content_for_llm(self, df: pd.DataFrame, output_format: str, include_stats: bool = True) -> str:
|
||||
"""Format extracted CSV content to be LLM-friendly.
|
||||
|
||||
Args:
|
||||
df: Pandas DataFrame with CSV data
|
||||
output_format: Desired output format
|
||||
include_stats: Whether to include statistical summary
|
||||
|
||||
Returns:
|
||||
Formatted content string
|
||||
"""
|
||||
if output_format.lower() == "markdown":
|
||||
# Convert to markdown table
|
||||
content = df.to_markdown(index=False, tablefmt="github")
|
||||
|
||||
if include_stats:
|
||||
# Add statistical summary
|
||||
stats_content = "\n\n## Data Summary\n\n"
|
||||
stats_content += f"- **Rows**: {len(df)}\n"
|
||||
stats_content += f"- **Columns**: {len(df.columns)}\n"
|
||||
stats_content += f"- **Column Names**: {', '.join(df.columns)}\n\n"
|
||||
|
||||
# Add data types info
|
||||
stats_content += "### Column Data Types\n\n"
|
||||
for col, dtype in df.dtypes.items():
|
||||
stats_content += f"- **{col}**: {dtype}\n"
|
||||
|
||||
# Add basic statistics for numeric columns
|
||||
numeric_cols = df.select_dtypes(include=["number"]).columns
|
||||
if len(numeric_cols) > 0:
|
||||
stats_content += "\n### Numeric Column Statistics\n\n"
|
||||
stats_df = df[numeric_cols].describe()
|
||||
stats_content += stats_df.to_markdown(tablefmt="github")
|
||||
|
||||
content += stats_content
|
||||
|
||||
elif output_format.lower() == "json":
|
||||
# Convert to JSON with metadata
|
||||
data_dict = {
|
||||
"data": df.to_dict(orient="records"),
|
||||
"metadata": {
|
||||
"rows": len(df),
|
||||
"columns": len(df.columns),
|
||||
"column_names": list(df.columns),
|
||||
"data_types": {col: str(dtype) for col, dtype in df.dtypes.items()},
|
||||
},
|
||||
}
|
||||
if include_stats:
|
||||
numeric_cols = df.select_dtypes(include=["number"]).columns
|
||||
if len(numeric_cols) > 0:
|
||||
data_dict["statistics"] = df[numeric_cols].describe().to_dict()
|
||||
|
||||
content = json.dumps(data_dict, indent=2, default=str)
|
||||
|
||||
elif output_format.lower() == "html":
|
||||
# Convert to HTML table
|
||||
content = df.to_html(index=False, classes="table table-striped")
|
||||
|
||||
else:
|
||||
# Plain text format
|
||||
content = df.to_string(index=False)
|
||||
|
||||
return content
|
||||
|
||||
def mcp_extract_csv_content(
|
||||
self,
|
||||
file_path: str = Field(description="Path to the CSV document file to extract content from"),
|
||||
output_format: Literal["markdown", "json", "html", "text"] = Field(
|
||||
default="markdown", description="Output format: 'markdown', 'json', 'html', or 'text'"
|
||||
),
|
||||
max_rows: int | None = Field(default=None, description="Maximum number of rows to read (None for all rows)"),
|
||||
include_statistics: bool = Field(default=True, description="Whether to include statistical summary in output"),
|
||||
generate_visualizations: bool = Field(
|
||||
default=False, description="Whether to generate and save data visualizations"
|
||||
),
|
||||
encoding: str | None = Field(default=None, description="File encoding (auto-detected if None)"),
|
||||
delimiter: str | None = Field(default=None, description="CSV delimiter (auto-detected if None)"),
|
||||
) -> ActionResponse:
|
||||
"""Extract content from CSV documents using pandas.
|
||||
|
||||
This tool provides comprehensive CSV document content extraction with support for:
|
||||
- CSV, TSV, and delimited text files
|
||||
- Automatic encoding and delimiter detection
|
||||
- Statistical analysis and data profiling
|
||||
- Multiple output formats (Markdown, JSON, HTML, Text)
|
||||
- Optional data visualizations
|
||||
- Memory-efficient processing for large files
|
||||
|
||||
Args:
|
||||
file_path: Path to the CSV file
|
||||
output_format: Desired output format
|
||||
max_rows: Maximum rows to process (None for all)
|
||||
include_statistics: Include statistical summary
|
||||
generate_visualizations: Generate data visualizations
|
||||
encoding: File encoding (auto-detected if None)
|
||||
delimiter: CSV delimiter (auto-detected if None)
|
||||
|
||||
Returns:
|
||||
ActionResponse with extracted content, metadata, and optional visualizations
|
||||
"""
|
||||
try:
|
||||
# Handle FieldInfo objects from pydantic
|
||||
if isinstance(file_path, FieldInfo):
|
||||
file_path = file_path.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
if isinstance(max_rows, FieldInfo):
|
||||
max_rows = max_rows.default
|
||||
if isinstance(include_statistics, FieldInfo):
|
||||
include_statistics = include_statistics.default
|
||||
if isinstance(generate_visualizations, FieldInfo):
|
||||
generate_visualizations = generate_visualizations.default
|
||||
if isinstance(encoding, FieldInfo):
|
||||
encoding = encoding.default
|
||||
if isinstance(delimiter, FieldInfo):
|
||||
delimiter = delimiter.default
|
||||
|
||||
# Validate input file
|
||||
file_path: Path = self._validate_file_path(file_path)
|
||||
self._color_log(f"Processing CSV file: {file_path.name}", Color.cyan)
|
||||
|
||||
# Extract CSV content
|
||||
extraction_result = self._extract_csv_content(
|
||||
file_path, max_rows=max_rows, encoding=encoding, delimiter=delimiter
|
||||
)
|
||||
|
||||
df: pd.DataFrame = extraction_result["dataframe"]
|
||||
|
||||
# Format content for LLM consumption
|
||||
formatted_content = self._format_content_for_llm(df, output_format, include_stats=include_statistics)
|
||||
|
||||
# Prepare metadata
|
||||
file_stats = file_path.stat()
|
||||
document_metadata = DocumentMetadata(
|
||||
file_name=file_path.name,
|
||||
file_size=file_stats.st_size,
|
||||
file_type=file_path.suffix.lower(),
|
||||
absolute_path=str(file_path.absolute()),
|
||||
page_count=None, # Not applicable for CSV
|
||||
processing_time=extraction_result["processing_time"],
|
||||
extracted_images=[], # CSV files don't contain images
|
||||
extracted_media=None,
|
||||
output_format=output_format,
|
||||
llm_enhanced=False,
|
||||
ocr_applied=False,
|
||||
)
|
||||
|
||||
# Add CSV-specific metadata
|
||||
csv_metadata = {
|
||||
"total_rows": extraction_result["total_rows"],
|
||||
"total_columns": extraction_result["total_columns"],
|
||||
"rows_processed": len(df),
|
||||
"columns_processed": len(df.columns),
|
||||
"column_names": extraction_result["columns"],
|
||||
"data_types": {k: str(v) for k, v in extraction_result["data_types"].items()},
|
||||
"encoding": extraction_result["encoding"],
|
||||
"delimiter": extraction_result["delimiter"],
|
||||
"memory_usage_bytes": int(extraction_result["memory_usage"]),
|
||||
}
|
||||
|
||||
# Merge metadata
|
||||
final_metadata = {**document_metadata.model_dump(), **csv_metadata}
|
||||
|
||||
self._color_log(
|
||||
f"Successfully extracted CSV content from {file_path.name} "
|
||||
f"({extraction_result['total_rows']} rows, {extraction_result['total_columns']} columns",
|
||||
Color.green,
|
||||
)
|
||||
|
||||
return ActionResponse(success=True, message=formatted_content, metadata=final_metadata)
|
||||
|
||||
except FileNotFoundError as e:
|
||||
self.logger.error(f"File not found: {str(e)}")
|
||||
return ActionResponse(
|
||||
success=False, message=f"File not found: {str(e)}", metadata={"error_type": "file_not_found"}
|
||||
)
|
||||
except ValueError as e:
|
||||
self.logger.error(f"Invalid input: {str(e)}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}",
|
||||
metadata={"error_type": "invalid_input"},
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"CSV extraction failed: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"CSV extraction failed: {str(e)}",
|
||||
metadata={"error_type": "extraction_error"},
|
||||
)
|
||||
|
||||
def mcp_list_supported_formats(self) -> ActionResponse:
|
||||
"""List all supported CSV formats for extraction.
|
||||
|
||||
Returns:
|
||||
ActionResponse with list of supported file formats and their descriptions
|
||||
"""
|
||||
supported_formats = {
|
||||
"CSV": "Comma-Separated Values files (.csv)",
|
||||
"TSV": "Tab-Separated Values files (.tsv)",
|
||||
"TXT": "Delimited text files (.txt)",
|
||||
}
|
||||
|
||||
format_list = "\n".join(
|
||||
[f"**{format_name}**: {description}" for format_name, description in supported_formats.items()]
|
||||
)
|
||||
|
||||
return ActionResponse(
|
||||
success=True,
|
||||
message=f"Supported CSV formats:\n\n{format_list}",
|
||||
metadata={"supported_formats": list(supported_formats.keys()), "total_formats": len(supported_formats)},
|
||||
)
|
||||
|
||||
|
||||
# Example usage and entry point
|
||||
if __name__ == "__main__":
|
||||
load_dotenv()
|
||||
|
||||
# Default arguments for testing
|
||||
args = ActionArguments(
|
||||
name="csv_extraction_service",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
|
||||
)
|
||||
|
||||
# Initialize and run the CSV extraction service
|
||||
try:
|
||||
service = CSVExtractionCollection(args)
|
||||
service.run()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}: {traceback.format_exc()}")
|
||||
@@ -0,0 +1,617 @@
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
from docx import Document
|
||||
from docx.document import Document as DocumentType
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import Field
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from aworld.logs.util import Color
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
from examples.gaia.mcp_collections.documents.models import DocumentMetadata
|
||||
|
||||
|
||||
class DOCXExtractionCollection(ActionCollection):
|
||||
"""MCP service for DOCX/DOC document content extraction using python-docx.
|
||||
|
||||
Supports extraction from DOCX and DOC files with comprehensive content parsing.
|
||||
Provides LLM-friendly text output with structured metadata and media file handling.
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
self._media_output_dir = self.workspace / "extracted_media"
|
||||
self._media_output_dir.mkdir(exist_ok=True)
|
||||
|
||||
self.supported_extensions = {".docx", ".doc"}
|
||||
|
||||
self._color_log("DOCX Extraction Service initialized", Color.green, "debug")
|
||||
self._color_log(f"Media output directory: {self._media_output_dir}", Color.blue, "debug")
|
||||
|
||||
def _extract_images_from_docx(self, file_path: Path, file_stem: str) -> list[dict[str, str]]:
|
||||
"""Extract embedded images from DOCX file.
|
||||
|
||||
Args:
|
||||
file_path: Path to the DOCX file
|
||||
file_stem: Base name for saving files
|
||||
|
||||
Returns:
|
||||
List of dictionaries containing image file paths and metadata
|
||||
"""
|
||||
saved_media = []
|
||||
|
||||
try:
|
||||
# DOCX files are ZIP archives
|
||||
with zipfile.ZipFile(file_path, "r") as docx_zip:
|
||||
# Look for media files in the word/media/ directory
|
||||
media_files = [f for f in docx_zip.namelist() if f.startswith("word/media/")]
|
||||
|
||||
for idx, media_file in enumerate(media_files):
|
||||
try:
|
||||
# Extract file extension and create appropriate filename
|
||||
original_name = Path(media_file).name
|
||||
file_extension = Path(media_file).suffix
|
||||
|
||||
# Generate unique filename
|
||||
media_filename = f"{file_stem}_media_{idx}_{original_name}"
|
||||
media_path = self._media_output_dir / media_filename
|
||||
|
||||
# Extract and save the media file
|
||||
with docx_zip.open(media_file) as source:
|
||||
with open(media_path, "wb") as target:
|
||||
target.write(source.read())
|
||||
|
||||
# Determine media type based on extension
|
||||
media_type = "image"
|
||||
if file_extension.lower() in [".mp3", ".wav", ".m4a", ".ogg"]:
|
||||
media_type = "audio"
|
||||
elif file_extension.lower() in [".mp4", ".avi", ".mov", ".wmv"]:
|
||||
media_type = "video"
|
||||
elif file_extension.lower() in [".png", ".jpg", ".jpeg", ".gif", ".bmp", ".tiff"]:
|
||||
media_type = "image"
|
||||
else:
|
||||
media_type = "other"
|
||||
|
||||
saved_media.append(
|
||||
{
|
||||
"type": media_type,
|
||||
"path": str(media_path),
|
||||
"filename": media_filename,
|
||||
"original_name": original_name,
|
||||
"size_bytes": media_path.stat().st_size,
|
||||
}
|
||||
)
|
||||
|
||||
self._color_log(f"Extracted {media_type}: {media_filename}", Color.blue)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to extract media file {media_file}: {e}")
|
||||
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Could not extract media from DOCX: {e}")
|
||||
|
||||
return saved_media
|
||||
|
||||
def _extract_document_structure(self, doc: DocumentType) -> dict[str, Any]:
|
||||
"""Extract document structure and metadata.
|
||||
|
||||
Args:
|
||||
doc: python-docx Document object
|
||||
|
||||
Returns:
|
||||
Dictionary containing document structure information
|
||||
"""
|
||||
structure = {
|
||||
"paragraphs_count": len(doc.paragraphs),
|
||||
"tables_count": len(doc.tables),
|
||||
"sections_count": len(doc.sections),
|
||||
"styles_used": [],
|
||||
"has_headers_footers": False,
|
||||
"page_count": None, # Not directly available in python-docx
|
||||
}
|
||||
|
||||
# Collect unique styles used in the document
|
||||
styles_used = set()
|
||||
for paragraph in doc.paragraphs:
|
||||
if paragraph.style and paragraph.style.name:
|
||||
styles_used.add(paragraph.style.name)
|
||||
|
||||
structure["styles_used"] = list(styles_used)
|
||||
|
||||
# Check for headers and footers
|
||||
for section in doc.sections:
|
||||
if (section.header.paragraphs and any(p.text.strip() for p in section.header.paragraphs)) or (
|
||||
section.footer.paragraphs and any(p.text.strip() for p in section.footer.paragraphs)
|
||||
):
|
||||
structure["has_headers_footers"] = True
|
||||
break
|
||||
|
||||
return structure
|
||||
|
||||
def _extract_tables_content(self, doc: DocumentType) -> list[dict[str, Any]]:
|
||||
"""Extract content from all tables in the document.
|
||||
|
||||
Args:
|
||||
doc: python-docx Document object
|
||||
|
||||
Returns:
|
||||
List of dictionaries containing table data
|
||||
"""
|
||||
tables_data = []
|
||||
|
||||
for table_idx, table in enumerate(doc.tables):
|
||||
table_data = {
|
||||
"table_index": table_idx,
|
||||
"rows_count": len(table.rows),
|
||||
"columns_count": len(table.columns) if table.rows else 0,
|
||||
"data": [],
|
||||
}
|
||||
|
||||
# Extract table content
|
||||
for _, row in enumerate(table.rows):
|
||||
row_data = []
|
||||
for cell in row.cells:
|
||||
cell_text = cell.text.strip()
|
||||
row_data.append(cell_text)
|
||||
table_data["data"].append(row_data)
|
||||
|
||||
tables_data.append(table_data)
|
||||
|
||||
return tables_data
|
||||
|
||||
def _convert_doc_to_docx(self, doc_path: Path) -> Path:
|
||||
"""Convert .doc file to .docx using LibreOffice.
|
||||
|
||||
Args:
|
||||
doc_path: Path to the .doc file
|
||||
|
||||
Returns:
|
||||
Path to the converted .docx file
|
||||
"""
|
||||
import subprocess
|
||||
|
||||
output_dir = self._media_output_dir
|
||||
docx_path = output_dir / f"{doc_path.stem}.docx"
|
||||
|
||||
# Check if already converted
|
||||
if docx_path.exists():
|
||||
self._color_log(f"Using existing converted file: {docx_path.name}", Color.blue)
|
||||
return docx_path
|
||||
|
||||
self._color_log(f"Converting .doc to .docx: {doc_path.name}", Color.yellow)
|
||||
|
||||
cmd = ["libreoffice", "--headless", "--convert-to", "docx", "--outdir", str(output_dir), str(doc_path)]
|
||||
|
||||
try:
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, check=True, timeout=60) # pylint: disable=W0612
|
||||
if docx_path.exists():
|
||||
self._color_log(f"Conversion successful: {docx_path.name}", Color.green)
|
||||
return docx_path
|
||||
else:
|
||||
raise RuntimeError("Conversion completed but output file not found")
|
||||
|
||||
except subprocess.CalledProcessError as e:
|
||||
self.logger.error(f"LibreOffice conversion failed: {e.stderr}")
|
||||
raise RuntimeError(f"Failed to convert .doc to .docx: {e.stderr}") from e
|
||||
except subprocess.TimeoutExpired as e:
|
||||
self.logger.error("LibreOffice conversion timed out")
|
||||
raise RuntimeError("Conversion timed out after 60 seconds") from e
|
||||
|
||||
# Modify the _extract_content_from_docx method to handle .doc files
|
||||
def _extract_content_from_docx(
|
||||
self, file_path: Path, extract_tables: bool = True, extract_headers_footers: bool = True
|
||||
) -> dict[str, Any]:
|
||||
"""Extract content from DOCX file using python-docx.
|
||||
|
||||
Args:
|
||||
file_path: Path to the DOCX/DOC file
|
||||
extract_tables: Whether to extract table content
|
||||
extract_headers_footers: Whether to extract headers and footers
|
||||
|
||||
Returns:
|
||||
Dictionary containing extracted content and metadata
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Convert .doc to .docx if needed
|
||||
if file_path.suffix.lower() == ".doc":
|
||||
file_path = self._convert_doc_to_docx(file_path)
|
||||
|
||||
try:
|
||||
# Load the document
|
||||
doc = Document(str(file_path))
|
||||
|
||||
# Extract main text content
|
||||
paragraphs = []
|
||||
for paragraph in doc.paragraphs:
|
||||
if paragraph.text.strip(): # Skip empty paragraphs
|
||||
para_data = {"text": paragraph.text, "style": paragraph.style.name if paragraph.style else "Normal"}
|
||||
paragraphs.append(para_data)
|
||||
|
||||
# Extract document structure
|
||||
structure = self._extract_document_structure(doc)
|
||||
|
||||
# Extract tables if requested
|
||||
tables = []
|
||||
if extract_tables:
|
||||
tables = self._extract_tables_content(doc)
|
||||
|
||||
# Extract headers and footers if requested
|
||||
headers_footers = []
|
||||
if extract_headers_footers:
|
||||
for section_idx, section in enumerate(doc.sections):
|
||||
# Extract header content
|
||||
header_text = []
|
||||
for para in section.header.paragraphs:
|
||||
if para.text.strip():
|
||||
header_text.append(para.text)
|
||||
|
||||
# Extract footer content
|
||||
footer_text = []
|
||||
for para in section.footer.paragraphs:
|
||||
if para.text.strip():
|
||||
footer_text.append(para.text)
|
||||
|
||||
if header_text or footer_text:
|
||||
headers_footers.append(
|
||||
{"section_index": section_idx, "header": header_text, "footer": footer_text}
|
||||
)
|
||||
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
return {
|
||||
"paragraphs": paragraphs,
|
||||
"tables": tables,
|
||||
"headers_footers": headers_footers,
|
||||
"structure": structure,
|
||||
"processing_time": processing_time,
|
||||
"word_count": sum(len(p["text"].split()) for p in paragraphs),
|
||||
"character_count": sum(len(p["text"]) for p in paragraphs),
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to extract content from DOCX: {e}")
|
||||
raise
|
||||
|
||||
def _format_content_for_llm(
|
||||
self, extraction_result: dict[str, Any], output_format: str, include_structure: bool = True
|
||||
) -> str:
|
||||
"""Format extracted DOCX content to be LLM-friendly.
|
||||
|
||||
Args:
|
||||
extraction_result: Dictionary containing extracted content
|
||||
output_format: Desired output format
|
||||
include_structure: Whether to include document structure information
|
||||
|
||||
Returns:
|
||||
Formatted content string
|
||||
"""
|
||||
if output_format.lower() == "markdown":
|
||||
content_parts = []
|
||||
|
||||
# Add document structure info if requested
|
||||
if include_structure:
|
||||
structure = extraction_result["structure"]
|
||||
content_parts.append("# Document Structure\n")
|
||||
content_parts.append(f"- **Paragraphs**: {structure['paragraphs_count']}\n")
|
||||
content_parts.append(f"- **Tables**: {structure['tables_count']}\n")
|
||||
content_parts.append(f"- **Sections**: {structure['sections_count']}\n")
|
||||
content_parts.append(f"- **Word Count**: {extraction_result['word_count']}\n")
|
||||
content_parts.append(f"- **Character Count**: {extraction_result['character_count']}\n")
|
||||
if structure["styles_used"]:
|
||||
content_parts.append(f"- **Styles Used**: {', '.join(structure['styles_used'])}\n")
|
||||
content_parts.append("\n---\n\n")
|
||||
|
||||
# Add main content
|
||||
content_parts.append("# Document Content\n\n")
|
||||
|
||||
# Add paragraphs
|
||||
for para in extraction_result["paragraphs"]:
|
||||
# Format based on style
|
||||
text = para["text"]
|
||||
style = para["style"]
|
||||
|
||||
if "Heading" in style:
|
||||
# Convert heading styles to markdown headers
|
||||
if "Heading 1" in style:
|
||||
content_parts.append(f"# {text}\n\n")
|
||||
elif "Heading 2" in style:
|
||||
content_parts.append(f"## {text}\n\n")
|
||||
elif "Heading 3" in style:
|
||||
content_parts.append(f"### {text}\n\n")
|
||||
else:
|
||||
content_parts.append(f"#### {text}\n\n")
|
||||
else:
|
||||
content_parts.append(f"{text}\n\n")
|
||||
|
||||
# Add tables
|
||||
if extraction_result["tables"]:
|
||||
content_parts.append("\n## Tables\n\n")
|
||||
for table_idx, table in enumerate(extraction_result["tables"]):
|
||||
content_parts.append(f"### Table {table_idx + 1}\n\n")
|
||||
|
||||
if table["data"]:
|
||||
# Create markdown table
|
||||
headers = table["data"][0] if table["data"] else []
|
||||
if headers:
|
||||
content_parts.append("| " + " | ".join(headers) + " |\n")
|
||||
content_parts.append("|" + "---|" * len(headers) + "\n")
|
||||
|
||||
for row in table["data"][1:]:
|
||||
content_parts.append("| " + " | ".join(row) + " |\n")
|
||||
content_parts.append("\n")
|
||||
|
||||
# Add headers and footers
|
||||
if extraction_result["headers_footers"]:
|
||||
content_parts.append("\n## Headers and Footers\n\n")
|
||||
for hf in extraction_result["headers_footers"]:
|
||||
if hf["header"]:
|
||||
content_parts.append(f"**Header (Section {hf['section_index'] + 1}):**\n")
|
||||
for header_line in hf["header"]:
|
||||
content_parts.append(f"{header_line}\n")
|
||||
content_parts.append("\n")
|
||||
|
||||
if hf["footer"]:
|
||||
content_parts.append(f"**Footer (Section {hf['section_index'] + 1}):**\n")
|
||||
for footer_line in hf["footer"]:
|
||||
content_parts.append(f"{footer_line}\n")
|
||||
content_parts.append("\n")
|
||||
|
||||
return "".join(content_parts)
|
||||
|
||||
elif output_format.lower() == "json":
|
||||
# Return structured JSON
|
||||
return json.dumps(extraction_result, indent=2, ensure_ascii=False)
|
||||
|
||||
elif output_format.lower() == "html":
|
||||
# Convert to HTML
|
||||
html_parts = ["<html><body>"]
|
||||
|
||||
if include_structure:
|
||||
html_parts.append("<h1>Document Structure</h1>")
|
||||
structure = extraction_result["structure"]
|
||||
html_parts.append(f"<p><strong>Paragraphs:</strong> {structure['paragraphs_count']}</p>")
|
||||
html_parts.append(f"<p><strong>Tables:</strong> {structure['tables_count']}</p>")
|
||||
html_parts.append(f"<p><strong>Word Count:</strong> {extraction_result['word_count']}</p>")
|
||||
html_parts.append("<hr>")
|
||||
|
||||
html_parts.append("<h1>Document Content</h1>")
|
||||
|
||||
# Add paragraphs
|
||||
for para in extraction_result["paragraphs"]:
|
||||
text = para["text"].replace("&", "&").replace("<", "<").replace(">", ">")
|
||||
style = para["style"]
|
||||
|
||||
if "Heading" in style:
|
||||
if "Heading 1" in style:
|
||||
html_parts.append(f"<h1>{text}</h1>")
|
||||
elif "Heading 2" in style:
|
||||
html_parts.append(f"<h2>{text}</h2>")
|
||||
elif "Heading 3" in style:
|
||||
html_parts.append(f"<h3>{text}</h3>")
|
||||
else:
|
||||
html_parts.append(f"<h4>{text}</h4>")
|
||||
else:
|
||||
html_parts.append(f"<p>{text}</p>")
|
||||
|
||||
# Add tables
|
||||
if extraction_result["tables"]:
|
||||
html_parts.append("<h2>Tables</h2>")
|
||||
for table_idx, table in enumerate(extraction_result["tables"]):
|
||||
html_parts.append(f"<h3>Table {table_idx + 1}</h3>")
|
||||
html_parts.append("<table border='1'>")
|
||||
|
||||
for row_idx, row in enumerate(table["data"]):
|
||||
html_parts.append("<tr>")
|
||||
tag = "th" if row_idx == 0 else "td"
|
||||
for cell in row:
|
||||
cell_text = cell.replace("&", "&").replace("<", "<").replace(">", ">")
|
||||
html_parts.append(f"<{tag}>{cell_text}</{tag}>")
|
||||
html_parts.append("</tr>")
|
||||
|
||||
html_parts.append("</table>")
|
||||
|
||||
html_parts.append("</body></html>")
|
||||
return "".join(html_parts)
|
||||
|
||||
else:
|
||||
# Plain text format
|
||||
text_parts = []
|
||||
|
||||
# Add main content
|
||||
for para in extraction_result["paragraphs"]:
|
||||
text_parts.append(para["text"])
|
||||
|
||||
# Add tables as plain text
|
||||
if extraction_result["tables"]:
|
||||
text_parts.append("\n\n=== TABLES ===\n")
|
||||
for table_idx, table in enumerate(extraction_result["tables"]):
|
||||
text_parts.append(f"\nTable {table_idx + 1}:\n")
|
||||
for row in table["data"]:
|
||||
text_parts.append("\t".join(row) + "\n")
|
||||
|
||||
return "\n\n".join(text_parts)
|
||||
|
||||
def mcp_extract_docx_content(
|
||||
self,
|
||||
file_path: str = Field(description="Path to the DOCX/DOC document file to extract content from"),
|
||||
output_format: Literal["markdown", "json", "html", "text"] = Field(
|
||||
default="markdown", description="Output format: 'markdown', 'json', 'html', or 'text'"
|
||||
),
|
||||
extract_images: bool = Field(default=True, description="Whether to extract and save embedded images and media"),
|
||||
extract_tables: bool = Field(default=True, description="Whether to extract table content"),
|
||||
extract_headers_footers: bool = Field(default=True, description="Whether to extract headers and footers"),
|
||||
include_structure: bool = Field(
|
||||
default=True, description="Whether to include document structure information in output"
|
||||
),
|
||||
) -> ActionResponse:
|
||||
"""Extract content from DOCX/DOC documents using python-docx.
|
||||
|
||||
This tool provides comprehensive DOCX/DOC document content extraction with support for:
|
||||
- DOCX and DOC files
|
||||
- Text extraction with style preservation
|
||||
- Table extraction and formatting
|
||||
- Headers and footers extraction
|
||||
- Embedded media extraction (images, audio, etc.)
|
||||
- Multiple output formats (Markdown, JSON, HTML, Text)
|
||||
- Document structure analysis
|
||||
|
||||
Args:
|
||||
file_path: Path to the DOCX/DOC file
|
||||
output_format: Desired output format
|
||||
extract_images: Extract embedded media files
|
||||
extract_tables: Extract table content
|
||||
extract_headers_footers: Extract headers and footers
|
||||
include_structure: Include document structure info
|
||||
|
||||
Returns:
|
||||
ActionResponse with extracted content, metadata, and media file paths
|
||||
"""
|
||||
try:
|
||||
# Handle FieldInfo objects from pydantic
|
||||
if isinstance(file_path, FieldInfo):
|
||||
file_path = file_path.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
if isinstance(extract_images, FieldInfo):
|
||||
extract_images = extract_images.default
|
||||
if isinstance(extract_tables, FieldInfo):
|
||||
extract_tables = extract_tables.default
|
||||
if isinstance(extract_headers_footers, FieldInfo):
|
||||
extract_headers_footers = extract_headers_footers.default
|
||||
if isinstance(include_structure, FieldInfo):
|
||||
include_structure = include_structure.default
|
||||
|
||||
# Validate input file
|
||||
file_path: Path = self._validate_file_path(file_path)
|
||||
self._color_log(f"Processing DOCX document: {file_path.name}", Color.cyan)
|
||||
|
||||
# Extract embedded media if requested
|
||||
saved_media = []
|
||||
if extract_images and file_path.suffix.lower() == ".docx":
|
||||
saved_media = self._extract_images_from_docx(file_path, file_path.stem)
|
||||
|
||||
# Extract document content
|
||||
extraction_result = self._extract_content_from_docx(
|
||||
file_path, extract_tables=extract_tables, extract_headers_footers=extract_headers_footers
|
||||
)
|
||||
|
||||
# Format content for LLM consumption
|
||||
formatted_content = self._format_content_for_llm(
|
||||
extraction_result, output_format, include_structure=include_structure
|
||||
)
|
||||
|
||||
# Prepare metadata
|
||||
file_stats = file_path.stat()
|
||||
document_metadata = DocumentMetadata(
|
||||
file_name=file_path.name,
|
||||
file_size=file_stats.st_size,
|
||||
file_type=file_path.suffix.lower(),
|
||||
absolute_path=str(file_path.absolute()),
|
||||
page_count=None, # Not directly available for DOCX
|
||||
processing_time=extraction_result["processing_time"],
|
||||
extracted_images=[media["path"] for media in saved_media if media["type"] == "image"],
|
||||
extracted_media=saved_media,
|
||||
output_format=output_format,
|
||||
llm_enhanced=False,
|
||||
ocr_applied=False,
|
||||
)
|
||||
|
||||
# Add DOCX-specific metadata
|
||||
docx_metadata = {
|
||||
"paragraphs_count": extraction_result["structure"]["paragraphs_count"],
|
||||
"tables_count": extraction_result["structure"]["tables_count"],
|
||||
"sections_count": extraction_result["structure"]["sections_count"],
|
||||
"word_count": extraction_result["word_count"],
|
||||
"character_count": extraction_result["character_count"],
|
||||
"styles_used": extraction_result["structure"]["styles_used"],
|
||||
"has_headers_footers": extraction_result["structure"]["has_headers_footers"],
|
||||
"has_embedded_media": len(saved_media) > 0,
|
||||
"media_files_count": len(saved_media),
|
||||
}
|
||||
|
||||
# Merge metadata
|
||||
final_metadata = {**document_metadata.model_dump(), **docx_metadata}
|
||||
|
||||
self._color_log(
|
||||
f"Successfully extracted DOCX content from {file_path.name} "
|
||||
f"({extraction_result['word_count']} words, {extraction_result['structure']['tables_count']} tables, "
|
||||
f"{len(saved_media)} media files)",
|
||||
Color.green,
|
||||
)
|
||||
|
||||
return ActionResponse(success=True, message=formatted_content, metadata=final_metadata)
|
||||
|
||||
except FileNotFoundError as e:
|
||||
self.logger.error(f"File not found: {str(e)}")
|
||||
return ActionResponse(
|
||||
success=False, message=f"File not found: {str(e)}", metadata={"error_type": "file_not_found"}
|
||||
)
|
||||
except ValueError as e:
|
||||
self.logger.error(f"Invalid input: {str(e)}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}",
|
||||
metadata={"error_type": "invalid_input"},
|
||||
)
|
||||
except ImportError as e:
|
||||
self.logger.error(f"Missing dependency: {str(e)}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Missing dependency: {str(e)}. Please install python-docx: pip install python-docx",
|
||||
metadata={"error_type": "missing_dependency"},
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"DOCX extraction failed: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"DOCX extraction failed: {str(e)}",
|
||||
metadata={"error_type": "extraction_error"},
|
||||
)
|
||||
|
||||
def mcp_list_supported_formats(self) -> ActionResponse:
|
||||
"""List all supported document formats for extraction.
|
||||
|
||||
Returns:
|
||||
ActionResponse with list of supported file formats and their descriptions
|
||||
"""
|
||||
supported_formats = {
|
||||
"DOCX": "Microsoft Word Open XML Document (.docx)",
|
||||
"DOC": "Microsoft Word Document (.doc) - limited support",
|
||||
}
|
||||
|
||||
format_list = "\n".join(
|
||||
[f"**{format_name}**: {description}" for format_name, description in supported_formats.items()]
|
||||
)
|
||||
|
||||
return ActionResponse(
|
||||
success=True,
|
||||
message=f"Supported document formats:\n\n{format_list}\n\n"
|
||||
"**Note**: DOC format support is limited. For best results, convert to DOCX format.",
|
||||
metadata={"supported_formats": list(supported_formats.keys()), "total_formats": len(supported_formats)},
|
||||
)
|
||||
|
||||
|
||||
# Example usage and entry point
|
||||
if __name__ == "__main__":
|
||||
load_dotenv()
|
||||
|
||||
# Default arguments for testing
|
||||
args = ActionArguments(
|
||||
name="docx_extraction_service",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
|
||||
)
|
||||
|
||||
# Initialize and run the DOCX extraction service
|
||||
try:
|
||||
service = DOCXExtractionCollection(args)
|
||||
service.run()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}: {traceback.format_exc()}")
|
||||
@@ -0,0 +1,545 @@
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from pptx import Presentation
|
||||
from pptx.presentation import Presentation as PresentationType
|
||||
from pydantic import Field
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from aworld.logs.util import Color
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
from examples.gaia.mcp_collections.documents.models import DocumentMetadata
|
||||
|
||||
|
||||
class PPTXExtractionCollection(ActionCollection):
|
||||
"""MCP service for PPTX/PPT presentation content extraction using python-pptx.
|
||||
|
||||
Supports extraction from PPTX and PPT files with comprehensive content parsing.
|
||||
Provides LLM-friendly text output with structured metadata and media file handling.
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
self._media_output_dir = self.workspace / "extracted_media"
|
||||
self._media_output_dir.mkdir(exist_ok=True)
|
||||
|
||||
self.supported_extensions = {".pptx", ".ppt"}
|
||||
|
||||
self._color_log("PPTX Extraction Service initialized", Color.green, "debug")
|
||||
self._color_log(f"Media output directory: {self._media_output_dir}", Color.blue, "debug")
|
||||
|
||||
def _extract_images_from_pptx(self, file_path: Path, file_stem: str) -> list[dict[str, str]]:
|
||||
"""Extract embedded images from PPTX file.
|
||||
|
||||
Args:
|
||||
file_path: Path to the PPTX file
|
||||
file_stem: Base name for saving files
|
||||
|
||||
Returns:
|
||||
List of dictionaries containing image file paths and metadata
|
||||
"""
|
||||
saved_media = []
|
||||
|
||||
try:
|
||||
# PPTX files are ZIP archives
|
||||
with zipfile.ZipFile(file_path, "r") as zip_file:
|
||||
# Find media files in the archive
|
||||
media_files = [f for f in zip_file.namelist() if f.startswith("ppt/media/")]
|
||||
|
||||
for idx, media_file in enumerate(media_files):
|
||||
try:
|
||||
# Extract file extension
|
||||
original_ext = Path(media_file).suffix
|
||||
if not original_ext:
|
||||
original_ext = ".png" # Default extension
|
||||
|
||||
# Generate unique filename
|
||||
media_filename = f"{file_stem}_media_{idx}{original_ext}"
|
||||
media_path = self._media_output_dir / media_filename
|
||||
|
||||
# Extract and save media file
|
||||
with zip_file.open(media_file) as source:
|
||||
with open(media_path, "wb") as target:
|
||||
target.write(source.read())
|
||||
|
||||
# Determine media type
|
||||
media_type = (
|
||||
"image" if original_ext.lower() in {".png", ".jpg", ".jpeg", ".gif", ".bmp"} else "media"
|
||||
)
|
||||
|
||||
saved_media.append(
|
||||
{
|
||||
"type": media_type,
|
||||
"path": str(media_path),
|
||||
"filename": media_filename,
|
||||
"original_path": media_file,
|
||||
}
|
||||
)
|
||||
|
||||
self._color_log(f"Saved media: {media_filename}", Color.blue)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to extract media {media_file}: {e}")
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to extract media from PPTX: {e}")
|
||||
|
||||
return saved_media
|
||||
|
||||
def _extract_slide_structure(self, presentation: PresentationType) -> dict[str, Any]:
|
||||
"""Extract presentation structure information.
|
||||
|
||||
Args:
|
||||
presentation: python-pptx Presentation object
|
||||
|
||||
Returns:
|
||||
Dictionary containing structure metadata
|
||||
"""
|
||||
structure = {
|
||||
"slide_count": len(presentation.slides),
|
||||
"slide_layouts": [],
|
||||
"has_notes": False,
|
||||
"has_comments": False,
|
||||
"slide_sizes": None,
|
||||
}
|
||||
|
||||
# Get slide size information
|
||||
if hasattr(presentation.slide_width, "inches") and hasattr(presentation.slide_height, "inches"):
|
||||
structure["slide_sizes"] = {
|
||||
"width_inches": presentation.slide_width.inches,
|
||||
"height_inches": presentation.slide_height.inches,
|
||||
}
|
||||
|
||||
# Analyze slide layouts and content
|
||||
for slide_idx, slide in enumerate(presentation.slides):
|
||||
layout_info = {
|
||||
"slide_index": slide_idx,
|
||||
"layout_name": slide.slide_layout.name if hasattr(slide.slide_layout, "name") else "Unknown",
|
||||
"shape_count": len(slide.shapes),
|
||||
"has_title": False,
|
||||
"has_content": False,
|
||||
}
|
||||
|
||||
# Check for title and content
|
||||
for shape in slide.shapes:
|
||||
if hasattr(shape, "text") and shape.text.strip():
|
||||
if (
|
||||
"title" in str(shape.placeholder_format.type).lower()
|
||||
if hasattr(shape, "placeholder_format")
|
||||
else False
|
||||
):
|
||||
layout_info["has_title"] = True
|
||||
else:
|
||||
layout_info["has_content"] = True
|
||||
|
||||
# Check for notes
|
||||
if hasattr(slide, "notes_slide") and slide.notes_slide:
|
||||
if hasattr(slide.notes_slide, "notes_text_frame") and slide.notes_slide.notes_text_frame.text.strip():
|
||||
structure["has_notes"] = True
|
||||
|
||||
structure["slide_layouts"].append(layout_info)
|
||||
|
||||
return structure
|
||||
|
||||
def _extract_content_from_pptx(self, file_path: Path, extract_notes: bool = True) -> dict[str, Any]:
|
||||
"""Extract content from PPTX file using python-pptx.
|
||||
|
||||
Args:
|
||||
file_path: Path to the PPTX file
|
||||
extract_notes: Whether to extract speaker notes
|
||||
|
||||
Returns:
|
||||
Dictionary containing extracted content and metadata
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# Load the presentation
|
||||
presentation = Presentation(str(file_path))
|
||||
|
||||
# Extract slide content
|
||||
slides_content = []
|
||||
total_text_length = 0
|
||||
|
||||
for slide_idx, slide in enumerate(presentation.slides):
|
||||
slide_data = {
|
||||
"slide_number": slide_idx + 1,
|
||||
"title": "",
|
||||
"content": [],
|
||||
"notes": "",
|
||||
"shapes_count": len(slide.shapes),
|
||||
}
|
||||
|
||||
# Extract text from all shapes
|
||||
for shape in slide.shapes:
|
||||
if hasattr(shape, "text") and shape.text.strip():
|
||||
text_content = shape.text.strip()
|
||||
|
||||
# Try to identify if this is a title
|
||||
if (
|
||||
hasattr(shape, "placeholder_format")
|
||||
and "title" in str(shape.placeholder_format.type).lower()
|
||||
):
|
||||
slide_data["title"] = text_content
|
||||
else:
|
||||
slide_data["content"].append(text_content)
|
||||
|
||||
total_text_length += len(text_content)
|
||||
|
||||
# Extract text from tables if present
|
||||
if hasattr(shape, "table"):
|
||||
table_text = []
|
||||
for row in shape.table.rows:
|
||||
row_text = []
|
||||
for cell in row.cells:
|
||||
if cell.text.strip():
|
||||
row_text.append(cell.text.strip())
|
||||
if row_text:
|
||||
table_text.append(" | ".join(row_text))
|
||||
if table_text:
|
||||
slide_data["content"].append("\n".join(table_text))
|
||||
|
||||
# Extract speaker notes if requested
|
||||
if extract_notes and hasattr(slide, "notes_slide"):
|
||||
try:
|
||||
if (
|
||||
hasattr(slide.notes_slide, "notes_text_frame")
|
||||
and slide.notes_slide.notes_text_frame.text.strip()
|
||||
):
|
||||
slide_data["notes"] = slide.notes_slide.notes_text_frame.text.strip()
|
||||
total_text_length += len(slide_data["notes"])
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to extract notes from slide {slide_idx + 1}: {e}")
|
||||
|
||||
slides_content.append(slide_data)
|
||||
|
||||
# Extract presentation structure
|
||||
structure = self._extract_slide_structure(presentation)
|
||||
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
return {
|
||||
"slides": slides_content,
|
||||
"structure": structure,
|
||||
"processing_time": processing_time,
|
||||
"total_text_length": total_text_length,
|
||||
"slide_count": len(slides_content),
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to extract content from PPTX: {e}")
|
||||
raise
|
||||
|
||||
def _format_content_for_llm(
|
||||
self, extraction_result: dict[str, Any], output_format: str, include_structure: bool = True
|
||||
) -> str:
|
||||
"""Format extracted PPTX content to be LLM-friendly.
|
||||
|
||||
Args:
|
||||
extraction_result: Dictionary containing extracted content
|
||||
output_format: Desired output format
|
||||
include_structure: Whether to include presentation structure information
|
||||
|
||||
Returns:
|
||||
Formatted content string
|
||||
"""
|
||||
slides = extraction_result["slides"]
|
||||
structure = extraction_result["structure"]
|
||||
|
||||
if output_format.lower() == "markdown":
|
||||
content_parts = []
|
||||
|
||||
if include_structure:
|
||||
content_parts.append("# Presentation Overview")
|
||||
content_parts.append(f"- **Total Slides**: {structure['slide_count']}")
|
||||
if structure.get("slide_sizes"):
|
||||
sizes = structure["slide_sizes"]
|
||||
content_parts.append(
|
||||
f'- **Slide Size**: {sizes["width_inches"]:.1f}" × {sizes["height_inches"]:.1f}"'
|
||||
)
|
||||
content_parts.append(f"- **Has Speaker Notes**: {'Yes' if structure['has_notes'] else 'No'}")
|
||||
content_parts.append("")
|
||||
|
||||
# Format each slide
|
||||
for slide in slides:
|
||||
content_parts.append(f"## Slide {slide['slide_number']}")
|
||||
|
||||
if slide["title"]:
|
||||
content_parts.append(f"**Title**: {slide['title']}")
|
||||
content_parts.append("")
|
||||
|
||||
if slide["content"]:
|
||||
content_parts.append("**Content**:")
|
||||
for content_item in slide["content"]:
|
||||
# Format multi-line content properly
|
||||
for line in content_item.split("\n"):
|
||||
if line.strip():
|
||||
content_parts.append(f"- {line.strip()}")
|
||||
content_parts.append("")
|
||||
|
||||
if slide["notes"]:
|
||||
content_parts.append("**Speaker Notes**:")
|
||||
content_parts.append(slide["notes"])
|
||||
content_parts.append("")
|
||||
|
||||
content_parts.append("---")
|
||||
content_parts.append("")
|
||||
|
||||
return "\n".join(content_parts)
|
||||
|
||||
elif output_format.lower() == "json":
|
||||
return json.dumps(
|
||||
{
|
||||
"presentation_structure": structure if include_structure else None,
|
||||
"slides": slides,
|
||||
"metadata": {
|
||||
"total_slides": len(slides),
|
||||
"total_text_length": extraction_result["total_text_length"],
|
||||
"processing_time": extraction_result["processing_time"],
|
||||
},
|
||||
},
|
||||
indent=2,
|
||||
)
|
||||
|
||||
elif output_format.lower() == "html":
|
||||
html_parts = ["<div class='presentation'>"]
|
||||
|
||||
if include_structure:
|
||||
html_parts.append("<div class='presentation-overview'>")
|
||||
html_parts.append("<h1>Presentation Overview</h1>")
|
||||
html_parts.append(f"<p><strong>Total Slides:</strong> {structure['slide_count']}</p>")
|
||||
if structure.get("slide_sizes"):
|
||||
sizes = structure["slide_sizes"]
|
||||
html_parts.append(
|
||||
"<p>"
|
||||
"<strong>Slide Size:</strong> "
|
||||
f"{sizes['width_inches']:.1f} × {sizes['height_inches']:.1f}"
|
||||
"</p>"
|
||||
)
|
||||
html_parts.append(
|
||||
f"<p><strong>Has Speaker Notes:</strong> {'Yes' if structure['has_notes'] else 'No'}</p>"
|
||||
)
|
||||
html_parts.append("</div>")
|
||||
|
||||
for slide in slides:
|
||||
html_parts.append(f"<div class='slide' data-slide='{slide['slide_number']}'>")
|
||||
html_parts.append(f"<h2>Slide {slide['slide_number']}</h2>")
|
||||
|
||||
if slide["title"]:
|
||||
html_parts.append(f"<h3>{slide['title']}</h3>")
|
||||
|
||||
if slide["content"]:
|
||||
html_parts.append("<div class='slide-content'>")
|
||||
for content_item in slide["content"]:
|
||||
html_parts.append(f"<p>{content_item.replace(chr(10), '<br>')}</p>")
|
||||
html_parts.append("</div>")
|
||||
|
||||
if slide["notes"]:
|
||||
html_parts.append(
|
||||
"<div class='speaker-notes'>"
|
||||
"<strong>Speaker Notes:</strong>"
|
||||
f"<br>{slide['notes'].replace(chr(10), '<br>')}"
|
||||
"</div>"
|
||||
)
|
||||
|
||||
html_parts.append("</div>")
|
||||
|
||||
html_parts.append("</div>")
|
||||
return "\n".join(html_parts)
|
||||
|
||||
else: # Plain text
|
||||
text_parts = []
|
||||
|
||||
if include_structure:
|
||||
text_parts.append("PRESENTATION OVERVIEW")
|
||||
text_parts.append(f"Total Slides: {structure['slide_count']}")
|
||||
if structure.get("slide_sizes"):
|
||||
sizes = structure["slide_sizes"]
|
||||
text_parts.append(f'Slide Size: {sizes["width_inches"]:.1f}" × {sizes["height_inches"]:.1f}"')
|
||||
text_parts.append(f"Has Speaker Notes: {'Yes' if structure['has_notes'] else 'No'}")
|
||||
text_parts.append("\n" + "=" * 50 + "\n")
|
||||
|
||||
for slide in slides:
|
||||
text_parts.append(f"SLIDE {slide['slide_number']}")
|
||||
|
||||
if slide["title"]:
|
||||
text_parts.append(f"Title: {slide['title']}")
|
||||
|
||||
if slide["content"]:
|
||||
text_parts.append("Content:")
|
||||
for content_item in slide["content"]:
|
||||
text_parts.append(f" {content_item}")
|
||||
|
||||
if slide["notes"]:
|
||||
text_parts.append(f"Speaker Notes: {slide['notes']}")
|
||||
|
||||
text_parts.append("\n" + "-" * 30 + "\n")
|
||||
|
||||
return "\n".join(text_parts)
|
||||
|
||||
def mcp_extract_pptx_content(
|
||||
self,
|
||||
file_path: str = Field(description="Path to the PPTX/PPT presentation file to extract content from"),
|
||||
output_format: Literal["markdown", "json", "html", "text"] = Field(
|
||||
default="markdown", description="Output format: 'markdown', 'json', 'html', or 'text'"
|
||||
),
|
||||
extract_images: bool = Field(
|
||||
default=True, description="Whether to extract and save images from the presentation"
|
||||
),
|
||||
extract_notes: bool = Field(default=True, description="Whether to extract speaker notes"),
|
||||
include_structure: bool = Field(
|
||||
default=True, description="Whether to include presentation structure information"
|
||||
),
|
||||
) -> ActionResponse:
|
||||
"""Extract content from PPTX/PPT presentations using python-pptx.
|
||||
|
||||
This tool provides comprehensive PowerPoint presentation content extraction with support for:
|
||||
- PPTX and PPT files
|
||||
- Text extraction from slides, titles, and content
|
||||
- Speaker notes extraction
|
||||
- Image and media extraction
|
||||
- Presentation structure analysis
|
||||
- Multiple output formats (Markdown, JSON, HTML, Text)
|
||||
- LLM-optimized formatting
|
||||
|
||||
Args:
|
||||
file_path: Path to the PPTX/PPT presentation file
|
||||
output_format: Desired output format
|
||||
extract_images: Extract embedded media files
|
||||
extract_notes: Extract speaker notes
|
||||
include_structure: Include presentation structure info
|
||||
|
||||
Returns:
|
||||
ActionResponse with extracted content, metadata, and media file paths
|
||||
"""
|
||||
try:
|
||||
# Handle FieldInfo objects from pydantic
|
||||
if isinstance(file_path, FieldInfo):
|
||||
file_path = file_path.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
if isinstance(extract_images, FieldInfo):
|
||||
extract_images = extract_images.default
|
||||
if isinstance(extract_notes, FieldInfo):
|
||||
extract_notes = extract_notes.default
|
||||
if isinstance(include_structure, FieldInfo):
|
||||
include_structure = include_structure.default
|
||||
|
||||
# Validate input file
|
||||
file_path: Path = self._validate_file_path(file_path)
|
||||
self._color_log(f"Processing PPTX presentation: {file_path.name}", Color.cyan)
|
||||
|
||||
# Extract embedded media if requested
|
||||
saved_media = []
|
||||
if extract_images and file_path.suffix.lower() == ".pptx":
|
||||
saved_media = self._extract_images_from_pptx(file_path, file_path.stem)
|
||||
|
||||
# Extract presentation content
|
||||
extraction_result = self._extract_content_from_pptx(file_path, extract_notes=extract_notes)
|
||||
|
||||
# Format content for LLM consumption
|
||||
formatted_content = self._format_content_for_llm(
|
||||
extraction_result, output_format, include_structure=include_structure
|
||||
)
|
||||
|
||||
# Prepare metadata
|
||||
file_stats = file_path.stat()
|
||||
document_metadata = DocumentMetadata(
|
||||
file_name=file_path.name,
|
||||
file_size=file_stats.st_size,
|
||||
file_type=file_path.suffix.lower(),
|
||||
absolute_path=str(file_path.absolute()),
|
||||
page_count=extraction_result["slide_count"], # Use slide count as page count
|
||||
processing_time=extraction_result["processing_time"],
|
||||
extracted_images=[media["path"] for media in saved_media if media["type"] == "image"],
|
||||
extracted_media=saved_media,
|
||||
output_format=output_format,
|
||||
llm_enhanced=False,
|
||||
ocr_applied=False,
|
||||
)
|
||||
|
||||
# Add PPTX-specific metadata
|
||||
pptx_metadata = {
|
||||
"slide_count": extraction_result["slide_count"],
|
||||
"total_text_length": extraction_result["total_text_length"],
|
||||
"has_speaker_notes": extraction_result["structure"]["has_notes"],
|
||||
"slide_layouts": extraction_result["structure"]["slide_layouts"],
|
||||
"presentation_size": extraction_result["structure"].get("slide_sizes"),
|
||||
"media_files_count": len(saved_media),
|
||||
}
|
||||
|
||||
# Merge metadata
|
||||
final_metadata = {**document_metadata.model_dump(), **pptx_metadata}
|
||||
|
||||
self._color_log(
|
||||
f"Successfully extracted content from {file_path.name} "
|
||||
f"({extraction_result['slide_count']} slides, "
|
||||
f"{len(formatted_content)} characters, {len(saved_media)} media files)",
|
||||
Color.green,
|
||||
)
|
||||
|
||||
return ActionResponse(success=True, message=formatted_content, metadata=final_metadata)
|
||||
|
||||
except FileNotFoundError as e:
|
||||
self.logger.error(f"File not found: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False, message=f"File not found: {str(e)}", metadata={"error_type": "file_not_found"}
|
||||
)
|
||||
except ValueError as e:
|
||||
self.logger.error(f"Invalid input: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False, message=f"Invalid input: {str(e)}", metadata={"error_type": "invalid_input"}
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"PPTX extraction failed: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"PPTX extraction failed: {str(e)}",
|
||||
metadata={"error_type": "extraction_error"},
|
||||
)
|
||||
|
||||
def mcp_list_supported_formats(self) -> ActionResponse:
|
||||
"""List all supported presentation formats for extraction.
|
||||
|
||||
Returns:
|
||||
ActionResponse with list of supported file formats and their descriptions
|
||||
"""
|
||||
supported_formats = {
|
||||
"PPTX": "Microsoft PowerPoint Presentation (.pptx) - full support",
|
||||
"PPT": "Microsoft PowerPoint Presentation (.ppt) - limited support",
|
||||
}
|
||||
|
||||
format_list = "\n".join(
|
||||
[f"**{format_name}**: {description}" for format_name, description in supported_formats.items()]
|
||||
)
|
||||
|
||||
return ActionResponse(
|
||||
success=True,
|
||||
message=f"Supported presentation formats:\n\n{format_list}",
|
||||
metadata={"supported_formats": list(supported_formats.keys()), "total_formats": len(supported_formats)},
|
||||
)
|
||||
|
||||
|
||||
# Example usage and entry point
|
||||
if __name__ == "__main__":
|
||||
load_dotenv()
|
||||
|
||||
# Default arguments for testing
|
||||
args = ActionArguments(
|
||||
name="pptx_extraction_service",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
|
||||
)
|
||||
|
||||
# Initialize and run the PPTX extraction service
|
||||
try:
|
||||
service = PPTXExtractionCollection(args)
|
||||
service.run()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}: {traceback.format_exc()}")
|
||||
@@ -0,0 +1,680 @@
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import traceback
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
import pandas as pd
|
||||
from dotenv import load_dotenv
|
||||
from openpyxl import load_workbook
|
||||
from pydantic import Field
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from aworld.logs.util import Color
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
from examples.gaia.mcp_collections.documents.models import DocumentMetadata
|
||||
|
||||
|
||||
class XLSXExtractionCollection(ActionCollection):
|
||||
"""MCP service for Excel document content extraction using xlrd and pandas.
|
||||
|
||||
Supports extraction from XLSX and XLS files.
|
||||
Provides LLM-friendly text output with structured metadata and media file handling.
|
||||
Extracts worksheets, formulas, charts, and embedded images.
|
||||
Includes screenshot functionality for visual representation of Excel data.
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
self._media_output_dir = self.workspace / "extracted_media"
|
||||
self._media_output_dir.mkdir(exist_ok=True)
|
||||
|
||||
# Create screenshots directory
|
||||
self._screenshots_dir = self.workspace / "excel_screenshots"
|
||||
self._screenshots_dir.mkdir(exist_ok=True)
|
||||
|
||||
self.supported_extensions: set = {
|
||||
".xlsx",
|
||||
".xls",
|
||||
}
|
||||
|
||||
self._color_log("Excel Extraction Service initialized", Color.green, "debug")
|
||||
self._color_log(f"Media output directory: {self._media_output_dir}", Color.blue, "debug")
|
||||
self._color_log(f"Screenshots directory: {self._screenshots_dir}", Color.blue, "debug")
|
||||
|
||||
def _create_excel_screenshot(self, file_path: Path, sheet_name: str = None) -> str:
|
||||
"""Create a JPEG screenshot of the valid Excel area using pyautogui.
|
||||
|
||||
Args:
|
||||
file_path: Path to the Excel file
|
||||
sheet_name: Specific sheet to screenshot (None for first sheet)
|
||||
|
||||
Returns:
|
||||
Path to the generated JPEG screenshot
|
||||
"""
|
||||
try:
|
||||
import pyautogui
|
||||
|
||||
# Generate unique filename
|
||||
timestamp = int(time.time())
|
||||
screenshot_filename = f"{file_path.stem}_{sheet_name or 'sheet'}_{timestamp}.jpg"
|
||||
screenshot_path = self._screenshots_dir / screenshot_filename
|
||||
|
||||
# Open Excel file with default application
|
||||
if sys.platform == "darwin": # macOS
|
||||
subprocess.run(["open", str(file_path)], check=True)
|
||||
elif sys.platform == "win32": # Windows
|
||||
subprocess.run(["start", str(file_path)], shell=True, check=True)
|
||||
else: # Linux
|
||||
subprocess.run(["xdg-open", str(file_path)], check=True)
|
||||
|
||||
# Wait for Excel to open
|
||||
time.sleep(3)
|
||||
|
||||
# Take screenshot of the entire screen
|
||||
screenshot = pyautogui.screenshot()
|
||||
|
||||
# Convert RGBA to RGB before saving as JPEG
|
||||
if screenshot.mode == "RGBA":
|
||||
screenshot = screenshot.convert("RGB")
|
||||
|
||||
screenshot.save(screenshot_path, "JPEG", quality=95)
|
||||
|
||||
self._color_log(f"Created Excel screenshot: {screenshot_filename}", Color.green)
|
||||
return str(screenshot_path)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to create Excel screenshot with pyautogui: {str(e)}")
|
||||
raise
|
||||
|
||||
def _extract_embedded_media_xlsx(self, file_path: Path) -> list[dict[str, str]]:
|
||||
"""Extract embedded media from XLSX files.
|
||||
|
||||
Args:
|
||||
file_path: Path to the XLSX file
|
||||
|
||||
Returns:
|
||||
List of dictionaries containing media information
|
||||
"""
|
||||
saved_media = []
|
||||
|
||||
try:
|
||||
# Load workbook to extract images
|
||||
workbook = load_workbook(file_path, data_only=False)
|
||||
|
||||
for sheet_name in workbook.sheetnames:
|
||||
worksheet = workbook[sheet_name]
|
||||
|
||||
# Extract images from worksheet
|
||||
if hasattr(worksheet, "_images"):
|
||||
for idx, image in enumerate(worksheet._images):
|
||||
try:
|
||||
# Generate unique filename
|
||||
image_filename = f"{file_path.stem}_{sheet_name}_img_{idx}.png"
|
||||
image_path = self._media_output_dir / image_filename
|
||||
|
||||
# Save image
|
||||
if hasattr(image, "ref"):
|
||||
# Extract image data
|
||||
img_data = image._data()
|
||||
if img_data:
|
||||
with open(image_path, "wb") as f:
|
||||
f.write(img_data)
|
||||
|
||||
saved_media.append(
|
||||
{
|
||||
"type": "image",
|
||||
"path": str(image_path),
|
||||
"sheet": sheet_name,
|
||||
"filename": image_filename,
|
||||
}
|
||||
)
|
||||
|
||||
self._color_log(f"Saved image: {image_filename}", Color.blue)
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to extract image {idx} from sheet {sheet_name}: {str(e)}")
|
||||
|
||||
# Also try to extract from ZIP structure for additional media
|
||||
with zipfile.ZipFile(file_path, "r") as zip_file:
|
||||
media_files = [f for f in zip_file.namelist() if f.startswith("xl/media/")]
|
||||
|
||||
for media_file in media_files:
|
||||
try:
|
||||
media_data = zip_file.read(media_file)
|
||||
media_filename = f"{file_path.stem}_{Path(media_file).name}"
|
||||
media_path = self._media_output_dir / media_filename
|
||||
|
||||
with open(media_path, "wb") as f:
|
||||
f.write(media_data)
|
||||
|
||||
# Determine media type based on extension
|
||||
media_ext = Path(media_file).suffix.lower()
|
||||
if media_ext in [".png", ".jpg", ".jpeg", ".gif", ".bmp"]:
|
||||
media_type = "image"
|
||||
elif media_ext in [".mp3", ".wav", ".m4a"]:
|
||||
media_type = "audio"
|
||||
elif media_ext in [".mp4", ".avi", ".mov"]:
|
||||
media_type = "video"
|
||||
else:
|
||||
media_type = "other"
|
||||
|
||||
saved_media.append(
|
||||
{
|
||||
"type": media_type,
|
||||
"path": str(media_path),
|
||||
"filename": media_filename,
|
||||
"original_path": media_file,
|
||||
}
|
||||
)
|
||||
|
||||
self._color_log(f"Saved media: {media_filename}", Color.blue)
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to extract media {media_file}: {str(e)}")
|
||||
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to extract media from XLSX: {str(e)}")
|
||||
|
||||
return saved_media
|
||||
|
||||
def _extract_excel_content(self, file_path: Path, sheet_names: list[str] | None = None) -> dict[str, Any]:
|
||||
"""Extract content from Excel files using pandas and xlrd.
|
||||
|
||||
Args:
|
||||
file_path: Path to the Excel file
|
||||
sheet_names: Specific sheets to process (None for all sheets)
|
||||
|
||||
Returns:
|
||||
Dictionary containing extracted content and metadata
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# Determine file type and read accordingly
|
||||
if file_path.suffix.lower() == ".xlsx":
|
||||
# Use openpyxl engine for XLSX files
|
||||
excel_file = pd.ExcelFile(file_path, engine="openpyxl")
|
||||
else:
|
||||
# Use xlrd engine for XLS files
|
||||
excel_file = pd.ExcelFile(file_path, engine="xlrd")
|
||||
|
||||
# Get all sheet names if not specified
|
||||
if sheet_names is None:
|
||||
sheet_names = excel_file.sheet_names
|
||||
|
||||
sheets_data = {}
|
||||
total_rows = 0
|
||||
total_cols = 0
|
||||
|
||||
# Extract data from each sheet
|
||||
for sheet_name in sheet_names:
|
||||
if sheet_name in excel_file.sheet_names:
|
||||
try:
|
||||
# Read sheet data
|
||||
df = pd.read_excel(excel_file, sheet_name=sheet_name, header=None)
|
||||
|
||||
# Remove completely empty rows and columns
|
||||
df = df.dropna(how="all").dropna(axis=1, how="all")
|
||||
|
||||
if not df.empty:
|
||||
sheets_data[sheet_name] = {
|
||||
"data": df,
|
||||
"shape": df.shape,
|
||||
"columns": df.columns.tolist(),
|
||||
"non_empty_cells": df.count().sum(),
|
||||
}
|
||||
|
||||
total_rows += df.shape[0]
|
||||
total_cols = max(total_cols, df.shape[1])
|
||||
else:
|
||||
sheets_data[sheet_name] = {"data": df, "shape": (0, 0), "columns": [], "non_empty_cells": 0}
|
||||
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to read sheet '{sheet_name}': {str(e)}")
|
||||
sheets_data[sheet_name] = {
|
||||
"error": str(e),
|
||||
"shape": (0, 0),
|
||||
"columns": [],
|
||||
"non_empty_cells": 0,
|
||||
}
|
||||
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
return {
|
||||
"sheets_data": sheets_data,
|
||||
"sheet_names": list(sheets_data.keys()),
|
||||
"total_sheets": len(sheets_data),
|
||||
"total_rows": total_rows,
|
||||
"total_columns": total_cols,
|
||||
"processing_time": processing_time,
|
||||
"file_engine": "openpyxl" if file_path.suffix.lower() == ".xlsx" else "xlrd",
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to extract Excel content: {str(e)}")
|
||||
raise
|
||||
|
||||
def _format_content_for_llm(
|
||||
self, extraction_result: dict[str, Any], output_format: str, include_empty_cells: bool = False
|
||||
) -> str:
|
||||
"""Format extracted Excel content to be LLM-friendly.
|
||||
|
||||
Args:
|
||||
extraction_result: Result from _extract_excel_content
|
||||
output_format: Desired output format
|
||||
include_empty_cells: Whether to include empty cells in output
|
||||
|
||||
Returns:
|
||||
Formatted content string
|
||||
"""
|
||||
sheets_data = extraction_result["sheets_data"]
|
||||
|
||||
if output_format.lower() == "markdown":
|
||||
content_parts = []
|
||||
content_parts.append("# Excel Document Content\n")
|
||||
content_parts.append(f"**Total Sheets:** {extraction_result['total_sheets']}\n")
|
||||
content_parts.append(f"**Processing Engine:** {extraction_result['file_engine']}\n\n")
|
||||
|
||||
for sheet_name, sheet_info in sheets_data.items():
|
||||
content_parts.append(f"## Sheet: {sheet_name}\n")
|
||||
|
||||
if "error" in sheet_info:
|
||||
content_parts.append(f"**Error:** {sheet_info['error']}\n\n")
|
||||
continue
|
||||
|
||||
df: pd.DataFrame = sheet_info["data"]
|
||||
shape = sheet_info["shape"]
|
||||
|
||||
content_parts.append(f"**Dimensions:** {shape[0]} rows × {shape[1]} columns\n")
|
||||
content_parts.append(f"**Non-empty cells:** {sheet_info['non_empty_cells']}\n\n")
|
||||
|
||||
if not df.empty:
|
||||
# Convert DataFrame to markdown table
|
||||
if include_empty_cells:
|
||||
# Fill NaN values with empty string for display
|
||||
df_display = df.fillna("")
|
||||
else:
|
||||
# Keep NaN values as they are
|
||||
df_display = df
|
||||
|
||||
# Convert to markdown table
|
||||
try:
|
||||
markdown_table = df_display.to_markdown(index=False, tablefmt="pipe")
|
||||
content_parts.append(f"### Data:\n{markdown_table}\n\n")
|
||||
except Exception:
|
||||
# Fallback to string representation
|
||||
content_parts.append(f"### Data (text format):\n```\n{df_display.to_string()}\n```\n\n")
|
||||
else:
|
||||
content_parts.append("*Sheet is empty*\n\n")
|
||||
|
||||
return "".join(content_parts)
|
||||
|
||||
elif output_format.lower() == "json":
|
||||
json_data = {
|
||||
"document_info": {
|
||||
"total_sheets": extraction_result["total_sheets"],
|
||||
"total_rows": extraction_result["total_rows"],
|
||||
"total_columns": extraction_result["total_columns"],
|
||||
"processing_engine": extraction_result["file_engine"],
|
||||
},
|
||||
"sheets": {},
|
||||
}
|
||||
|
||||
for sheet_name, sheet_info in sheets_data.items():
|
||||
if "error" in sheet_info:
|
||||
json_data["sheets"][sheet_name] = {"error": sheet_info["error"], "shape": sheet_info["shape"]}
|
||||
continue
|
||||
|
||||
df = sheet_info["data"]
|
||||
|
||||
if not df.empty:
|
||||
# Convert DataFrame to records
|
||||
if include_empty_cells:
|
||||
df_records = df.fillna("").to_dict("records")
|
||||
else:
|
||||
df_records = df.to_dict("records")
|
||||
|
||||
json_data["sheets"][sheet_name] = {
|
||||
"shape": sheet_info["shape"],
|
||||
"non_empty_cells": sheet_info["non_empty_cells"],
|
||||
"data": df_records,
|
||||
}
|
||||
else:
|
||||
json_data["sheets"][sheet_name] = {"shape": sheet_info["shape"], "non_empty_cells": 0, "data": []}
|
||||
|
||||
return json.dumps(json_data, indent=2, default=str)
|
||||
|
||||
elif output_format.lower() == "html":
|
||||
html_parts = []
|
||||
html_parts.append("<html><body>")
|
||||
html_parts.append("<h1>Excel Document Content</h1>")
|
||||
html_parts.append(f"<p><strong>Total Sheets:</strong> {extraction_result['total_sheets']}</p>")
|
||||
html_parts.append(f"<p><strong>Processing Engine:</strong> {extraction_result['file_engine']}</p>")
|
||||
|
||||
for sheet_name, sheet_info in sheets_data.items():
|
||||
html_parts.append(f"<h2>Sheet: {sheet_name}</h2>")
|
||||
|
||||
if "error" in sheet_info:
|
||||
html_parts.append(f"<p><strong>Error:</strong> {sheet_info['error']}</p>")
|
||||
continue
|
||||
|
||||
df = sheet_info["data"]
|
||||
shape = sheet_info["shape"]
|
||||
|
||||
html_parts.append(f"<p><strong>Dimensions:</strong> {shape[0]} rows × {shape[1]} columns</p>")
|
||||
html_parts.append(f"<p><strong>Non-empty cells:</strong> {sheet_info['non_empty_cells']}</p>")
|
||||
|
||||
if not df.empty:
|
||||
# Convert DataFrame to HTML table
|
||||
if include_empty_cells:
|
||||
df_display = df.fillna("")
|
||||
else:
|
||||
df_display = df
|
||||
|
||||
html_table = df_display.to_html(index=False, escape=False, table_id=f"sheet_{sheet_name}")
|
||||
html_parts.append(html_table)
|
||||
else:
|
||||
html_parts.append("<p><em>Sheet is empty</em></p>")
|
||||
|
||||
html_parts.append("</body></html>")
|
||||
return "".join(html_parts)
|
||||
|
||||
else: # text format
|
||||
content_parts = []
|
||||
content_parts.append(f"Excel Document Content\n{'=' * 50}\n")
|
||||
content_parts.append(f"Total Sheets: {extraction_result['total_sheets']}\n")
|
||||
content_parts.append(f"Processing Engine: {extraction_result['file_engine']}\n\n")
|
||||
|
||||
for sheet_name, sheet_info in sheets_data.items():
|
||||
content_parts.append(f"Sheet: {sheet_name}\n{'-' * 30}\n")
|
||||
|
||||
if "error" in sheet_info:
|
||||
content_parts.append(f"Error: {sheet_info['error']}\n\n")
|
||||
continue
|
||||
|
||||
df = sheet_info["data"]
|
||||
shape = sheet_info["shape"]
|
||||
|
||||
content_parts.append(f"Dimensions: {shape[0]} rows × {shape[1]} columns\n")
|
||||
content_parts.append(f"Non-empty cells: {sheet_info['non_empty_cells']}\n\n")
|
||||
|
||||
if not df.empty:
|
||||
if include_empty_cells:
|
||||
df_display = df.fillna("")
|
||||
else:
|
||||
df_display = df
|
||||
|
||||
content_parts.append(f"Data:\n{df_display.to_string()}\n\n")
|
||||
else:
|
||||
content_parts.append("Sheet is empty\n\n")
|
||||
|
||||
return "".join(content_parts)
|
||||
|
||||
def mcp_extract_excel_content(
|
||||
self,
|
||||
file_path: str = Field(description="Path to the Excel document file to extract content from"),
|
||||
output_format: Literal["markdown", "json", "html", "text"] = Field(
|
||||
default="markdown", description="Output format: 'markdown', 'json', 'html', or 'text'"
|
||||
),
|
||||
extract_images: bool = Field(default=True, description="Whether to extract and save images from the document"),
|
||||
create_screenshot: bool = Field(
|
||||
default=False, description="Whether to create a JPEG screenshot of the Excel data"
|
||||
),
|
||||
sheet_names: str | None = Field(
|
||||
default=None, description="Comma-separated list of specific sheet names to process (None for all sheets)"
|
||||
),
|
||||
include_empty_cells: bool = Field(default=False, description="Whether to include empty cells in the output"),
|
||||
screenshot_max_rows: int = Field(default=50, description="Maximum rows to include in screenshot"),
|
||||
screenshot_max_cols: int = Field(default=20, description="Maximum columns to include in screenshot"),
|
||||
) -> ActionResponse:
|
||||
"""Extract content from Excel documents using pandas and xlrd.
|
||||
|
||||
This tool provides comprehensive Excel document content extraction with support for:
|
||||
- XLSX and XLS files
|
||||
- Multiple worksheets
|
||||
- Text and numeric data extraction
|
||||
- Image and media extraction (XLSX only)
|
||||
- JPEG screenshot generation of Excel data
|
||||
- Metadata collection
|
||||
- LLM-optimized output formatting
|
||||
|
||||
Args:
|
||||
file_path: Path to the Excel file
|
||||
output_format: Desired output format
|
||||
extract_images: Whether to extract embedded images
|
||||
create_screenshot: Whether to create a JPEG screenshot
|
||||
sheet_names: Specific sheets to process
|
||||
include_empty_cells: Whether to include empty cells
|
||||
screenshot_max_rows: Maximum rows in screenshot
|
||||
screenshot_max_cols: Maximum columns in screenshot
|
||||
|
||||
Returns:
|
||||
ActionResponse with extracted content, metadata, media file paths, and screenshot path
|
||||
"""
|
||||
try:
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(file_path, FieldInfo):
|
||||
file_path = file_path.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
if isinstance(extract_images, FieldInfo):
|
||||
extract_images = extract_images.default
|
||||
if isinstance(create_screenshot, FieldInfo):
|
||||
create_screenshot = create_screenshot.default
|
||||
if isinstance(sheet_names, FieldInfo):
|
||||
sheet_names = sheet_names.default
|
||||
if isinstance(include_empty_cells, FieldInfo):
|
||||
include_empty_cells = include_empty_cells.default
|
||||
if isinstance(screenshot_max_rows, FieldInfo):
|
||||
screenshot_max_rows = screenshot_max_rows.default
|
||||
if isinstance(screenshot_max_cols, FieldInfo):
|
||||
screenshot_max_cols = screenshot_max_cols.default
|
||||
|
||||
# Validate input file
|
||||
file_path: Path = self._validate_file_path(file_path)
|
||||
self._color_log(f"Processing Excel document: {file_path.name}", Color.cyan)
|
||||
|
||||
# Parse sheet names if provided
|
||||
target_sheets = None
|
||||
if sheet_names:
|
||||
target_sheets = [name.strip() for name in sheet_names.split(",")]
|
||||
|
||||
# Extract content from Excel file
|
||||
extraction_result = self._extract_excel_content(file_path, target_sheets)
|
||||
|
||||
# Extract embedded media if requested (XLSX only)
|
||||
saved_media = []
|
||||
if extract_images and file_path.suffix.lower() == ".xlsx":
|
||||
saved_media = self._extract_embedded_media_xlsx(file_path)
|
||||
elif extract_images and file_path.suffix.lower() == ".xls":
|
||||
self._color_log("Image extraction not supported for XLS files", Color.yellow)
|
||||
|
||||
# Create screenshot if requested
|
||||
screenshot_path = None
|
||||
if create_screenshot:
|
||||
target_sheet = target_sheets[0] if target_sheets else None
|
||||
screenshot_path = self._create_excel_screenshot(file_path, target_sheet)
|
||||
|
||||
# Format content for LLM consumption
|
||||
formatted_content = self._format_content_for_llm(extraction_result, output_format, include_empty_cells)
|
||||
|
||||
# Prepare metadata
|
||||
file_stats = file_path.stat()
|
||||
|
||||
# Create Excel-specific metadata
|
||||
excel_metadata = {
|
||||
"sheet_count": extraction_result["total_sheets"],
|
||||
"sheet_names": extraction_result["sheet_names"],
|
||||
"total_rows": extraction_result["total_rows"],
|
||||
"total_columns": extraction_result["total_columns"],
|
||||
"processing_engine": extraction_result["file_engine"],
|
||||
"extracted_images": [media["path"] for media in saved_media if media["type"] == "image"],
|
||||
"extracted_media": saved_media,
|
||||
"screenshot_path": screenshot_path,
|
||||
"include_empty_cells": include_empty_cells,
|
||||
"processed_sheets": target_sheets or extraction_result["sheet_names"],
|
||||
}
|
||||
|
||||
document_metadata = DocumentMetadata(
|
||||
file_name=file_path.name,
|
||||
file_size=file_stats.st_size,
|
||||
file_type=file_path.suffix.lower(),
|
||||
absolute_path=str(file_path.absolute()),
|
||||
page_count=extraction_result["total_sheets"], # Use sheet count as "page" count
|
||||
processing_time=extraction_result["processing_time"],
|
||||
extracted_images=[media["path"] for media in saved_media if media["type"] == "image"],
|
||||
extracted_media=saved_media,
|
||||
output_format=output_format,
|
||||
llm_enhanced=False,
|
||||
ocr_applied=False,
|
||||
)
|
||||
|
||||
# Combine standard and Excel-specific metadata
|
||||
combined_metadata = document_metadata.model_dump()
|
||||
combined_metadata.update(excel_metadata)
|
||||
|
||||
success_message = (
|
||||
f"Successfully extracted content from {file_path.name} "
|
||||
f"({len(formatted_content)} characters, {extraction_result['total_sheets']} sheets, "
|
||||
f"{len(saved_media)} media files"
|
||||
)
|
||||
|
||||
if screenshot_path:
|
||||
success_message += f", screenshot saved to: {screenshot_path}"
|
||||
|
||||
self._color_log(success_message, Color.green)
|
||||
|
||||
return ActionResponse(success=True, message=formatted_content, metadata=combined_metadata)
|
||||
|
||||
except FileNotFoundError as e:
|
||||
self.logger.error(f"File not found: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False, message=f"File not found: {str(e)}", metadata={"error_type": "file_not_found"}
|
||||
)
|
||||
except ValueError as e:
|
||||
self.logger.error(f"Invalid input: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}",
|
||||
metadata={"error_type": "invalid_input"},
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Excel extraction failed: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Excel extraction failed: {str(e)}",
|
||||
metadata={"error_type": "extraction_error"},
|
||||
)
|
||||
|
||||
def mcp_create_excel_screenshot(
|
||||
self,
|
||||
file_path: str = Field(description="Path to the Excel document file"),
|
||||
sheet_name: str | None = Field(
|
||||
default=None, description="Specific sheet name to screenshot (None for first sheet)"
|
||||
),
|
||||
max_rows: int = Field(default=50, description="Maximum number of rows to include"),
|
||||
max_cols: int = Field(default=20, description="Maximum number of columns to include"),
|
||||
) -> ActionResponse:
|
||||
"""Create a JPEG screenshot of the valid Excel area.
|
||||
|
||||
This tool creates a visual representation of Excel data as a JPEG image,
|
||||
useful for further image processing or visual analysis.
|
||||
|
||||
Args:
|
||||
file_path: Path to the Excel file
|
||||
sheet_name: Specific sheet to screenshot
|
||||
max_rows: Maximum rows to include in screenshot
|
||||
max_cols: Maximum columns to include in screenshot
|
||||
|
||||
Returns:
|
||||
ActionResponse with screenshot file path and metadata
|
||||
"""
|
||||
try:
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(file_path, FieldInfo):
|
||||
file_path = file_path.default
|
||||
if isinstance(sheet_name, FieldInfo):
|
||||
sheet_name = sheet_name.default
|
||||
if isinstance(max_rows, FieldInfo):
|
||||
max_rows = max_rows.default
|
||||
if isinstance(max_cols, FieldInfo):
|
||||
max_cols = max_cols.default
|
||||
|
||||
# Validate input file
|
||||
file_path: Path = self._validate_file_path(file_path)
|
||||
self._color_log(f"Creating screenshot for Excel document: {file_path.name}", Color.cyan)
|
||||
|
||||
# Create screenshot
|
||||
screenshot_path = self._create_excel_screenshot(file_path, sheet_name)
|
||||
|
||||
# Prepare metadata
|
||||
file_stats = file_path.stat()
|
||||
screenshot_stats = Path(screenshot_path).stat()
|
||||
|
||||
metadata = {
|
||||
"source_file": str(file_path.absolute()),
|
||||
"source_file_size": file_stats.st_size,
|
||||
"screenshot_path": screenshot_path,
|
||||
"screenshot_size": screenshot_stats.st_size,
|
||||
"sheet_name": sheet_name,
|
||||
"max_rows_displayed": max_rows,
|
||||
"max_cols_displayed": max_cols,
|
||||
"format": "JPEG",
|
||||
}
|
||||
|
||||
return ActionResponse(
|
||||
success=True,
|
||||
message=f"Excel screenshot created successfully. File saved to: {screenshot_path}",
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Screenshot creation failed: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Screenshot creation failed: {str(e)}",
|
||||
metadata={"error_type": "screenshot_error"},
|
||||
)
|
||||
|
||||
def mcp_list_supported_formats(self) -> ActionResponse:
|
||||
"""List all supported Excel formats for extraction.
|
||||
|
||||
Returns:
|
||||
ActionResponse with list of supported file formats and their descriptions
|
||||
"""
|
||||
supported_formats = {
|
||||
"XLSX": "Excel 2007+ format files (.xlsx) - Full support including images",
|
||||
"XLS": "Excel 97-2003 format files (.xls) - Text and data only",
|
||||
}
|
||||
|
||||
format_list = "\n".join(
|
||||
[f"**{format_name}**: {description}" for format_name, description in supported_formats.items()]
|
||||
)
|
||||
|
||||
return ActionResponse(
|
||||
success=True,
|
||||
message=f"Supported Excel formats:\n\n{format_list}",
|
||||
metadata={"supported_formats": list(supported_formats.keys()), "total_formats": len(supported_formats)},
|
||||
)
|
||||
|
||||
|
||||
# Example usage and entry point
|
||||
if __name__ == "__main__":
|
||||
load_dotenv()
|
||||
|
||||
# Default arguments for testing
|
||||
args = ActionArguments(
|
||||
name="excel_extraction_service",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "/tmp"),
|
||||
)
|
||||
|
||||
# Initialize and run the Excel extraction service
|
||||
try:
|
||||
service = XLSXExtractionCollection(args)
|
||||
service.run()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}: {traceback.format_exc()}")
|
||||
@@ -0,0 +1,341 @@
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal, Optional # Added Optional
|
||||
|
||||
import markdown
|
||||
from dotenv import load_dotenv
|
||||
from marker.converters.pdf import PdfConverter
|
||||
from marker.models import create_model_dict
|
||||
from marker.output import text_from_rendered
|
||||
from marker.settings import settings
|
||||
from pydantic import Field
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from aworld.logs.util import Color
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
from examples.gaia.mcp_collections.documents.models import DocumentMetadata
|
||||
|
||||
|
||||
class DocumentExtractionCollection(ActionCollection):
|
||||
"""MCP service for PDF document content extraction using marker package.
|
||||
|
||||
Supports extraction from PDF files only.
|
||||
Provides LLM-friendly text output with structured metadata and media file handling.
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
self._models_loaded = False
|
||||
self._marker_models = None
|
||||
self._media_output_dir = self.workspace / "extracted_media"
|
||||
self._media_output_dir.mkdir(exist_ok=True)
|
||||
self._extracted_texts_dir = self.workspace / "extracted_texts" # New directory for text files
|
||||
self._extracted_texts_dir.mkdir(exist_ok=True)
|
||||
|
||||
self.supported_extensions = {".pdf"}
|
||||
|
||||
self._color_log("PDF Extraction Service initialized", Color.green, "debug")
|
||||
self._color_log(f"Media output directory: {self._media_output_dir}", Color.blue, "debug")
|
||||
|
||||
def _load_marker_models(self) -> None:
|
||||
"""Load marker models for document processing.
|
||||
|
||||
Lazy loading to avoid unnecessary resource consumption.
|
||||
"""
|
||||
if not self._models_loaded:
|
||||
try:
|
||||
self._color_log("Loading marker models...", Color.yellow)
|
||||
self._marker_models = create_model_dict()
|
||||
self._models_loaded = True
|
||||
self._color_log("Marker models loaded successfully", Color.green)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to load marker models: {str(e)}")
|
||||
raise
|
||||
|
||||
def _extract_content_with_marker(
|
||||
self, file_path: Path, page_range: str | None, force_ocr: bool = False
|
||||
) -> dict[str, Any]:
|
||||
"""Extract content using marker package.
|
||||
|
||||
Args:
|
||||
file_path: Path to the document file
|
||||
page_range: Specific pages to process (e.g., '0,5-10,20')
|
||||
force_ocr: Use OCR to extract text from images if available
|
||||
|
||||
Returns:
|
||||
Dictionary containing extracted content and metadata
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Prepare marker arguments
|
||||
marker_args = {
|
||||
"fname": str(file_path),
|
||||
"model_lst": self._marker_models,
|
||||
"max_pages": None,
|
||||
"langs": None,
|
||||
"batch_multiplier": 1,
|
||||
"force_ocr": force_ocr,
|
||||
}
|
||||
|
||||
# Handle page range
|
||||
if page_range:
|
||||
# Parse page range string (e.g., "0,5-10,20")
|
||||
pages = []
|
||||
for part in page_range.split(","):
|
||||
if "-" in part:
|
||||
start, end = map(int, part.split("-"))
|
||||
pages.extend(range(start, end + 1))
|
||||
else:
|
||||
pages.append(int(part))
|
||||
marker_args["page_range"] = pages
|
||||
converter: PdfConverter = PdfConverter(artifact_dict=self._marker_models)
|
||||
rendered = converter(str(file_path))
|
||||
text, _, images = text_from_rendered(rendered)
|
||||
text = text.encode(settings.OUTPUT_ENCODING, errors="replace").decode(settings.OUTPUT_ENCODING)
|
||||
|
||||
processing_time = time.time() - start_time
|
||||
return {
|
||||
"content": text,
|
||||
"images": images or {},
|
||||
"metadata": defaultdict(),
|
||||
"processing_time": processing_time,
|
||||
}
|
||||
|
||||
def _save_extracted_media(self, images: dict[str, Any], file_stem: str) -> list[dict[str, str]]:
|
||||
"""Save extracted images and return their paths.
|
||||
|
||||
Args:
|
||||
images: Dictionary of extracted images from marker
|
||||
file_stem: Base name for saving files
|
||||
|
||||
Returns:
|
||||
list of dictionaries containing media type and file paths
|
||||
"""
|
||||
saved_media = []
|
||||
|
||||
for idx, (page_num, image_data) in enumerate(images.items()):
|
||||
try:
|
||||
# Generate unique filename
|
||||
image_filename = f"{file_stem}_page_{page_num}_img_{idx}.png"
|
||||
image_path = self._media_output_dir / image_filename
|
||||
|
||||
# Save image data
|
||||
if hasattr(image_data, "save"):
|
||||
# PIL Image object
|
||||
image_data.save(image_path)
|
||||
elif isinstance(image_data, bytes):
|
||||
# Raw image bytes
|
||||
with open(image_path, "wb") as f:
|
||||
f.write(image_data)
|
||||
else:
|
||||
# Handle other formats
|
||||
self.logger.warning(f"Unknown image data type for page {page_num}: {type(image_data)}")
|
||||
continue
|
||||
|
||||
saved_media.append(
|
||||
{"type": "image", "path": str(image_path), "page": str(page_num), "filename": image_filename}
|
||||
)
|
||||
|
||||
self._color_log(f"Saved image: {image_filename}", Color.blue)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to save image from page {page_num}: {str(e)}")
|
||||
|
||||
return saved_media
|
||||
|
||||
def _format_content_for_llm(self, content: str, output_format: str) -> str:
|
||||
"""Format extracted content to be LLM-friendly.
|
||||
|
||||
Args:
|
||||
content: Raw extracted content
|
||||
output_format: Desired output format
|
||||
|
||||
Returns:
|
||||
Formatted content string
|
||||
"""
|
||||
if output_format.lower() == "markdown":
|
||||
# Content is already in markdown format from marker
|
||||
return content
|
||||
elif output_format.lower() == "json":
|
||||
# Structure content as JSON
|
||||
|
||||
return json.dumps({"content": content, "format": "structured_text"}, indent=2)
|
||||
elif output_format.lower() == "html":
|
||||
# Convert markdown to HTML if needed
|
||||
try:
|
||||
return markdown.markdown(content)
|
||||
except ImportError:
|
||||
self.logger.warning("markdown package not available, returning raw content")
|
||||
return content
|
||||
else:
|
||||
return content
|
||||
|
||||
def mcp_extract_document_content(
|
||||
self,
|
||||
file_path: str = Field(description="Path to the PDF document file to extract content from"),
|
||||
output_format: Literal["markdown", "json", "html"] = Field(
|
||||
default="markdown", description="Output format: 'markdown', 'json', or 'html'"
|
||||
),
|
||||
extract_images: bool = Field(default=True, description="Whether to extract and save images from the document"),
|
||||
save_extracted_text_to_file: bool = Field(
|
||||
default=False, description="Save extracted text to a local file"
|
||||
), # New parameter
|
||||
use_llm: bool = Field(default=False, description="Use LLM for enhanced accuracy (requires additional setup)"),
|
||||
page_range: str | None = Field(default=None, description="Specific pages to process (e.g., '0,5-10,20')"),
|
||||
force_ocr: bool = Field(default=False, description="Force OCR processing on the entire document"),
|
||||
format_lines: bool = Field(
|
||||
default=False, description="Reformat lines using local OCR model for better quality"
|
||||
),
|
||||
) -> ActionResponse:
|
||||
"""Extract content from PDF documents using marker package.
|
||||
|
||||
This tool provides comprehensive PDF document content extraction with support for:
|
||||
- PDF files
|
||||
- Text extraction with proper formatting
|
||||
- Image and media extraction
|
||||
- Metadata collection
|
||||
- LLM-optimized output formatting
|
||||
|
||||
Args:
|
||||
args: Document extraction arguments including file path and options
|
||||
|
||||
Returns:
|
||||
ActionResponse with extracted content, metadata, and media file paths
|
||||
"""
|
||||
try:
|
||||
if isinstance(file_path, FieldInfo):
|
||||
file_path = file_path.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
if isinstance(extract_images, FieldInfo):
|
||||
extract_images = extract_images.default
|
||||
if isinstance(save_extracted_text_to_file, FieldInfo): # Handle new parameter
|
||||
save_extracted_text_to_file = save_extracted_text_to_file.default
|
||||
if isinstance(page_range, FieldInfo):
|
||||
page_range = page_range.default
|
||||
if isinstance(use_llm, FieldInfo):
|
||||
use_llm = use_llm.default
|
||||
if isinstance(force_ocr, FieldInfo):
|
||||
force_ocr = force_ocr.default
|
||||
if isinstance(format_lines, FieldInfo):
|
||||
format_lines = format_lines.default
|
||||
|
||||
# Validate input file
|
||||
file_path: Path = self._validate_file_path(file_path)
|
||||
self._color_log(f"Processing document: {file_path.name}", Color.cyan)
|
||||
|
||||
# Load marker models if needed
|
||||
self._load_marker_models()
|
||||
|
||||
# Extract content using marker
|
||||
extraction_result = self._extract_content_with_marker(file_path, page_range, force_ocr)
|
||||
|
||||
# Save extracted media if requested
|
||||
saved_media = []
|
||||
if extract_images and extraction_result["images"]:
|
||||
saved_media = self._save_extracted_media(extraction_result["images"], file_path.stem)
|
||||
|
||||
# Format content for LLM consumption
|
||||
formatted_content = self._format_content_for_llm(extraction_result["content"], output_format)
|
||||
|
||||
# Save extracted text to file if requested
|
||||
saved_text_path_str: Optional[str] = None
|
||||
if save_extracted_text_to_file:
|
||||
text_file_name = f"{file_path.stem}_extracted_text.txt"
|
||||
saved_text_path = self._extracted_texts_dir / text_file_name
|
||||
try:
|
||||
with open(saved_text_path, "w", encoding="utf-8") as f:
|
||||
f.write(formatted_content)
|
||||
saved_text_path_str = str(saved_text_path.absolute())
|
||||
self._color_log(f"Saved extracted text to: {saved_text_path_str}", Color.blue)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to save extracted text to {saved_text_path}: {str(e)}")
|
||||
# Optionally, you might want to reflect this failure in the response
|
||||
|
||||
# Prepare metadata
|
||||
file_stats = file_path.stat()
|
||||
document_metadata = DocumentMetadata(
|
||||
file_name=file_path.name,
|
||||
file_size=file_stats.st_size,
|
||||
file_type=file_path.suffix.lower(),
|
||||
absolute_path=str(file_path.absolute()),
|
||||
page_count=extraction_result["metadata"].get("page_count"),
|
||||
processing_time=extraction_result["processing_time"],
|
||||
extracted_images=[media["path"] for media in saved_media if media["type"] == "image"],
|
||||
extracted_media=saved_media,
|
||||
output_format=output_format,
|
||||
llm_enhanced=use_llm,
|
||||
ocr_applied=force_ocr or format_lines,
|
||||
extracted_text_file_path=saved_text_path_str,
|
||||
)
|
||||
|
||||
self._color_log(
|
||||
f"Successfully extracted content from {file_path.name} "
|
||||
f"({len(formatted_content)} characters, {len(saved_media)} media files)",
|
||||
Color.green,
|
||||
)
|
||||
|
||||
return ActionResponse(success=True, message=formatted_content, metadata=document_metadata.model_dump())
|
||||
|
||||
except FileNotFoundError as e:
|
||||
self.logger.error(f"File not found: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False, message=f"File not found: {str(e)}", metadata={"error_type": "file_not_found"}
|
||||
)
|
||||
except ValueError as e:
|
||||
self.logger.error(f"Invalid input: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}: {traceback.format_exc()}",
|
||||
metadata={"error_type": "invalid_input"},
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Document extraction failed: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Document extraction failed: {str(e)}",
|
||||
metadata={"error_type": "extraction_error"},
|
||||
)
|
||||
|
||||
def mcp_list_supported_formats(self) -> ActionResponse:
|
||||
"""list all supported document formats for extraction.
|
||||
|
||||
Returns:
|
||||
ActionResponse with list of supported file formats and their descriptions
|
||||
"""
|
||||
supported_formats = {
|
||||
"PDF": "Portable Document Format files (.pdf)",
|
||||
}
|
||||
|
||||
format_list = "\n".join(
|
||||
[f"**{format_name}**: {description}" for format_name, description in supported_formats.items()]
|
||||
)
|
||||
|
||||
return ActionResponse(
|
||||
success=True,
|
||||
message=f"Supported document formats:\n\n{format_list}",
|
||||
metadata={"supported_formats": list(supported_formats.keys()), "total_formats": len(supported_formats)},
|
||||
)
|
||||
|
||||
|
||||
# Example usage and entry point
|
||||
if __name__ == "__main__":
|
||||
load_dotenv()
|
||||
|
||||
# Default arguments for testing
|
||||
args = ActionArguments(
|
||||
name="document_extraction_service",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
|
||||
)
|
||||
|
||||
# Initialize and run the document extraction service
|
||||
try:
|
||||
service = DocumentExtractionCollection(args)
|
||||
service.run()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}: {traceback.format_exc()}")
|
||||
@@ -0,0 +1,648 @@
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
import chardet
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import Field
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from aworld.logs.util import Color
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
from examples.gaia.mcp_collections.documents.models import DocumentMetadata
|
||||
from examples.gaia.mcp_collections.utils import get_mime_type
|
||||
|
||||
|
||||
class TextExtractionCollection(ActionCollection):
|
||||
"""MCP service for text document content extraction.
|
||||
|
||||
Supports extraction from TXT and other raw text format files.
|
||||
Provides LLM-friendly text output with structured metadata and encoding detection.
|
||||
Handles various text encodings and provides comprehensive file analysis.
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
self._media_output_dir = self.workspace / "extracted_media"
|
||||
self._media_output_dir.mkdir(exist_ok=True)
|
||||
|
||||
self.supported_extensions: set = {
|
||||
".txt",
|
||||
".text",
|
||||
".log",
|
||||
".md",
|
||||
".markdown",
|
||||
".rst",
|
||||
".rtf",
|
||||
".csv",
|
||||
".tsv",
|
||||
".json",
|
||||
".xml",
|
||||
".yaml",
|
||||
".yml",
|
||||
".ini",
|
||||
".cfg",
|
||||
".conf",
|
||||
".properties",
|
||||
".sql",
|
||||
".py",
|
||||
".js",
|
||||
".html",
|
||||
".htm",
|
||||
".css",
|
||||
".java",
|
||||
".cpp",
|
||||
".c",
|
||||
".h",
|
||||
".php",
|
||||
".rb",
|
||||
".go",
|
||||
".rs",
|
||||
".sh",
|
||||
".bat",
|
||||
".ps1",
|
||||
".r",
|
||||
".m",
|
||||
".swift",
|
||||
".kt",
|
||||
".scala",
|
||||
".pl",
|
||||
".lua",
|
||||
".vim",
|
||||
".tex",
|
||||
".bib",
|
||||
}
|
||||
|
||||
self._color_log("Text Extraction Service initialized", Color.green, "debug")
|
||||
self._color_log(f"Media output directory: {self._media_output_dir}", Color.blue, "debug")
|
||||
|
||||
def _validate_file_path(self, file_path: str) -> Path:
|
||||
"""Validate and resolve file path.
|
||||
|
||||
Args:
|
||||
file_path: Path to the text document file
|
||||
|
||||
Returns:
|
||||
Resolved Path object
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If file doesn't exist
|
||||
ValueError: If file type is not supported
|
||||
"""
|
||||
path = Path(file_path)
|
||||
if not path.is_absolute():
|
||||
path = self.workspace / path
|
||||
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"File not found: {path}")
|
||||
|
||||
# Also check MIME type for files without extensions or unknown extensions
|
||||
mime_type = get_mime_type(str(path), default_mime="text/plain")
|
||||
is_text_mime = mime_type and mime_type.startswith("text/")
|
||||
|
||||
if path.suffix.lower() not in self.supported_extensions and not is_text_mime:
|
||||
# Try to detect if it's a text file by reading a small sample
|
||||
try:
|
||||
with open(path, "rb") as f:
|
||||
sample = f.read(1024)
|
||||
# Check if the sample contains mostly printable characters
|
||||
if self._is_likely_text(sample):
|
||||
self._color_log(f"Detected text file without standard extension: {path.suffix}", Color.yellow)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported file type: {path.suffix}. "
|
||||
f"Supported types: {', '.join(sorted(self.supported_extensions))} or text MIME types"
|
||||
)
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Cannot determine if file is text: {str(e)}. "
|
||||
f"Supported types: {', '.join(sorted(self.supported_extensions))}"
|
||||
) from e
|
||||
|
||||
return path
|
||||
|
||||
def _is_likely_text(self, data: bytes) -> bool:
|
||||
"""Check if binary data is likely to be text.
|
||||
|
||||
Args:
|
||||
data: Binary data sample
|
||||
|
||||
Returns:
|
||||
True if data appears to be text
|
||||
"""
|
||||
if not data:
|
||||
return True
|
||||
|
||||
# Check for null bytes (common in binary files)
|
||||
if b"\x00" in data:
|
||||
return False
|
||||
|
||||
# Try to decode as text
|
||||
try:
|
||||
data.decode("utf-8")
|
||||
return True
|
||||
except UnicodeDecodeError:
|
||||
pass
|
||||
|
||||
# Check if most bytes are printable ASCII
|
||||
printable_count = sum(1 for byte in data if 32 <= byte <= 126 or byte in [9, 10, 13])
|
||||
return printable_count / len(data) > 0.7
|
||||
|
||||
def _detect_encoding(self, file_path: Path) -> dict[str, Any]:
|
||||
"""Detect file encoding and other characteristics.
|
||||
|
||||
Args:
|
||||
file_path: Path to the text file
|
||||
|
||||
Returns:
|
||||
Dictionary containing encoding information
|
||||
"""
|
||||
encoding_info = {
|
||||
"detected_encoding": None,
|
||||
"confidence": 0.0,
|
||||
"bom_detected": False,
|
||||
"line_endings": None,
|
||||
"is_binary": False,
|
||||
}
|
||||
|
||||
try:
|
||||
# Read file in binary mode for encoding detection
|
||||
with open(file_path, "rb") as f:
|
||||
raw_data = f.read()
|
||||
|
||||
if not raw_data:
|
||||
encoding_info["detected_encoding"] = "utf-8"
|
||||
encoding_info["confidence"] = 1.0
|
||||
return encoding_info
|
||||
|
||||
# Check for BOM (Byte Order Mark)
|
||||
if raw_data.startswith(b"\xef\xbb\xbf"):
|
||||
encoding_info["bom_detected"] = True
|
||||
encoding_info["detected_encoding"] = "utf-8-sig"
|
||||
encoding_info["confidence"] = 1.0
|
||||
elif raw_data.startswith(b"\xff\xfe"):
|
||||
encoding_info["bom_detected"] = True
|
||||
encoding_info["detected_encoding"] = "utf-16-le"
|
||||
encoding_info["confidence"] = 1.0
|
||||
elif raw_data.startswith(b"\xfe\xff"):
|
||||
encoding_info["bom_detected"] = True
|
||||
encoding_info["detected_encoding"] = "utf-16-be"
|
||||
encoding_info["confidence"] = 1.0
|
||||
else:
|
||||
# Use chardet for encoding detection
|
||||
detection_result = chardet.detect(raw_data)
|
||||
encoding_info["detected_encoding"] = detection_result.get("encoding", "utf-8")
|
||||
encoding_info["confidence"] = detection_result.get("confidence", 0.0)
|
||||
|
||||
# Detect line endings
|
||||
if b"\r\n" in raw_data:
|
||||
encoding_info["line_endings"] = "CRLF (Windows)"
|
||||
elif b"\n" in raw_data:
|
||||
encoding_info["line_endings"] = "LF (Unix/Linux/Mac)"
|
||||
elif b"\r" in raw_data:
|
||||
encoding_info["line_endings"] = "CR (Classic Mac)"
|
||||
else:
|
||||
encoding_info["line_endings"] = "None detected"
|
||||
|
||||
# Check if file appears to be binary
|
||||
encoding_info["is_binary"] = not self._is_likely_text(raw_data[:1024])
|
||||
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to detect encoding: {str(e)}")
|
||||
encoding_info["detected_encoding"] = "utf-8"
|
||||
encoding_info["confidence"] = 0.0
|
||||
|
||||
return encoding_info
|
||||
|
||||
def _extract_text_content(self, file_path: Path, encoding: str | None = None) -> dict[str, Any]:
|
||||
"""Extract content from text files.
|
||||
|
||||
Args:
|
||||
file_path: Path to the text file
|
||||
encoding: Specific encoding to use (None for auto-detection)
|
||||
|
||||
Returns:
|
||||
Dictionary containing extracted content and metadata
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Detect encoding if not specified
|
||||
encoding_info = self._detect_encoding(file_path)
|
||||
|
||||
if encoding:
|
||||
# Use specified encoding
|
||||
target_encoding = encoding
|
||||
self._color_log(f"Using specified encoding: {encoding}", Color.blue)
|
||||
else:
|
||||
# Use detected encoding
|
||||
target_encoding = encoding_info["detected_encoding"]
|
||||
self._color_log(
|
||||
f"Detected encoding: {target_encoding} (confidence: {encoding_info['confidence']:.2f})", Color.blue
|
||||
)
|
||||
|
||||
try:
|
||||
# Read file content
|
||||
with open(file_path, "r", encoding=target_encoding, errors="replace") as f:
|
||||
content = f.read()
|
||||
|
||||
# Analyze content
|
||||
lines = content.splitlines()
|
||||
|
||||
# Calculate statistics
|
||||
char_count = len(content)
|
||||
line_count = len(lines)
|
||||
word_count = len(content.split()) if content.strip() else 0
|
||||
|
||||
# Find longest and shortest lines
|
||||
line_lengths = [len(line) for line in lines]
|
||||
max_line_length = max(line_lengths) if line_lengths else 0
|
||||
min_line_length = min(line_lengths) if line_lengths else 0
|
||||
avg_line_length = sum(line_lengths) / len(line_lengths) if line_lengths else 0
|
||||
|
||||
# Count empty lines
|
||||
empty_lines = sum(1 for line in lines if not line.strip())
|
||||
|
||||
# Detect file type based on content patterns
|
||||
content_type = self._detect_content_type(content, file_path)
|
||||
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
return {
|
||||
"content": content,
|
||||
"encoding_info": encoding_info,
|
||||
"statistics": {
|
||||
"character_count": char_count,
|
||||
"line_count": line_count,
|
||||
"word_count": word_count,
|
||||
"empty_lines": empty_lines,
|
||||
"max_line_length": max_line_length,
|
||||
"min_line_length": min_line_length,
|
||||
"avg_line_length": round(avg_line_length, 2),
|
||||
},
|
||||
"content_type": content_type,
|
||||
"processing_time": processing_time,
|
||||
"used_encoding": target_encoding,
|
||||
}
|
||||
|
||||
except UnicodeDecodeError as e:
|
||||
self.logger.error(f"Failed to decode file with encoding {target_encoding}: {str(e)}")
|
||||
# Try with fallback encodings
|
||||
fallback_encodings = ["utf-8", "latin-1", "cp1252", "iso-8859-1"]
|
||||
|
||||
for fallback_encoding in fallback_encodings:
|
||||
if fallback_encoding != target_encoding:
|
||||
try:
|
||||
with open(file_path, "r", encoding=fallback_encoding, errors="replace") as f:
|
||||
content = f.read()
|
||||
|
||||
self._color_log(f"Successfully read with fallback encoding: {fallback_encoding}", Color.yellow)
|
||||
|
||||
# Recalculate with fallback encoding
|
||||
lines = content.splitlines()
|
||||
char_count = len(content)
|
||||
line_count = len(lines)
|
||||
word_count = len(content.split()) if content.strip() else 0
|
||||
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
return {
|
||||
"content": content,
|
||||
"encoding_info": encoding_info,
|
||||
"statistics": {
|
||||
"character_count": char_count,
|
||||
"line_count": line_count,
|
||||
"word_count": word_count,
|
||||
"empty_lines": sum(1 for line in lines if not line.strip()),
|
||||
"max_line_length": max(len(line) for line in lines) if lines else 0,
|
||||
"min_line_length": min(len(line) for line in lines) if lines else 0,
|
||||
"avg_line_length": round(sum(len(line) for line in lines) / len(lines), 2)
|
||||
if lines
|
||||
else 0,
|
||||
},
|
||||
"content_type": self._detect_content_type(content, file_path),
|
||||
"processing_time": processing_time,
|
||||
"used_encoding": fallback_encoding,
|
||||
"encoding_fallback": True,
|
||||
}
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
raise ValueError("Unable to decode file with any supported encoding") from e
|
||||
|
||||
def _detect_content_type(self, content: str, file_path: Path) -> str:
|
||||
"""Detect the type of content based on file extension and content patterns.
|
||||
|
||||
Args:
|
||||
content: File content
|
||||
file_path: Path to the file
|
||||
|
||||
Returns:
|
||||
Detected content type
|
||||
"""
|
||||
extension = file_path.suffix.lower()
|
||||
|
||||
# Map extensions to content types
|
||||
extension_map = {
|
||||
".py": "Python source code",
|
||||
".js": "JavaScript source code",
|
||||
".html": "HTML document",
|
||||
".htm": "HTML document",
|
||||
".css": "CSS stylesheet",
|
||||
".json": "JSON data",
|
||||
".xml": "XML document",
|
||||
".yaml": "YAML configuration",
|
||||
".yml": "YAML configuration",
|
||||
".md": "Markdown document",
|
||||
".markdown": "Markdown document",
|
||||
".rst": "reStructuredText document",
|
||||
".csv": "CSV data",
|
||||
".tsv": "TSV data",
|
||||
".sql": "SQL script",
|
||||
".log": "Log file",
|
||||
".ini": "Configuration file",
|
||||
".cfg": "Configuration file",
|
||||
".conf": "Configuration file",
|
||||
}
|
||||
|
||||
if extension in extension_map:
|
||||
return extension_map[extension]
|
||||
|
||||
# Content-based detection
|
||||
content_lower = content.lower().strip()
|
||||
|
||||
if content_lower.startswith("<?xml"):
|
||||
return "XML document"
|
||||
elif content_lower.startswith("{") and content_lower.endswith("}"):
|
||||
return "JSON-like data"
|
||||
elif content_lower.startswith("[") and content_lower.endswith("]"):
|
||||
return "JSON array or configuration"
|
||||
elif "#!/" in content[:50]:
|
||||
return "Script file"
|
||||
elif content.count(",") > content.count("\n") * 2:
|
||||
return "CSV-like data"
|
||||
else:
|
||||
return "Plain text"
|
||||
|
||||
def _format_content_for_llm(
|
||||
self, extraction_result: dict[str, Any], output_format: str, max_length: int | None = None
|
||||
) -> str:
|
||||
"""Format extracted text content to be LLM-friendly.
|
||||
|
||||
Args:
|
||||
extraction_result: Result from _extract_text_content
|
||||
output_format: Desired output format
|
||||
max_length: Maximum length of content to include (None for no limit)
|
||||
|
||||
Returns:
|
||||
Formatted content string
|
||||
"""
|
||||
content = extraction_result["content"]
|
||||
stats = extraction_result["statistics"]
|
||||
content_type = extraction_result["content_type"]
|
||||
|
||||
# Truncate content if needed
|
||||
if max_length and len(content) > max_length:
|
||||
content = (
|
||||
content[:max_length]
|
||||
+ f"\n\n[Content truncated - showing first {max_length} characters of {stats['character_count']} total]"
|
||||
)
|
||||
|
||||
if output_format.lower() == "markdown":
|
||||
formatted_parts = []
|
||||
formatted_parts.append("# Text Document Content\n")
|
||||
formatted_parts.append(f"**File Type:** {content_type}\n")
|
||||
formatted_parts.append(f"**Encoding:** {extraction_result['used_encoding']}\n")
|
||||
formatted_parts.append("**Statistics:**\n")
|
||||
formatted_parts.append(f"- Characters: {stats['character_count']:,}\n")
|
||||
formatted_parts.append(f"- Lines: {stats['line_count']:,}\n")
|
||||
formatted_parts.append(f"- Words: {stats['word_count']:,}\n")
|
||||
formatted_parts.append(f"- Empty lines: {stats['empty_lines']:,}\n")
|
||||
formatted_parts.append(f"- Average line length: {stats['avg_line_length']} characters\n\n")
|
||||
|
||||
formatted_parts.append(f"## Content\n\n```\n{content}\n```")
|
||||
|
||||
return "".join(formatted_parts)
|
||||
|
||||
elif output_format.lower() == "json":
|
||||
json_data = {
|
||||
"document_info": {
|
||||
"content_type": content_type,
|
||||
"encoding": extraction_result["used_encoding"],
|
||||
"statistics": stats,
|
||||
},
|
||||
"content": content,
|
||||
}
|
||||
|
||||
return json.dumps(json_data, indent=2, ensure_ascii=False)
|
||||
|
||||
elif output_format.lower() == "html":
|
||||
html_parts = []
|
||||
html_parts.append("<html><head><meta charset='utf-8'></head><body>")
|
||||
html_parts.append("<h1>Text Document Content</h1>")
|
||||
html_parts.append(f"<p><strong>File Type:</strong> {content_type}</p>")
|
||||
html_parts.append(f"<p><strong>Encoding:</strong> {extraction_result['used_encoding']}</p>")
|
||||
html_parts.append("<h2>Statistics</h2>")
|
||||
html_parts.append("<ul>")
|
||||
html_parts.append(f"<li>Characters: {stats['character_count']:,}</li>")
|
||||
html_parts.append(f"<li>Lines: {stats['line_count']:,}</li>")
|
||||
html_parts.append(f"<li>Words: {stats['word_count']:,}</li>")
|
||||
html_parts.append(f"<li>Empty lines: {stats['empty_lines']:,}</li>")
|
||||
html_parts.append(f"<li>Average line length: {stats['avg_line_length']} characters</li>")
|
||||
html_parts.append("</ul>")
|
||||
html_parts.append("<h2>Content</h2>")
|
||||
html_parts.append(f"<pre><code>{content}</code></pre>")
|
||||
html_parts.append("</body></html>")
|
||||
|
||||
return "".join(html_parts)
|
||||
|
||||
else: # text format
|
||||
text_parts = []
|
||||
text_parts.append(f"Text Document Content\n{'=' * 50}\n")
|
||||
text_parts.append(f"File Type: {content_type}\n")
|
||||
text_parts.append(f"Encoding: {extraction_result['used_encoding']}\n")
|
||||
text_parts.append("\nStatistics:\n")
|
||||
text_parts.append(f" Characters: {stats['character_count']:,}\n")
|
||||
text_parts.append(f" Lines: {stats['line_count']:,}\n")
|
||||
text_parts.append(f" Words: {stats['word_count']:,}\n")
|
||||
text_parts.append(f" Empty lines: {stats['empty_lines']:,}\n")
|
||||
text_parts.append(f" Average line length: {stats['avg_line_length']} characters\n")
|
||||
text_parts.append(f"\nContent:\n{'-' * 30}\n{content}")
|
||||
|
||||
return "".join(text_parts)
|
||||
|
||||
def mcp_extract_text_content(
|
||||
self,
|
||||
file_path: str = Field(description="Path to the text document file to extract content from"),
|
||||
output_format: Literal["markdown", "json", "html", "text"] = Field(
|
||||
default="markdown", description="Output format: 'markdown', 'json', 'html', or 'text'"
|
||||
),
|
||||
encoding: str | None = Field(default=None, description="Specific encoding to use (None for auto-detection)"),
|
||||
max_content_length: int | None = Field(
|
||||
default=None, description="Maximum length of content to include in output (None for no limit)"
|
||||
),
|
||||
) -> ActionResponse:
|
||||
"""Extract content from text documents with encoding detection and analysis.
|
||||
|
||||
This tool provides comprehensive text document content extraction with support for:
|
||||
- Various text file formats (TXT, MD, CSV, JSON, XML, source code, etc.)
|
||||
- Automatic encoding detection with fallback options
|
||||
- Content type detection and analysis
|
||||
- Comprehensive text statistics
|
||||
- LLM-optimized output formatting
|
||||
- Binary file detection and handling
|
||||
|
||||
Args:
|
||||
file_path: Path to the text file
|
||||
output_format: Desired output format
|
||||
encoding: Specific encoding to use
|
||||
max_content_length: Maximum content length to include
|
||||
|
||||
Returns:
|
||||
ActionResponse with extracted content, metadata, and file analysis
|
||||
"""
|
||||
try:
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(file_path, FieldInfo):
|
||||
file_path = file_path.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
if isinstance(encoding, FieldInfo):
|
||||
encoding = encoding.default
|
||||
if isinstance(max_content_length, FieldInfo):
|
||||
max_content_length = max_content_length.default
|
||||
|
||||
# Validate input file
|
||||
file_path: Path = self._validate_file_path(file_path)
|
||||
self._color_log(f"Processing text document: {file_path.name}", Color.cyan)
|
||||
|
||||
# Extract content from text file
|
||||
extraction_result = self._extract_text_content(file_path, encoding)
|
||||
|
||||
# Check if file appears to be binary
|
||||
if extraction_result["encoding_info"]["is_binary"]:
|
||||
self._color_log("Warning: File appears to contain binary data", Color.yellow)
|
||||
|
||||
# Format content for LLM consumption
|
||||
formatted_content = self._format_content_for_llm(extraction_result, output_format, max_content_length)
|
||||
|
||||
# Prepare metadata
|
||||
file_stats = file_path.stat()
|
||||
|
||||
# Create text-specific metadata
|
||||
text_metadata = {
|
||||
"content_type": extraction_result["content_type"],
|
||||
"encoding_info": extraction_result["encoding_info"],
|
||||
"text_statistics": extraction_result["statistics"],
|
||||
"used_encoding": extraction_result["used_encoding"],
|
||||
"encoding_fallback": extraction_result.get("encoding_fallback", False),
|
||||
"content_truncated": max_content_length and len(extraction_result["content"]) > max_content_length,
|
||||
"original_content_length": extraction_result["statistics"]["character_count"],
|
||||
}
|
||||
|
||||
document_metadata = DocumentMetadata(
|
||||
file_name=file_path.name,
|
||||
file_size=file_stats.st_size,
|
||||
file_type=file_path.suffix.lower() or ".txt",
|
||||
absolute_path=str(file_path.absolute()),
|
||||
page_count=extraction_result["statistics"]["line_count"], # Use line count as "page" count
|
||||
processing_time=extraction_result["processing_time"],
|
||||
extracted_images=[], # Text files don't contain images
|
||||
extracted_media=[], # Text files don't contain media
|
||||
output_format=output_format,
|
||||
llm_enhanced=False,
|
||||
ocr_applied=False,
|
||||
)
|
||||
|
||||
# Combine standard and text-specific metadata
|
||||
combined_metadata = document_metadata.model_dump()
|
||||
combined_metadata.update(text_metadata)
|
||||
|
||||
self._color_log(
|
||||
f"Successfully extracted content from {file_path.name} "
|
||||
f"({extraction_result['statistics']['character_count']:,} characters, "
|
||||
f"{extraction_result['statistics']['line_count']:,} lines, "
|
||||
f"encoding: {extraction_result['used_encoding']})",
|
||||
Color.green,
|
||||
)
|
||||
|
||||
return ActionResponse(success=True, message=formatted_content, metadata=combined_metadata)
|
||||
|
||||
except FileNotFoundError as e:
|
||||
self.logger.error(f"File not found: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False, message=f"File not found: {str(e)}", metadata={"error_type": "file_not_found"}
|
||||
)
|
||||
except ValueError as e:
|
||||
self.logger.error(f"Invalid input: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}",
|
||||
metadata={"error_type": "invalid_input"},
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Text extraction failed: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Text extraction failed: {str(e)}",
|
||||
metadata={"error_type": "extraction_error"},
|
||||
)
|
||||
|
||||
def mcp_list_supported_formats(self) -> ActionResponse:
|
||||
"""List all supported text formats for extraction.
|
||||
|
||||
Returns:
|
||||
ActionResponse with list of supported file formats and their descriptions
|
||||
"""
|
||||
supported_formats = {
|
||||
"TXT": "Plain text files (.txt, .text)",
|
||||
"Markdown": "Markdown documents (.md, .markdown)",
|
||||
"CSV/TSV": "Comma/Tab separated values (.csv, .tsv)",
|
||||
"JSON": "JSON data files (.json)",
|
||||
"XML": "XML documents (.xml)",
|
||||
"YAML": "YAML configuration files (.yaml, .yml)",
|
||||
"Source Code": "Programming language files (.py, .js, .html, .css, etc.)",
|
||||
"Configuration": "Config files (.ini, .cfg, .conf, .properties)",
|
||||
"Logs": "Log files (.log)",
|
||||
"Documentation": "Documentation files (.rst, .rtf)",
|
||||
"Scripts": "Script files (.sh, .bat, .ps1)",
|
||||
"Other Text": "Any file with text MIME type or detectable text content",
|
||||
}
|
||||
|
||||
format_list = "\n".join(
|
||||
[f"**{format_name}**: {description}" for format_name, description in supported_formats.items()]
|
||||
)
|
||||
|
||||
return ActionResponse(
|
||||
success=True,
|
||||
message=f"Supported text formats:\n\n{format_list}\n\n"
|
||||
"**Note:** The service automatically detects encoding and "
|
||||
"can handle files without standard extensions if they contain readable text.",
|
||||
metadata={
|
||||
"supported_formats": list(supported_formats.keys()),
|
||||
"total_formats": len(supported_formats),
|
||||
"encoding_detection": True,
|
||||
"binary_detection": True,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# Example usage and entry point
|
||||
if __name__ == "__main__":
|
||||
load_dotenv()
|
||||
|
||||
# Default arguments for testing
|
||||
args = ActionArguments(
|
||||
name="text_extraction_service",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
|
||||
)
|
||||
|
||||
# Initialize and run the text extraction service
|
||||
try:
|
||||
service = TextExtractionCollection(args)
|
||||
service.run()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}: {traceback.format_exc()}")
|
||||
@@ -0,0 +1,349 @@
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from aworld.config.conf import AgentConfig
|
||||
from aworld.logs.util import Color
|
||||
from aworld.models.llm import call_llm_model, get_llm_model
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
|
||||
|
||||
class CodeGenerationMetadata(BaseModel):
|
||||
"""Metadata for code generation results."""
|
||||
|
||||
model_name: str | None = None
|
||||
code_style: str | None = None
|
||||
code_length: int | None = None
|
||||
line_count: int | None = None
|
||||
processing_time_seconds: float | None = None
|
||||
temperature: float | None = None
|
||||
has_requirements: bool | None = None
|
||||
has_context: bool | None = None
|
||||
saved_file_path: str | None = None
|
||||
file_save_error: str | None = None
|
||||
error_type: str | None = None
|
||||
error_message: str | None = None
|
||||
|
||||
|
||||
class CodeCollection(ActionCollection):
|
||||
"""MCP service for generating executable Python code snippets using LLM.
|
||||
|
||||
Supports code generation for:
|
||||
- Data processing and analysis tasks
|
||||
- Algorithm implementations
|
||||
- Utility functions and scripts
|
||||
- Problem-solving code snippets
|
||||
- Educational programming examples
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
|
||||
# Initialize code generation model configuration
|
||||
self._llm_config = AgentConfig(
|
||||
llm_provider="openai",
|
||||
llm_model_name=os.getenv("CODE_LLM_MODEL_NAME", "anthropic/claude-sonnet-4"),
|
||||
llm_api_key=os.getenv("CODE_LLM_API_KEY"),
|
||||
llm_base_url=os.getenv("CODE_LLM_BASE_URL"),
|
||||
)
|
||||
|
||||
self._color_log("Code Generation Service initialized", Color.green, "debug")
|
||||
self._color_log(f"Using model: {self._llm_config.llm_model_name}", Color.blue, "debug")
|
||||
|
||||
def _prepare_code_prompt(self, task_description: str, requirements: str = "", context: str = "") -> str:
|
||||
"""Prepare the code generation prompt with task description and optional requirements.
|
||||
|
||||
Args:
|
||||
task_description: The main task for code generation
|
||||
requirements: Optional specific requirements or constraints
|
||||
context: Optional additional context or background information
|
||||
|
||||
Returns:
|
||||
Formatted prompt string
|
||||
"""
|
||||
prompt_parts = [f"Task: {task_description}"]
|
||||
|
||||
if requirements:
|
||||
prompt_parts.append(f"Requirements: {requirements}")
|
||||
|
||||
if context:
|
||||
prompt_parts.append(f"Context: {context}")
|
||||
|
||||
return "\n\n".join(prompt_parts)
|
||||
|
||||
def _call_code_model(self, prompt: str, temperature: float = 0.1) -> str:
|
||||
"""Call the code generation model with the prepared prompt.
|
||||
|
||||
Args:
|
||||
prompt: The formatted prompt for code generation
|
||||
temperature: Model temperature for response variability
|
||||
|
||||
Returns:
|
||||
Generated code from the model
|
||||
|
||||
Raises:
|
||||
Exception: If model call fails
|
||||
"""
|
||||
response = call_llm_model(
|
||||
llm_model=get_llm_model(conf=self._llm_config),
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"You are an expert Python programmer. Generate clean, efficient, and "
|
||||
"well-documented Python code that solves the given task. "
|
||||
"Include proper error handling and follow Python best practices. "
|
||||
"Return only executable Python code with minimal explanatory comments."
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
temperature=temperature,
|
||||
)
|
||||
|
||||
return response.content
|
||||
|
||||
def _extract_python_code(self, response: str) -> str:
|
||||
"""Extract Python code from the model response.
|
||||
|
||||
Args:
|
||||
response: Raw response from the model
|
||||
|
||||
Returns:
|
||||
Extracted Python code
|
||||
"""
|
||||
# Remove markdown code blocks if present
|
||||
lines = response.strip().split("\n")
|
||||
|
||||
# Find code block boundaries
|
||||
start_idx = 0
|
||||
end_idx = len(lines)
|
||||
|
||||
for i, line in enumerate(lines):
|
||||
if line.strip().startswith("```python") or line.strip().startswith("```"):
|
||||
start_idx = i + 1
|
||||
break
|
||||
|
||||
for i in range(len(lines) - 1, -1, -1):
|
||||
if lines[i].strip() == "```":
|
||||
end_idx = i
|
||||
break
|
||||
|
||||
# Extract the code
|
||||
code_lines = lines[start_idx:end_idx]
|
||||
return "\n".join(code_lines).strip()
|
||||
|
||||
def mcp_generate_python_code(
|
||||
self,
|
||||
task_description: str = Field(description="Description of the programming task or problem to solve"),
|
||||
requirements: str = Field(
|
||||
default="", description="Specific requirements, constraints, or specifications for the code"
|
||||
),
|
||||
context: str = Field(default="", description="Additional context or background information"),
|
||||
temperature: float = Field(
|
||||
default=0.1,
|
||||
description="Model temperature for code generation (0.0-1.0, lower = more deterministic)",
|
||||
ge=0.0,
|
||||
le=1.0,
|
||||
),
|
||||
code_style: Literal["minimal", "documented", "verbose"] = Field(
|
||||
default="documented",
|
||||
description="Style of generated code: minimal (concise), documented (with comments), verbose (detailed)",
|
||||
),
|
||||
save_to_file_path: str | None = Field(
|
||||
default=None,
|
||||
description="Optional. Path to save the generated Python snippet. e.g., 'output/generated_script.py'",
|
||||
),
|
||||
) -> ActionResponse:
|
||||
"""Generate executable Python code snippets based on task description.
|
||||
|
||||
This tool provides comprehensive code generation capabilities for:
|
||||
- Solve simple math tasks and validations
|
||||
- Data processing and analysis tasks
|
||||
- Algorithm implementations and optimizations
|
||||
- Utility functions and helper scripts
|
||||
- Problem-solving code snippets
|
||||
- Educational programming examples
|
||||
- API integrations and automation scripts
|
||||
|
||||
Strengths:
|
||||
- Generates clean, executable Python code
|
||||
- Follows modern Python best practices (>=3.11)
|
||||
- Includes proper error handling
|
||||
- Supports various coding styles and complexity levels
|
||||
|
||||
Limitations:
|
||||
- Cannot execute or test the generated code
|
||||
- May require manual adjustments for specific environments
|
||||
- Limited to Python programming language
|
||||
|
||||
Args:
|
||||
task_description: Clear description of the programming task
|
||||
requirements: Specific requirements or constraints
|
||||
context: Additional context or background information
|
||||
temperature: Model temperature controlling randomness
|
||||
code_style: Style preference for the generated code
|
||||
save_to_file_path: Optional. If provided, saves the generated code to this path within the workspace.
|
||||
|
||||
Returns:
|
||||
ActionResponse with generated Python code and metadata
|
||||
"""
|
||||
try:
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(task_description, FieldInfo):
|
||||
task_description = task_description.default
|
||||
if isinstance(requirements, FieldInfo):
|
||||
requirements = requirements.default
|
||||
if isinstance(context, FieldInfo):
|
||||
context = context.default
|
||||
if isinstance(temperature, FieldInfo):
|
||||
temperature = temperature.default
|
||||
if isinstance(code_style, FieldInfo):
|
||||
code_style = code_style.default
|
||||
if isinstance(save_to_file_path, FieldInfo):
|
||||
save_to_file_path = save_to_file_path.default
|
||||
|
||||
# Validate input
|
||||
if not task_description or not task_description.strip():
|
||||
raise ValueError("Task description is required for code generation")
|
||||
|
||||
self._color_log(f"Generating code for: {task_description[:100]}...", Color.cyan)
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Prepare the code generation prompt
|
||||
prompt = self._prepare_code_prompt(task_description, requirements, context)
|
||||
|
||||
# Enhance prompt based on code style
|
||||
if code_style == "minimal":
|
||||
prompt += "\n\nGenerate concise, minimal code without extensive comments."
|
||||
elif code_style == "verbose":
|
||||
prompt += "\n\nGenerate detailed code with comprehensive comments and explanations."
|
||||
elif code_style == "documented":
|
||||
prompt += "\n\nGenerate well-documented code with clear comments and docstrings."
|
||||
|
||||
# Call the code generation model
|
||||
raw_response = self._call_code_model(prompt, temperature)
|
||||
|
||||
# Extract clean Python code
|
||||
generated_code = self._extract_python_code(raw_response)
|
||||
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
# Populate metadata fields
|
||||
metadata = CodeGenerationMetadata(
|
||||
model_name=self._llm_config.llm_model_name,
|
||||
code_style=code_style,
|
||||
code_length=len(generated_code),
|
||||
line_count=len(generated_code.split("\n")),
|
||||
processing_time_seconds=round(processing_time, 2),
|
||||
temperature=temperature,
|
||||
has_requirements=bool(requirements.strip()),
|
||||
has_context=bool(context.strip()),
|
||||
)
|
||||
|
||||
# Save the generated code to a file if path is provided
|
||||
if save_to_file_path:
|
||||
try:
|
||||
# Use _validate_file_path to ensure path is within workspace and get absolute path
|
||||
# The check_existence=False allows creating a new file.
|
||||
output_file_path_obj = Path(self._validate_file_path(save_to_file_path))
|
||||
|
||||
# Ensure parent directories exist
|
||||
output_file_path_obj.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with open(output_file_path_obj, "w", encoding="utf-8") as f:
|
||||
f.write(generated_code)
|
||||
|
||||
metadata.saved_file_path = str(output_file_path_obj)
|
||||
self._color_log(f"Generated code also saved to: {output_file_path_obj}", Color.blue)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to save code to file '{save_to_file_path}': {str(e)}")
|
||||
metadata.file_save_error = str(e)
|
||||
|
||||
self._color_log(
|
||||
f"Successfully generated code ({metadata.code_length} characters, "
|
||||
f"{metadata.processing_time_seconds:.2f}s)",
|
||||
Color.green,
|
||||
)
|
||||
|
||||
return ActionResponse(success=True, message=generated_code, metadata=metadata.model_dump(exclude_none=True))
|
||||
|
||||
except ValueError as e:
|
||||
self.logger.error(f"Invalid input: {str(e)}")
|
||||
metadata.error_type = "invalid_input"
|
||||
metadata.error_message = str(e)
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}",
|
||||
metadata=metadata.model_dump(exclude_none=True),
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Code generation failed: {str(e)}: {traceback.format_exc()}")
|
||||
metadata.error_type = "generation_error"
|
||||
metadata.error_message = str(e)
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Code generation failed: {str(e)}",
|
||||
metadata=metadata.model_dump(exclude_none=True),
|
||||
)
|
||||
|
||||
def mcp_get_code_capabilities(self) -> ActionResponse:
|
||||
"""Get information about the code generation service capabilities.
|
||||
|
||||
Returns:
|
||||
ActionResponse with service capabilities and configuration
|
||||
"""
|
||||
capabilities = {
|
||||
"Data Processing": "Generate code for data manipulation, analysis, and visualization",
|
||||
"Algorithm Implementation": "Create efficient algorithms and data structures",
|
||||
"Utility Functions": "Build helper functions and reusable code components",
|
||||
"Problem Solving": "Generate solutions for programming challenges and tasks",
|
||||
"API Integration": "Create code for working with APIs and web services",
|
||||
"Automation Scripts": "Build scripts for task automation and workflow optimization",
|
||||
}
|
||||
|
||||
capability_list = "\n".join(
|
||||
[f"**{capability}**: {description}" for capability, description in capabilities.items()]
|
||||
)
|
||||
|
||||
metadata = {
|
||||
"model_name": self._llm_config.llm_model_name,
|
||||
"provider": self._llm_config.llm_provider,
|
||||
"supported_capabilities": list(capabilities.keys()),
|
||||
"total_capabilities": len(capabilities),
|
||||
"code_styles": ["minimal", "documented", "verbose"],
|
||||
"python_version": ">=3.11",
|
||||
"supported_language": "Python",
|
||||
}
|
||||
|
||||
return ActionResponse(
|
||||
success=True,
|
||||
message=f"Code Generation Service Capabilities:\n\n{capability_list}",
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
|
||||
# Example usage and entry point
|
||||
if __name__ == "__main__":
|
||||
load_dotenv()
|
||||
|
||||
# Default arguments for testing
|
||||
args = ActionArguments(
|
||||
name="code_generation_service",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
|
||||
)
|
||||
|
||||
# Initialize and run the code generation service
|
||||
try:
|
||||
service = CodeCollection(args)
|
||||
service.run()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}: {traceback.format_exc()}")
|
||||
@@ -0,0 +1,260 @@
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from typing import Literal
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import Field
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from aworld.config.conf import AgentConfig
|
||||
from aworld.logs.util import Color
|
||||
from aworld.models.llm import call_llm_model, get_llm_model
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
|
||||
|
||||
class GuardCollection(ActionCollection):
|
||||
"""MCP service for diagnosing and correcting (if necessary) the reasoning/thinking process already existed in the currect context, or avoid the potential loopholes in the future, towards solving the complex problem correctly, through powerful guarding model with sophisticated experience.
|
||||
The MUST Choice for the Thinking Process Reviewing phase, good at diagnosing the reasoning process in the context or giving valuable suggestions in advance.
|
||||
|
||||
Supports advanced guarding for reasoning process:
|
||||
- Identify potential loopholes or oversights in the reasoning process already existed in the currect context, while solving the complex problem.
|
||||
- If necessary, provide the corresponding supplements or guidance to the reasoning process in advance, to maneuver the reasoning/thinking process towards solving the complex problem in a proper direction.
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
env_path = "/Users/zhitianxie/PycharmProjects/AWorld_gaia_July/AWorld/examples/gaia/cmd/agent_deploy/gaia_agent/.env"
|
||||
load_dotenv(env_path, override=True, verbose=True)
|
||||
|
||||
# Initialize guarding model configuration
|
||||
self._llm_config = AgentConfig(
|
||||
llm_provider="openai",
|
||||
# llm_model_name="google/gemini-2.5-flash-preview-05-20:thinking",
|
||||
llm_model_name=os.getenv("GUARD_LLM_MODEL_NAME", "deepseek/deepseek-r1-0528:free"),
|
||||
llm_api_key=os.getenv("GUARD_LLM_API_KEY"),
|
||||
llm_base_url=os.getenv("GUARD_LLM_BASE_URL"),
|
||||
)
|
||||
|
||||
self._color_log("Intelligence Guard Service initialized", Color.green, "debug")
|
||||
self._color_log(f"Using model: {self._llm_config.llm_model_name}", Color.blue, "debug")
|
||||
|
||||
def _prepare_guarding_prompt(self, question: str, original_task: str = "") -> str:
|
||||
"""Prepare the guarding prompt with question and optional context.
|
||||
|
||||
Args:
|
||||
question: The main question for guarding the reasoning process, such as 'is there any potential loopholes or oversights in the reasoning process?'
|
||||
original_task: Optional original task description for context
|
||||
|
||||
Returns:
|
||||
Formatted prompt string
|
||||
"""
|
||||
if original_task:
|
||||
return f"Original Task: {original_task}\n\nQuestion: {question}"
|
||||
return f"Question: {question}"
|
||||
|
||||
def _call_guarding_model(self, prompt: str, temperature: float = 0.1) -> str:
|
||||
"""Call the guarding model with the prepared prompt.
|
||||
|
||||
Args:
|
||||
prompt: The formatted prompt for guarding the reasoning process
|
||||
temperature: Model temperature for response variability
|
||||
|
||||
Returns:
|
||||
guarding result from the model
|
||||
|
||||
Raises:
|
||||
Exception: If model call fails
|
||||
"""
|
||||
response = call_llm_model(
|
||||
llm_model=get_llm_model(conf=self._llm_config),
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"## Your Role\n"
|
||||
"You are an expert at identifying the potential loopholes or oversights"
|
||||
"of the current reasoning process while solving the complex problem.\n\n "
|
||||
"## Your Task: \n"
|
||||
"Based on the gathered information retrieved from the internet, and the reasoning process already"
|
||||
"generated towards solving a complex task, you need to do the following 1 or 2 things, to guarntee the quality of the reasoning process, and a clear final answer: \n"
|
||||
" 1. Provide your diagnosing result on the generated reasoning process and the corresponding the correction if necessary;\n"
|
||||
" 2. Provide your insight and supplements in advance to avoid the potential loopholes or oversights in the future;\n\n"
|
||||
"## Requirements: \n"
|
||||
" 1. If the reasoning process already generated is complete and correct in your opinion, just say 'No loopholes or oversights found'. \n"
|
||||
" 2. If the reasoning process already generated contains the materials that may lead to the potential logic mistake or lack of some important guardrails in your opinion, you may give a hint to the current reasoning process, with the necessary supplements.\n"
|
||||
" 3. If the reasoning process already generated is seriously incorrect in your opinion, you may give the turn signal to the reasoning process, to maneuver the reasoning process towards solving the complex problem correctly. \n\n"
|
||||
"## Restriction: \n"
|
||||
" 1. Please do not make judgments about the authenticity of externally sourced information obtained through searches, as this is not part of your job responsibilities;\n"
|
||||
" 2. Do not make additional inferences or assumptions about the content of such information itself.\n"
|
||||
" 3. If the question lacks necessary details/data/clues in your opinion, you may ask for more details.\n\n"
|
||||
"## Example 1:\n"
|
||||
" Question: Is my reasoning process correct?\n"
|
||||
" Reasoning Process: (nothing specified)\n"
|
||||
" Your Identification Result: Your question lacks some information, please provide me more details so I can help you.\n\n"
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
temperature=temperature,
|
||||
)
|
||||
|
||||
return response.content
|
||||
|
||||
def mcp_guarding_reasoning_process(
|
||||
self,
|
||||
question: str = Field(
|
||||
description="The input question for diagnosing the completeness/correctness of the reasoning process.\n"
|
||||
"For example: based on the staged/phased information/data concluded as 1.xxxx 2. xxxx 3. xxxx...., is there any faults in the current reasoning process aaaaaa"
|
||||
"that should be corrected? Or is there any loopholes or oversights that should be emphized in advance, towards solving the bbbbb problem?\n"
|
||||
"Requirement: This input question should include a clear question with the necessary details/data/clues from the previous context"
|
||||
"(such as the key information/data retrieved from the internet), to present more clues to help the diagnosing process."
|
||||
"The more exact details/data contained in this input question, the better the diagnosing result will be."
|
||||
),
|
||||
original_task: str = Field(default="", description="The original task. This field is required and cannot be simplied, has to be true to the original task."),
|
||||
temperature: float = Field(
|
||||
default=0.1,
|
||||
description="Model temperature for response variability (0.0-1.0)",
|
||||
ge=0.0,
|
||||
le=1.0,
|
||||
),
|
||||
guarding_style: Literal["detailed", "concise", "step-by-step"] = Field(
|
||||
default="detailed",
|
||||
description="Style of guarding output: detailed analysis, concise summary, or step-by-step breakdown",
|
||||
),
|
||||
) -> ActionResponse:
|
||||
"""This tool provides advanced logic diagonsing and correcting ability, to improve the quality of the reasoning process that already exists in the current context, while solving the complex question:
|
||||
- Identify potential loopholes or oversights in the current reasoning process while solving the complex problem.
|
||||
- Providing the guidance, suggestions to the reasoning process, to correct the loopholes or oversights if identified in this shot.
|
||||
|
||||
Invoke Timing: During Thinking Process Reviewing, while diagnosing the reasoning process in the context or give valuable suggestions in advance, this tool is a reliable selection.
|
||||
|
||||
Strengths:
|
||||
- Be relatively sensitive to common logical traps in some mathematics or logic problems.
|
||||
|
||||
Weakness:
|
||||
- Inability to process media types: image, audio, or video.
|
||||
- Inability to check the correctness of the retrieved information from the internet.
|
||||
- Require precise description of problem context and settings, including the reasoning process, retrieved data and the complex task itself.
|
||||
|
||||
Args:
|
||||
question: The input question that invokes this tool to diagnose and/or correct the suspected reasoning process in the context
|
||||
original_task: Optional original task description for additional context
|
||||
temperature: Model temperature controlling response variability
|
||||
guarding_style: Style of guarding output format
|
||||
|
||||
Returns:
|
||||
ActionResponse with guarding result and processing metadata
|
||||
"""
|
||||
try:
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(question, FieldInfo):
|
||||
question = question.default
|
||||
if isinstance(original_task, FieldInfo):
|
||||
original_task = original_task.default
|
||||
if isinstance(temperature, FieldInfo):
|
||||
temperature = temperature.default
|
||||
if isinstance(guarding_style, FieldInfo):
|
||||
guarding_style = guarding_style.default
|
||||
|
||||
# Validate input
|
||||
if not question or not question.strip():
|
||||
raise ValueError("Question is required for guarding the complex problem reasoning process")
|
||||
|
||||
self._color_log(f"Processing guarding request: {question[:100]}...", Color.cyan)
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Prepare the guarding prompt
|
||||
prompt = self._prepare_guarding_prompt(question, original_task) ## 简单的原始问题+分配给mcp server的问题
|
||||
|
||||
# Enhance prompt based on guarding style
|
||||
if guarding_style == "step-by-step":
|
||||
prompt += "\n\nPlease provide a clear step-by-step breakdown of your reviewing process of the reasoning process."
|
||||
elif guarding_style == "concise":
|
||||
prompt += "\n\nPlease provide a concise and final guarding answer."
|
||||
elif guarding_style == "detailed":
|
||||
prompt += "\n\nPlease provide detailed reviewing analysis with comprehensive guarding."
|
||||
|
||||
# Call the guarding model
|
||||
guarding_result = self._call_guarding_model(prompt, temperature)
|
||||
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
# Prepare metadata
|
||||
metadata = {
|
||||
"model_name": self._llm_config.llm_model_name,
|
||||
"guarding_style": guarding_style,
|
||||
"response_length": len(guarding_result),
|
||||
}
|
||||
|
||||
self._color_log(
|
||||
f"Successfully completed guarding ({len(guarding_result)} characters, {processing_time:.2f}s)",
|
||||
Color.green,
|
||||
)
|
||||
|
||||
return ActionResponse(success=True, message=guarding_result, metadata=metadata)
|
||||
|
||||
except ValueError as e:
|
||||
self.logger.error(f"Invalid input: {str(e)}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}",
|
||||
metadata={"error_type": "invalid_input", "error_message": str(e)},
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Guarding failed: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Guarding failed: {str(e)}",
|
||||
metadata={"error_type": "guarding_error", "error_message": str(e)},
|
||||
)
|
||||
|
||||
def mcp_get_guarding_capabilities(self) -> ActionResponse:
|
||||
"""Get information about the guarding reasoning process service capabilities.
|
||||
|
||||
Returns:
|
||||
ActionResponse with service capabilities and configuration
|
||||
"""
|
||||
capabilities = {
|
||||
"Logic Loopholes Detecting": "Identifying the logic loopholes in the reasoning process already generated previsouly",
|
||||
"Detected Loopholes Correcting": "Correcting the logic loopholes identified in the reasoning process already generated previously",
|
||||
"Oversights Prevention": "Providing necessary supplements as hints to the currect reasoning process, to prevent the possible oversights in the future",
|
||||
}
|
||||
|
||||
capability_list = "\n".join(
|
||||
[f"**{capability}**: {description}" for capability, description in capabilities.items()]
|
||||
)
|
||||
|
||||
metadata = {
|
||||
"model_name": self._llm_config.llm_model_name,
|
||||
"provider": self._llm_config.llm_provider,
|
||||
"supported_capabilities": list(capabilities.keys()),
|
||||
"total_capabilities": len(capabilities),
|
||||
"guarding_styles": ["detailed", "concise", "step-by-step"],
|
||||
}
|
||||
|
||||
return ActionResponse(
|
||||
success=True,
|
||||
message=f"Intelligence guarding Service Capabilities:\n\n{capability_list}",
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
|
||||
# Example usage and entry point
|
||||
if __name__ == "__main__":
|
||||
load_dotenv()
|
||||
|
||||
# Default arguments for testing
|
||||
args = ActionArguments(
|
||||
name="intelligence_guarding_service",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
|
||||
)
|
||||
|
||||
# Initialize and run the intelligence guarding service
|
||||
try:
|
||||
service = GuardCollection(args)
|
||||
service.run()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}: {traceback.format_exc()}")
|
||||
@@ -0,0 +1,237 @@
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from typing import Literal
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import Field
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from aworld.config.conf import AgentConfig
|
||||
from aworld.logs.util import Color
|
||||
from aworld.models.llm import call_llm_model, get_llm_model
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
|
||||
|
||||
class ThinkCollection(ActionCollection):
|
||||
"""MCP service for complex problem reasoning using powerful reasoning models.
|
||||
|
||||
Supports advanced reasoning for:
|
||||
- Mathematical problems and proofs
|
||||
- Code contest and programming challenges
|
||||
- Logic puzzles and riddles
|
||||
- Competition-level STEM problems
|
||||
- Multi-step analytical reasoning
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
|
||||
# Initialize reasoning model configuration
|
||||
self._llm_config = AgentConfig(
|
||||
llm_provider="openai",
|
||||
# llm_model_name="google/gemini-2.5-flash-preview-05-20:thinking",
|
||||
llm_model_name=os.getenv("THINK_LLM_MODEL_NAME", "deepseek/deepseek-r1-0528:free"),
|
||||
llm_api_key=os.getenv("THINK_LLM_API_KEY"),
|
||||
llm_base_url=os.getenv("THINK_LLM_BASE_URL"),
|
||||
)
|
||||
|
||||
self._color_log("Intelligence Reasoning Service initialized", Color.green, "debug")
|
||||
self._color_log(f"Using model: {self._llm_config.llm_model_name}", Color.blue, "debug")
|
||||
|
||||
def _prepare_reasoning_prompt(self, question: str, original_task: str = "") -> str:
|
||||
"""Prepare the reasoning prompt with question and optional context.
|
||||
|
||||
Args:
|
||||
question: The main question for reasoning
|
||||
original_task: Optional original task description for context
|
||||
|
||||
Returns:
|
||||
Formatted prompt string
|
||||
"""
|
||||
if original_task:
|
||||
return f"Original Task: {original_task}\n\nQuestion: {question}"
|
||||
return f"Question: {question}"
|
||||
|
||||
def _call_reasoning_model(self, prompt: str, temperature: float = 0.3) -> str:
|
||||
"""Call the reasoning model with the prepared prompt.
|
||||
|
||||
Args:
|
||||
prompt: The formatted prompt for reasoning
|
||||
temperature: Model temperature for response variability
|
||||
|
||||
Returns:
|
||||
Reasoning result from the model
|
||||
|
||||
Raises:
|
||||
Exception: If model call fails
|
||||
"""
|
||||
response = call_llm_model(
|
||||
llm_model=get_llm_model(conf=self._llm_config),
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"You are an expert at solving complex problems including math, "
|
||||
"code contests, riddles, and puzzles. "
|
||||
"Provide detailed step-by-step reasoning and a clear final answer."
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
temperature=temperature,
|
||||
)
|
||||
|
||||
return response.content
|
||||
|
||||
def mcp_complex_problem_reasoning(
|
||||
self,
|
||||
question: str = Field(
|
||||
description="The input question for complex problem reasoning, such as math and code contest problems"
|
||||
),
|
||||
original_task: str = Field(default="", description="The original task description."),
|
||||
temperature: float = Field(
|
||||
default=0.3,
|
||||
description="Model temperature for response variability (0.0-1.0)",
|
||||
ge=0.0,
|
||||
le=1.0,
|
||||
),
|
||||
reasoning_style: Literal["detailed", "concise", "step-by-step"] = Field(
|
||||
default="detailed",
|
||||
description="Style of reasoning output: detailed analysis, concise summary, or step-by-step breakdown",
|
||||
),
|
||||
) -> ActionResponse:
|
||||
"""This tool provides comprehensive reasoning capabilities for:
|
||||
- Mathematical problems and proofs
|
||||
- Programming and algorithm challenges
|
||||
- Logic puzzles, brain teasers, and fun riddles
|
||||
- Competition-level STEM problems
|
||||
- Multi-step analytical reasoning tasks
|
||||
|
||||
Weakness:
|
||||
- Inability to process media types: image, audio, or video
|
||||
- Require precise description of problem context and settings
|
||||
|
||||
Args:
|
||||
question: The input question requiring complex reasoning
|
||||
original_task: Optional original task description for additional context
|
||||
temperature: Model temperature controlling response variability
|
||||
reasoning_style: Style of reasoning output format
|
||||
|
||||
Returns:
|
||||
ActionResponse with reasoning result and processing metadata
|
||||
"""
|
||||
try:
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(question, FieldInfo):
|
||||
question = question.default
|
||||
if isinstance(original_task, FieldInfo):
|
||||
original_task = original_task.default
|
||||
if isinstance(temperature, FieldInfo):
|
||||
temperature = temperature.default
|
||||
if isinstance(reasoning_style, FieldInfo):
|
||||
reasoning_style = reasoning_style.default
|
||||
|
||||
# Validate input
|
||||
if not question or not question.strip():
|
||||
raise ValueError("Question is required for complex problem reasoning")
|
||||
|
||||
self._color_log(f"Processing reasoning request: {question[:100]}...", Color.cyan)
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Prepare the reasoning prompt
|
||||
prompt = self._prepare_reasoning_prompt(question, original_task)
|
||||
|
||||
# Enhance prompt based on reasoning style
|
||||
if reasoning_style == "step-by-step":
|
||||
prompt += "\n\nPlease provide a clear step-by-step breakdown of your reasoning process."
|
||||
elif reasoning_style == "concise":
|
||||
prompt += "\n\nPlease provide a concise but complete reasoning and final answer."
|
||||
elif reasoning_style == "detailed":
|
||||
prompt += "\n\nPlease provide detailed analysis with comprehensive reasoning."
|
||||
|
||||
# Call the reasoning model
|
||||
reasoning_result = self._call_reasoning_model(prompt, temperature)
|
||||
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
# Prepare metadata
|
||||
metadata = {
|
||||
"model_name": self._llm_config.llm_model_name,
|
||||
"reasoning_style": reasoning_style,
|
||||
"response_length": len(reasoning_result),
|
||||
}
|
||||
|
||||
self._color_log(
|
||||
f"Successfully completed reasoning ({len(reasoning_result)} characters, {processing_time:.2f}s)",
|
||||
Color.green,
|
||||
)
|
||||
|
||||
return ActionResponse(success=True, message=reasoning_result, metadata=metadata)
|
||||
|
||||
except ValueError as e:
|
||||
self.logger.error(f"Invalid input: {str(e)}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}",
|
||||
metadata={"error_type": "invalid_input", "error_message": str(e)},
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Reasoning failed: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Reasoning failed: {str(e)}",
|
||||
metadata={"error_type": "reasoning_error", "error_message": str(e)},
|
||||
)
|
||||
|
||||
def mcp_get_reasoning_capabilities(self) -> ActionResponse:
|
||||
"""Get information about the reasoning service capabilities.
|
||||
|
||||
Returns:
|
||||
ActionResponse with service capabilities and configuration
|
||||
"""
|
||||
capabilities = {
|
||||
"Mathematical Problems": "Advanced mathematical reasoning, proofs, and calculations",
|
||||
"Code Contests": "Programming challenges, algorithm design, and optimization",
|
||||
"Logic Puzzles": "Brain teasers, riddles, and logical reasoning problems",
|
||||
"STEM Problems": "Competition-level science, technology, engineering, and math",
|
||||
"Multi-step Analysis": "Complex analytical reasoning with multiple interconnected steps",
|
||||
}
|
||||
|
||||
capability_list = "\n".join(
|
||||
[f"**{capability}**: {description}" for capability, description in capabilities.items()]
|
||||
)
|
||||
|
||||
metadata = {
|
||||
"model_name": self._llm_config.llm_model_name,
|
||||
"provider": self._llm_config.llm_provider,
|
||||
"supported_capabilities": list(capabilities.keys()),
|
||||
"total_capabilities": len(capabilities),
|
||||
"reasoning_styles": ["detailed", "concise", "step-by-step"],
|
||||
}
|
||||
|
||||
return ActionResponse(
|
||||
success=True,
|
||||
message=f"Intelligence Reasoning Service Capabilities:\n\n{capability_list}",
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
|
||||
# Example usage and entry point
|
||||
if __name__ == "__main__":
|
||||
load_dotenv()
|
||||
|
||||
# Default arguments for testing
|
||||
args = ActionArguments(
|
||||
name="intelligence_reasoning_service",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
|
||||
)
|
||||
|
||||
# Initialize and run the intelligence reasoning service
|
||||
try:
|
||||
service = ThinkCollection(args)
|
||||
service.run()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}: {traceback.format_exc()}")
|
||||
@@ -0,0 +1,555 @@
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from openai import OpenAI
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from aworld.logs.util import Color
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
|
||||
|
||||
class AudioMetadata(BaseModel):
|
||||
"""Metadata extracted from audio processing."""
|
||||
|
||||
file_name: str = Field(description="Original audio file name")
|
||||
file_size: int = Field(description="File size in bytes")
|
||||
file_type: str = Field(description="Audio file type/extension")
|
||||
absolute_path: str = Field(description="Absolute path to the audio file")
|
||||
duration: float | None = Field(default=None, description="Duration of audio in seconds")
|
||||
sample_rate: int | None = Field(default=None, description="Audio sample rate in Hz")
|
||||
channels: int | None = Field(default=None, description="Number of audio channels")
|
||||
bitrate: int | None = Field(default=None, description="Audio bitrate in kbps")
|
||||
codec: str | None = Field(default=None, description="Audio codec used")
|
||||
processing_time: float = Field(description="Time taken to process the audio in seconds")
|
||||
output_files: list[str] = Field(default_factory=list, description="Paths to generated output files")
|
||||
transcription: str | None = Field(default=None, description="Transcribed text from audio")
|
||||
word_count: int | None = Field(default=None, description="Number of words in transcription")
|
||||
output_format: str = Field(description="Format of the processed output")
|
||||
|
||||
|
||||
class AudioCollection(ActionCollection):
|
||||
"""MCP service for comprehensive audio processing using ffmpeg.
|
||||
|
||||
Supports various audio operations including:
|
||||
- Audio format conversion
|
||||
- Audio transcription (speech-to-text)
|
||||
- Audio metadata extraction
|
||||
- Audio quality enhancement
|
||||
- Audio trimming and editing
|
||||
- Audio analysis and feature extraction
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
self._audio_output_dir = self.workspace / "processed_audio"
|
||||
self._audio_output_dir.mkdir(exist_ok=True)
|
||||
|
||||
# Supported audio formats
|
||||
self.supported_extensions = {
|
||||
".mp3",
|
||||
".wav",
|
||||
".flac",
|
||||
".aac",
|
||||
".ogg",
|
||||
".m4a",
|
||||
".wma",
|
||||
".opus",
|
||||
".aiff",
|
||||
".au",
|
||||
".ra",
|
||||
".amr",
|
||||
}
|
||||
|
||||
self._color_log("Audio Processing Service initialized", Color.green, "debug")
|
||||
self._color_log(f"Audio output directory: {self._audio_output_dir}", Color.blue, "debug")
|
||||
|
||||
# Check ffmpeg availability
|
||||
self._check_ffmpeg_availability()
|
||||
|
||||
def _check_ffmpeg_availability(self) -> bool:
|
||||
"""Check if ffmpeg is available in the system.
|
||||
|
||||
Returns:
|
||||
bool: True if ffmpeg is available, False otherwise
|
||||
"""
|
||||
try:
|
||||
result = subprocess.run(["ffmpeg", "-version"], capture_output=True, text=True, timeout=10, check=False)
|
||||
if result.returncode == 0:
|
||||
self._color_log("FFmpeg is available", Color.green, "debug")
|
||||
else:
|
||||
self._color_log("FFmpeg not found in system PATH", Color.red)
|
||||
return result.returncode == 0
|
||||
except (subprocess.TimeoutExpired, FileNotFoundError):
|
||||
self._color_log(
|
||||
"FFmpeg not available or timeout, Please refer to https://github.com/inclusionAI/AWorld/blob/main/examples/gaia/README.md#system-tools-setup for more details",
|
||||
Color.red,
|
||||
)
|
||||
return False
|
||||
|
||||
def _prepare_audio_for_transcription(self, file_path: Path) -> Path:
|
||||
"""Prepare audio file for transcription by converting to optimal format.
|
||||
|
||||
Args:
|
||||
file_path: Path to the original audio file
|
||||
|
||||
Returns:
|
||||
Path to the prepared audio file (WAV format, 16kHz)
|
||||
"""
|
||||
output_path = self._audio_output_dir / f"{file_path.stem}_for_transcription.wav"
|
||||
|
||||
# Convert to WAV format with 16kHz sample rate for optimal transcription
|
||||
cmd = [
|
||||
"ffmpeg",
|
||||
"-i",
|
||||
str(file_path),
|
||||
"-ar",
|
||||
"16000", # 16kHz sample rate
|
||||
"-ac",
|
||||
"1", # Mono channel
|
||||
"-c:a",
|
||||
"pcm_s16le", # 16-bit PCM
|
||||
"-y",
|
||||
str(output_path),
|
||||
]
|
||||
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=300, check=False)
|
||||
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(f"Audio preparation for transcription failed: {result.stderr}")
|
||||
|
||||
return output_path
|
||||
|
||||
def _transcribe_with_whisper(self, audio_path: Path) -> dict[str, Any]:
|
||||
"""Transcribe audio using OpenAI Whisper.
|
||||
|
||||
Args:
|
||||
audio_path: Path to the audio file
|
||||
|
||||
Returns:
|
||||
Dictionary containing transcription results
|
||||
"""
|
||||
try:
|
||||
client: OpenAI = OpenAI(api_key=os.getenv("AUDIO_LLM_API_KEY"), base_url=os.getenv("AUDIO_LLM_BASE_URL"))
|
||||
|
||||
# Use the file for transcription
|
||||
with open(audio_path, "rb") as audio_file:
|
||||
transcription: str = client.audio.transcriptions.create(
|
||||
file=audio_file,
|
||||
model=os.getenv("AUDIO_LLM_MODEL_NAME"),
|
||||
response_format="text",
|
||||
)
|
||||
|
||||
return {"text": transcription.strip() if transcription else ""}
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Audio transcription failed: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
def _get_audio_metadata(self, file_path: Path) -> dict[str, Any]:
|
||||
"""Extract audio metadata using ffprobe.
|
||||
|
||||
Args:
|
||||
file_path: Path to the audio file
|
||||
|
||||
Returns:
|
||||
Dictionary containing audio metadata
|
||||
"""
|
||||
try:
|
||||
cmd = ["ffprobe", "-v", "quiet", "-print_format", "json", "-show_format", "-show_streams", str(file_path)]
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=30, check=False)
|
||||
|
||||
if result.returncode == 0:
|
||||
metadata = json.loads(result.stdout)
|
||||
|
||||
# Extract relevant audio information
|
||||
format_info = metadata.get("format", {})
|
||||
streams = metadata.get("streams", [])
|
||||
audio_stream = next((s for s in streams if s.get("codec_type") == "audio"), {})
|
||||
|
||||
return {
|
||||
"duration": float(format_info.get("duration", 0)),
|
||||
"sample_rate": int(audio_stream.get("sample_rate", 0)) if audio_stream.get("sample_rate") else None,
|
||||
"channels": int(audio_stream.get("channels", 0)) if audio_stream.get("channels") else None,
|
||||
"bitrate": int(format_info.get("bit_rate", 0)) // 1000 if format_info.get("bit_rate") else None,
|
||||
"codec": audio_stream.get("codec_name"),
|
||||
}
|
||||
else:
|
||||
self.logger.warning(f"Failed to extract metadata: {result.stderr}")
|
||||
return {}
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error extracting audio metadata: {str(e)}")
|
||||
return {}
|
||||
|
||||
def _trim_audio(self, input_path: Path, start_time: float, duration: float | None = None) -> Path:
|
||||
"""Trim audio file to specified time range.
|
||||
|
||||
Args:
|
||||
input_path: Path to input audio file
|
||||
start_time: Start time in seconds
|
||||
duration: Duration in seconds (if None, trim to end)
|
||||
|
||||
Returns:
|
||||
Path to trimmed audio file
|
||||
"""
|
||||
output_path = self._audio_output_dir / f"{input_path.stem}_trimmed{input_path.suffix}"
|
||||
|
||||
cmd = ["ffmpeg", "-i", str(input_path), "-ss", str(start_time), "-y"]
|
||||
|
||||
if duration is not None:
|
||||
cmd.extend(["-t", str(duration)])
|
||||
|
||||
cmd.extend(["-c", "copy", str(output_path)])
|
||||
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=300, check=False)
|
||||
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(f"Audio trimming failed: {result.stderr}")
|
||||
|
||||
return output_path
|
||||
|
||||
def mcp_transcribe_audio(
|
||||
self,
|
||||
file_path: str = Field(description="Path to the audio file to transcribe"),
|
||||
model_size: Literal["tiny", "base", "small", "medium", "large"] = Field(
|
||||
default="base",
|
||||
description="Whisper model size: tiny (fastest), base (balanced), small, medium, large (most accurate)",
|
||||
),
|
||||
output_format: Literal["text", "detailed", "segments"] = Field(
|
||||
default="text",
|
||||
description="Output format: 'text' (plain text), 'detailed' (with metadata), 'segments' (timestamped)",
|
||||
),
|
||||
) -> ActionResponse:
|
||||
"""Transcribe audio file to text using OpenAI Whisper.
|
||||
|
||||
This tool converts speech in audio files to text with high accuracy.
|
||||
Supports multiple languages and provides various output formats including
|
||||
timestamped segments for detailed analysis.
|
||||
|
||||
Args:
|
||||
file_path: Path to the audio file to transcribe
|
||||
model_size: Whisper model size affecting speed vs accuracy trade-off
|
||||
output_format: Format of transcription output
|
||||
|
||||
Returns:
|
||||
ActionResponse with transcribed text and detailed metadata
|
||||
"""
|
||||
try:
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(file_path, FieldInfo):
|
||||
file_path = file_path.default
|
||||
if isinstance(model_size, FieldInfo):
|
||||
model_size = model_size.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Validate input file
|
||||
file_path: Path = self._validate_file_path(file_path)
|
||||
self._color_log(f"Transcribing audio: {file_path.name}", Color.cyan)
|
||||
|
||||
# Get original metadata
|
||||
original_metadata = self._get_audio_metadata(file_path)
|
||||
|
||||
# Prepare audio for transcription
|
||||
prepared_audio = self._prepare_audio_for_transcription(file_path)
|
||||
|
||||
# Perform transcription
|
||||
transcription_result = self._transcribe_with_whisper(prepared_audio)
|
||||
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
# Prepare file statistics
|
||||
file_stats = file_path.stat()
|
||||
|
||||
# Count words in transcription
|
||||
word_count = len(transcription_result["text"].split()) if transcription_result["text"] else 0
|
||||
|
||||
# Create metadata object
|
||||
audio_metadata = AudioMetadata(
|
||||
file_name=file_path.name,
|
||||
file_size=file_stats.st_size,
|
||||
file_type=file_path.suffix.lower(),
|
||||
absolute_path=str(file_path.absolute()),
|
||||
duration=original_metadata.get("duration"),
|
||||
sample_rate=original_metadata.get("sample_rate"),
|
||||
channels=original_metadata.get("channels"),
|
||||
bitrate=original_metadata.get("bitrate"),
|
||||
codec=original_metadata.get("codec"),
|
||||
processing_time=processing_time,
|
||||
output_files=[str(prepared_audio)],
|
||||
transcription=transcription_result["text"],
|
||||
word_count=word_count,
|
||||
output_format=f"transcription_{output_format}",
|
||||
)
|
||||
|
||||
# Format output based on requested format
|
||||
if output_format == "text":
|
||||
result_message = transcription_result["text"]
|
||||
elif output_format == "detailed":
|
||||
confidence_str = (
|
||||
f"{transcription_result['confidence']:.2f}" if transcription_result["confidence"] else "N/A"
|
||||
)
|
||||
result_message = (
|
||||
f"Transcription Results for {file_path.name}:\n\n"
|
||||
f"**Text:** {transcription_result['text']}\n\n"
|
||||
f"**Confidence:** {confidence_str}\n"
|
||||
f"**Word Count:** {word_count}\n"
|
||||
f"**Duration:** {original_metadata.get('duration', 0):.2f} seconds\n"
|
||||
f"**Model:** {model_size}\n"
|
||||
f"**Processing Time:** {processing_time:.2f} seconds"
|
||||
)
|
||||
elif output_format == "segments":
|
||||
segments_text = "\n".join(
|
||||
[
|
||||
f"[{seg.get('start', 0):.2f}s - {seg.get('end', 0):.2f}s]: {seg.get('text', '').strip()}"
|
||||
for seg in transcription_result.get("segments", [])
|
||||
]
|
||||
)
|
||||
result_message = (
|
||||
f"Timestamped Transcription for {file_path.name}:\n\n"
|
||||
f"{segments_text}\n\n"
|
||||
f"**Full Text:** {transcription_result['text']}"
|
||||
)
|
||||
else:
|
||||
result_message = transcription_result["text"]
|
||||
|
||||
# Clean up temporary file
|
||||
try:
|
||||
prepared_audio.unlink()
|
||||
except Exception:
|
||||
pass # Ignore cleanup errors
|
||||
|
||||
self._color_log(f"Transcription completed: {word_count} words, {processing_time:.2f}s", Color.green)
|
||||
|
||||
return ActionResponse(success=True, message=result_message, metadata=audio_metadata.model_dump())
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Audio transcription failed: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Audio transcription failed: {str(e)}",
|
||||
metadata={"error_type": "transcription_error"},
|
||||
)
|
||||
|
||||
def mcp_extract_audio_metadata(
|
||||
self,
|
||||
file_path: str = Field(description="Path to the audio file to analyze"),
|
||||
) -> ActionResponse:
|
||||
"""Extract comprehensive metadata from audio files.
|
||||
|
||||
This tool analyzes audio files and extracts detailed metadata including
|
||||
duration, sample rate, channels, bitrate, codec, and other technical information.
|
||||
|
||||
Args:
|
||||
file_path: Path to the audio file to analyze
|
||||
|
||||
Returns:
|
||||
ActionResponse with detailed audio metadata
|
||||
"""
|
||||
try:
|
||||
if isinstance(file_path, FieldInfo):
|
||||
file_path = file_path.default
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Validate input file
|
||||
file_path: Path = self._validate_file_path(file_path)
|
||||
self._color_log(f"Extracting metadata from: {file_path.name}", Color.cyan)
|
||||
|
||||
# Extract metadata
|
||||
metadata = self._get_audio_metadata(file_path)
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
# Prepare file statistics
|
||||
file_stats = file_path.stat()
|
||||
|
||||
# Create metadata object
|
||||
audio_metadata = AudioMetadata(
|
||||
file_name=file_path.name,
|
||||
file_size=file_stats.st_size,
|
||||
file_type=file_path.suffix.lower(),
|
||||
absolute_path=str(file_path.absolute()),
|
||||
duration=metadata.get("duration"),
|
||||
sample_rate=metadata.get("sample_rate"),
|
||||
channels=metadata.get("channels"),
|
||||
bitrate=metadata.get("bitrate"),
|
||||
codec=metadata.get("codec"),
|
||||
processing_time=processing_time,
|
||||
output_files=[],
|
||||
output_format="metadata",
|
||||
)
|
||||
|
||||
# Format metadata for LLM consumption
|
||||
result_message = (
|
||||
f"Audio Metadata for {file_path.name}:\n"
|
||||
f"Duration: {metadata.get('duration', 'Unknown'):.2f} seconds\n"
|
||||
f"Sample Rate: {metadata.get('sample_rate', 'Unknown')} Hz\n"
|
||||
f"Channels: {metadata.get('channels', 'Unknown')}\n"
|
||||
f"Bitrate: {metadata.get('bitrate', 'Unknown')} kbps\n"
|
||||
f"Codec: {metadata.get('codec', 'Unknown')}\n"
|
||||
f"File Size: {file_stats.st_size / 1024 / 1024:.2f} MB\n"
|
||||
f"Format: {file_path.suffix.upper()}"
|
||||
)
|
||||
|
||||
self._color_log(f"Metadata extraction completed in {processing_time:.2f}s", Color.green)
|
||||
|
||||
return ActionResponse(success=True, message=result_message, metadata=audio_metadata.model_dump())
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Metadata extraction failed: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Metadata extraction failed: {str(e)}",
|
||||
metadata={"error_type": "metadata_error"},
|
||||
)
|
||||
|
||||
def mcp_trim_audio(
|
||||
self,
|
||||
file_path: str = Field(description="Path to the audio file to trim"),
|
||||
start_time: float = Field(description="Start time in seconds"),
|
||||
duration: float | None = Field(default=None, description="Duration in seconds (if None, trim to end)"),
|
||||
) -> ActionResponse:
|
||||
"""Trim audio file to specified time range.
|
||||
|
||||
This tool cuts audio files to extract specific segments based on start time
|
||||
and duration. Useful for creating clips or removing unwanted sections.
|
||||
|
||||
Args:
|
||||
file_path: Path to the source audio file
|
||||
start_time: Start time in seconds for trimming
|
||||
duration: Duration of the trimmed segment (optional)
|
||||
|
||||
Returns:
|
||||
ActionResponse with trimmed audio file path and metadata
|
||||
"""
|
||||
try:
|
||||
if isinstance(file_path, FieldInfo):
|
||||
file_path = file_path.default
|
||||
if isinstance(start_time, FieldInfo):
|
||||
start_time = start_time.default
|
||||
if isinstance(duration, FieldInfo):
|
||||
duration = duration.default
|
||||
|
||||
process_start = time.time()
|
||||
|
||||
# Validate input file
|
||||
file_path: Path = self._validate_file_path(file_path)
|
||||
self._color_log(f"Trimming audio: {file_path.name}", Color.cyan)
|
||||
|
||||
# Get original metadata
|
||||
original_metadata = self._get_audio_metadata(file_path)
|
||||
|
||||
# Validate time parameters
|
||||
if start_time < 0:
|
||||
raise ValueError("Start time cannot be negative")
|
||||
if duration is not None and duration <= 0:
|
||||
raise ValueError("Duration must be positive")
|
||||
if original_metadata.get("duration") and start_time >= original_metadata["duration"]:
|
||||
raise ValueError("Start time exceeds audio duration")
|
||||
|
||||
# Trim audio
|
||||
output_path = self._trim_audio(file_path, start_time, duration)
|
||||
|
||||
# Get trimmed file metadata
|
||||
trimmed_metadata = self._get_audio_metadata(output_path)
|
||||
processing_time = time.time() - process_start
|
||||
|
||||
# Prepare file statistics
|
||||
file_stats = file_path.stat()
|
||||
|
||||
# Create metadata object
|
||||
audio_metadata = AudioMetadata(
|
||||
file_name=file_path.name,
|
||||
file_size=file_stats.st_size,
|
||||
file_type=file_path.suffix.lower(),
|
||||
absolute_path=str(file_path.absolute()),
|
||||
duration=trimmed_metadata.get("duration"),
|
||||
sample_rate=trimmed_metadata.get("sample_rate"),
|
||||
channels=trimmed_metadata.get("channels"),
|
||||
bitrate=trimmed_metadata.get("bitrate"),
|
||||
codec=trimmed_metadata.get("codec"),
|
||||
processing_time=processing_time,
|
||||
output_files=[str(output_path)],
|
||||
output_format="trimmed_audio",
|
||||
)
|
||||
|
||||
end_time = start_time + (duration or (original_metadata.get("duration", 0) - start_time))
|
||||
result_message = (
|
||||
f"Successfully trimmed {file_path.name}\n"
|
||||
f"Original duration: {original_metadata.get('duration', 0):.2f} seconds\n"
|
||||
f"Trimmed segment: {start_time:.2f}s - {end_time:.2f}s\n"
|
||||
f"New duration: {trimmed_metadata.get('duration', 0):.2f} seconds\n"
|
||||
f"Output file: {output_path.name}"
|
||||
)
|
||||
|
||||
self._color_log(f"Audio trimming completed in {processing_time:.2f}s", Color.green)
|
||||
|
||||
return ActionResponse(success=True, message=result_message, metadata=audio_metadata.model_dump())
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Audio trimming failed: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False, message=f"Audio trimming failed: {str(e)}", metadata={"error_type": "trimming_error"}
|
||||
)
|
||||
|
||||
def mcp_list_supported_formats(self) -> ActionResponse:
|
||||
"""List all supported audio formats for processing.
|
||||
|
||||
Returns:
|
||||
ActionResponse with list of supported audio formats and their descriptions
|
||||
"""
|
||||
supported_formats = {
|
||||
"MP3": "MPEG Audio Layer III (.mp3) - Most common compressed format",
|
||||
"WAV": "Waveform Audio File Format (.wav) - Uncompressed, high quality",
|
||||
"FLAC": "Free Lossless Audio Codec (.flac) - Lossless compression",
|
||||
"AAC": "Advanced Audio Coding (.aac) - Efficient compression",
|
||||
"OGG": "Ogg Vorbis (.ogg) - Open source compressed format",
|
||||
"M4A": "MPEG-4 Audio (.m4a) - Apple's preferred format",
|
||||
"WMA": "Windows Media Audio (.wma) - Microsoft format",
|
||||
"OPUS": "Opus Audio (.opus) - Modern, efficient codec",
|
||||
"AIFF": "Audio Interchange File Format (.aiff) - Apple's uncompressed format",
|
||||
"AU": "Sun Audio (.au) - Unix audio format",
|
||||
"RA": "RealAudio (.ra) - Streaming audio format",
|
||||
"AMR": "Adaptive Multi-Rate (.amr) - Mobile audio format",
|
||||
}
|
||||
|
||||
format_list = "\n".join(
|
||||
[f"**{format_name}**: {description}" for format_name, description in supported_formats.items()]
|
||||
)
|
||||
|
||||
return ActionResponse(
|
||||
success=True,
|
||||
message=f"Supported audio formats:\n\n{format_list}",
|
||||
metadata={
|
||||
"supported_formats": list(supported_formats.keys()),
|
||||
"total_formats": len(supported_formats),
|
||||
"ffmpeg_available": self._check_ffmpeg_availability(),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# Example usage and entry point
|
||||
if __name__ == "__main__":
|
||||
load_dotenv()
|
||||
|
||||
# Default arguments for testing
|
||||
args = ActionArguments(
|
||||
name="audio_processing_service",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
|
||||
)
|
||||
|
||||
# Initialize and run the audio processing service
|
||||
try:
|
||||
service = AudioCollection(args)
|
||||
service.run()
|
||||
print("Audio processing service started")
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}: {traceback.format_exc()}")
|
||||
@@ -0,0 +1,484 @@
|
||||
import base64
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytesseract
|
||||
from dotenv import load_dotenv
|
||||
from PIL import Image, ImageEnhance, ImageFilter
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from aworld.config.conf import AgentConfig
|
||||
from aworld.logs.util import Color
|
||||
from aworld.models.llm import call_llm_model, get_llm_model
|
||||
from aworld.models.model_response import ModelResponse
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
|
||||
|
||||
class ImageMetadata(BaseModel):
|
||||
"""Metadata extracted from image processing."""
|
||||
|
||||
file_name: str = Field(description="Original image file name")
|
||||
file_size: int = Field(description="File size in bytes")
|
||||
file_type: str = Field(description="Image file type/extension")
|
||||
absolute_path: str = Field(description="Absolute path to the image file")
|
||||
width: int | None = Field(default=None, description="Image width in pixels")
|
||||
height: int | None = Field(default=None, description="Image height in pixels")
|
||||
mode: str | None = Field(default=None, description="Image color mode (RGB, RGBA, L, etc.)")
|
||||
format: str | None = Field(default=None, description="Image format (JPEG, PNG, etc.)")
|
||||
has_transparency: bool = Field(default=False, description="Whether image has transparency")
|
||||
processing_time: float = Field(description="Time taken to process the image in seconds", exclude=True)
|
||||
output_files: list[str] = Field(default_factory=list, description="Paths to generated output files")
|
||||
extracted_text: str | None = Field(default=None, description="Text extracted via OCR")
|
||||
analysis_result: str | None = Field(default=None, description="AI analysis result")
|
||||
compression_ratio: float | None = Field(default=None, description="Compression ratio if optimized")
|
||||
output_format: str = Field(description="Format of the processed output")
|
||||
|
||||
|
||||
class ImageCollection(ActionCollection):
|
||||
"""MCP service for comprehensive image processing and analysis.
|
||||
|
||||
Supports various image operations including:
|
||||
- Metadata extraction
|
||||
- OCR (Optical Character Recognition)
|
||||
- AI-powered image analysis and reasoning
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
self._image_output_dir = self.workspace / "processed_images"
|
||||
self._image_output_dir.mkdir(exist_ok=True)
|
||||
|
||||
# Supported image formats
|
||||
self.supported_extensions = {
|
||||
".jpg",
|
||||
".jpeg",
|
||||
".png",
|
||||
".gif",
|
||||
".webp",
|
||||
".bmp",
|
||||
".tiff",
|
||||
".tif",
|
||||
".ico",
|
||||
".svg",
|
||||
}
|
||||
|
||||
self._llm_config = AgentConfig(
|
||||
llm_provider="openai",
|
||||
llm_model_name=os.getenv("IMAGE_LLM_MODEL_NAME", "gpt-4o"),
|
||||
llm_api_key=os.getenv("IMAGE_LLM_API_KEY"),
|
||||
llm_base_url=os.getenv("IMAGE_LLM_BASE_URL"),
|
||||
)
|
||||
|
||||
self._color_log("Image Processing Service initialized", Color.green, "debug")
|
||||
self._color_log(f"Image output directory: {self._image_output_dir}", Color.blue, "debug")
|
||||
|
||||
def _load_image(self, file_path: Path) -> Image.Image:
|
||||
"""Load image from file path.
|
||||
|
||||
Args:
|
||||
file_path: Path to the image file
|
||||
|
||||
Returns:
|
||||
PIL Image object
|
||||
"""
|
||||
try:
|
||||
return Image.open(file_path)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to load image {file_path}: {str(e)}") from e
|
||||
|
||||
def _get_image_metadata(self, image: Image.Image, file_path: Path) -> dict[str, Any]:
|
||||
"""Extract metadata from PIL Image object.
|
||||
|
||||
Args:
|
||||
image: PIL Image object
|
||||
file_path: Path to the original file
|
||||
|
||||
Returns:
|
||||
Dictionary containing image metadata
|
||||
"""
|
||||
return {
|
||||
"width": image.width,
|
||||
"height": image.height,
|
||||
"mode": image.mode,
|
||||
"format": image.format or file_path.suffix.upper().lstrip("."),
|
||||
"has_transparency": image.mode in ("RGBA", "LA") or "transparency" in image.info,
|
||||
}
|
||||
|
||||
def _optimize_image(self, image: Image.Image, max_size: tuple[int, int] | None = None) -> Image.Image:
|
||||
"""Optimize image for size and quality.
|
||||
|
||||
Args:
|
||||
image: PIL Image object
|
||||
max_size: Maximum dimensions (width, height)
|
||||
|
||||
Returns:
|
||||
Optimized PIL Image object
|
||||
"""
|
||||
optimized = image.copy()
|
||||
|
||||
# Resize if max_size specified
|
||||
if max_size:
|
||||
optimized.thumbnail(max_size, Image.Resampling.LANCZOS)
|
||||
|
||||
return optimized
|
||||
|
||||
def _perform_ocr(self, image: Image.Image) -> str:
|
||||
"""Perform OCR on image to extract text.
|
||||
|
||||
Args:
|
||||
image: PIL Image object
|
||||
|
||||
Returns:
|
||||
Extracted text string
|
||||
"""
|
||||
try:
|
||||
return pytesseract.image_to_string(image).strip()
|
||||
except ImportError:
|
||||
return "OCR not available - pytesseract not installed"
|
||||
except Exception as e:
|
||||
return f"OCR failed: {str(e)}"
|
||||
|
||||
def _analyze_with_ai(self, image_base64: str, task: str) -> str:
|
||||
"""Analyze image using AI model.
|
||||
|
||||
Args:
|
||||
image_base64: Base64 encoded image
|
||||
task: Analysis task description
|
||||
|
||||
Returns:
|
||||
AI analysis result
|
||||
"""
|
||||
try:
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"role": "text", "content": task},
|
||||
{"type": "image_url", "image_url": {"url": image_base64}},
|
||||
],
|
||||
},
|
||||
]
|
||||
response: ModelResponse = call_llm_model(
|
||||
llm_model=get_llm_model(conf=self._llm_config),
|
||||
messages=messages,
|
||||
temperature=float(os.getenv("LLM_TEMPERATURE", "1.0")),
|
||||
)
|
||||
self._color_log(f"{response.content=}", Color.green)
|
||||
return response.content
|
||||
except Exception as e:
|
||||
return f"AI analysis failed: {str(e)}"
|
||||
|
||||
def _image_to_base64(self, image: Image.Image, output_format: str = "JPEG") -> str:
|
||||
"""Convert PIL Image to base64 string.
|
||||
|
||||
Args:
|
||||
image: PIL Image object
|
||||
output_format: Output format (JPEG, PNG, etc.)
|
||||
|
||||
Returns:
|
||||
Base64 encoded image string
|
||||
"""
|
||||
buffer = BytesIO()
|
||||
|
||||
# Convert RGBA or P to RGB for JPEG
|
||||
if output_format.upper() == "JPEG":
|
||||
if image.mode in ("RGBA", "LA"):
|
||||
background = Image.new("RGB", image.size, (255, 255, 255))
|
||||
background.paste(image, mask=image.split()[-1] if image.mode == "RGBA" else None)
|
||||
image = background
|
||||
elif image.mode == "P":
|
||||
image = image.convert("RGB")
|
||||
|
||||
image.save(buffer, format=output_format, quality=85 if output_format.upper() == "JPEG" else None)
|
||||
|
||||
mime_type = f"image/{output_format.lower()}"
|
||||
img_base64 = base64.b64encode(buffer.getvalue()).decode()
|
||||
|
||||
return f"data:{mime_type};base64,{img_base64}"
|
||||
|
||||
def mcp_extract_text_ocr(
|
||||
self,
|
||||
file_path: str = Field(description="Path to the image file for OCR"),
|
||||
language: str = Field(default="eng", description="OCR language code (e.g., 'eng', 'spa', 'fra')"),
|
||||
preprocess: bool = Field(default=True, description="Whether to preprocess image for better OCR"),
|
||||
) -> ActionResponse:
|
||||
"""Extract text from images using Optical Character Recognition (OCR).
|
||||
|
||||
This tool uses Tesseract OCR to extract text content from images,
|
||||
with optional preprocessing to improve recognition accuracy.
|
||||
|
||||
Args:
|
||||
file_path: Path to the image file
|
||||
language: OCR language for better recognition
|
||||
preprocess: Whether to enhance image for OCR
|
||||
|
||||
Returns:
|
||||
ActionResponse with extracted text and metadata
|
||||
"""
|
||||
try:
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(file_path, FieldInfo):
|
||||
file_path = file_path.default
|
||||
if isinstance(language, FieldInfo):
|
||||
language = language.default
|
||||
if isinstance(preprocess, FieldInfo):
|
||||
preprocess = preprocess.default
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Validate input file
|
||||
file_path: Path = self._validate_file_path(file_path)
|
||||
self._color_log(f"Performing OCR on: {file_path.name}", Color.cyan)
|
||||
|
||||
# Load image
|
||||
image = self._load_image(file_path)
|
||||
original_metadata = self._get_image_metadata(image, file_path)
|
||||
|
||||
# Preprocess image for better OCR if requested
|
||||
if preprocess:
|
||||
# Convert to grayscale
|
||||
if image.mode != "L":
|
||||
image = image.convert("L")
|
||||
|
||||
# Enhance contrast
|
||||
enhancer = ImageEnhance.Contrast(image)
|
||||
image = enhancer.enhance(2.0)
|
||||
|
||||
# Apply slight sharpening
|
||||
image = image.filter(ImageFilter.SHARPEN)
|
||||
|
||||
# Perform OCR
|
||||
extracted_text = self._perform_ocr(image)
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
# Count words and characters
|
||||
word_count = len(extracted_text.split()) if extracted_text else 0
|
||||
char_count = len(extracted_text) if extracted_text else 0
|
||||
|
||||
# Create metadata object
|
||||
image_metadata = ImageMetadata(
|
||||
file_name=file_path.name,
|
||||
file_size=file_path.stat().st_size,
|
||||
file_type=file_path.suffix.lower(),
|
||||
absolute_path=str(file_path.absolute()),
|
||||
width=original_metadata["width"],
|
||||
height=original_metadata["height"],
|
||||
mode=original_metadata["mode"],
|
||||
format=original_metadata["format"],
|
||||
has_transparency=original_metadata["has_transparency"],
|
||||
processing_time=processing_time,
|
||||
output_files=[],
|
||||
extracted_text=extracted_text,
|
||||
output_format="ocr_text",
|
||||
)
|
||||
|
||||
if extracted_text:
|
||||
result_message = (
|
||||
f"OCR Results for {file_path.name}:\n\n"
|
||||
f"**Extracted Text:**\n{extracted_text}\n\n"
|
||||
f"**Statistics:**\n"
|
||||
f"- Words: {word_count}\n"
|
||||
f"- Characters: {char_count}\n"
|
||||
f"- Language: {language}\n"
|
||||
f"- Processing time: {processing_time:.2f}s"
|
||||
)
|
||||
else:
|
||||
result_message = (
|
||||
f"No text detected in {file_path.name}."
|
||||
" The image may not contain readable text or OCR preprocessing may be needed."
|
||||
)
|
||||
|
||||
self._color_log(f"OCR completed: {word_count} words extracted in {processing_time:.2f}s", Color.green)
|
||||
|
||||
return ActionResponse(success=True, message=result_message, metadata=image_metadata.model_dump())
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"OCR failed: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"OCR failed: {str(e)}",
|
||||
metadata={"error_type": "ocr_error"},
|
||||
)
|
||||
|
||||
def mcp_analyze_image_ai(
|
||||
self,
|
||||
file_path: str = Field(description="Path to the image file for AI analysis"),
|
||||
task: str = Field(
|
||||
default="Describe what you see in this image",
|
||||
description="Specific analysis task or question about the image",
|
||||
),
|
||||
) -> ActionResponse:
|
||||
"""Analyze image content using AI vision models.
|
||||
|
||||
This tool uses advanced AI models to analyze and describe image content,
|
||||
answer questions about images, or perform specific visual reasoning tasks.
|
||||
|
||||
Args:
|
||||
file_path: Path to the image file
|
||||
task: Specific analysis task or question
|
||||
|
||||
Returns:
|
||||
ActionResponse with AI analysis results and metadata
|
||||
"""
|
||||
try:
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(file_path, FieldInfo):
|
||||
file_path = file_path.default
|
||||
if isinstance(task, FieldInfo):
|
||||
task = task.default
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Validate input file
|
||||
file_path: Path = self._validate_file_path(file_path)
|
||||
self._color_log(f"Analyzing image with AI: {file_path.name}", Color.cyan)
|
||||
|
||||
# Load image
|
||||
image = self._load_image(file_path)
|
||||
original_metadata = self._get_image_metadata(image, file_path)
|
||||
|
||||
# Convert to base64 for AI analysis
|
||||
image_base64 = self._image_to_base64(image, "JPEG")
|
||||
|
||||
# Perform AI analysis
|
||||
analysis_result = self._analyze_with_ai(image_base64, task)
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
# Create metadata object
|
||||
metadata_dict = {
|
||||
"file_name": file_path.name,
|
||||
"file_size": file_path.stat().st_size,
|
||||
"file_type": file_path.suffix.lower(),
|
||||
"absolute_path": str(file_path.absolute()),
|
||||
"width": original_metadata["width"],
|
||||
"height": original_metadata["height"],
|
||||
"mode": original_metadata["mode"],
|
||||
"format": original_metadata["format"],
|
||||
"has_transparency": original_metadata["has_transparency"],
|
||||
"processing_time": processing_time,
|
||||
"output_files": [],
|
||||
"analysis_result": analysis_result,
|
||||
"output_format": "ai_analysis",
|
||||
}
|
||||
|
||||
image_metadata = ImageMetadata(**metadata_dict)
|
||||
|
||||
result_message = (
|
||||
f"AI Analysis Results for {file_path.name}:\n\n"
|
||||
f"**Task:** {task}\n\n"
|
||||
f"**Analysis:**\n{analysis_result}\n\n"
|
||||
f"**Image Info:**\n"
|
||||
f"- Dimensions: {original_metadata['width']}x{original_metadata['height']}\n"
|
||||
f"- Format: {original_metadata['format']}\n"
|
||||
f"- Processing time: {processing_time:.2f}s"
|
||||
)
|
||||
|
||||
self._color_log(f"AI analysis completed in {processing_time:.2f}s", Color.green)
|
||||
|
||||
return ActionResponse(success=True, message=result_message, metadata=image_metadata.model_dump())
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"AI image analysis failed: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"AI image analysis failed: {str(e)}",
|
||||
metadata={"error_type": "ai_analysis_error"},
|
||||
)
|
||||
|
||||
def mcp_get_image_metadata(
|
||||
self,
|
||||
file_path: str = Field(description="Path to the image file to analyze"),
|
||||
) -> ActionResponse:
|
||||
"""Extract comprehensive metadata from image files.
|
||||
|
||||
This tool analyzes image files and extracts detailed metadata including
|
||||
dimensions, format, color mode, file size, and other technical information.
|
||||
|
||||
Args:
|
||||
file_path: Path to the image file to analyze
|
||||
|
||||
Returns:
|
||||
ActionResponse with detailed image metadata
|
||||
"""
|
||||
try:
|
||||
if isinstance(file_path, FieldInfo):
|
||||
file_path = file_path.default
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Validate input file
|
||||
file_path: Path = self._validate_file_path(file_path)
|
||||
self._color_log(f"Extracting metadata from: {file_path.name}", Color.cyan)
|
||||
|
||||
# Load image
|
||||
image = self._load_image(file_path)
|
||||
metadata = self._get_image_metadata(image, file_path)
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
# Get file statistics
|
||||
file_stats = file_path.stat()
|
||||
|
||||
# Create metadata object
|
||||
image_metadata = ImageMetadata(
|
||||
file_name=file_path.name,
|
||||
file_size=file_stats.st_size,
|
||||
file_type=file_path.suffix.lower(),
|
||||
absolute_path=str(file_path.absolute()),
|
||||
width=metadata["width"],
|
||||
height=metadata["height"],
|
||||
mode=metadata["mode"],
|
||||
format=metadata["format"],
|
||||
has_transparency=metadata["has_transparency"],
|
||||
processing_time=processing_time,
|
||||
output_files=[],
|
||||
output_format="metadata",
|
||||
)
|
||||
|
||||
# Calculate additional info
|
||||
aspect_ratio = metadata["width"] / metadata["height"] if metadata["height"] > 0 else 0
|
||||
megapixels = (metadata["width"] * metadata["height"]) / 1_000_000
|
||||
|
||||
result_message = (
|
||||
f"Image Metadata for {file_path.name}:\n"
|
||||
f"Dimensions: {metadata['width']}x{metadata['height']} pixels\n"
|
||||
f"Aspect Ratio: {aspect_ratio:.2f}:1\n"
|
||||
f"Megapixels: {megapixels:.2f} MP\n"
|
||||
f"Color Mode: {metadata['mode']}\n"
|
||||
f"Format: {metadata['format']}\n"
|
||||
f"Has Transparency: {metadata['has_transparency']}\n"
|
||||
f"File Size: {file_stats.st_size / 1024 / 1024:.2f} MB\n"
|
||||
f"File Type: {file_path.suffix.upper()}"
|
||||
)
|
||||
|
||||
self._color_log(f"Metadata extraction completed in {processing_time:.2f}s", Color.green)
|
||||
|
||||
return ActionResponse(success=True, message=result_message, metadata=image_metadata.model_dump())
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Metadata extraction failed: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Metadata extraction failed: {str(e)}",
|
||||
metadata={"error_type": "metadata_error"},
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
load_dotenv()
|
||||
# Default arguments for testing
|
||||
args = ActionArguments(
|
||||
name="image_analysis_service",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
|
||||
)
|
||||
# Initialize and run the image analysis service
|
||||
try:
|
||||
service = ImageCollection(args)
|
||||
service.run()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {str(e)}")
|
||||
@@ -0,0 +1,992 @@
|
||||
"""
|
||||
Video MCP Service
|
||||
|
||||
This module provides MCP service functionality for video operations including:
|
||||
- Video content analysis with AI-powered insights
|
||||
- Video summarization and key point extraction
|
||||
- Keyframe extraction with scene detection
|
||||
- Subtitle extraction from video content
|
||||
|
||||
It handles various video formats with proper validation, error handling,
|
||||
and progress tracking while providing LLM-friendly formatted results.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from aworld.config.conf import AgentConfig
|
||||
from aworld.logs.util import Color
|
||||
from aworld.models.llm import call_llm_model, get_llm_model
|
||||
from aworld.models.model_response import ModelResponse
|
||||
from examples.gaia.mcp_collections.base import (
|
||||
ActionArguments,
|
||||
ActionCollection,
|
||||
ActionResponse,
|
||||
)
|
||||
|
||||
from ..utils import get_file_from_source
|
||||
|
||||
|
||||
class VideoAnalysisResult(BaseModel):
|
||||
"""Video analysis result model with structured data"""
|
||||
|
||||
video_source: str
|
||||
analysis_result: str
|
||||
frame_count: int
|
||||
duration_analyzed: float
|
||||
success: bool
|
||||
error: str | None = None
|
||||
|
||||
|
||||
class VideoSummaryResult(BaseModel):
|
||||
"""Video summary result model with structured data"""
|
||||
|
||||
video_source: str
|
||||
summary: str
|
||||
frame_count: int
|
||||
duration_analyzed: float
|
||||
success: bool
|
||||
error: str | None = None
|
||||
|
||||
|
||||
class KeyframeResult(BaseModel):
|
||||
"""Keyframe extraction result model with file information"""
|
||||
|
||||
frame_paths: list[str]
|
||||
frame_timestamps: list[float]
|
||||
output_directory: str
|
||||
frame_count: int
|
||||
target_time: float
|
||||
window_size: float
|
||||
success: bool
|
||||
error: str | None = None
|
||||
|
||||
|
||||
class VideoMetadata(BaseModel):
|
||||
"""Metadata for video operation results"""
|
||||
|
||||
operation: str
|
||||
video_source: str | None = None
|
||||
sample_rate: int | None = None
|
||||
start_time: float | None = None
|
||||
end_time: float | None = None
|
||||
target_time: float | None = None
|
||||
window_size: float | None = None
|
||||
output_directory: str | None = None
|
||||
frame_count: int | None = None
|
||||
execution_time: float | None = None
|
||||
error_type: str | None = None
|
||||
|
||||
|
||||
class VideoCollection(ActionCollection):
|
||||
"""MCP service for video operations with AI-powered analysis.
|
||||
|
||||
Provides video processing capabilities including:
|
||||
- AI-powered video content analysis
|
||||
- Video summarization with key insights
|
||||
- Keyframe extraction with scene detection
|
||||
- Subtitle extraction from video content
|
||||
- LLM-friendly result formatting
|
||||
- Error handling and logging
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
|
||||
# Initialize supported video extensions
|
||||
self.supported_extensions = {".mp4", ".avi", ".mov", ".mkv", ".webm", ".flv"}
|
||||
|
||||
# Video analysis prompts
|
||||
self.video_analyze_prompt = (
|
||||
"Input is a sequence of video frames. Given user's task: {task}. "
|
||||
"analyze the video content following these steps:\n"
|
||||
"1. Temporal sequence understanding\n"
|
||||
"2. Motion and action analysis\n"
|
||||
"3. Scene context interpretation\n"
|
||||
"4. Object and person tracking\n"
|
||||
)
|
||||
|
||||
self.video_summarize_prompt = (
|
||||
"Input is a sequence of video frames. "
|
||||
"Summarize the main content of the video. "
|
||||
"Include key points, main topics, and important visual elements. "
|
||||
)
|
||||
|
||||
self._color_log("Video service initialized", Color.green, "debug")
|
||||
|
||||
def _get_video_frames(
|
||||
self,
|
||||
video_source: str,
|
||||
sample_rate: int = 2,
|
||||
start_time: float = 0,
|
||||
end_time: float | None = None,
|
||||
) -> list[dict[str, any]]:
|
||||
"""Extract frames from video with given sample rate.
|
||||
|
||||
Args:
|
||||
video_source: Path or URL to the video file
|
||||
sample_rate: Number of frames to sample per second
|
||||
start_time: Start time of the video segment in seconds
|
||||
end_time: End time of the video segment in seconds
|
||||
|
||||
Returns:
|
||||
List of dictionaries containing frame data and timestamp
|
||||
|
||||
Raises:
|
||||
ValueError: When video file cannot be opened or is not valid
|
||||
"""
|
||||
try:
|
||||
# Get file with validation (only video files allowed)
|
||||
file_path, _, _ = get_file_from_source(
|
||||
video_source,
|
||||
max_size_mb=2500.0, # 2500MB limit for videos
|
||||
)
|
||||
|
||||
# Open video file
|
||||
video = cv2.VideoCapture(file_path) # pylint: disable=E1101
|
||||
if not video.isOpened():
|
||||
raise ValueError(f"Could not open video file: {file_path}")
|
||||
|
||||
fps = video.get(cv2.CAP_PROP_FPS) # pylint: disable=E1101
|
||||
frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) # pylint: disable=E1101
|
||||
video_duration = frame_count / fps
|
||||
|
||||
if end_time is None:
|
||||
end_time = video_duration
|
||||
|
||||
if start_time > end_time:
|
||||
raise ValueError("Start time cannot be greater than end time.")
|
||||
|
||||
if start_time < 0:
|
||||
start_time = 0
|
||||
|
||||
if end_time > video_duration:
|
||||
end_time = video_duration
|
||||
|
||||
start_frame = int(start_time * fps)
|
||||
end_frame = int(end_time * fps)
|
||||
|
||||
all_frames = []
|
||||
frames = []
|
||||
|
||||
# Calculate frame interval based on sample rate
|
||||
frame_interval = max(1, int(fps / sample_rate))
|
||||
|
||||
# Set the video capture to the start frame
|
||||
video.set(cv2.CAP_PROP_POS_FRAMES, start_frame) # pylint: disable=E1101
|
||||
|
||||
for i in range(start_frame, end_frame):
|
||||
ret, frame = video.read()
|
||||
if not ret:
|
||||
break
|
||||
|
||||
# Convert frame to JPEG format
|
||||
_, buffer = cv2.imencode(".jpg", frame) # pylint: disable=E1101
|
||||
frame_data = base64.b64encode(buffer).decode("utf-8")
|
||||
|
||||
# Add data URL prefix for JPEG image
|
||||
frame_data = f"data:image/jpeg;base64,{frame_data}"
|
||||
|
||||
all_frames.append({"data": frame_data, "time": i / fps})
|
||||
|
||||
for i in range(0, len(all_frames), frame_interval):
|
||||
frames.append(all_frames[i])
|
||||
|
||||
video.release()
|
||||
|
||||
# Clean up temporary file if it was created for a URL
|
||||
if (
|
||||
file_path != str(Path(video_source).resolve())
|
||||
and Path(file_path).exists()
|
||||
):
|
||||
Path(file_path).unlink()
|
||||
|
||||
if not frames:
|
||||
raise ValueError(
|
||||
f"Could not extract any frames from video: {video_source}"
|
||||
)
|
||||
|
||||
return frames
|
||||
|
||||
except Exception as e:
|
||||
self._color_log(
|
||||
f"Error extracting frames from {video_source}: {str(e)}", Color.red
|
||||
)
|
||||
raise
|
||||
|
||||
def _create_video_content(
|
||||
self, prompt: str, video_frames: list[dict[str, any]]
|
||||
) -> list[dict[str, any]]:
|
||||
"""Create uniform video format for querying LLM."""
|
||||
content = [{"type": "text", "text": prompt}]
|
||||
content.extend(
|
||||
[
|
||||
{"type": "image_url", "image_url": {"url": frame["data"]}}
|
||||
for frame in video_frames
|
||||
]
|
||||
)
|
||||
return content
|
||||
|
||||
def _format_analysis_output(
|
||||
self, result: VideoAnalysisResult, format_type: str = "markdown"
|
||||
) -> str:
|
||||
"""Format video analysis results for LLM consumption.
|
||||
|
||||
Args:
|
||||
result: Video analysis result
|
||||
format_type: Output format ('markdown', 'json', 'text')
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if not result.success:
|
||||
return f"Failed to analyze video: {result.error}"
|
||||
|
||||
if format_type == "json":
|
||||
return result.model_dump_json(indent=2)
|
||||
|
||||
elif format_type == "text":
|
||||
output_parts = [
|
||||
"Video Analysis Results",
|
||||
f"Source: {result.video_source}",
|
||||
f"Frames Analyzed: {result.frame_count}",
|
||||
f"Duration: {result.duration_analyzed:.2f} seconds",
|
||||
"",
|
||||
"Analysis:",
|
||||
result.analysis_result,
|
||||
]
|
||||
return "\n".join(output_parts)
|
||||
|
||||
else: # markdown (default)
|
||||
output_parts = [
|
||||
"# Video Analysis Results ✅",
|
||||
"",
|
||||
"## Video Information",
|
||||
f"**Source:** `{result.video_source}`",
|
||||
f"**Frames Analyzed:** {result.frame_count}",
|
||||
f"**Duration:** {result.duration_analyzed:.2f} seconds",
|
||||
"",
|
||||
"## Analysis Results",
|
||||
result.analysis_result,
|
||||
]
|
||||
return "\n".join(output_parts)
|
||||
|
||||
def _format_summary_output(
|
||||
self, result: VideoSummaryResult, format_type: str = "markdown"
|
||||
) -> str:
|
||||
"""Format video summary results for LLM consumption.
|
||||
|
||||
Args:
|
||||
result: Video summary result
|
||||
format_type: Output format ('markdown', 'json', 'text')
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if not result.success:
|
||||
return f"Failed to summarize video: {result.error}"
|
||||
|
||||
if format_type == "json":
|
||||
return result.model_dump_json(indent=2)
|
||||
|
||||
elif format_type == "text":
|
||||
output_parts = [
|
||||
"Video Summary",
|
||||
f"Source: {result.video_source}",
|
||||
f"Frames Analyzed: {result.frame_count}",
|
||||
f"Duration: {result.duration_analyzed:.2f} seconds",
|
||||
"",
|
||||
"Summary:",
|
||||
result.summary,
|
||||
]
|
||||
return "\n".join(output_parts)
|
||||
|
||||
else: # markdown (default)
|
||||
output_parts = [
|
||||
"# Video Summary ✅",
|
||||
"",
|
||||
"## Video Information",
|
||||
f"**Source:** `{result.video_source}`",
|
||||
f"**Frames Analyzed:** {result.frame_count}",
|
||||
f"**Duration:** {result.duration_analyzed:.2f} seconds",
|
||||
"",
|
||||
"## Summary",
|
||||
result.summary,
|
||||
]
|
||||
return "\n".join(output_parts)
|
||||
|
||||
def _format_keyframe_output(
|
||||
self, result: KeyframeResult, format_type: str = "markdown"
|
||||
) -> str:
|
||||
"""Format keyframe extraction results for LLM consumption.
|
||||
|
||||
Args:
|
||||
result: Keyframe extraction result
|
||||
format_type: Output format ('markdown', 'json', 'text')
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if not result.success:
|
||||
return f"Failed to extract keyframes: {result.error}"
|
||||
|
||||
if format_type == "json":
|
||||
return result.model_dump_json(indent=2)
|
||||
|
||||
elif format_type == "text":
|
||||
output_parts = [
|
||||
"Keyframe Extraction Results",
|
||||
f"Target Time: {result.target_time}s",
|
||||
f"Window Size: {result.window_size}s",
|
||||
f"Frames Extracted: {result.frame_count}",
|
||||
f"Output Directory: {result.output_directory}",
|
||||
"",
|
||||
"Frame Files:",
|
||||
]
|
||||
for i, (path, timestamp) in enumerate(
|
||||
zip(result.frame_paths, result.frame_timestamps), 1
|
||||
):
|
||||
output_parts.append(f"{i}. {path} (at {timestamp:.2f}s)")
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
else: # markdown (default)
|
||||
output_parts = [
|
||||
"# Keyframe Extraction Results ✅",
|
||||
"",
|
||||
"## Extraction Parameters",
|
||||
f"**Target Time:** {result.target_time}s",
|
||||
f"**Window Size:** {result.window_size}s",
|
||||
f"**Frames Extracted:** {result.frame_count}",
|
||||
f"**Output Directory:** `{result.output_directory}`",
|
||||
"",
|
||||
"## Extracted Frames",
|
||||
]
|
||||
|
||||
for i, (path, timestamp) in enumerate(
|
||||
zip(result.frame_paths, result.frame_timestamps), 1
|
||||
):
|
||||
output_parts.append(f"{i}. `{path}` (at {timestamp:.2f}s)")
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
def _analyze_frame_chunk(
|
||||
self, chunk_data: tuple[int, list, str]
|
||||
) -> tuple[int, str]:
|
||||
"""Analyze a chunk of video frames using LLM.
|
||||
|
||||
Args:
|
||||
chunk_data: Tuple containing (chunk_index, frames, question)
|
||||
|
||||
Returns:
|
||||
Tuple of (chunk_index, analysis_result)
|
||||
"""
|
||||
chunk_index, frames, question = chunk_data
|
||||
|
||||
try:
|
||||
content = self._create_video_content(
|
||||
self.video_analyze_prompt.format(task=question), frames
|
||||
)
|
||||
inputs = [{"role": "user", "content": content}]
|
||||
|
||||
response: ModelResponse = call_llm_model(
|
||||
get_llm_model(
|
||||
conf=AgentConfig(
|
||||
llm_provider="openai",
|
||||
llm_model_name=os.getenv("VIDEO_LLM_MODEL_NAME"),
|
||||
llm_api_key=os.getenv("VIDEO_LLM_API_KEY"),
|
||||
llm_base_url=os.getenv("VIDEO_LLM_BASE_URL"),
|
||||
)
|
||||
),
|
||||
inputs,
|
||||
temperature=float(os.getenv("VIDEO_LLM_TEMPERATURE", "1.0")),
|
||||
)
|
||||
analysis_result = response.content
|
||||
self._color_log(
|
||||
f"✅ Completed analysis for chunk {chunk_index + 1}", Color.green
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self._color_log(
|
||||
f"❌ LLM analysis error for chunk {chunk_index + 1}: {str(e)}",
|
||||
Color.yellow,
|
||||
)
|
||||
analysis_result = (
|
||||
f"Analysis failed for video segment {chunk_index + 1}: {str(e)}"
|
||||
)
|
||||
|
||||
return chunk_index, analysis_result
|
||||
|
||||
async def mcp_analyze_video(
|
||||
self,
|
||||
video_url: str = Field(description="Path or URL to the video file to analyze"),
|
||||
question: str = Field(description="Question or task for video analysis"),
|
||||
sample_rate: float = Field(
|
||||
default=1.0, description="Frame sampling rate (frames per second)"
|
||||
),
|
||||
start_time: float = Field(default=0.0, description="Start time in seconds"),
|
||||
end_time: float | None = Field(
|
||||
default=None, description="End time in seconds (None for full video)"
|
||||
),
|
||||
output_format: str = Field(
|
||||
default="markdown",
|
||||
description="Output format: 'markdown', 'json', or 'text'",
|
||||
),
|
||||
max_workers: int = Field(
|
||||
default=4, description="Maximum number of parallel workers for analysis"
|
||||
),
|
||||
) -> ActionResponse:
|
||||
"""Analyze video content using AI with parallel processing.
|
||||
|
||||
This tool provides comprehensive video analysis capabilities including:
|
||||
- Content understanding and description
|
||||
- Object and scene detection
|
||||
- Action and movement analysis
|
||||
- Temporal event tracking
|
||||
- Question-answering about video content
|
||||
- Parallel processing for faster analysis
|
||||
|
||||
Args:
|
||||
video_url: Path or URL to the video file
|
||||
question: Specific question or analysis task
|
||||
sample_rate: Frame sampling rate for analysis
|
||||
start_time: Start time of the video segment in seconds
|
||||
end_time: End time of the video segment in seconds
|
||||
output_format: Format for the response output
|
||||
max_workers: Maximum number of parallel workers
|
||||
|
||||
Returns:
|
||||
ActionResponse with video analysis results and metadata
|
||||
"""
|
||||
start_exec_time = time.time()
|
||||
|
||||
try:
|
||||
# Validate video file
|
||||
video_path = self._validate_file_path(video_url)
|
||||
|
||||
self._color_log(f"🎬 Analyzing video: {video_url}", Color.cyan)
|
||||
self._color_log(f"📋 Question: {question}", Color.blue)
|
||||
|
||||
# Extract video frames
|
||||
video_frames = self._get_video_frames(
|
||||
str(video_path), sample_rate, start_time, end_time
|
||||
)
|
||||
self._color_log(f"📸 Extracted {len(video_frames)} frames", Color.blue)
|
||||
|
||||
# Process frames in chunks of 64 frames for parallel analysis
|
||||
chunk_size = 128
|
||||
chunks = []
|
||||
|
||||
# Create chunks of `chunk_size` continuous frames
|
||||
for i in range(0, len(video_frames), chunk_size):
|
||||
chunk_frames = video_frames[i : i + chunk_size]
|
||||
chunks.append((i // chunk_size, chunk_frames, question))
|
||||
|
||||
self._color_log(
|
||||
f"🔄 Processing {len(chunks)} chunks with {max_workers} parallel workers",
|
||||
Color.blue,
|
||||
)
|
||||
|
||||
# Process chunks in parallel
|
||||
all_results = [None] * len(chunks) # Pre-allocate to maintain order
|
||||
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
# Submit all chunk analysis tasks
|
||||
future_to_chunk = {
|
||||
executor.submit(self._analyze_frame_chunk, chunk_data): chunk_data[
|
||||
0
|
||||
]
|
||||
for chunk_data in chunks
|
||||
}
|
||||
|
||||
# Collect results as they complete
|
||||
for future in as_completed(future_to_chunk):
|
||||
try:
|
||||
chunk_index, result = future.result()
|
||||
all_results[chunk_index] = (
|
||||
f"Result of video part {chunk_index + 1}: {result}"
|
||||
)
|
||||
except Exception as e:
|
||||
chunk_index = future_to_chunk[future]
|
||||
self._color_log(
|
||||
f"❌ Error processing chunk {chunk_index + 1}: {str(e)}",
|
||||
Color.red,
|
||||
)
|
||||
all_results[chunk_index] = (
|
||||
f"Result of video part {chunk_index + 1}: Analysis failed - {str(e)}"
|
||||
)
|
||||
|
||||
# Filter out None results and join
|
||||
analysis_result = "\n".join(
|
||||
[result for result in all_results if result is not None]
|
||||
)
|
||||
duration_analyzed = (
|
||||
end_time - start_time if end_time else len(video_frames) / sample_rate
|
||||
)
|
||||
|
||||
# Create result
|
||||
result = VideoAnalysisResult(
|
||||
video_source=video_url,
|
||||
analysis_result=analysis_result,
|
||||
frame_count=len(video_frames),
|
||||
duration_analyzed=duration_analyzed,
|
||||
success=True,
|
||||
error=None,
|
||||
)
|
||||
|
||||
# Format output for LLM
|
||||
message = self._format_analysis_output(result, output_format)
|
||||
execution_time = time.time() - start_exec_time
|
||||
|
||||
# Create metadata
|
||||
metadata = {
|
||||
"video_source": video_url,
|
||||
"frame_count": len(video_frames),
|
||||
"chunks_processed": len(chunks),
|
||||
"chunk_size": chunk_size,
|
||||
"parallel_workers": max_workers,
|
||||
"duration_analyzed": duration_analyzed,
|
||||
"sample_rate": sample_rate,
|
||||
"start_time": start_time,
|
||||
"end_time": end_time,
|
||||
"execution_time": execution_time,
|
||||
"output_format": output_format,
|
||||
"success": True,
|
||||
}
|
||||
|
||||
self._color_log(
|
||||
f"✅ Video analysis completed in {execution_time:.2f}s "
|
||||
f"({len(chunks)} chunks, {len(video_frames)} frames)",
|
||||
Color.green,
|
||||
)
|
||||
|
||||
return ActionResponse(success=True, message=message, metadata=metadata)
|
||||
|
||||
except Exception as e:
|
||||
execution_time = time.time() - start_exec_time
|
||||
error_msg = f"Video analysis failed: {str(e)}"
|
||||
self._color_log(f"❌ {error_msg}", Color.red)
|
||||
self.logger.error(f"{error_msg}: {traceback.format_exc()}")
|
||||
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=error_msg,
|
||||
metadata={
|
||||
"video_source": video_url,
|
||||
"execution_time": execution_time,
|
||||
"error": str(e),
|
||||
"success": False,
|
||||
},
|
||||
)
|
||||
|
||||
async def mcp_summarize_video(
|
||||
self,
|
||||
video_url: str = Field(
|
||||
description="The input video filepath or URL to summarize."
|
||||
),
|
||||
sample_rate: int = Field(
|
||||
default=1, description="Sample n frames per second (default: 1)."
|
||||
),
|
||||
start_time: float = Field(
|
||||
default=0,
|
||||
description="Start time of the video segment in seconds (default: 0).",
|
||||
),
|
||||
end_time: float | None = Field(
|
||||
default=None,
|
||||
description="End time of the video segment in seconds (default: None).",
|
||||
),
|
||||
output_format: str = Field(
|
||||
default="markdown",
|
||||
description="Output format: 'markdown', 'json', or 'text' (default: markdown).",
|
||||
),
|
||||
) -> ActionResponse:
|
||||
"""Summarize the main content of a video using AI analysis.
|
||||
|
||||
This tool provides AI-powered video summarization with:
|
||||
- Key point extraction
|
||||
- Main topic identification
|
||||
- Important visual element recognition
|
||||
- LLM-optimized result formatting
|
||||
|
||||
Args:
|
||||
video_url: The input video filepath or URL to summarize
|
||||
sample_rate: Sample n frames per second
|
||||
start_time: Start time of the video segment in seconds
|
||||
end_time: End time of the video segment in seconds
|
||||
output_format: Format for the response output
|
||||
|
||||
Returns:
|
||||
ActionResponse with video summary results and metadata
|
||||
"""
|
||||
start_exec_time = time.time()
|
||||
|
||||
try:
|
||||
# Validate video file
|
||||
video_path = self._validate_file_path(video_url)
|
||||
|
||||
self._color_log(f"🎬 Summarizing video: {video_url}", Color.cyan)
|
||||
|
||||
# Extract video frames
|
||||
video_frames = self._get_video_frames(
|
||||
str(video_path), sample_rate, start_time, end_time
|
||||
)
|
||||
self._color_log(f"📸 Extracted {len(video_frames)} frames", Color.blue)
|
||||
|
||||
# Process frames in larger chunks for summarization
|
||||
interval = 490
|
||||
frame_nums = 500
|
||||
all_results = []
|
||||
|
||||
for i in range(0, len(video_frames), interval):
|
||||
cur_frames = video_frames[i : i + frame_nums]
|
||||
content = self._create_video_content(
|
||||
self.video_summarize_prompt, cur_frames
|
||||
)
|
||||
inputs = [{"role": "user", "content": content}]
|
||||
|
||||
try:
|
||||
response: ModelResponse = call_llm_model(
|
||||
get_llm_model(
|
||||
conf=AgentConfig(
|
||||
llm_provider="openai",
|
||||
llm_model_name=os.getenv(
|
||||
"VIDEO_LLM_MODEL_NAME", "gpt-4o"
|
||||
),
|
||||
llm_api_key=os.getenv("VIDEO_LLM_API_KEY"),
|
||||
llm_base_url=os.getenv("VIDEO_LLM_BASE_URL"),
|
||||
)
|
||||
),
|
||||
inputs,
|
||||
temperature=float(os.getenv("VIDEO_LLM_TEMPERATURE", "1.0")),
|
||||
)
|
||||
cur_summary = response.content
|
||||
except Exception as e:
|
||||
self._color_log(
|
||||
f"LLM summary error for chunk {i // interval + 1}: {str(e)}",
|
||||
Color.yellow,
|
||||
)
|
||||
cur_summary = (
|
||||
f"Summary failed for video segment {i // interval + 1}"
|
||||
)
|
||||
|
||||
all_results.append(
|
||||
f"Summary of video part {i // interval + 1}: {cur_summary}"
|
||||
)
|
||||
|
||||
if i + frame_nums >= len(video_frames):
|
||||
break
|
||||
|
||||
summary_result = "\n".join(all_results)
|
||||
duration_analyzed = (
|
||||
end_time - start_time if end_time else len(video_frames) / sample_rate
|
||||
)
|
||||
|
||||
# Create result
|
||||
result = VideoSummaryResult(
|
||||
video_source=video_url,
|
||||
summary=summary_result,
|
||||
frame_count=len(video_frames),
|
||||
duration_analyzed=duration_analyzed,
|
||||
success=True,
|
||||
error=None,
|
||||
)
|
||||
|
||||
# Format output for LLM
|
||||
message = self._format_summary_output(result, output_format)
|
||||
execution_time = time.time() - start_exec_time
|
||||
|
||||
# Create metadata
|
||||
metadata = VideoMetadata(
|
||||
operation="summarize",
|
||||
video_source=video_url,
|
||||
sample_rate=sample_rate,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
frame_count=len(video_frames),
|
||||
execution_time=execution_time,
|
||||
).model_dump()
|
||||
|
||||
self._color_log(
|
||||
"✅ Video summarization completed successfully", Color.green
|
||||
)
|
||||
return ActionResponse(success=True, message=message, metadata=metadata)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
self._color_log(
|
||||
f"❌ Video summarization error: {traceback.format_exc()}", Color.red
|
||||
)
|
||||
|
||||
# Format error for LLM
|
||||
message = f"Failed to summarize video: {error_msg}"
|
||||
execution_time = time.time() - start_exec_time
|
||||
|
||||
# Create metadata
|
||||
metadata = VideoMetadata(
|
||||
operation="summarize",
|
||||
video_source=video_url,
|
||||
error_type="summarization_failure",
|
||||
execution_time=execution_time,
|
||||
).model_dump()
|
||||
|
||||
return ActionResponse(success=False, message=message, metadata=metadata)
|
||||
|
||||
async def mcp_extract_keyframes(
|
||||
self,
|
||||
video_path: str = Field(description="The input video filepath or URL."),
|
||||
target_time: int = Field(
|
||||
description="The specific time point for extraction (in seconds), centered within the window_size."
|
||||
),
|
||||
window_size: int = Field(
|
||||
default=5,
|
||||
description="The window size for extraction (in seconds, default: 5).",
|
||||
),
|
||||
output_dir: str = Field(
|
||||
default=None,
|
||||
description="Directory where extracted frames will be saved (default: workspace/keyframes).",
|
||||
),
|
||||
output_format: str = Field(
|
||||
default="markdown",
|
||||
description="Output format: 'markdown', 'json', or 'text' (default: markdown).",
|
||||
),
|
||||
) -> ActionResponse:
|
||||
"""Extract key frames around a target time with scene detection.
|
||||
|
||||
This tool provides keyframe extraction with:
|
||||
- Scene detection for significant frame changes
|
||||
- Configurable time windows
|
||||
- Automatic output directory management
|
||||
- LLM-optimized result formatting
|
||||
|
||||
Args:
|
||||
video_path: The input video filepath or URL
|
||||
target_time: Specific time point (in seconds) to extract frames around
|
||||
window_size: Time window (in seconds) centered on target_time
|
||||
cleanup: Whether to delete the original video file after processing
|
||||
output_dir: Directory where extracted frames will be saved
|
||||
output_format: Format for the response output
|
||||
|
||||
Returns:
|
||||
ActionResponse with keyframe extraction results and metadata
|
||||
"""
|
||||
start_exec_time = time.time()
|
||||
|
||||
try:
|
||||
# Validate video file
|
||||
validated_path = self._validate_file_path(video_path)
|
||||
|
||||
# Set default output directory
|
||||
output_dir = (
|
||||
str(self.workspace / "keyframes") if output_dir is None else output_dir
|
||||
)
|
||||
|
||||
output_path = Path(output_dir)
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
self._color_log(f"🎬 Extracting keyframes from: {video_path}", Color.cyan)
|
||||
self._color_log(
|
||||
f"🎯 Target time: {target_time}s, Window: {window_size}s", Color.blue
|
||||
)
|
||||
|
||||
# Extract keyframes with scene detection
|
||||
frames, frame_times = self._extract_keyframes_with_scene_detection(
|
||||
str(validated_path), target_time, window_size
|
||||
)
|
||||
|
||||
# Save frames to disk
|
||||
frame_paths, frame_timestamps = self._save_keyframes(
|
||||
frames, frame_times, str(output_path)
|
||||
)
|
||||
|
||||
# Cleanup if requested
|
||||
# if cleanup and validated_path.exists():
|
||||
# validated_path.unlink()
|
||||
# self._color_log("🗑️ Cleaned up original video file", Color.yellow)
|
||||
|
||||
# Create result
|
||||
result = KeyframeResult(
|
||||
frame_paths=frame_paths,
|
||||
frame_timestamps=frame_timestamps,
|
||||
output_directory=str(output_path),
|
||||
frame_count=len(frame_paths),
|
||||
target_time=float(target_time),
|
||||
window_size=float(window_size),
|
||||
success=True,
|
||||
error=None,
|
||||
)
|
||||
|
||||
# Format output for LLM
|
||||
message = self._format_keyframe_output(result, output_format)
|
||||
execution_time = time.time() - start_exec_time
|
||||
|
||||
# Create metadata
|
||||
metadata = VideoMetadata(
|
||||
operation="extract_keyframes",
|
||||
video_source=video_path,
|
||||
target_time=float(target_time),
|
||||
window_size=float(window_size),
|
||||
output_directory=str(output_path),
|
||||
frame_count=len(frame_paths),
|
||||
execution_time=execution_time,
|
||||
).model_dump()
|
||||
|
||||
self._color_log(
|
||||
f"✅ Extracted {len(frame_paths)} keyframes successfully", Color.green
|
||||
)
|
||||
return ActionResponse(success=True, message=message, metadata=metadata)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
self._color_log(
|
||||
f"❌ Keyframe extraction error: {traceback.format_exc()}", Color.red
|
||||
)
|
||||
|
||||
# Format error for LLM
|
||||
message = f"Failed to extract keyframes: {error_msg}"
|
||||
execution_time = time.time() - start_exec_time
|
||||
|
||||
# Create metadata
|
||||
metadata = VideoMetadata(
|
||||
operation="extract_keyframes",
|
||||
video_source=video_path,
|
||||
target_time=float(target_time) if target_time else None,
|
||||
window_size=float(window_size) if window_size else None,
|
||||
error_type="keyframe_extraction_failure",
|
||||
execution_time=execution_time,
|
||||
).model_dump()
|
||||
|
||||
return ActionResponse(success=False, message=message, metadata=metadata)
|
||||
|
||||
def _extract_keyframes_with_scene_detection(
|
||||
self, video_path: str, target_time: int, window_size: int
|
||||
) -> tuple[list[any], list[float]]:
|
||||
"""Extract key frames around the target time with scene detection.
|
||||
|
||||
Args:
|
||||
video_path: Path to the video file
|
||||
target_time: Target time in seconds
|
||||
window_size: Window size in seconds
|
||||
|
||||
Returns:
|
||||
Tuple of (frames, frame_times)
|
||||
"""
|
||||
cap = cv2.VideoCapture(video_path) # pylint: disable=E1101
|
||||
fps = cap.get(cv2.CAP_PROP_FPS) # pylint: disable=E1101
|
||||
|
||||
# Calculate frame numbers for the time window
|
||||
start_frame = int((target_time - window_size / 2) * fps)
|
||||
end_frame = int((target_time + window_size / 2) * fps)
|
||||
total_frames_in_window = end_frame - start_frame
|
||||
|
||||
max_frames = 384 # Maximum allowed frames to prevent memory issues
|
||||
|
||||
# Calculate sampling interval for even distribution
|
||||
if total_frames_in_window <= max_frames:
|
||||
# If total frames is within limit, use scene detection normally
|
||||
frame_interval = 1
|
||||
use_scene_detection = True
|
||||
else:
|
||||
# If exceeds limit, sample evenly across the window
|
||||
frame_interval = total_frames_in_window // max_frames
|
||||
use_scene_detection = False # Skip scene detection for even sampling
|
||||
|
||||
frames = []
|
||||
frame_times = []
|
||||
|
||||
# Set video position to start_frame
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, max(0, start_frame)) # pylint: disable=E1101
|
||||
|
||||
prev_frame = None
|
||||
frame_count = 0
|
||||
|
||||
while cap.isOpened() and len(frames) < max_frames:
|
||||
frame_pos = cap.get(cv2.CAP_PROP_POS_FRAMES) # pylint: disable=E1101
|
||||
if frame_pos >= end_frame:
|
||||
break
|
||||
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
break
|
||||
|
||||
if use_scene_detection:
|
||||
# Use original scene detection logic
|
||||
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) # pylint: disable=E1101
|
||||
|
||||
# If this is the first frame, save it
|
||||
if prev_frame is None:
|
||||
frames.append(frame)
|
||||
frame_times.append(frame_pos / fps)
|
||||
else:
|
||||
# Calculate difference between current and previous frame
|
||||
diff = cv2.absdiff(gray, prev_frame) # pylint: disable=E1101
|
||||
mean_diff = np.mean(diff)
|
||||
|
||||
# If significant change detected, save frame
|
||||
if mean_diff > 20: # Threshold for scene change
|
||||
frames.append(frame)
|
||||
frame_times.append(frame_pos / fps)
|
||||
|
||||
prev_frame = gray
|
||||
else:
|
||||
# Use even sampling for large windows
|
||||
if frame_count % frame_interval == 0:
|
||||
frames.append(frame)
|
||||
frame_times.append(frame_pos / fps)
|
||||
|
||||
frame_count += 1
|
||||
|
||||
cap.release()
|
||||
return frames, frame_times
|
||||
|
||||
def _save_keyframes(
|
||||
self, frames: list[any], frame_times: list[float], output_dir: str
|
||||
) -> tuple[list[str], list[float]]:
|
||||
"""Save extracted frames to disk.
|
||||
|
||||
Args:
|
||||
frames: List of frame objects
|
||||
frame_times: List of frame timestamps
|
||||
output_dir: Output directory path
|
||||
|
||||
Returns:
|
||||
Tuple of (saved_paths, saved_timestamps)
|
||||
"""
|
||||
output_path = Path(output_dir)
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
saved_paths = []
|
||||
saved_timestamps = []
|
||||
|
||||
for _, (frame, timestamp) in enumerate(zip(frames, frame_times)):
|
||||
filename = output_path / f"frame_{timestamp:.2f}s.jpg"
|
||||
cv2.imwrite(str(filename), frame) # pylint: disable=E1101
|
||||
saved_paths.append(str(filename))
|
||||
saved_timestamps.append(timestamp)
|
||||
|
||||
return saved_paths, saved_timestamps
|
||||
|
||||
|
||||
# Default arguments for testing
|
||||
if __name__ == "__main__":
|
||||
load_dotenv()
|
||||
|
||||
arguments = ActionArguments(
|
||||
name="video",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
|
||||
)
|
||||
|
||||
try:
|
||||
service = VideoCollection(arguments)
|
||||
service.run()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {str(e)}")
|
||||
@@ -0,0 +1,374 @@
|
||||
"""
|
||||
Browser MCP Server
|
||||
|
||||
This module provides MCP server functionality for browser automation and interaction.
|
||||
It handles tasks such as web scraping, form submission, and automated browsing using browser-use package.
|
||||
|
||||
Main functions:
|
||||
- mcp_browser_use: Performs browser automation tasks with LLM-friendly output
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
import traceback
|
||||
|
||||
try:
|
||||
from browser_use import Agent, AgentHistoryList, BrowserProfile
|
||||
from browser_use.llm import ChatOpenAI
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from aworld.logs.util import Color
|
||||
|
||||
from ..base import ActionArguments, ActionCollection, ActionResponse
|
||||
except Exception as e:
|
||||
print(f"Failed to import browser tool: {traceback.format_exc()}")
|
||||
raise e
|
||||
|
||||
|
||||
print(f"Browser tool sys.path: {sys.path}")
|
||||
|
||||
|
||||
class BrowserMetadata(BaseModel):
|
||||
"""Metadata for browser automation results."""
|
||||
|
||||
task: str
|
||||
execution_successful: bool
|
||||
steps_taken: int | None = None
|
||||
downloaded_files: list[str] = Field(default_factory=list)
|
||||
visited_urls: list[str] = Field(default_factory=list)
|
||||
execution_time: float | None = None
|
||||
error_type: str | None = None
|
||||
trace_log_path: str | None = None
|
||||
|
||||
|
||||
class BrowserActionCollection(ActionCollection):
|
||||
"""MCP service for browser automation using browser-use package.
|
||||
|
||||
Provides comprehensive web automation capabilities including:
|
||||
- Web scraping and content extraction
|
||||
- Form submission and interaction
|
||||
- File downloads and media handling
|
||||
- LLM-enhanced browsing with memory
|
||||
- Robot detection and paywall handling
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
|
||||
# Load environment variables
|
||||
load_dotenv()
|
||||
|
||||
# Extended system prompt for browser automation
|
||||
self.extended_browser_system_prompt = """
|
||||
10. URL ends with .pdf
|
||||
- If the go_to_url function with `https://any_url/any_file_name.pdf` as the parameter, just report the url link and hint the user to download using `download` mcp tool or `curl`, then execute `done` action.
|
||||
|
||||
11. Robot Detection:
|
||||
- If the page is a robot detection page, abort immediately. Then navigate to the most authoritative source for similar information instead
|
||||
|
||||
# Efficiency Guidelines
|
||||
0. if download option is available, always **DOWNLOAD** as possible! Also, report the download url link in your result.
|
||||
1. Use specific search queries with key terms from the task
|
||||
2. Avoid getting distracted by tangential information
|
||||
3. If blocked by paywalls, try archive.org or similar alternatives
|
||||
4. Document each significant finding clearly and concisely
|
||||
5. Precisely extract the necessary information with minimal browsing steps.
|
||||
"""
|
||||
|
||||
# Initialize LLM configuration
|
||||
self.llm_config = ChatOpenAI(
|
||||
model=os.getenv("LLM_MODEL_NAME"),
|
||||
api_key=os.getenv("LLM_API_KEY"),
|
||||
base_url=os.getenv("LLM_BASE_URL"),
|
||||
temperature=1.0,
|
||||
)
|
||||
self._color_log(f"Browser llm_config: {self.llm_config}", Color.green)
|
||||
|
||||
# Browser profile configuration
|
||||
self.browser_profile = BrowserProfile(
|
||||
cookies_file=os.getenv("COOKIES_FILE_PATH"),
|
||||
downloads_dir=str(self.workspace),
|
||||
downloads_path=str(self.workspace),
|
||||
save_recording_path=str(self.workspace),
|
||||
save_downloads_path=str(self.workspace),
|
||||
chromium_sandbox=False,
|
||||
headless=True,
|
||||
)
|
||||
self._color_log(f"Browser browser_profile: {self.browser_profile}", Color.green)
|
||||
|
||||
# Log configuration
|
||||
self.trace_log_dir = str(self.workspace / "logs")
|
||||
os.makedirs(f"{self.trace_log_dir}/browser_log", exist_ok=True)
|
||||
|
||||
self._color_log("Browser automation service initialized", Color.green)
|
||||
self._color_log(
|
||||
f"Downloads directory: {self.browser_profile.downloads_path}", Color.blue
|
||||
)
|
||||
self._color_log(
|
||||
f"Trace logs directory: {self.trace_log_dir}/browser_log", Color.blue
|
||||
)
|
||||
|
||||
def _create_browser_agent(self, task: str) -> Agent:
|
||||
"""Create a browser agent instance with configured settings.
|
||||
|
||||
Args:
|
||||
task: The task description for the browser agent
|
||||
|
||||
Returns:
|
||||
Configured Agent instance
|
||||
"""
|
||||
return Agent(
|
||||
task=task,
|
||||
llm=self.llm_config,
|
||||
extend_system_message=self.extended_browser_system_prompt,
|
||||
use_vision=True,
|
||||
enable_memory=False,
|
||||
browser_profile=self.browser_profile,
|
||||
save_conversation_path=f"{self.trace_log_dir}/browser_log/trace.log",
|
||||
)
|
||||
|
||||
def _extract_visited_urls(self, extracted_content: list[str]) -> list[str]:
|
||||
"""Inner method to extract URLs from content using regex.
|
||||
|
||||
Args:
|
||||
content_list: List of content strings to search for URLs
|
||||
|
||||
Returns:
|
||||
List of unique URLs found in the content
|
||||
"""
|
||||
url_pattern = r'https?://[^\s<>"\[\]{}|\\^`]+'
|
||||
visited_urls = set()
|
||||
|
||||
for content in extracted_content:
|
||||
if content and isinstance(content, str):
|
||||
urls = re.findall(url_pattern, content)
|
||||
visited_urls.update(urls)
|
||||
|
||||
return list(visited_urls)
|
||||
|
||||
def _format_extracted_content(self, extracted_content: list[str]) -> str:
|
||||
"""Format extracted content to be LLM-friendly.
|
||||
|
||||
Args:
|
||||
extracted_content: List of extracted content strings from browser execution
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if not extracted_content:
|
||||
return "No content extracted from browser execution."
|
||||
|
||||
# Handle list of strings
|
||||
if len(extracted_content) == 1:
|
||||
# Single item - return it directly with formatting
|
||||
return f"**Extracted Content:**\n{extracted_content[0]}"
|
||||
else:
|
||||
# Multiple items - format as numbered list
|
||||
formatted_parts = ["**Extracted Content:**"]
|
||||
for i, content in enumerate(extracted_content, 1):
|
||||
if content.strip(): # Only include non-empty content
|
||||
formatted_parts.append(f"{i}. {content}")
|
||||
|
||||
return (
|
||||
"\n".join(formatted_parts)
|
||||
if len(formatted_parts) > 1
|
||||
else "No meaningful content extracted from browser execution."
|
||||
)
|
||||
|
||||
async def mcp_browser_use(
|
||||
self,
|
||||
task: str = Field(
|
||||
description="The task to perform using the browser automation agent"
|
||||
),
|
||||
max_steps: int = Field(
|
||||
default=50, description="Maximum number of steps for browser execution"
|
||||
),
|
||||
extract_format: str = Field(
|
||||
default="markdown",
|
||||
description="Format for extracted content: 'markdown', 'json', or 'text'",
|
||||
),
|
||||
) -> ActionResponse:
|
||||
"""Perform browser automation tasks using the browser-use package.
|
||||
|
||||
This tool provides comprehensive browser automation capabilities including:
|
||||
- Web scraping and content extraction
|
||||
- Form submission and automated interactions
|
||||
- File downloads and media handling
|
||||
- LLM-enhanced browsing with memory and vision
|
||||
- Automatic handling of robot detection and paywalls
|
||||
|
||||
Args:
|
||||
task: Description of the browser automation task to perform
|
||||
max_steps: Maximum number of execution steps (default: 50)
|
||||
extract_format: Output format for extracted content
|
||||
|
||||
Returns:
|
||||
ActionResponse with LLM-friendly extracted content and execution metadata
|
||||
"""
|
||||
try:
|
||||
self._color_log(f"🎯 Starting browser task: {task}", Color.cyan)
|
||||
|
||||
# Create browser agent
|
||||
agent = self._create_browser_agent(task)
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
browser_execution: AgentHistoryList = await agent.run(max_steps=max_steps)
|
||||
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
if (
|
||||
browser_execution is not None
|
||||
and browser_execution.is_done()
|
||||
and browser_execution.is_successful()
|
||||
):
|
||||
# Extract and format content
|
||||
extracted_content = browser_execution.extracted_content()
|
||||
final_result = browser_execution.final_result()
|
||||
|
||||
# Format content based on requested format
|
||||
if extract_format.lower() == "json":
|
||||
formatted_content = json.dumps(
|
||||
{"summary": final_result, "extracted_data": extracted_content},
|
||||
indent=2,
|
||||
)
|
||||
elif extract_format.lower() == "text":
|
||||
formatted_content = f"{final_result}\n\n{self._format_extracted_content(extracted_content)}"
|
||||
else: # markdown (default)
|
||||
formatted_content = (
|
||||
f"## Browser Automation Result\n\n**Summary:** {final_result}\n\n"
|
||||
f"{self._format_extracted_content(extracted_content)}"
|
||||
)
|
||||
|
||||
# Prepare metadata
|
||||
metadata = BrowserMetadata(
|
||||
task=task,
|
||||
execution_successful=True,
|
||||
steps_taken=(
|
||||
len(browser_execution.history)
|
||||
if hasattr(browser_execution, "history")
|
||||
else None
|
||||
),
|
||||
downloaded_files=[],
|
||||
visited_urls=self._extract_visited_urls(extracted_content),
|
||||
execution_time=execution_time,
|
||||
trace_log_path=f"{self.trace_log_dir}/browser_log/trace.log",
|
||||
)
|
||||
|
||||
self._color_log(f"🗒️ Detail: {extracted_content}", Color.lightgrey)
|
||||
self._color_log(f"🌏 Result: {final_result}", Color.green)
|
||||
|
||||
return ActionResponse(
|
||||
success=True,
|
||||
message=formatted_content,
|
||||
metadata=metadata.model_dump(),
|
||||
)
|
||||
|
||||
else:
|
||||
# Handle execution failure
|
||||
error_msg = "Browser execution failed or was not completed successfully"
|
||||
|
||||
metadata = BrowserMetadata(
|
||||
task=task,
|
||||
execution_successful=False,
|
||||
execution_time=execution_time,
|
||||
error_type="execution_failure",
|
||||
trace_log_path=f"{self.trace_log_dir}/browser_log/trace.log",
|
||||
)
|
||||
|
||||
self._color_log(f"❌ {error_msg}", Color.red)
|
||||
|
||||
return ActionResponse(
|
||||
success=False, message=error_msg, metadata=metadata.model_dump()
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Browser automation failed: {str(e)}"
|
||||
error_trace = traceback.format_exc()
|
||||
|
||||
self.logger.error(f"Browser execution error: {error_trace}")
|
||||
|
||||
metadata = BrowserMetadata(
|
||||
task=task,
|
||||
execution_successful=False,
|
||||
error_type="exception",
|
||||
trace_log_path=f"{self.trace_log_dir}/browser_log/trace.log",
|
||||
)
|
||||
|
||||
self._color_log(f"❌ {error_msg}", Color.red)
|
||||
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"{error_msg}\n\nError details: {error_trace}",
|
||||
metadata=metadata.model_dump(),
|
||||
)
|
||||
|
||||
def mcp_get_browser_capabilities(self) -> ActionResponse:
|
||||
"""Get information about browser automation capabilities and configuration.
|
||||
|
||||
Returns:
|
||||
ActionResponse with browser service capabilities and current configuration
|
||||
"""
|
||||
capabilities = {
|
||||
"automation_features": [
|
||||
"Web scraping and content extraction",
|
||||
"Form submission and interaction",
|
||||
"File downloads and media handling",
|
||||
"LLM-enhanced browsing with vision",
|
||||
"Memory-enabled browsing sessions",
|
||||
"Robot detection and paywall handling",
|
||||
],
|
||||
"supported_formats": ["markdown", "json", "text"],
|
||||
"configuration": {
|
||||
"llm_model": os.getenv("LLM_MODEL_NAME", "Not configured"),
|
||||
"downloads_directory": self.browser_profile.downloads_path,
|
||||
"cookies_enabled": bool(os.getenv("COOKIES_FILE_PATH")),
|
||||
"trace_logging": True,
|
||||
"vision_enabled": True,
|
||||
"headless": True,
|
||||
},
|
||||
}
|
||||
|
||||
formatted_info = f"""# Browser Automation Service Capabilities
|
||||
|
||||
## Features
|
||||
{chr(10).join(f"- {feature}" for feature in capabilities["automation_features"])}
|
||||
|
||||
## Supported Output Formats
|
||||
{chr(10).join(f"- {fmt}" for fmt in capabilities["supported_formats"])}
|
||||
|
||||
## Current Configuration
|
||||
- **LLM Model:** {capabilities["configuration"]["llm_model"]}
|
||||
- **Downloads Directory:** {capabilities["configuration"]["downloads_directory"]}
|
||||
- **Cookies Enabled:** {capabilities["configuration"]["cookies_enabled"]}
|
||||
- **Vision Enabled:** {capabilities["configuration"]["vision_enabled"]}
|
||||
- **Memory Enabled:** {capabilities["configuration"]["memory_enabled"]}
|
||||
- **Trace Logging:** {capabilities["configuration"]["trace_logging"]}
|
||||
"""
|
||||
|
||||
return ActionResponse(
|
||||
success=True, message=formatted_info, metadata=capabilities
|
||||
)
|
||||
|
||||
|
||||
# Example usage and entry point
|
||||
if __name__ == "__main__":
|
||||
load_dotenv()
|
||||
|
||||
# Default arguments for testing
|
||||
args = ActionArguments(
|
||||
name="browser_automation_service",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
|
||||
)
|
||||
|
||||
# Initialize and run the browser automation service
|
||||
try:
|
||||
service = BrowserActionCollection(args)
|
||||
service.run()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}: {traceback.format_exc()}")
|
||||
@@ -0,0 +1,479 @@
|
||||
"""
|
||||
Download MCP Server
|
||||
|
||||
This module provides MCP server functionality for downloading files from URLs.
|
||||
It supports HTTP/HTTPS downloads with configurable options and returns LLM-friendly formatted results.
|
||||
|
||||
Key features:
|
||||
- Download files from HTTP/HTTPS URLs
|
||||
- Configurable timeout and overwrite options
|
||||
- Custom headers support for authentication
|
||||
- LLM-optimized output formatting
|
||||
- Comprehensive error handling and logging
|
||||
- Path validation and directory creation
|
||||
|
||||
Main functions:
|
||||
- mcp_download_file: Download files from URLs with comprehensive options
|
||||
- mcp_get_download_capabilities: Get download service capabilities
|
||||
"""
|
||||
|
||||
import json
|
||||
import shutil
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import requests
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from aworld.logs.util import Color
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
|
||||
|
||||
class DownloadResult(BaseModel):
|
||||
"""Individual download operation result with structured data."""
|
||||
|
||||
url: str
|
||||
file_path: str
|
||||
success: bool
|
||||
file_size: int | None = None
|
||||
duration: str
|
||||
timestamp: str
|
||||
error_message: str | None = None
|
||||
|
||||
|
||||
class DownloadMetadata(BaseModel):
|
||||
"""Metadata for download operation results."""
|
||||
|
||||
url: str
|
||||
output_path: str
|
||||
timeout_seconds: int
|
||||
overwrite_enabled: bool
|
||||
execution_time: float | None = None
|
||||
file_size_bytes: int | None = None
|
||||
content_type: str | None = None
|
||||
status_code: int | None = None
|
||||
error_type: str | None = None
|
||||
headers_used: bool = False
|
||||
|
||||
|
||||
class DownloadCollection(ActionCollection):
|
||||
"""MCP service for file download operations with comprehensive controls.
|
||||
|
||||
Provides secure file download capabilities including:
|
||||
- HTTP/HTTPS URL support
|
||||
- Configurable timeout controls
|
||||
- Custom headers for authentication
|
||||
- Path validation and directory creation
|
||||
- LLM-friendly result formatting
|
||||
- Error handling and logging
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
|
||||
# Configuration
|
||||
self.default_timeout = 60 * 3 # 3 minutes timeout
|
||||
self.max_file_size = 1024 * 1024 * 1024 # 1GB limit
|
||||
self.supported_schemes = {"http", "https"}
|
||||
|
||||
self.headers = {
|
||||
"User-Agent": (
|
||||
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
|
||||
"AppleWebKit/537.36 (KHTML, like Gecko) "
|
||||
"Chrome/91.0.4472.124 Safari/537.36"
|
||||
),
|
||||
"Accept-Language": "en-US,en;q=0.9",
|
||||
}
|
||||
|
||||
self._color_log("Download service initialized", Color.green, "debug")
|
||||
self._color_log(f"Workspace: {self.workspace}", Color.blue, "debug")
|
||||
|
||||
def _validate_url(self, url: str) -> tuple[bool, str | None]:
|
||||
"""Validate URL format and scheme.
|
||||
|
||||
Args:
|
||||
url: URL to validate
|
||||
|
||||
Returns:
|
||||
Tuple of (is_valid, error_message)
|
||||
"""
|
||||
try:
|
||||
parsed = urlparse(url)
|
||||
|
||||
if not parsed.scheme:
|
||||
return False, "URL must include a scheme (http:// or https://)"
|
||||
|
||||
if parsed.scheme.lower() not in self.supported_schemes:
|
||||
return False, f"Unsupported URL scheme: {parsed.scheme}. Supported: {', '.join(self.supported_schemes)}"
|
||||
|
||||
if not parsed.netloc:
|
||||
return False, "URL must include a valid domain"
|
||||
|
||||
return True, None
|
||||
|
||||
except Exception as e:
|
||||
return False, f"Invalid URL format: {str(e)}"
|
||||
|
||||
def _resolve_output_path(self, output_path: str) -> Path:
|
||||
"""Resolve and validate output file path.
|
||||
|
||||
Args:
|
||||
output_path: Output file path (absolute or relative)
|
||||
|
||||
Returns:
|
||||
Resolved Path object
|
||||
"""
|
||||
path = Path(output_path).expanduser()
|
||||
|
||||
if not path.is_absolute():
|
||||
path = self.workspace / path
|
||||
|
||||
# Ensure parent directory exists
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
return path.resolve()
|
||||
|
||||
def _format_download_output(self, result: DownloadResult, output_format: str = "markdown") -> str:
|
||||
"""Format download results for LLM consumption.
|
||||
|
||||
Args:
|
||||
result: Download execution result
|
||||
output_format: Format type ('markdown', 'json', 'text')
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if output_format == "json":
|
||||
return json.dumps(result.model_dump(), indent=2)
|
||||
|
||||
elif output_format == "text":
|
||||
output_parts = [
|
||||
f"URL: {result.url}",
|
||||
f"File Path: {result.file_path}",
|
||||
f"Status: {'SUCCESS' if result.success else 'FAILED'}",
|
||||
f"Duration: {result.duration}",
|
||||
f"Timestamp: {result.timestamp}",
|
||||
]
|
||||
|
||||
if result.file_size is not None:
|
||||
output_parts.append(f"File Size: {result.file_size:,} bytes")
|
||||
|
||||
if result.error_message:
|
||||
output_parts.append(f"Error: {result.error_message}")
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
else: # markdown (default)
|
||||
status_emoji = "✅" if result.success else "❌"
|
||||
|
||||
output_parts = [
|
||||
f"# File Download {status_emoji}",
|
||||
f"**URL:** `{result.url}`",
|
||||
f"**File Path:** `{result.file_path}`",
|
||||
f"**Status:** {'SUCCESS' if result.success else 'FAILED'}",
|
||||
f"**Duration:** {result.duration}",
|
||||
f"**Timestamp:** {result.timestamp}",
|
||||
]
|
||||
|
||||
if result.file_size is not None:
|
||||
size_mb = result.file_size / (1024 * 1024)
|
||||
output_parts.append(f"**File Size:** {result.file_size:,} bytes ({size_mb:.2f} MB)")
|
||||
|
||||
if result.error_message:
|
||||
output_parts.extend(["\n## Error Details", f"```\n{result.error_message}\n```"])
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
async def _download_file_async(
|
||||
self, url: str, output_path: Path, timeout: int, headers: dict[str, str] | None
|
||||
) -> DownloadResult:
|
||||
"""Download file asynchronously with comprehensive error handling.
|
||||
|
||||
Args:
|
||||
url: URL to download from
|
||||
output_path: Local path to save file
|
||||
timeout: Request timeout in seconds
|
||||
headers: Optional custom headers
|
||||
|
||||
Returns:
|
||||
DownloadResult with execution details
|
||||
"""
|
||||
start_time = datetime.now()
|
||||
|
||||
try:
|
||||
self._color_log(f"📥 Starting download: {url}", Color.cyan)
|
||||
|
||||
with requests.get(url, stream=True, timeout=timeout, headers=headers) as response:
|
||||
response.raise_for_status()
|
||||
|
||||
# Check content length if available
|
||||
content_length = response.headers.get("content-length")
|
||||
if content_length and int(content_length) > self.max_file_size:
|
||||
raise ValueError(f"File too large: {content_length} bytes (max: {self.max_file_size})")
|
||||
|
||||
# Download file
|
||||
with open(output_path, "wb") as f:
|
||||
shutil.copyfileobj(response.raw, f)
|
||||
|
||||
file_size = output_path.stat().st_size
|
||||
duration = str(datetime.now() - start_time)
|
||||
|
||||
self._color_log(f"✅ Download completed: {file_size:,} bytes", Color.green)
|
||||
|
||||
return DownloadResult(
|
||||
url=url,
|
||||
file_path=str(output_path),
|
||||
success=True,
|
||||
file_size=file_size,
|
||||
duration=duration,
|
||||
timestamp=start_time.isoformat(),
|
||||
)
|
||||
|
||||
except requests.exceptions.Timeout:
|
||||
duration = str(datetime.now() - start_time)
|
||||
error_msg = f"Download timed out after {timeout} seconds"
|
||||
self._color_log(f"⏰ {error_msg}", Color.red)
|
||||
|
||||
return DownloadResult(
|
||||
url=url,
|
||||
file_path=str(output_path),
|
||||
success=False,
|
||||
duration=duration,
|
||||
timestamp=start_time.isoformat(),
|
||||
error_message=error_msg,
|
||||
)
|
||||
|
||||
except requests.exceptions.RequestException as e:
|
||||
duration = str(datetime.now() - start_time)
|
||||
error_msg = f"Request failed: {str(e)}"
|
||||
self._color_log(f"❌ {error_msg}", Color.red)
|
||||
|
||||
return DownloadResult(
|
||||
url=url,
|
||||
file_path=str(output_path),
|
||||
success=False,
|
||||
duration=duration,
|
||||
timestamp=start_time.isoformat(),
|
||||
error_message=error_msg,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
duration = str(datetime.now() - start_time)
|
||||
error_msg = f"Unexpected error: {str(e)}"
|
||||
self._color_log(f"💥 {error_msg}", Color.red)
|
||||
|
||||
return DownloadResult(
|
||||
url=url,
|
||||
file_path=str(output_path),
|
||||
success=False,
|
||||
duration=duration,
|
||||
timestamp=start_time.isoformat(),
|
||||
error_message=error_msg,
|
||||
)
|
||||
|
||||
async def mcp_download_file(
|
||||
self,
|
||||
url: str = Field(description="HTTP/HTTPS URL of the file to download"),
|
||||
output_file_path: str = Field(
|
||||
description="Local path where the file should be saved (absolute or relative to workspace)"
|
||||
),
|
||||
overwrite: bool = Field(default=False, description="Whether to overwrite existing files (default: False)"),
|
||||
timeout: int = Field(default=60, description="Download timeout in seconds (default: 60)"),
|
||||
output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'"),
|
||||
) -> ActionResponse:
|
||||
"""Download a file from a URL with comprehensive options and controls.
|
||||
|
||||
This tool provides secure file download capabilities with:
|
||||
- HTTP/HTTPS URL support
|
||||
- Configurable timeout controls
|
||||
- Path validation and directory creation
|
||||
- File size limits and safety checks
|
||||
- LLM-optimized result formatting
|
||||
|
||||
Args:
|
||||
url: The HTTP/HTTPS URL of the file to download
|
||||
output_file_path: Local path to save the downloaded file
|
||||
overwrite: Whether to overwrite existing files
|
||||
timeout: Maximum download time in seconds
|
||||
output_format: Format for the response output
|
||||
|
||||
Returns:
|
||||
ActionResponse with download results and metadata
|
||||
"""
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(url, FieldInfo):
|
||||
url = url.default
|
||||
if isinstance(output_file_path, FieldInfo):
|
||||
output_file_path = output_file_path.default
|
||||
if isinstance(overwrite, FieldInfo):
|
||||
overwrite = overwrite.default
|
||||
if isinstance(timeout, FieldInfo):
|
||||
timeout = timeout.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
try:
|
||||
# Validate URL
|
||||
url_valid, url_error = self._validate_url(url)
|
||||
if not url_valid:
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid URL: {url_error}",
|
||||
metadata=DownloadMetadata(
|
||||
url=url,
|
||||
output_path=output_file_path,
|
||||
timeout_seconds=timeout,
|
||||
overwrite_enabled=overwrite,
|
||||
error_type="invalid_url",
|
||||
).model_dump(),
|
||||
)
|
||||
|
||||
# Resolve output path
|
||||
output_path = self._resolve_output_path(output_file_path)
|
||||
|
||||
# Check if file exists and overwrite setting
|
||||
if output_path.exists() and not overwrite:
|
||||
existing_size = output_path.stat().st_size
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"File already exists at {output_path} ({existing_size:,} bytes) and overwrite is disabled",
|
||||
metadata=DownloadMetadata(
|
||||
url=url,
|
||||
output_path=str(output_path),
|
||||
timeout_seconds=timeout,
|
||||
overwrite_enabled=overwrite,
|
||||
file_size_bytes=existing_size,
|
||||
error_type="file_exists",
|
||||
).model_dump(),
|
||||
)
|
||||
|
||||
# Perform download
|
||||
start_time = time.time()
|
||||
result = await self._download_file_async(url, output_path, timeout, self.headers)
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
# Format output
|
||||
formatted_output = self._format_download_output(result, output_format)
|
||||
|
||||
# Create metadata
|
||||
metadata = DownloadMetadata(
|
||||
url=url,
|
||||
output_path=str(output_path),
|
||||
timeout_seconds=timeout,
|
||||
overwrite_enabled=overwrite,
|
||||
execution_time=execution_time,
|
||||
file_size_bytes=result.file_size,
|
||||
headers_used=self.headers is not None,
|
||||
)
|
||||
|
||||
if not result.success:
|
||||
metadata.error_type = "download_failure"
|
||||
|
||||
return ActionResponse(
|
||||
success=result.success,
|
||||
message=formatted_output,
|
||||
metadata=metadata.model_dump(),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to download file: {str(e)}"
|
||||
self.logger.error(f"Download error: {traceback.format_exc()}")
|
||||
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=error_msg,
|
||||
metadata=DownloadMetadata(
|
||||
url=url,
|
||||
output_path=output_file_path,
|
||||
timeout_seconds=timeout,
|
||||
overwrite_enabled=overwrite,
|
||||
error_type="internal_error",
|
||||
).model_dump(),
|
||||
)
|
||||
|
||||
def mcp_get_download_capabilities(self) -> ActionResponse:
|
||||
"""Get information about download service capabilities and configuration.
|
||||
|
||||
Returns:
|
||||
ActionResponse with download service capabilities and current configuration
|
||||
"""
|
||||
capabilities = {
|
||||
"requests_available": requests is not None,
|
||||
"supported_schemes": list(self.supported_schemes),
|
||||
"supported_features": [
|
||||
"HTTP/HTTPS URL downloads",
|
||||
"Configurable timeout controls",
|
||||
"Custom headers support",
|
||||
"Path validation and directory creation",
|
||||
"File size limits and safety checks",
|
||||
"Multiple output formats (markdown, json, text)",
|
||||
"LLM-optimized result formatting",
|
||||
"Comprehensive error handling",
|
||||
],
|
||||
"supported_formats": ["markdown", "json", "text"],
|
||||
"configuration": {
|
||||
"default_timeout": self.default_timeout,
|
||||
"max_file_size_bytes": self.max_file_size,
|
||||
"workspace": str(self.workspace),
|
||||
},
|
||||
"safety_features": [
|
||||
"URL validation",
|
||||
"File size limits",
|
||||
"Timeout controls",
|
||||
"Path validation",
|
||||
"Overwrite protection",
|
||||
"Error handling and logging",
|
||||
],
|
||||
}
|
||||
|
||||
max_size_mb = self.max_file_size / (1024 * 1024)
|
||||
formatted_info = f"""# Download Service Capabilities
|
||||
|
||||
## Status
|
||||
- **Workspace:** `{self.workspace}`
|
||||
|
||||
## Supported Features
|
||||
{chr(10).join(f"- {feature}" for feature in capabilities["supported_features"])}
|
||||
|
||||
## Supported URL Schemes
|
||||
{chr(10).join(f"- {scheme}://" for scheme in capabilities["supported_schemes"])}
|
||||
|
||||
## Supported Output Formats
|
||||
{chr(10).join(f"- {fmt}" for fmt in capabilities["supported_formats"])}
|
||||
|
||||
## Configuration
|
||||
- **Default Timeout:** {capabilities["configuration"]["default_timeout"]} seconds
|
||||
- **Max File Size:** {self.max_file_size:,} bytes ({max_size_mb:.1f} MB)
|
||||
|
||||
## Safety Features
|
||||
{chr(10).join(f"- {feature}" for feature in capabilities["safety_features"])}
|
||||
"""
|
||||
|
||||
return ActionResponse(
|
||||
success=True,
|
||||
message=formatted_info,
|
||||
metadata=capabilities,
|
||||
)
|
||||
|
||||
|
||||
# Default arguments for testing
|
||||
if __name__ == "__main__":
|
||||
import os
|
||||
|
||||
load_dotenv()
|
||||
|
||||
arguments = ActionArguments(
|
||||
name="download",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
|
||||
)
|
||||
try:
|
||||
service = DownloadCollection(arguments)
|
||||
service.run()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}: {traceback.format_exc()}")
|
||||
@@ -0,0 +1,723 @@
|
||||
"""
|
||||
ArXiv MCP Server
|
||||
|
||||
This module provides MCP server functionality for ArXiv academic paper operations.
|
||||
It supports paper search, metadata extraction, and content retrieval with LLM-friendly formatting.
|
||||
|
||||
Key features:
|
||||
- Search ArXiv papers by query, author, category, or ID
|
||||
- Extract paper metadata (title, authors, abstract, etc.)
|
||||
- Download and process paper PDFs
|
||||
- LLM-optimized content formatting
|
||||
- Comprehensive error handling and logging
|
||||
|
||||
Main functions:
|
||||
- mcp_search_papers: Search ArXiv papers with flexible criteria
|
||||
- mcp_get_paper_details: Get detailed information about specific papers
|
||||
- mcp_download_paper: Download paper PDF and extract text content
|
||||
- mcp_get_categories: Get available ArXiv subject categories
|
||||
- mcp_get_arxiv_capabilities: Get service capabilities and configuration
|
||||
"""
|
||||
|
||||
import json
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
|
||||
import arxiv
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from aworld.logs.util import Color
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
|
||||
|
||||
class PaperResult(BaseModel):
|
||||
"""Individual paper search result with structured data."""
|
||||
|
||||
entry_id: str
|
||||
title: str
|
||||
authors: list[str]
|
||||
summary: str
|
||||
published: str
|
||||
updated: str | None = None
|
||||
categories: list[str]
|
||||
primary_category: str
|
||||
pdf_url: str | None = None
|
||||
doi: str | None = None
|
||||
journal_ref: str | None = None
|
||||
comment: str | None = None
|
||||
|
||||
|
||||
class ArxivMetadata(BaseModel):
|
||||
"""Metadata for ArXiv operation results."""
|
||||
|
||||
operation: str
|
||||
query: str | None = None
|
||||
max_results: int | None = None
|
||||
sort_by: str | None = None
|
||||
sort_order: str | None = None
|
||||
total_results: int | None = None
|
||||
execution_time: float | None = None
|
||||
error_type: str | None = None
|
||||
paper_id: str | None = None
|
||||
download_path: str | None = None
|
||||
file_size: int | None = None
|
||||
|
||||
|
||||
class ArxivActionCollection(ActionCollection):
|
||||
"""MCP service for ArXiv academic paper operations.
|
||||
|
||||
Provides comprehensive ArXiv functionality including:
|
||||
- Paper search with flexible criteria (query, author, category, ID)
|
||||
- Detailed paper metadata extraction
|
||||
- PDF download and text content extraction
|
||||
- Subject category information
|
||||
- LLM-optimized result formatting
|
||||
- Error handling and logging
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
|
||||
# Initialize supported file extensions for PDF processing
|
||||
self.supported_extensions = {".pdf"}
|
||||
|
||||
# ArXiv client configuration
|
||||
self.client = arxiv.Client(
|
||||
page_size=100,
|
||||
delay_seconds=3.0, # Be respectful to ArXiv servers
|
||||
num_retries=3,
|
||||
)
|
||||
|
||||
# Create downloads directory
|
||||
self._downloads_dir = self.workspace / "arxiv_downloads"
|
||||
self._downloads_dir.mkdir(exist_ok=True)
|
||||
|
||||
# ArXiv subject categories mapping
|
||||
self.subject_categories = {
|
||||
"cs": "Computer Science",
|
||||
"math": "Mathematics",
|
||||
"physics": "Physics",
|
||||
"astro-ph": "Astrophysics",
|
||||
"cond-mat": "Condensed Matter",
|
||||
"gr-qc": "General Relativity and Quantum Cosmology",
|
||||
"hep-ex": "High Energy Physics - Experiment",
|
||||
"hep-lat": "High Energy Physics - Lattice",
|
||||
"hep-ph": "High Energy Physics - Phenomenology",
|
||||
"hep-th": "High Energy Physics - Theory",
|
||||
"math-ph": "Mathematical Physics",
|
||||
"nlin": "Nonlinear Sciences",
|
||||
"nucl-ex": "Nuclear Experiment",
|
||||
"nucl-th": "Nuclear Theory",
|
||||
"quant-ph": "Quantum Physics",
|
||||
"q-bio": "Quantitative Biology",
|
||||
"q-fin": "Quantitative Finance",
|
||||
"stat": "Statistics",
|
||||
"econ": "Economics",
|
||||
"eess": "Electrical Engineering and Systems Science",
|
||||
}
|
||||
|
||||
self._color_log("ArXiv service initialized", Color.green, "debug")
|
||||
self._color_log(f"Downloads directory: {self._downloads_dir}", Color.blue, "debug")
|
||||
|
||||
def _format_paper_result(self, paper: arxiv.Result) -> PaperResult:
|
||||
"""Convert arxiv.Result to structured PaperResult.
|
||||
|
||||
Args:
|
||||
paper: ArXiv paper result object
|
||||
|
||||
Returns:
|
||||
Structured PaperResult object
|
||||
"""
|
||||
return PaperResult(
|
||||
entry_id=paper.entry_id,
|
||||
title=paper.title.strip(),
|
||||
authors=[author.name for author in paper.authors],
|
||||
summary=paper.summary.strip(),
|
||||
published=paper.published.isoformat(),
|
||||
updated=paper.updated.isoformat() if paper.updated else None,
|
||||
categories=paper.categories,
|
||||
primary_category=paper.primary_category,
|
||||
pdf_url=paper.pdf_url,
|
||||
doi=paper.doi,
|
||||
journal_ref=paper.journal_ref,
|
||||
comment=paper.comment,
|
||||
)
|
||||
|
||||
def _format_search_results(self, results: list[PaperResult], output_format: str = "markdown") -> str:
|
||||
"""Format paper search results for LLM consumption.
|
||||
|
||||
Args:
|
||||
results: List of paper results
|
||||
output_format: Format type ('markdown', 'json', 'text')
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if not results:
|
||||
return "No papers found matching the search criteria."
|
||||
|
||||
if output_format == "json":
|
||||
return json.dumps([result.model_dump() for result in results], indent=2)
|
||||
|
||||
elif output_format == "text":
|
||||
output_parts = [f"Found {len(results)} papers:\n"]
|
||||
|
||||
for i, paper in enumerate(results, 1):
|
||||
authors_str = ", ".join(paper.authors[:3])
|
||||
if len(paper.authors) > 3:
|
||||
authors_str += f" et al. ({len(paper.authors)} total)"
|
||||
|
||||
output_parts.extend(
|
||||
[
|
||||
f"{i}. {paper.title}",
|
||||
f" Authors: {authors_str}",
|
||||
f" Published: {paper.published[:10]}",
|
||||
f" Categories: {', '.join(paper.categories)}",
|
||||
f" ArXiv ID: {paper.entry_id.split('/')[-1]}",
|
||||
f" Abstract: {paper.summary[:200]}...",
|
||||
"",
|
||||
]
|
||||
)
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
else: # markdown (default)
|
||||
output_parts = [f"# ArXiv Search Results\n\nFound **{len(results)}** papers:\n"]
|
||||
|
||||
for i, paper in enumerate(results, 1):
|
||||
authors_str = ", ".join(paper.authors[:3])
|
||||
if len(paper.authors) > 3:
|
||||
authors_str += f" *et al.* ({len(paper.authors)} total)"
|
||||
|
||||
arxiv_id = paper.entry_id.split("/")[-1]
|
||||
|
||||
output_parts.extend(
|
||||
[
|
||||
f"## {i}. {paper.title}",
|
||||
f"**Authors:** {authors_str}",
|
||||
f"**Published:** {paper.published[:10]}",
|
||||
f"**Categories:** {', '.join(paper.categories)}",
|
||||
f"**ArXiv ID:** `{arxiv_id}`",
|
||||
f"**PDF:** [Download]({paper.pdf_url})" if paper.pdf_url else "",
|
||||
"",
|
||||
f"**Abstract:** {paper.summary}",
|
||||
"",
|
||||
"---",
|
||||
"",
|
||||
]
|
||||
)
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
def _format_paper_details(self, paper: PaperResult, output_format: str = "markdown") -> str:
|
||||
"""Format detailed paper information for LLM consumption.
|
||||
|
||||
Args:
|
||||
paper: Paper result object
|
||||
output_format: Format type ('markdown', 'json', 'text')
|
||||
|
||||
Returns:
|
||||
Formatted string with detailed paper information
|
||||
"""
|
||||
if output_format == "json":
|
||||
return json.dumps(paper.model_dump(), indent=2)
|
||||
|
||||
elif output_format == "text":
|
||||
output_parts = [
|
||||
f"Title: {paper.title}",
|
||||
f"Authors: {', '.join(paper.authors)}",
|
||||
f"Published: {paper.published}",
|
||||
f"Updated: {paper.updated or 'N/A'}",
|
||||
f"Primary Category: {paper.primary_category}",
|
||||
f"All Categories: {', '.join(paper.categories)}",
|
||||
f"ArXiv ID: {paper.entry_id.split('/')[-1]}",
|
||||
f"PDF URL: {paper.pdf_url or 'N/A'}",
|
||||
f"DOI: {paper.doi or 'N/A'}",
|
||||
f"Journal Reference: {paper.journal_ref or 'N/A'}",
|
||||
f"Comment: {paper.comment or 'N/A'}",
|
||||
"",
|
||||
"Abstract:",
|
||||
paper.summary,
|
||||
]
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
else: # markdown (default)
|
||||
arxiv_id = paper.entry_id.split("/")[-1]
|
||||
|
||||
output_parts = [
|
||||
f"# {paper.title}",
|
||||
"",
|
||||
f"**Authors:** {', '.join(paper.authors)}",
|
||||
f"**Published:** {paper.published[:10]}",
|
||||
f"**Updated:** {paper.updated[:10] if paper.updated else 'N/A'}",
|
||||
f"**Primary Category:** {paper.primary_category}",
|
||||
f"**All Categories:** {', '.join(paper.categories)}",
|
||||
f"**ArXiv ID:** `{arxiv_id}`",
|
||||
f"**PDF:** [Download]({paper.pdf_url})" if paper.pdf_url else "**PDF:** N/A",
|
||||
f"**DOI:** {paper.doi}" if paper.doi else "**DOI:** N/A",
|
||||
f"**Journal Reference:** {paper.journal_ref}" if paper.journal_ref else "**Journal Reference:** N/A",
|
||||
f"**Comment:** {paper.comment}" if paper.comment else "**Comment:** N/A",
|
||||
"",
|
||||
"## Abstract",
|
||||
"",
|
||||
paper.summary,
|
||||
]
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
async def mcp_search_papers(
|
||||
self,
|
||||
query: str = Field(description="Search query (keywords, title, author, etc.)"),
|
||||
sort_by: str = Field(
|
||||
default="relevance", description="Sort by: 'relevance', 'lastUpdatedDate', 'submittedDate'"
|
||||
),
|
||||
sort_order: str = Field(default="descending", description="Sort order: 'ascending' or 'descending'"),
|
||||
category: str | None = Field(default=None, description="Filter by ArXiv category (e.g., 'cs.AI', 'math.CO')"),
|
||||
output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'"),
|
||||
) -> ActionResponse:
|
||||
"""Search ArXiv papers with flexible criteria.
|
||||
|
||||
This tool provides comprehensive ArXiv paper search with:
|
||||
- Keyword, title, and author search capabilities
|
||||
- Category filtering for specific subject areas
|
||||
- Flexible sorting options (relevance, date)
|
||||
- Configurable result limits
|
||||
- LLM-optimized result formatting
|
||||
|
||||
Args:
|
||||
query: Search terms (can include keywords, titles, author names)
|
||||
sort_by: Sorting criteria for results
|
||||
sort_order: Order of sorting (ascending/descending)
|
||||
category: Optional category filter (e.g., 'cs.AI' for AI papers)
|
||||
output_format: Format for the response output
|
||||
|
||||
Returns:
|
||||
ActionResponse with search results and metadata
|
||||
"""
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(query, FieldInfo):
|
||||
query = query.default
|
||||
if isinstance(sort_by, FieldInfo):
|
||||
sort_by = sort_by.default
|
||||
if isinstance(sort_order, FieldInfo):
|
||||
sort_order = sort_order.default
|
||||
if isinstance(category, FieldInfo):
|
||||
category = category.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
try:
|
||||
self._color_log(f"🔍 Searching ArXiv for: {query}", Color.cyan)
|
||||
|
||||
start_time = datetime.now()
|
||||
|
||||
# Build search query
|
||||
search_query = query
|
||||
if category:
|
||||
search_query = f"cat:{category} AND ({query})"
|
||||
|
||||
# Configure sort criteria
|
||||
sort_criterion = arxiv.SortCriterion.Relevance
|
||||
if sort_by == "lastUpdatedDate":
|
||||
sort_criterion = arxiv.SortCriterion.LastUpdatedDate
|
||||
elif sort_by == "submittedDate":
|
||||
sort_criterion = arxiv.SortCriterion.SubmittedDate
|
||||
|
||||
sort_order_enum = arxiv.SortOrder.Descending
|
||||
if sort_order == "ascending":
|
||||
sort_order_enum = arxiv.SortOrder.Ascending
|
||||
|
||||
# Perform search
|
||||
search = arxiv.Search(
|
||||
query=search_query, max_results=300000, sort_by=sort_criterion, sort_order=sort_order_enum
|
||||
)
|
||||
|
||||
# Execute search and collect results
|
||||
results = []
|
||||
for paper in self.client.results(search):
|
||||
results.append(self._format_paper_result(paper))
|
||||
|
||||
execution_time = (datetime.now() - start_time).total_seconds()
|
||||
|
||||
# Format output
|
||||
formatted_output = self._format_search_results(results, output_format)
|
||||
|
||||
# Create metadata
|
||||
metadata = ArxivMetadata(
|
||||
operation="search_papers",
|
||||
query=query,
|
||||
sort_by=sort_by,
|
||||
sort_order=sort_order,
|
||||
total_results=len(results),
|
||||
execution_time=execution_time,
|
||||
)
|
||||
|
||||
self._color_log(f"✅ Found {len(results)} papers in {execution_time:.2f}s", Color.green)
|
||||
|
||||
return ActionResponse(success=True, message=formatted_output, metadata=metadata.model_dump())
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to search ArXiv papers: {str(e)}"
|
||||
self.logger.error(f"ArXiv search error: {traceback.format_exc()}")
|
||||
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=error_msg,
|
||||
metadata=ArxivMetadata(operation="search_papers", query=query, error_type="search_error").model_dump(),
|
||||
)
|
||||
|
||||
async def mcp_get_paper_details(
|
||||
self,
|
||||
paper_id: str = Field(description="ArXiv paper ID (e.g., '2301.07041' or 'arxiv:2301.07041')"),
|
||||
output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'"),
|
||||
) -> ActionResponse:
|
||||
"""Get detailed information about a specific ArXiv paper.
|
||||
|
||||
Args:
|
||||
paper_id: ArXiv paper identifier
|
||||
output_format: Format for the response output
|
||||
|
||||
Returns:
|
||||
ActionResponse with detailed paper information and metadata
|
||||
"""
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(paper_id, FieldInfo):
|
||||
paper_id = paper_id.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
try:
|
||||
# Clean paper ID
|
||||
clean_id = paper_id.replace("arxiv:", "").strip()
|
||||
|
||||
self._color_log(f"📄 Getting details for paper: {clean_id}", Color.cyan)
|
||||
|
||||
start_time = datetime.now()
|
||||
|
||||
# Search for the specific paper
|
||||
search = arxiv.Search(id_list=[clean_id])
|
||||
|
||||
paper = next(self.client.results(search), None)
|
||||
if not paper:
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Paper not found: {clean_id}",
|
||||
metadata=ArxivMetadata(
|
||||
operation="get_paper_details", paper_id=clean_id, error_type="paper_not_found"
|
||||
).model_dump(),
|
||||
)
|
||||
|
||||
execution_time = (datetime.now() - start_time).total_seconds()
|
||||
|
||||
# Format paper details
|
||||
paper_result = self._format_paper_result(paper)
|
||||
formatted_output = self._format_paper_details(paper_result, output_format)
|
||||
|
||||
# Create metadata
|
||||
metadata = ArxivMetadata(operation="get_paper_details", paper_id=clean_id, execution_time=execution_time)
|
||||
|
||||
self._color_log(f"✅ Retrieved paper details in {execution_time:.2f}s", Color.green)
|
||||
|
||||
return ActionResponse(success=True, message=formatted_output, metadata=metadata.model_dump())
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to get paper details: {str(e)}"
|
||||
self.logger.error(f"Paper details error: {traceback.format_exc()}")
|
||||
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=error_msg,
|
||||
metadata=ArxivMetadata(
|
||||
operation="get_paper_details", paper_id=paper_id, error_type="retrieval_error"
|
||||
).model_dump(),
|
||||
)
|
||||
|
||||
async def mcp_download_paper(
|
||||
self,
|
||||
paper_id: str = Field(description="ArXiv paper ID (e.g., '2301.07041' or 'arxiv:2301.07041')"),
|
||||
extract_text: bool = Field(default=True, description="Whether to extract text content from PDF"),
|
||||
output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'"),
|
||||
) -> ActionResponse:
|
||||
"""Download ArXiv paper PDF and optionally extract text content.
|
||||
|
||||
Args:
|
||||
paper_id: ArXiv paper identifier
|
||||
extract_text: Whether to extract and return text content
|
||||
output_format: Format for the response output
|
||||
|
||||
Returns:
|
||||
ActionResponse with download status and optional text content
|
||||
"""
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(paper_id, FieldInfo):
|
||||
paper_id = paper_id.default
|
||||
if isinstance(extract_text, FieldInfo):
|
||||
extract_text = extract_text.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
try:
|
||||
# Clean paper ID
|
||||
clean_id = paper_id.replace("arxiv:", "").strip()
|
||||
|
||||
self._color_log(f"📥 Downloading paper: {clean_id}", Color.cyan)
|
||||
|
||||
start_time = datetime.now()
|
||||
|
||||
# Search for the paper
|
||||
search = arxiv.Search(id_list=[clean_id])
|
||||
paper = next(self.client.results(search), None)
|
||||
|
||||
if not paper:
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Paper not found: {clean_id}",
|
||||
metadata=ArxivMetadata(
|
||||
operation="download_paper", paper_id=clean_id, error_type="paper_not_found"
|
||||
).model_dump(),
|
||||
)
|
||||
|
||||
# Download PDF
|
||||
filename = f"{clean_id.replace('/', '_')}.pdf"
|
||||
download_path = self._downloads_dir / filename
|
||||
|
||||
paper.download_pdf(dirpath=str(self._downloads_dir), filename=filename)
|
||||
|
||||
execution_time = (datetime.now() - start_time).total_seconds()
|
||||
file_size = download_path.stat().st_size if download_path.exists() else 0
|
||||
|
||||
# Prepare response message
|
||||
if output_format == "json":
|
||||
response_data = {
|
||||
"paper_id": clean_id,
|
||||
"title": paper.title,
|
||||
"download_path": str(download_path),
|
||||
"file_size": file_size,
|
||||
"download_time": execution_time,
|
||||
}
|
||||
|
||||
if extract_text:
|
||||
try:
|
||||
# Basic text extraction (would need additional libraries like PyPDF2 or pdfplumber)
|
||||
response_data["text_extraction"] = (
|
||||
"Text extraction requires additional PDF processing libraries"
|
||||
)
|
||||
except Exception:
|
||||
response_data["text_extraction"] = "Text extraction failed"
|
||||
|
||||
formatted_output = json.dumps(response_data, indent=2)
|
||||
|
||||
elif output_format == "text":
|
||||
output_parts = [
|
||||
"Paper Downloaded Successfully",
|
||||
f"Paper ID: {clean_id}",
|
||||
f"Title: {paper.title}",
|
||||
f"Download Path: {download_path}",
|
||||
f"File Size: {file_size:,} bytes",
|
||||
f"Download Time: {execution_time:.2f} seconds",
|
||||
]
|
||||
|
||||
if extract_text:
|
||||
output_parts.append("\nNote: Text extraction requires additional PDF processing libraries")
|
||||
|
||||
formatted_output = "\n".join(output_parts)
|
||||
|
||||
else: # markdown (default)
|
||||
output_parts = [
|
||||
"# 📥 Paper Download Complete",
|
||||
"",
|
||||
f"**Paper ID:** `{clean_id}`",
|
||||
f"**Title:** {paper.title}",
|
||||
f"**Download Path:** `{download_path}`",
|
||||
f"**File Size:** {file_size:,} bytes",
|
||||
f"**Download Time:** {execution_time:.2f} seconds",
|
||||
]
|
||||
|
||||
if extract_text:
|
||||
output_parts.extend(
|
||||
[
|
||||
"",
|
||||
"## 📄 Text Extraction",
|
||||
(
|
||||
"*Note: Text extraction requires additional "
|
||||
"PDF processing libraries like PyPDF2 or pdfplumber*"
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
formatted_output = "\n".join(output_parts)
|
||||
|
||||
# Create metadata
|
||||
metadata = ArxivMetadata(
|
||||
operation="download_paper",
|
||||
paper_id=clean_id,
|
||||
download_path=str(download_path),
|
||||
file_size=file_size,
|
||||
execution_time=execution_time,
|
||||
)
|
||||
|
||||
self._color_log(f"✅ Downloaded paper in {execution_time:.2f}s ({file_size:,} bytes)", Color.green)
|
||||
|
||||
return ActionResponse(success=True, message=formatted_output, metadata=metadata.model_dump())
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to download paper: {str(e)}"
|
||||
self.logger.error(f"Paper download error: {traceback.format_exc()}")
|
||||
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=error_msg,
|
||||
metadata=ArxivMetadata(
|
||||
operation="download_paper", paper_id=paper_id, error_type="download_error"
|
||||
).model_dump(),
|
||||
)
|
||||
|
||||
def mcp_get_categories(
|
||||
self,
|
||||
output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'"),
|
||||
) -> ActionResponse:
|
||||
"""Get available ArXiv subject categories.
|
||||
|
||||
Args:
|
||||
output_format: Format for the response output
|
||||
|
||||
Returns:
|
||||
ActionResponse with category information
|
||||
"""
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
try:
|
||||
if output_format == "json":
|
||||
formatted_output = json.dumps(self.subject_categories, indent=2)
|
||||
|
||||
elif output_format == "text":
|
||||
output_parts = ["ArXiv Subject Categories:\n"]
|
||||
for code, name in self.subject_categories.items():
|
||||
output_parts.append(f"{code}: {name}")
|
||||
formatted_output = "\n".join(output_parts)
|
||||
|
||||
else: # markdown (default)
|
||||
output_parts = [
|
||||
"# ArXiv Subject Categories",
|
||||
"",
|
||||
"Available categories for filtering search results:",
|
||||
"",
|
||||
]
|
||||
|
||||
for code, name in self.subject_categories.items():
|
||||
output_parts.append(f"- **`{code}`**: {name}")
|
||||
|
||||
output_parts.extend(
|
||||
[
|
||||
"",
|
||||
"## Usage Examples",
|
||||
"- `cs.AI` - Artificial Intelligence",
|
||||
"- `cs.LG` - Machine Learning",
|
||||
"- `math.CO` - Combinatorics",
|
||||
"- `physics.gen-ph` - General Physics",
|
||||
]
|
||||
)
|
||||
|
||||
formatted_output = "\n".join(output_parts)
|
||||
|
||||
metadata = ArxivMetadata(operation="get_categories", total_results=len(self.subject_categories))
|
||||
|
||||
return ActionResponse(success=True, message=formatted_output, metadata=metadata.model_dump())
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to get categories: {str(e)}"
|
||||
self.logger.error(f"Categories error: {traceback.format_exc()}")
|
||||
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=error_msg,
|
||||
metadata=ArxivMetadata(operation="get_categories", error_type="internal_error").model_dump(),
|
||||
)
|
||||
|
||||
def mcp_get_arxiv_capabilities(self) -> ActionResponse:
|
||||
"""Get information about ArXiv service capabilities and configuration.
|
||||
|
||||
Returns:
|
||||
ActionResponse with service capabilities and current configuration
|
||||
"""
|
||||
capabilities = {
|
||||
"supported_operations": [
|
||||
"Paper search with flexible criteria",
|
||||
"Detailed paper metadata retrieval",
|
||||
"PDF download and storage",
|
||||
"Subject category filtering",
|
||||
"Multiple output formats (markdown, json, text)",
|
||||
"LLM-optimized result formatting",
|
||||
],
|
||||
"search_capabilities": [
|
||||
"Keyword and phrase search",
|
||||
"Author name search",
|
||||
"Title search",
|
||||
"Category filtering",
|
||||
"Date-based sorting",
|
||||
"Relevance-based sorting",
|
||||
],
|
||||
"supported_formats": ["markdown", "json", "text"],
|
||||
"configuration": {
|
||||
"downloads_directory": str(self._downloads_dir),
|
||||
"client_page_size": 100,
|
||||
"client_delay_seconds": 3.0,
|
||||
"client_num_retries": 3,
|
||||
"supported_categories_count": len(self.subject_categories),
|
||||
},
|
||||
"rate_limiting": {
|
||||
"delay_between_requests": "3.0 seconds",
|
||||
"retry_attempts": 3,
|
||||
"respectful_usage": "Configured for ArXiv server guidelines",
|
||||
},
|
||||
}
|
||||
|
||||
formatted_info = f"""# ArXiv Service Capabilities
|
||||
|
||||
## Supported Operations
|
||||
{chr(10).join(f"- {op}" for op in capabilities["supported_operations"])}
|
||||
|
||||
## Search Capabilities
|
||||
{chr(10).join(f"- {cap}" for cap in capabilities["search_capabilities"])}
|
||||
|
||||
## Supported Output Formats
|
||||
{chr(10).join(f"- {fmt}" for fmt in capabilities["supported_formats"])}
|
||||
|
||||
## Current Configuration
|
||||
- **Downloads Directory:** {capabilities["configuration"]["downloads_directory"]}
|
||||
- **Client Page Size:** {capabilities["configuration"]["client_page_size"]}
|
||||
- **Request Delay:** {capabilities["configuration"]["client_delay_seconds"]} seconds
|
||||
- **Retry Attempts:** {capabilities["configuration"]["client_num_retries"]}
|
||||
- **Available Categories:** {capabilities["configuration"]["supported_categories_count"]}
|
||||
|
||||
## Rate Limiting & Ethics
|
||||
- **Delay Between Requests:** {capabilities["rate_limiting"]["delay_between_requests"]}
|
||||
- **Retry Policy:** {capabilities["rate_limiting"]["retry_attempts"]} attempts
|
||||
- **Server Respect:** {capabilities["rate_limiting"]["respectful_usage"]}
|
||||
"""
|
||||
|
||||
return ActionResponse(success=True, message=formatted_info, metadata=capabilities)
|
||||
|
||||
|
||||
# Default arguments for testing
|
||||
if __name__ == "__main__":
|
||||
import os
|
||||
|
||||
load_dotenv()
|
||||
|
||||
arguments = ActionArguments(
|
||||
name="arxiv",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
|
||||
)
|
||||
|
||||
try:
|
||||
service = ArxivActionCollection(arguments)
|
||||
service.run()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}: {traceback.format_exc()}")
|
||||
@@ -0,0 +1,535 @@
|
||||
"""
|
||||
Chess MCP Server
|
||||
|
||||
This module provides MCP server functionality for chess game operations and analysis.
|
||||
It utilizes the 'python-chess' library to support various chess-related tasks.
|
||||
|
||||
Key features:
|
||||
- Manage chess game states (new game, load FEN, make moves)
|
||||
- Validate and execute moves in UCI or SAN format
|
||||
- Get legal moves for the current position
|
||||
- Check game status (checkmate, stalemate, draw, etc.)
|
||||
- Basic board visualization (ASCII)
|
||||
- LLM-optimized output formatting for game states and analysis
|
||||
|
||||
Main functions:
|
||||
- mcp_new_game: Start a new chess game
|
||||
- mcp_load_fen: Load a game state from FEN notation
|
||||
- mcp_make_move: Make a move on the current board
|
||||
- mcp_get_legal_moves: List all legal moves in the current position
|
||||
- mcp_get_board_state: Get the current board state (FEN, ASCII, status)
|
||||
- mcp_get_game_status: Check the current game status (e.g., checkmate)
|
||||
- mcp_get_chess_capabilities: Get service capabilities
|
||||
"""
|
||||
|
||||
import json
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
import chess
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from aworld.logs.util import Color
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
|
||||
|
||||
class ChessBoardState(BaseModel):
|
||||
"""Structured representation of the chess board state."""
|
||||
|
||||
fen: str
|
||||
turn: str # 'white' or 'black'
|
||||
castling_rights: str
|
||||
ep_square: str | None = None # The target square if an en passant capture is possible *right now*
|
||||
halfmove_clock: int
|
||||
fullmove_number: int
|
||||
is_check: bool
|
||||
is_checkmate: bool
|
||||
is_stalemate: bool
|
||||
is_insufficient_material: bool
|
||||
is_seventyfive_moves: bool
|
||||
is_fivefold_repetition: bool
|
||||
is_game_over: bool
|
||||
ascii_board: str
|
||||
legal_moves_uci: list[str]
|
||||
legal_moves_san: list[str]
|
||||
is_en_passant_possible: bool # True if there is a legal en passant capture
|
||||
en_passant_capture_square: str | None = None # The square a pawn would move TO for en passant
|
||||
|
||||
|
||||
class ChessMoveResult(BaseModel):
|
||||
"""Result of making a chess move."""
|
||||
|
||||
move_uci: str
|
||||
move_san: str
|
||||
is_capture: bool
|
||||
is_check: bool
|
||||
is_kingside_castling: bool
|
||||
is_queenside_castling: bool
|
||||
board_after_move: ChessBoardState
|
||||
|
||||
|
||||
class ChessMetadata(BaseModel):
|
||||
"""Metadata for Chess operation results."""
|
||||
|
||||
operation: str
|
||||
fen_before: str | None = None
|
||||
fen_after: str | None = None
|
||||
move_played: str | None = None
|
||||
execution_time: float | None = None
|
||||
error_type: str | None = None
|
||||
engine_analysis_depth: int | None = None
|
||||
|
||||
|
||||
class ChessCollection(ActionCollection):
|
||||
"""MCP service for chess game operations and analysis.
|
||||
|
||||
Provides capabilities to manage chess games, make moves, analyze positions,
|
||||
and get game status, all formatted for LLM interaction.
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
self.board = chess.Board()
|
||||
# For more advanced analysis, you might initialize a chess engine here
|
||||
# Example: self.engine = chess.engine.SimpleEngine.popen_uci("/path/to/stockfish")
|
||||
# Ensure Stockfish or another UCI engine is installed and path is correct.
|
||||
self._color_log("Chess service initialized", Color.green, "debug")
|
||||
self._color_log(f"Initial board FEN: {self.board.fen()}", Color.blue, "debug")
|
||||
|
||||
def _get_current_board_state(self) -> ChessBoardState:
|
||||
"""Helper to get the current board state in a structured format."""
|
||||
legal_moves_uci = [move.uci() for move in self.board.legal_moves]
|
||||
legal_moves_san = []
|
||||
# Generating SAN for all moves can be slow, do it carefully or on demand
|
||||
# for move in self.board.legal_moves:
|
||||
# try:
|
||||
# legal_moves_san.append(self.board.san(move))
|
||||
# except Exception:
|
||||
# legal_moves_san.append(move.uci()) # Fallback to UCI if SAN fails
|
||||
|
||||
has_legal_ep = self.board.has_legal_en_passant()
|
||||
ep_sq_name = chess.square_name(self.board.ep_square) if self.board.ep_square else None
|
||||
|
||||
return ChessBoardState(
|
||||
fen=self.board.fen(),
|
||||
turn="white" if self.board.turn == chess.WHITE else "black",
|
||||
castling_rights=self.board.castling_xfen(),
|
||||
ep_square=ep_sq_name, # This is the target square from FEN, might not be a legal capture
|
||||
halfmove_clock=self.board.halfmove_clock,
|
||||
fullmove_number=self.board.fullmove_number,
|
||||
is_check=self.board.is_check(),
|
||||
is_checkmate=self.board.is_checkmate(),
|
||||
is_stalemate=self.board.is_stalemate(),
|
||||
is_insufficient_material=self.board.is_insufficient_material(),
|
||||
is_seventyfive_moves=self.board.is_seventyfive_moves(),
|
||||
is_fivefold_repetition=self.board.is_fivefold_repetition(),
|
||||
is_game_over=self.board.is_game_over(),
|
||||
ascii_board=str(self.board),
|
||||
legal_moves_uci=legal_moves_uci,
|
||||
legal_moves_san=legal_moves_san, # Populate if SAN generation is enabled
|
||||
is_en_passant_possible=has_legal_ep,
|
||||
en_passant_capture_square=ep_sq_name if has_legal_ep else None,
|
||||
)
|
||||
|
||||
def _format_board_state_output(self, state: ChessBoardState, output_format: str = "markdown") -> str:
|
||||
"""Format board state for LLM consumption."""
|
||||
if output_format == "json":
|
||||
return json.dumps(state.model_dump(), indent=2)
|
||||
|
||||
status_parts = []
|
||||
if state.is_checkmate:
|
||||
status_parts.append("Checkmate!")
|
||||
elif state.is_stalemate:
|
||||
status_parts.append("Stalemate!")
|
||||
elif state.is_insufficient_material:
|
||||
status_parts.append("Draw by insufficient material.")
|
||||
elif state.is_seventyfive_moves:
|
||||
status_parts.append("Draw by 75-move rule.")
|
||||
elif state.is_fivefold_repetition:
|
||||
status_parts.append("Draw by fivefold repetition.")
|
||||
elif state.is_check:
|
||||
status_parts.append("Check!")
|
||||
game_status = " ".join(status_parts) if status_parts else "Game in progress."
|
||||
|
||||
en_passant_info = "N/A"
|
||||
if state.is_en_passant_possible and state.en_passant_capture_square:
|
||||
en_passant_info = f"Yes, capture on {state.en_passant_capture_square}"
|
||||
elif state.ep_square: # FEN might list an ep_square even if no legal ep move
|
||||
en_passant_info = f"Target square {state.ep_square} (no legal en passant capture)"
|
||||
|
||||
if output_format == "text":
|
||||
return (
|
||||
f"Board FEN: {state.fen}\n"
|
||||
f"Turn: {state.turn.capitalize()}\n"
|
||||
f"Status: {game_status}\n"
|
||||
f"Castling: {state.castling_rights}\n"
|
||||
f"En Passant Possible: {en_passant_info}\n" # Updated line
|
||||
f"Halfmove Clock: {state.halfmove_clock}\n"
|
||||
f"Fullmove Number: {state.fullmove_number}\n"
|
||||
f"Game Over: {'Yes' if state.is_game_over else 'No'}\n"
|
||||
f"Legal Moves (UCI): {', '.join(state.legal_moves_uci[:10])}... ({len(state.legal_moves_uci)} total)\n"
|
||||
f"Board:\n{state.ascii_board}"
|
||||
)
|
||||
else: # markdown (default)
|
||||
return (
|
||||
f"### Chess Board State\n"
|
||||
f"**FEN:** `{state.fen}`\n"
|
||||
f"**Turn:** {state.turn.capitalize()}\n"
|
||||
f"**Status:** {game_status}\n"
|
||||
f"**Castling Rights:** {state.castling_rights}\n"
|
||||
f"**En Passant Possible:** {en_passant_info}\n" # Updated line
|
||||
f"**Game Over:** {'Yes' if state.is_game_over else 'No'}\n"
|
||||
f"**Legal Moves (UCI, sample):** `{', '.join(state.legal_moves_uci[:5])}`... ({len(state.legal_moves_uci)} total)\n"
|
||||
f"```\n{state.ascii_board}\n```"
|
||||
)
|
||||
|
||||
async def mcp_new_game(self) -> ActionResponse:
|
||||
"""Starts a new standard chess game, resetting the board.
|
||||
|
||||
Returns:
|
||||
ActionResponse with the initial board state.
|
||||
"""
|
||||
start_time = datetime.now()
|
||||
self.board.reset()
|
||||
self._color_log("🚀 New chess game started", Color.green)
|
||||
|
||||
current_state = self._get_current_board_state()
|
||||
formatted_output = self._format_board_state_output(current_state)
|
||||
execution_time = (datetime.now() - start_time).total_seconds()
|
||||
|
||||
metadata = ChessMetadata(
|
||||
operation="new_game", fen_after=self.board.fen(), execution_time=execution_time
|
||||
).model_dump()
|
||||
|
||||
return ActionResponse(success=True, message=formatted_output, metadata=metadata)
|
||||
|
||||
async def mcp_load_fen(
|
||||
self, fen_string: str = Field(description="FEN string representing the board state.")
|
||||
) -> ActionResponse:
|
||||
"""Loads a chess game from a FEN (Forsyth-Edwards Notation) string.
|
||||
|
||||
Args:
|
||||
fen_string: The FEN string to load.
|
||||
|
||||
Returns:
|
||||
ActionResponse with the board state after loading the FEN.
|
||||
"""
|
||||
# Handle FieldInfo
|
||||
if isinstance(fen_string, FieldInfo):
|
||||
fen_string = fen_string.default
|
||||
|
||||
start_time = datetime.now()
|
||||
try:
|
||||
self.board.set_fen(fen_string)
|
||||
self._color_log(f"🔄 Board loaded from FEN: {fen_string}", Color.blue)
|
||||
current_state = self._get_current_board_state()
|
||||
formatted_output = self._format_board_state_output(current_state)
|
||||
execution_time = (datetime.now() - start_time).total_seconds()
|
||||
metadata = ChessMetadata(
|
||||
operation="load_fen", fen_after=self.board.fen(), execution_time=execution_time
|
||||
).model_dump()
|
||||
return ActionResponse(success=True, message=formatted_output, metadata=metadata)
|
||||
except ValueError as e:
|
||||
error_msg = f"Invalid FEN string: {str(e)}"
|
||||
self.logger.error(f"FEN loading error: {traceback.format_exc()}")
|
||||
execution_time = (datetime.now() - start_time).total_seconds()
|
||||
metadata = ChessMetadata(
|
||||
operation="load_fen", error_type="invalid_fen", execution_time=execution_time
|
||||
).model_dump()
|
||||
return ActionResponse(success=False, message=error_msg, metadata=metadata)
|
||||
|
||||
async def mcp_make_move(
|
||||
self, move_str: str = Field(description="Move in UCI (e.g., 'e2e4') or SAN (e.g., 'Nf3') format.")
|
||||
) -> ActionResponse:
|
||||
"""Makes a move on the current chess board.
|
||||
|
||||
The move can be in UCI (Universal Chess Interface) format (e.g., 'g1f3')
|
||||
or SAN (Standard Algebraic Notation) format (e.g., 'Nf3').
|
||||
|
||||
Args:
|
||||
move_str: The move to make.
|
||||
|
||||
Returns:
|
||||
ActionResponse with the result of the move and new board state.
|
||||
"""
|
||||
# Handle FieldInfo
|
||||
if isinstance(move_str, FieldInfo):
|
||||
move_str = move_str.default
|
||||
|
||||
start_time = datetime.now()
|
||||
fen_before = self.board.fen()
|
||||
try:
|
||||
move = None
|
||||
# Try parsing as UCI first, then SAN
|
||||
try:
|
||||
move = self.board.parse_uci(move_str)
|
||||
except ValueError:
|
||||
try:
|
||||
move = self.board.parse_san(move_str)
|
||||
except ValueError as e_san:
|
||||
raise ValueError(f"Invalid move format. UCI error: N/A, SAN error: {e_san}") from e_san
|
||||
|
||||
if move not in self.board.legal_moves:
|
||||
raise ValueError(f"Illegal move: {move_str}")
|
||||
|
||||
move_san = self.board.san(move)
|
||||
is_capture = self.board.is_capture(move)
|
||||
is_kingside_castling = self.board.is_kingside_castling(move)
|
||||
is_queenside_castling = self.board.is_queenside_castling(move)
|
||||
|
||||
self.board.push(move)
|
||||
is_check_after_move = self.board.is_check()
|
||||
|
||||
self._color_log(f"♟️ Move made: {move_str} (UCI: {move.uci()}, SAN: {move_san})", Color.cyan)
|
||||
|
||||
current_state = self._get_current_board_state()
|
||||
move_result = ChessMoveResult(
|
||||
move_uci=move.uci(),
|
||||
move_san=move_san,
|
||||
is_capture=is_capture,
|
||||
is_check=is_check_after_move, # Check status *after* the move
|
||||
is_kingside_castling=is_kingside_castling,
|
||||
is_queenside_castling=is_queenside_castling,
|
||||
board_after_move=current_state,
|
||||
)
|
||||
|
||||
# Format output (can be customized)
|
||||
formatted_output = f"Move {move_result.move_san} (UCI: {move_result.move_uci}) played.\n"
|
||||
formatted_output += self._format_board_state_output(current_state)
|
||||
|
||||
execution_time = (datetime.now() - start_time).total_seconds()
|
||||
metadata = ChessMetadata(
|
||||
operation="make_move",
|
||||
fen_before=fen_before,
|
||||
fen_after=self.board.fen(),
|
||||
move_played=move.uci(),
|
||||
execution_time=execution_time,
|
||||
).model_dump()
|
||||
|
||||
return ActionResponse(success=True, message=formatted_output, metadata=metadata)
|
||||
|
||||
except ValueError as e:
|
||||
error_msg = f"Failed to make move '{move_str}': {str(e)}"
|
||||
self.logger.error(f"Move error: {traceback.format_exc()}")
|
||||
execution_time = (datetime.now() - start_time).total_seconds()
|
||||
metadata = ChessMetadata(
|
||||
operation="make_move",
|
||||
fen_before=fen_before,
|
||||
move_played=move_str,
|
||||
error_type="invalid_or_illegal_move",
|
||||
execution_time=execution_time,
|
||||
).model_dump()
|
||||
return ActionResponse(success=False, message=error_msg, metadata=metadata)
|
||||
|
||||
async def mcp_get_legal_moves(
|
||||
self,
|
||||
output_format: str = Field(
|
||||
default="markdown", description="Output format: 'uci_list', 'san_list', 'markdown', 'json'"
|
||||
),
|
||||
) -> ActionResponse:
|
||||
"""Gets all legal moves for the current board position.
|
||||
|
||||
Args:
|
||||
output_format: 'uci_list' (simple list of UCI moves),
|
||||
'san_list' (simple list of SAN moves),
|
||||
'markdown' (formatted list),
|
||||
'json' (structured list).
|
||||
|
||||
Returns:
|
||||
ActionResponse with the list of legal moves.
|
||||
"""
|
||||
# Handle FieldInfo
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
start_time = datetime.now()
|
||||
legal_moves_uci = [move.uci() for move in self.board.legal_moves]
|
||||
|
||||
message_content: Any
|
||||
if output_format == "uci_list":
|
||||
message_content = legal_moves_uci
|
||||
elif output_format == "san_list":
|
||||
try:
|
||||
message_content = [self.board.san(move) for move in self.board.legal_moves]
|
||||
except Exception as e:
|
||||
self.logger.warning(
|
||||
f"Could not generate all SAN moves: {e}. Falling back to UCI for problematic moves."
|
||||
)
|
||||
san_moves = []
|
||||
for move in self.board.legal_moves:
|
||||
try:
|
||||
san_moves.append(self.board.san(move))
|
||||
except Exception:
|
||||
san_moves.append(move.uci() + " (SAN failed)")
|
||||
message_content = san_moves
|
||||
elif output_format == "json":
|
||||
moves_data = []
|
||||
for move in self.board.legal_moves:
|
||||
try:
|
||||
san = self.board.san(move)
|
||||
except Exception:
|
||||
san = move.uci() + " (SAN failed)"
|
||||
moves_data.append({"uci": move.uci(), "san": san})
|
||||
message_content = json.dumps(moves_data, indent=2)
|
||||
else: # markdown
|
||||
if not legal_moves_uci:
|
||||
message_content = "No legal moves available (game might be over)."
|
||||
else:
|
||||
san_formatted_moves = []
|
||||
for move_uci in legal_moves_uci[:20]: # Display sample for markdown
|
||||
try:
|
||||
move_obj = self.board.parse_uci(move_uci)
|
||||
san_formatted_moves.append(f"`{self.board.san(move_obj)}` ({move_uci})")
|
||||
except Exception:
|
||||
san_formatted_moves.append(f"`{move_uci}` (SAN failed)")
|
||||
|
||||
header = f"### Legal Moves ({len(legal_moves_uci)} total)\n"
|
||||
moves_list_md = "\n".join([f"- {m}" for m in san_formatted_moves])
|
||||
if len(legal_moves_uci) > 20:
|
||||
moves_list_md += "\n- ... (and more)"
|
||||
message_content = header + moves_list_md
|
||||
|
||||
execution_time = (datetime.now() - start_time).total_seconds()
|
||||
metadata = ChessMetadata(
|
||||
operation="get_legal_moves", fen_before=self.board.fen(), execution_time=execution_time
|
||||
).model_dump()
|
||||
|
||||
return ActionResponse(success=True, message=message_content, metadata=metadata)
|
||||
|
||||
async def mcp_get_board_state(
|
||||
self, output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'")
|
||||
) -> ActionResponse:
|
||||
"""Gets the current state of the chess board.
|
||||
|
||||
Includes FEN, turn, game status, ASCII board, and legal moves.
|
||||
|
||||
Args:
|
||||
output_format: Desired format for the board state.
|
||||
|
||||
Returns:
|
||||
ActionResponse with the current board state.
|
||||
"""
|
||||
# Handle FieldInfo
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
start_time = datetime.now()
|
||||
current_state = self._get_current_board_state()
|
||||
formatted_output = self._format_board_state_output(current_state, output_format)
|
||||
execution_time = (datetime.now() - start_time).total_seconds()
|
||||
|
||||
metadata = ChessMetadata(
|
||||
operation="get_board_state",
|
||||
fen_before=self.board.fen(), # FEN is part of the state, so 'before' and 'after' are same here
|
||||
execution_time=execution_time,
|
||||
).model_dump()
|
||||
|
||||
return ActionResponse(success=True, message=formatted_output, metadata=metadata)
|
||||
|
||||
async def mcp_get_game_status(self) -> ActionResponse:
|
||||
"""Checks and returns the current game status (e.g., checkmate, stalemate).
|
||||
|
||||
Returns:
|
||||
ActionResponse with a human-readable game status and structured data.
|
||||
"""
|
||||
start_time = datetime.now()
|
||||
state = self._get_current_board_state()
|
||||
|
||||
status_message = "Game in progress."
|
||||
if state.is_checkmate:
|
||||
status_message = f"Checkmate! {state.turn.capitalize()} is mated."
|
||||
elif state.is_stalemate:
|
||||
status_message = "Stalemate! The game is a draw."
|
||||
elif state.is_insufficient_material:
|
||||
status_message = "Draw by insufficient material."
|
||||
elif state.is_seventyfive_moves:
|
||||
status_message = "Draw by 75-move rule."
|
||||
elif state.is_fivefold_repetition:
|
||||
status_message = "Draw by fivefold repetition."
|
||||
elif state.is_check:
|
||||
status_message = f"{state.turn.capitalize()} is in check."
|
||||
|
||||
status_data = {
|
||||
"status_message": status_message,
|
||||
"is_game_over": state.is_game_over,
|
||||
"is_check": state.is_check,
|
||||
"is_checkmate": state.is_checkmate,
|
||||
"is_stalemate": state.is_stalemate,
|
||||
"is_draw": state.is_stalemate
|
||||
or state.is_insufficient_material
|
||||
or state.is_seventyfive_moves
|
||||
or state.is_fivefold_repetition,
|
||||
"winner": None, # Could be determined if checkmate
|
||||
}
|
||||
if state.is_checkmate:
|
||||
status_data["winner"] = "black" if self.board.turn == chess.WHITE else "white"
|
||||
|
||||
execution_time = (datetime.now() - start_time).total_seconds()
|
||||
metadata = ChessMetadata(
|
||||
operation="get_game_status", fen_before=self.board.fen(), execution_time=execution_time
|
||||
).model_dump()
|
||||
metadata.update(status_data) # Add specific status flags to metadata
|
||||
|
||||
return ActionResponse(success=True, message=status_message, metadata=metadata)
|
||||
|
||||
def mcp_get_chess_capabilities(self) -> ActionResponse:
|
||||
"""Get information about the Chess service capabilities.
|
||||
|
||||
Returns:
|
||||
ActionResponse with service capabilities.
|
||||
"""
|
||||
capabilities_info = {
|
||||
"service_name": "Chess MCP Service",
|
||||
"library_used": "python-chess",
|
||||
"supported_operations": [
|
||||
"new_game: Start a new chess game.",
|
||||
"load_fen: Load game state from FEN string.",
|
||||
"make_move: Make a move (UCI or SAN).",
|
||||
"get_legal_moves: List legal moves.",
|
||||
"get_board_state: Get current board FEN, ASCII, status, etc.",
|
||||
"get_game_status: Check for checkmate, stalemate, draw conditions.",
|
||||
],
|
||||
"output_formats": ["markdown", "json", "text"],
|
||||
"move_input_formats": ["UCI (e.g., e2e4)", "SAN (e.g., Nf3)"],
|
||||
"fen_support": "Full FEN loading and generation.",
|
||||
"engine_integration": "Basic structure for UCI engine integration (not fully implemented by default).",
|
||||
}
|
||||
|
||||
formatted_message = "# Chess Service Capabilities\n\n"
|
||||
formatted_message += f"**Service Name:** {capabilities_info['service_name']}\n"
|
||||
formatted_message += f"**Core Library:** {capabilities_info['library_used']}\n\n"
|
||||
formatted_message += "**Supported Operations:**\n"
|
||||
for op in capabilities_info["supported_operations"]:
|
||||
formatted_message += f"- {op}\n"
|
||||
formatted_message += "\n**Supported Output Formats:** " + ", ".join(capabilities_info["output_formats"]) + "\n"
|
||||
formatted_message += "**Move Input Formats:** " + ", ".join(capabilities_info["move_input_formats"]) + "\n"
|
||||
|
||||
return ActionResponse(success=True, message=formatted_message, metadata=capabilities_info)
|
||||
|
||||
# Optional: Method to close engine if it was initialized
|
||||
# def __del__(self):
|
||||
# if hasattr(self, 'engine') and self.engine:
|
||||
# self.engine.quit()
|
||||
# self._color_log("Chess engine quit.", Color.yellow)
|
||||
|
||||
|
||||
# Default arguments for testing
|
||||
if __name__ == "__main__":
|
||||
import os
|
||||
|
||||
load_dotenv()
|
||||
|
||||
arguments = ActionArguments(
|
||||
name="chess_service",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
|
||||
)
|
||||
|
||||
try:
|
||||
service = ChessCollection(arguments)
|
||||
service.run()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}: {traceback.format_exc()}")
|
||||
@@ -0,0 +1,641 @@
|
||||
"""
|
||||
PubChem MCP Server
|
||||
|
||||
This module provides MCP server functionality for accessing PubChem database programmatically.
|
||||
It supports compound searches, property retrieval, and structure-based queries using PubChem's REST API.
|
||||
|
||||
Key features:
|
||||
- Compound search by name, CID, SMILES, or InChI
|
||||
- Property retrieval (molecular weight, formula, etc.)
|
||||
- Structure similarity searches
|
||||
- Bioactivity data access
|
||||
- 3D structure downloads
|
||||
- Rate limiting compliance (max 5 requests/second)
|
||||
|
||||
Main functions:
|
||||
- mcp_search_compounds: Search for compounds by various identifiers
|
||||
- mcp_get_compound_properties: Retrieve compound properties
|
||||
- mcp_get_compound_synonyms: Get compound names and synonyms
|
||||
- mcp_search_similar_compounds: Find structurally similar compounds
|
||||
- mcp_get_bioactivity_data: Retrieve bioactivity assay data
|
||||
- mcp_download_structure: Download 2D/3D structure files
|
||||
"""
|
||||
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from typing import Literal
|
||||
from urllib.parse import quote
|
||||
|
||||
import requests
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from aworld.logs.util import Color
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
|
||||
|
||||
# pylint: disable=C0301
|
||||
class CompoundData(BaseModel):
|
||||
"""Structured compound data from PubChem."""
|
||||
|
||||
cid: int | None = None
|
||||
name: str | None = None
|
||||
molecular_formula: str | None = None
|
||||
molecular_weight: float | None = None
|
||||
smiles: str | None = None
|
||||
inchi: str | None = None
|
||||
synonyms: list[str] = []
|
||||
|
||||
|
||||
class PubChemMetadata(BaseModel):
|
||||
"""Metadata for PubChem operation results."""
|
||||
|
||||
query_type: str
|
||||
query_value: str
|
||||
api_endpoint: str
|
||||
response_time: float
|
||||
total_results: int | None = None
|
||||
rate_limit_delay: float | None = None
|
||||
error_type: str | None = None
|
||||
timestamp: str
|
||||
|
||||
|
||||
class PubChemCollection(ActionCollection):
|
||||
"""MCP service for PubChem database access with comprehensive chemical data retrieval.
|
||||
|
||||
Provides access to PubChem's extensive chemical database including:
|
||||
- Compound identification and search capabilities
|
||||
- Chemical property and structure data
|
||||
- Bioactivity and assay information
|
||||
- Structure similarity searches
|
||||
- 2D/3D molecular structure downloads
|
||||
- Synonym and nomenclature data
|
||||
|
||||
Complies with PubChem usage policies:
|
||||
- Maximum 5 requests per second
|
||||
- Automatic rate limiting
|
||||
- Proper error handling for timeouts
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
|
||||
# PubChem API configuration
|
||||
self.base_url = "https://pubchem.ncbi.nlm.nih.gov/rest/pug"
|
||||
self.base_url_view = "https://pubchem.ncbi.nlm.nih.gov/rest/pug_view"
|
||||
self.request_delay = 0.2 # 200ms delay to stay under 5 req/sec limit
|
||||
self.last_request_time = 0.0
|
||||
|
||||
# Request timeout settings
|
||||
self.timeout = 30 # PubChem's 30-second limit
|
||||
|
||||
# Initialize request session with headers
|
||||
self.session = requests.Session()
|
||||
self.session.headers.update(
|
||||
{"User-Agent": "AWorld-PubChem-MCP/1.0 (https://github.com/aworld-framework)", "Accept": "application/json"}
|
||||
)
|
||||
|
||||
self._color_log("PubChem MCP Service initialized", Color.green, "debug")
|
||||
self._color_log(f"Base URL: {self.base_url}", Color.blue, "debug")
|
||||
|
||||
def _rate_limit(self) -> float:
|
||||
"""Enforce rate limiting to comply with PubChem usage policy.
|
||||
|
||||
Returns:
|
||||
Actual delay time applied
|
||||
"""
|
||||
current_time = time.time()
|
||||
time_since_last = current_time - self.last_request_time
|
||||
|
||||
if time_since_last < self.request_delay:
|
||||
delay = self.request_delay - time_since_last
|
||||
time.sleep(delay)
|
||||
self.last_request_time = time.time()
|
||||
return delay
|
||||
|
||||
self.last_request_time = current_time
|
||||
return 0.0
|
||||
|
||||
def _make_request(self, url: str, params: dict = None) -> tuple[dict | None, float]:
|
||||
"""Make a rate-limited request to PubChem API.
|
||||
|
||||
Args:
|
||||
url: API endpoint URL
|
||||
params: Query parameters
|
||||
|
||||
Returns:
|
||||
Tuple of (response_data, response_time)
|
||||
|
||||
Raises:
|
||||
requests.RequestException: For API request failures
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
response = self.session.get(url, params=params, timeout=self.timeout)
|
||||
response_time = time.time() - start_time
|
||||
|
||||
if response.status_code == 200:
|
||||
return response.json(), response_time
|
||||
elif response.status_code == 503:
|
||||
raise requests.RequestException("PubChem service temporarily unavailable (503)")
|
||||
else:
|
||||
raise requests.RequestException(f"HTTP {response.status_code}: {response.text}")
|
||||
|
||||
except requests.Timeout as e:
|
||||
response_time = time.time() - start_time
|
||||
raise requests.RequestException(f"Request timeout after {self.timeout}s") from e
|
||||
except requests.RequestException:
|
||||
response_time = time.time() - start_time
|
||||
raise
|
||||
|
||||
def mcp_search_compounds(
|
||||
self,
|
||||
query: str = Field(description="Search query (compound name, CID, SMILES, InChI, etc.)"),
|
||||
search_type: Literal["name", "cid", "smiles", "inchi", "formula"] = Field(
|
||||
default="name",
|
||||
description="Type of search: name (compound name), cid (PubChem ID), smiles, inchi, or formula",
|
||||
),
|
||||
max_results: int = Field(default=10, description="Maximum number of results to return (1-100)", ge=1, le=100),
|
||||
) -> ActionResponse:
|
||||
"""Search for chemical compounds in PubChem database.
|
||||
|
||||
Supports multiple search types:
|
||||
- Name: Search by common or IUPAC names
|
||||
- CID: Search by PubChem Compound ID
|
||||
- SMILES: Search by SMILES notation
|
||||
- InChI: Search by InChI identifier
|
||||
- Formula: Search by molecular formula
|
||||
|
||||
Args:
|
||||
query: Search term or identifier
|
||||
search_type: Type of search to perform
|
||||
max_results: Maximum number of compounds to return
|
||||
|
||||
Returns:
|
||||
ActionResponse with compound search results and metadata
|
||||
"""
|
||||
try:
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(query, FieldInfo):
|
||||
query = query.default
|
||||
if isinstance(search_type, FieldInfo):
|
||||
search_type = search_type.default
|
||||
if isinstance(max_results, FieldInfo):
|
||||
max_results = max_results.default
|
||||
|
||||
if not query or not query.strip():
|
||||
raise ValueError("Search query is required")
|
||||
|
||||
self._color_log(f"Searching PubChem for: {query} (type: {search_type})", Color.cyan)
|
||||
|
||||
# Build API URL based on search type
|
||||
if search_type == "cid":
|
||||
url = f"{self.base_url}/compound/cid/{quote(str(query))}/property/Title,MolecularFormula,MolecularWeight,CanonicalSMILES,InChI/JSON"
|
||||
elif search_type == "name":
|
||||
url = f"{self.base_url}/compound/name/{quote(query)}/property/Title,MolecularFormula,MolecularWeight,CanonicalSMILES,InChI/JSON"
|
||||
elif search_type == "smiles":
|
||||
url = f"{self.base_url}/compound/smiles/{quote(query)}/property/Title,MolecularFormula,MolecularWeight,CanonicalSMILES,InChI/JSON"
|
||||
elif search_type == "inchi":
|
||||
url = f"{self.base_url}/compound/inchi/{quote(query)}/property/Title,MolecularFormula,MolecularWeight,CanonicalSMILES,InChI/JSON"
|
||||
elif search_type == "formula":
|
||||
url = f"{self.base_url}/compound/formula/{quote(query)}/property/Title,MolecularFormula,MolecularWeight,CanonicalSMILES,InChI/JSON"
|
||||
else:
|
||||
raise ValueError(f"Unsupported search type: {search_type}")
|
||||
|
||||
# Make API request
|
||||
data, response_time = self._make_request(url)
|
||||
|
||||
# Parse results
|
||||
compounds = []
|
||||
if data and "PropertyTable" in data and "Properties" in data["PropertyTable"]:
|
||||
properties_list = data["PropertyTable"]["Properties"][:max_results]
|
||||
|
||||
for prop in properties_list:
|
||||
compound = CompoundData(
|
||||
cid=prop.get("CID"),
|
||||
name=prop.get("Title"),
|
||||
molecular_formula=prop.get("MolecularFormula"),
|
||||
molecular_weight=prop.get("MolecularWeight"),
|
||||
smiles=prop.get("CanonicalSMILES"),
|
||||
inchi=prop.get("InChI"),
|
||||
)
|
||||
compounds.append(compound)
|
||||
|
||||
# Format results for LLM
|
||||
if compounds:
|
||||
result_lines = [f"Found {len(compounds)} compound(s) for query '{query}':\n"]
|
||||
|
||||
for i, compound in enumerate(compounds, 1):
|
||||
result_lines.append(f"{i}. **{compound.name}** (CID: {compound.cid})")
|
||||
result_lines.append(f" - Formula: {compound.molecular_formula}")
|
||||
result_lines.append(f" - Molecular Weight: {compound.molecular_weight} g/mol")
|
||||
result_lines.append(f" - SMILES: {compound.smiles}")
|
||||
if compound.inchi:
|
||||
result_lines.append(
|
||||
f" - InChI: {compound.inchi[:100]}..."
|
||||
if len(compound.inchi) > 100
|
||||
else f" - InChI: {compound.inchi}"
|
||||
)
|
||||
result_lines.append("")
|
||||
|
||||
message = "\n".join(result_lines)
|
||||
else:
|
||||
message = f"No compounds found for query '{query}' using search type '{search_type}'"
|
||||
|
||||
# Prepare metadata
|
||||
metadata = PubChemMetadata(
|
||||
query_type=search_type,
|
||||
query_value=query,
|
||||
api_endpoint=url,
|
||||
response_time=response_time,
|
||||
total_results=len(compounds),
|
||||
timestamp=datetime.now().isoformat(),
|
||||
)
|
||||
|
||||
self._color_log(f"Found {len(compounds)} compounds ({response_time:.2f}s)", Color.green)
|
||||
|
||||
return ActionResponse(success=True, message=message, metadata=metadata.model_dump())
|
||||
|
||||
except ValueError as e:
|
||||
self.logger.error(f"Invalid input: {str(e)}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}",
|
||||
metadata={"error_type": "invalid_input", "error_message": str(e)},
|
||||
)
|
||||
except requests.RequestException as e:
|
||||
self.logger.error(f"PubChem API error: {str(e)}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"PubChem API error: {str(e)}",
|
||||
metadata={"error_type": "api_error", "error_message": str(e)},
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Search failed: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Search failed: {str(e)}",
|
||||
metadata={"error_type": "general_error", "error_message": str(e)},
|
||||
)
|
||||
|
||||
def mcp_get_compound_synonyms(
|
||||
self,
|
||||
cid: int = Field(description="PubChem Compound ID (CID)"),
|
||||
max_synonyms: int = Field(default=20, description="Maximum number of synonyms to return (1-100)", ge=1, le=100),
|
||||
) -> ActionResponse:
|
||||
"""Retrieve synonyms and alternative names for a PubChem compound.
|
||||
|
||||
Args:
|
||||
cid: PubChem Compound ID
|
||||
max_synonyms: Maximum number of synonyms to return
|
||||
|
||||
Returns:
|
||||
ActionResponse with compound synonyms and metadata
|
||||
"""
|
||||
try:
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(cid, FieldInfo):
|
||||
cid = cid.default
|
||||
if isinstance(max_synonyms, FieldInfo):
|
||||
max_synonyms = max_synonyms.default
|
||||
|
||||
if not cid or cid <= 0:
|
||||
raise ValueError("Valid PubChem CID is required")
|
||||
|
||||
self._color_log(f"Retrieving synonyms for CID: {cid}", Color.cyan)
|
||||
|
||||
# Build API URL for synonyms
|
||||
url = f"{self.base_url}/compound/cid/{cid}/synonyms/JSON"
|
||||
|
||||
# Make API request
|
||||
data, response_time = self._make_request(url)
|
||||
|
||||
# Parse synonyms
|
||||
synonyms = []
|
||||
if data and "InformationList" in data and "Information" in data["InformationList"]:
|
||||
info_list = data["InformationList"]["Information"]
|
||||
if info_list and "Synonym" in info_list[0]:
|
||||
synonyms = info_list[0]["Synonym"][:max_synonyms]
|
||||
|
||||
# Format results for LLM
|
||||
if synonyms:
|
||||
result_lines = [f"Found {len(synonyms)} synonym(s) for CID {cid}:\n"]
|
||||
|
||||
for i, synonym in enumerate(synonyms, 1):
|
||||
result_lines.append(f"{i}. {synonym}")
|
||||
|
||||
message = "\n".join(result_lines)
|
||||
else:
|
||||
message = f"No synonyms found for CID {cid}"
|
||||
|
||||
# Prepare metadata
|
||||
metadata = PubChemMetadata(
|
||||
query_type="synonyms",
|
||||
query_value=str(cid),
|
||||
api_endpoint=url,
|
||||
response_time=response_time,
|
||||
total_results=len(synonyms),
|
||||
timestamp=datetime.now().isoformat(),
|
||||
)
|
||||
|
||||
self._color_log(f"Retrieved {len(synonyms)} synonyms ({response_time:.2f}s)", Color.green)
|
||||
|
||||
return ActionResponse(success=True, message=message, metadata=metadata.model_dump())
|
||||
|
||||
except ValueError as e:
|
||||
self.logger.error(f"Invalid input: {str(e)}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}",
|
||||
metadata={"error_type": "invalid_input", "error_message": str(e)},
|
||||
)
|
||||
except requests.RequestException as e:
|
||||
self.logger.error(f"PubChem API error: {str(e)}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"PubChem API error: {str(e)}",
|
||||
metadata={"error_type": "api_error", "error_message": str(e)},
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Synonym retrieval failed: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Synonym retrieval failed: {str(e)}",
|
||||
metadata={"error_type": "general_error", "error_message": str(e)},
|
||||
)
|
||||
|
||||
def mcp_get_compound_properties(
|
||||
self,
|
||||
cid: int = Field(description="PubChem Compound ID (CID)"),
|
||||
properties: list[str] = Field(
|
||||
default=[
|
||||
"MolecularWeight",
|
||||
"MolecularFormula",
|
||||
"CanonicalSMILES",
|
||||
"InChI",
|
||||
"XLogP",
|
||||
"TPSA",
|
||||
"HBondDonorCount",
|
||||
"HBondAcceptorCount",
|
||||
],
|
||||
description="List of properties to retrieve (e.g., MolecularWeight, XLogP, TPSA)",
|
||||
),
|
||||
) -> ActionResponse:
|
||||
"""Retrieve detailed chemical properties for a PubChem compound.
|
||||
|
||||
Common properties include:
|
||||
- MolecularWeight: Molecular weight in g/mol
|
||||
- MolecularFormula: Chemical formula
|
||||
- CanonicalSMILES: SMILES notation
|
||||
- InChI: InChI identifier
|
||||
- XLogP: Partition coefficient
|
||||
- TPSA: Topological polar surface area
|
||||
- HBondDonorCount: Hydrogen bond donor count
|
||||
- HBondAcceptorCount: Hydrogen bond acceptor count
|
||||
|
||||
Args:
|
||||
cid: PubChem Compound ID
|
||||
properties: List of property names to retrieve
|
||||
|
||||
Returns:
|
||||
ActionResponse with compound properties and metadata
|
||||
"""
|
||||
try:
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(cid, FieldInfo):
|
||||
cid = cid.default
|
||||
if isinstance(properties, FieldInfo):
|
||||
properties = properties.default
|
||||
|
||||
if not cid or cid <= 0:
|
||||
raise ValueError("Valid PubChem CID is required")
|
||||
|
||||
if not properties:
|
||||
properties = ["MolecularWeight", "MolecularFormula", "CanonicalSMILES"]
|
||||
|
||||
self._color_log(f"Retrieving properties for CID: {cid}", Color.cyan)
|
||||
|
||||
# Build API URL for properties
|
||||
props_str = ",".join(properties)
|
||||
url = f"{self.base_url}/compound/cid/{cid}/property/{props_str}/JSON"
|
||||
|
||||
# Make API request
|
||||
data, response_time = self._make_request(url)
|
||||
|
||||
# Parse properties
|
||||
compound_props = {}
|
||||
if data and "PropertyTable" in data and "Properties" in data["PropertyTable"]:
|
||||
props_data = data["PropertyTable"]["Properties"][0]
|
||||
compound_props = {k: v for k, v in props_data.items() if k != "CID"}
|
||||
|
||||
# Format results for LLM
|
||||
if compound_props:
|
||||
result_lines = [f"Properties for PubChem CID {cid}:\n"]
|
||||
|
||||
for prop_name, prop_value in compound_props.items():
|
||||
if prop_name == "InChI" and isinstance(prop_value, str) and len(prop_value) > 100:
|
||||
result_lines.append(f"**{prop_name}**: {prop_value[:100]}...")
|
||||
else:
|
||||
result_lines.append(f"**{prop_name}**: {prop_value}")
|
||||
|
||||
message = "\n".join(result_lines)
|
||||
else:
|
||||
message = f"No properties found for CID {cid}"
|
||||
|
||||
# Prepare metadata
|
||||
metadata = PubChemMetadata(
|
||||
query_type="properties",
|
||||
query_value=str(cid),
|
||||
api_endpoint=url,
|
||||
response_time=response_time,
|
||||
total_results=len(compound_props),
|
||||
timestamp=datetime.now().isoformat(),
|
||||
)
|
||||
|
||||
self._color_log(f"Retrieved {len(compound_props)} properties ({response_time:.2f}s)", Color.green)
|
||||
|
||||
return ActionResponse(success=True, message=message, metadata=metadata.model_dump())
|
||||
|
||||
except ValueError as e:
|
||||
self.logger.error(f"Invalid input: {str(e)}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}",
|
||||
metadata={"error_type": "invalid_input", "error_message": str(e)},
|
||||
)
|
||||
except requests.RequestException as e:
|
||||
self.logger.error(f"PubChem API error: {str(e)}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"PubChem API error: {str(e)}",
|
||||
metadata={"error_type": "api_error", "error_message": str(e)},
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Property retrieval failed: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Property retrieval failed: {str(e)}",
|
||||
metadata={"error_type": "general_error", "error_message": str(e)},
|
||||
)
|
||||
|
||||
def mcp_search_similar_compounds(
|
||||
self,
|
||||
cid: int = Field(description="PubChem Compound ID to find similar compounds for"),
|
||||
similarity_threshold: float = Field(
|
||||
default=0.9, description="Similarity threshold (0.0-1.0, higher = more similar)", ge=0.0, le=1.0
|
||||
),
|
||||
max_results: int = Field(
|
||||
default=10, description="Maximum number of similar compounds to return (1-50)", ge=1, le=50
|
||||
),
|
||||
) -> ActionResponse:
|
||||
"""Find structurally similar compounds using PubChem's similarity search.
|
||||
|
||||
Uses Tanimoto similarity coefficient for 2D structure comparison.
|
||||
|
||||
Args:
|
||||
cid: Reference compound CID for similarity search
|
||||
similarity_threshold: Minimum similarity score (0.0-1.0)
|
||||
max_results: Maximum number of similar compounds to return
|
||||
|
||||
Returns:
|
||||
ActionResponse with similar compounds and metadata
|
||||
"""
|
||||
try:
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(cid, FieldInfo):
|
||||
cid = cid.default
|
||||
if isinstance(similarity_threshold, FieldInfo):
|
||||
similarity_threshold = similarity_threshold.default
|
||||
if isinstance(max_results, FieldInfo):
|
||||
max_results = max_results.default
|
||||
|
||||
if not cid or cid <= 0:
|
||||
raise ValueError("Valid PubChem CID is required")
|
||||
|
||||
self._color_log(f"Searching for compounds similar to CID: {cid}", Color.cyan)
|
||||
|
||||
# Build API URL for similarity search
|
||||
threshold_percent = int(similarity_threshold * 100)
|
||||
url = f"{self.base_url}/compound/fastsimilarity_2d/cid/{cid}/property/Title,MolecularFormula,MolecularWeight/JSON"
|
||||
params = {"Threshold": threshold_percent, "MaxRecords": max_results}
|
||||
|
||||
# Make API request
|
||||
data, response_time = self._make_request(url, params)
|
||||
|
||||
# Parse similar compounds
|
||||
similar_compounds: list[CompoundData] = []
|
||||
if data and "PropertyTable" in data and "Properties" in data["PropertyTable"]:
|
||||
properties_list = data["PropertyTable"]["Properties"]
|
||||
|
||||
for prop in properties_list:
|
||||
if prop.get("CID") != cid: # Exclude the query compound itself
|
||||
compound = CompoundData(
|
||||
cid=prop.get("CID"),
|
||||
name=prop.get("Title"),
|
||||
molecular_formula=prop.get("MolecularFormula"),
|
||||
molecular_weight=prop.get("MolecularWeight"),
|
||||
)
|
||||
similar_compounds.append(compound)
|
||||
|
||||
# Format results for LLM
|
||||
if similar_compounds:
|
||||
result_lines = [
|
||||
f"Found {len(similar_compounds)} compound(s) similar to CID {cid} (threshold: {similarity_threshold}):\n"
|
||||
]
|
||||
|
||||
for i, compound in enumerate(similar_compounds, 1):
|
||||
result_lines.append(f"{i}. **{compound.name}** (CID: {compound.cid})")
|
||||
result_lines.append(f" - Formula: {compound.molecular_formula}")
|
||||
result_lines.append(f" - Molecular Weight: {compound.molecular_weight} g/mol")
|
||||
result_lines.append("")
|
||||
|
||||
message = "\n".join(result_lines)
|
||||
else:
|
||||
message = f"No similar compounds found for CID {cid} with similarity threshold {similarity_threshold}"
|
||||
|
||||
# Prepare metadata
|
||||
metadata = PubChemMetadata(
|
||||
query_type="similarity",
|
||||
query_value=str(cid),
|
||||
api_endpoint=url,
|
||||
response_time=response_time,
|
||||
total_results=len(similar_compounds),
|
||||
timestamp=datetime.now().isoformat(),
|
||||
)
|
||||
|
||||
self._color_log(f"Found {len(similar_compounds)} similar compounds ({response_time:.2f}s)", Color.green)
|
||||
|
||||
return ActionResponse(success=True, message=message, metadata=metadata.model_dump())
|
||||
|
||||
except ValueError as e:
|
||||
self.logger.error(f"Invalid input: {str(e)}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}",
|
||||
metadata={"error_type": "invalid_input", "error_message": str(e)},
|
||||
)
|
||||
except requests.RequestException as e:
|
||||
self.logger.error(f"PubChem API error: {str(e)}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"PubChem API error: {str(e)}",
|
||||
metadata={"error_type": "api_error", "error_message": str(e)},
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Similarity search failed: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Similarity search failed: {str(e)}",
|
||||
metadata={"error_type": "general_error", "error_message": str(e)},
|
||||
)
|
||||
|
||||
def mcp_get_pubchem_capabilities(self) -> ActionResponse:
|
||||
"""Get information about the PubChem service capabilities.
|
||||
|
||||
Returns:
|
||||
ActionResponse with service capabilities and configuration
|
||||
"""
|
||||
capabilities = {
|
||||
"Compound Search": "Search by name, CID, SMILES, InChI, or molecular formula",
|
||||
"Property Retrieval": "Get molecular weight, formula, SMILES, physicochemical properties",
|
||||
"Synonym Lookup": "Retrieve alternative names and identifiers for compounds",
|
||||
"Similarity Search": "Find structurally similar compounds using Tanimoto similarity",
|
||||
"Rate Limiting": "Compliant with PubChem's 5 requests/second limit",
|
||||
"Data Formats": "JSON responses with structured compound data",
|
||||
}
|
||||
|
||||
capability_list = "\n".join(
|
||||
[f"**{capability}**: {description}" for capability, description in capabilities.items()]
|
||||
)
|
||||
|
||||
metadata = {
|
||||
"base_url": self.base_url,
|
||||
"rate_limit": "5 requests/second",
|
||||
"timeout": f"{self.timeout} seconds",
|
||||
"supported_capabilities": list(capabilities.keys()),
|
||||
"total_capabilities": len(capabilities),
|
||||
"search_types": ["name", "cid", "smiles", "inchi", "formula"],
|
||||
"data_source": "PubChem (NCBI)",
|
||||
}
|
||||
|
||||
return ActionResponse(
|
||||
success=True, message=f"PubChem MCP Service Capabilities:\n\n{capability_list}", metadata=metadata
|
||||
)
|
||||
|
||||
|
||||
# Example usage and entry point
|
||||
if __name__ == "__main__":
|
||||
load_dotenv()
|
||||
|
||||
# Default arguments for testing
|
||||
args = ActionArguments(name="pubchem_service", transport="stdio", workspace="~")
|
||||
|
||||
# Initialize and run the PubChem service
|
||||
try:
|
||||
service = PubChemCollection(args)
|
||||
service.run()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}: {traceback.format_exc()}")
|
||||
@@ -0,0 +1,387 @@
|
||||
"""
|
||||
Search MCP Server
|
||||
|
||||
This module provides MCP server functionality for performing web searches using various search engines.
|
||||
It supports structured queries and returns LLM-friendly formatted search results.
|
||||
|
||||
Key features:
|
||||
- Perform web searches using Google Custom Search API
|
||||
- Filter and format search results for LLM consumption
|
||||
- Validate and process search queries with metadata tracking
|
||||
|
||||
Main functions:
|
||||
- mcp_search_google: Searches the web using Google Custom Search API
|
||||
- mcp_get_search_capabilities: Returns information about search service capabilities
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
|
||||
import requests
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from aworld.logs.util import Color
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
|
||||
|
||||
class SearchResult(BaseModel):
|
||||
"""Individual search result with structured data."""
|
||||
|
||||
id: str
|
||||
title: str
|
||||
url: str
|
||||
snippet: str
|
||||
source: str
|
||||
display_link: str | None = None
|
||||
formatted_url: str | None = None
|
||||
|
||||
|
||||
class SearchMetadata(BaseModel):
|
||||
"""Metadata for search operation results."""
|
||||
|
||||
query: str
|
||||
search_engine: str
|
||||
total_results: int
|
||||
search_time: float | None = None
|
||||
language: str = "en"
|
||||
country: str = "us"
|
||||
safe_search: bool = True
|
||||
error_type: str | None = None
|
||||
api_quota_used: bool = False
|
||||
|
||||
|
||||
class SearchCollection(ActionCollection):
|
||||
"""MCP service for web search operations using various search engines.
|
||||
|
||||
Provides comprehensive web search capabilities including:
|
||||
- Google Custom Search API integration
|
||||
- LLM-friendly result formatting
|
||||
- Search result filtering and validation
|
||||
- Metadata tracking for search operations
|
||||
- Error handling and quota management
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
|
||||
# Load environment variables
|
||||
load_dotenv()
|
||||
|
||||
# Validate required API credentials
|
||||
self.google_api_key = os.getenv("GOOGLE_API_KEY")
|
||||
self.google_cse_id = os.getenv("GOOGLE_CSE_ID")
|
||||
|
||||
# Log initialization status
|
||||
self._color_log("Search service initialized", Color.green, "debug")
|
||||
|
||||
if self.google_api_key and self.google_cse_id:
|
||||
self._color_log("Google Search API credentials found", Color.blue, "debug")
|
||||
else:
|
||||
self._color_log("Google Search API credentials missing - some features may be unavailable", Color.yellow)
|
||||
|
||||
def _format_search_results_for_llm(self, results: list[SearchResult], query: str) -> str:
|
||||
"""Format search results to be LLM-friendly.
|
||||
|
||||
Args:
|
||||
results: List of search results
|
||||
query: Original search query
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if not results:
|
||||
return f"No search results found for query: '{query}'"
|
||||
|
||||
formatted_parts = [f"# Search Results for: '{query}'", f"Found {len(results)} results:\n"]
|
||||
|
||||
for i, result in enumerate(results, 1):
|
||||
result_section = [
|
||||
f"## Result {i}: {result.title}",
|
||||
f"**URL:** {result.url}",
|
||||
f"**Source:** {result.source}",
|
||||
]
|
||||
|
||||
if result.display_link:
|
||||
result_section.append(f"**Domain:** {result.display_link}")
|
||||
|
||||
result_section.append(f"**Summary:** {result.snippet}")
|
||||
result_section.append("") # Empty line for spacing
|
||||
|
||||
formatted_parts.append("\n".join(result_section))
|
||||
|
||||
return "\n".join(formatted_parts)
|
||||
|
||||
def _validate_search_parameters(self, query: str, num_results: int) -> tuple[str, int]:
|
||||
"""Validate and normalize search parameters.
|
||||
|
||||
Args:
|
||||
query: Search query string
|
||||
num_results: Number of results requested
|
||||
|
||||
Returns:
|
||||
Tuple of (validated_query, validated_num_results)
|
||||
|
||||
Raises:
|
||||
ValueError: If parameters are invalid
|
||||
"""
|
||||
if not query or not query.strip():
|
||||
raise ValueError("Search query cannot be empty")
|
||||
|
||||
# Normalize query
|
||||
validated_query = query.strip()
|
||||
|
||||
# Validate and clamp num_results
|
||||
validated_num_results = max(1, min(num_results, 10)) # Google CSE limit is 10
|
||||
|
||||
return validated_query, validated_num_results
|
||||
|
||||
def mcp_search_google(
|
||||
self,
|
||||
query: str = Field(description="The search query string to search for"),
|
||||
num_results: int = Field(default=5, description="Number of search results to return (1-10, default: 5)"),
|
||||
safe_search: bool = Field(default=True, description="Whether to enable safe search filtering"),
|
||||
language: str = Field(default="en", description="Language code for search results (e.g., 'en', 'es', 'fr')"),
|
||||
country: str = Field(default="us", description="Country code for search results (e.g., 'us', 'uk', 'ca')"),
|
||||
output_format: str = Field(default="json", description="Output format: 'markdown', 'json', or 'text'"),
|
||||
) -> ActionResponse:
|
||||
"""Search the web using Google Custom Search API.
|
||||
|
||||
This tool provides comprehensive web search capabilities with:
|
||||
- Google Custom Search API integration
|
||||
- Configurable result count and filtering
|
||||
- Safe search and localization options
|
||||
- LLM-optimized result formatting
|
||||
- Detailed metadata tracking
|
||||
|
||||
Args:
|
||||
query: The search query string
|
||||
num_results: Number of results to return (1-10)
|
||||
safe_search: Enable safe search filtering
|
||||
language: Language code for results
|
||||
country: Country code for results
|
||||
output_format: Format for the response
|
||||
|
||||
Returns:
|
||||
ActionResponse with formatted search results and metadata
|
||||
"""
|
||||
if isinstance(query, FieldInfo):
|
||||
query = query.default
|
||||
if isinstance(num_results, FieldInfo):
|
||||
num_results = num_results.default
|
||||
if isinstance(safe_search, FieldInfo):
|
||||
safe_search = safe_search.default
|
||||
if isinstance(language, FieldInfo):
|
||||
language = language.default
|
||||
if isinstance(country, FieldInfo):
|
||||
country = country.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
try:
|
||||
# Validate API credentials
|
||||
if not self.google_api_key or not self.google_cse_id:
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=(
|
||||
"Google Search API credentials not configured. "
|
||||
"Please set GOOGLE_API_KEY and GOOGLE_CSE_ID environment variables."
|
||||
),
|
||||
metadata={"error_type": "missing_credentials"},
|
||||
)
|
||||
|
||||
# Validate parameters
|
||||
validated_query, validated_num_results = self._validate_search_parameters(query, num_results)
|
||||
|
||||
self._color_log(f"🔍 Searching Google for: '{validated_query}'", Color.cyan)
|
||||
|
||||
# Prepare API request
|
||||
start_time = time.time()
|
||||
|
||||
url = "https://www.googleapis.com/customsearch/v1"
|
||||
params = {
|
||||
"key": self.google_api_key,
|
||||
"cx": self.google_cse_id,
|
||||
"q": validated_query,
|
||||
"num": validated_num_results,
|
||||
"safe": "active" if safe_search else "off",
|
||||
"hl": language,
|
||||
"gl": country,
|
||||
}
|
||||
|
||||
# Make API request
|
||||
response = requests.get(url, params=params, timeout=10)
|
||||
response.raise_for_status()
|
||||
|
||||
search_time = time.time() - start_time
|
||||
data = response.json()
|
||||
|
||||
# Parse search results
|
||||
search_results = []
|
||||
if "items" in data:
|
||||
for i, item in enumerate(data["items"]):
|
||||
result = SearchResult(
|
||||
id=f"google-{i}",
|
||||
title=item.get("title", ""),
|
||||
url=item.get("link", ""),
|
||||
snippet=item.get("snippet", ""),
|
||||
source="google",
|
||||
display_link=item.get("displayLink", ""),
|
||||
formatted_url=item.get("formattedUrl", ""),
|
||||
)
|
||||
search_results.append(result)
|
||||
|
||||
# Format results based on requested format
|
||||
if "json" == "json":
|
||||
formatted_content = {
|
||||
"query": validated_query,
|
||||
"results": [result.model_dump() for result in search_results],
|
||||
"count": len(search_results),
|
||||
}
|
||||
|
||||
message_content = formatted_content
|
||||
elif output_format.lower() == "text":
|
||||
if search_results:
|
||||
result_lines = []
|
||||
for i, result in enumerate(search_results, 1):
|
||||
result_lines.append(f"{i}. {result.title}")
|
||||
result_lines.append(f" URL: {result.url}")
|
||||
result_lines.append(f" Summary: {result.snippet}")
|
||||
result_lines.append("") # Empty line
|
||||
message_content = "\n".join(result_lines)
|
||||
else:
|
||||
message_content = f"No results found for: {validated_query}"
|
||||
else: # markdown (default)
|
||||
message_content = self._format_search_results_for_llm(search_results, validated_query)
|
||||
|
||||
# Prepare metadata
|
||||
metadata = SearchMetadata(
|
||||
query=validated_query,
|
||||
search_engine="google",
|
||||
total_results=len(search_results),
|
||||
search_time=search_time,
|
||||
language=language,
|
||||
country=country,
|
||||
safe_search=safe_search,
|
||||
api_quota_used=True,
|
||||
)
|
||||
|
||||
self._color_log(f"✅ Found {len(search_results)} results in {search_time:.2f}s", Color.green)
|
||||
|
||||
return ActionResponse(success=True, message=message_content, metadata=metadata.model_dump())
|
||||
|
||||
except requests.exceptions.RequestException as e:
|
||||
error_msg = f"Google Search API request failed: {str(e)}"
|
||||
self.logger.error(f"Search API error: {traceback.format_exc()}")
|
||||
|
||||
metadata = SearchMetadata(
|
||||
query=query, search_engine="google", total_results=0, error_type="api_request_failed"
|
||||
)
|
||||
|
||||
self._color_log(f"❌ {error_msg}", Color.red)
|
||||
|
||||
return ActionResponse(success=False, message=error_msg, metadata=metadata.model_dump())
|
||||
|
||||
except ValueError as e:
|
||||
error_msg = f"Invalid search parameters: {str(e)}"
|
||||
|
||||
metadata = SearchMetadata(
|
||||
query=query, search_engine="google", total_results=0, error_type="invalid_parameters"
|
||||
)
|
||||
|
||||
self._color_log(f"❌ {error_msg}", Color.red)
|
||||
|
||||
return ActionResponse(success=False, message=error_msg, metadata=metadata.model_dump())
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Search operation failed: {str(e)}"
|
||||
error_trace = traceback.format_exc()
|
||||
|
||||
self.logger.error(f"Unexpected search error: {error_trace}")
|
||||
|
||||
metadata = SearchMetadata(
|
||||
query=query, search_engine="google", total_results=0, error_type="unexpected_error"
|
||||
)
|
||||
|
||||
self._color_log(f"❌ {error_msg}", Color.red)
|
||||
|
||||
return ActionResponse(
|
||||
success=False, message=f"{error_msg}\n\nError details: {error_trace}", metadata=metadata.model_dump()
|
||||
)
|
||||
|
||||
def mcp_get_search_capabilities(self) -> ActionResponse:
|
||||
"""Get information about search service capabilities and configuration.
|
||||
|
||||
Returns:
|
||||
ActionResponse with search service capabilities and current configuration
|
||||
"""
|
||||
capabilities = {
|
||||
"search_engines": ["Google Custom Search API"],
|
||||
"supported_features": [
|
||||
"Web search with customizable result count",
|
||||
"Safe search filtering",
|
||||
"Language and country localization",
|
||||
"Multiple output formats (markdown, json, text)",
|
||||
"LLM-optimized result formatting",
|
||||
"Detailed metadata tracking",
|
||||
],
|
||||
"supported_formats": ["markdown", "json", "text"],
|
||||
"configuration": {
|
||||
"google_api_configured": bool(self.google_api_key and self.google_cse_id),
|
||||
"max_results_per_query": 10,
|
||||
"default_language": "en",
|
||||
"default_country": "us",
|
||||
"safe_search_default": True,
|
||||
},
|
||||
"limitations": [
|
||||
"Google CSE has daily quota limits",
|
||||
"Maximum 10 results per query",
|
||||
"Requires valid API credentials",
|
||||
],
|
||||
}
|
||||
|
||||
formatted_info = f"""# Search Service Capabilities
|
||||
|
||||
## Available Search Engines
|
||||
{chr(10).join(f"- {engine}" for engine in capabilities["search_engines"])}
|
||||
|
||||
## Features
|
||||
{chr(10).join(f"- {feature}" for feature in capabilities["supported_features"])}
|
||||
|
||||
## Supported Output Formats
|
||||
{chr(10).join(f"- {fmt}" for fmt in capabilities["supported_formats"])}
|
||||
|
||||
## Current Configuration
|
||||
- **Google API Configured:** {capabilities["configuration"]["google_api_configured"]}
|
||||
- **Max Results Per Query:** {capabilities["configuration"]["max_results_per_query"]}
|
||||
- **Default Language:** {capabilities["configuration"]["default_language"]}
|
||||
- **Default Country:** {capabilities["configuration"]["default_country"]}
|
||||
- **Safe Search Default:** {capabilities["configuration"]["safe_search_default"]}
|
||||
|
||||
## Limitations
|
||||
{chr(10).join(f"- {limitation}" for limitation in capabilities["limitations"])}
|
||||
"""
|
||||
|
||||
return ActionResponse(success=True, message=formatted_info, metadata=capabilities)
|
||||
|
||||
|
||||
# Example usage and entry point
|
||||
if __name__ == "__main__":
|
||||
load_dotenv()
|
||||
|
||||
# Default arguments for testing
|
||||
args = ActionArguments(
|
||||
name="search_service",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
|
||||
)
|
||||
|
||||
# Initialize and run the search service
|
||||
try:
|
||||
service = SearchCollection(args)
|
||||
service.run()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}: {traceback.format_exc()}")
|
||||
@@ -0,0 +1,512 @@
|
||||
"""
|
||||
Terminal MCP Server
|
||||
|
||||
This module provides MCP server functionality for executing terminal commands safely.
|
||||
It supports command execution with timeout controls and returns LLM-friendly formatted results.
|
||||
|
||||
Key features:
|
||||
- Execute terminal commands with configurable timeouts
|
||||
- Cross-platform command execution support
|
||||
- Command history tracking and retrieval
|
||||
- Safety checks for dangerous commands
|
||||
- LLM-optimized output formatting
|
||||
|
||||
Main functions:
|
||||
- mcp_execute_command: Execute terminal commands with safety checks
|
||||
- mcp_get_command_history: Retrieve recent command execution history
|
||||
- mcp_get_terminal_capabilities: Get terminal service capabilities
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from aworld.logs.util import Color
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
|
||||
# pylint: disable=C0301
|
||||
|
||||
|
||||
class CommandResult(BaseModel):
|
||||
"""Individual command execution result with structured data."""
|
||||
|
||||
command: str
|
||||
success: bool
|
||||
stdout: str
|
||||
stderr: str
|
||||
return_code: int
|
||||
duration: str
|
||||
timestamp: str
|
||||
|
||||
|
||||
class TerminalMetadata(BaseModel):
|
||||
"""Metadata for terminal operation results."""
|
||||
|
||||
command: str
|
||||
platform: str
|
||||
working_directory: str
|
||||
timeout_seconds: int
|
||||
execution_time: float | None = None
|
||||
return_code: int | None = None
|
||||
safety_check_passed: bool = True
|
||||
error_type: str | None = None
|
||||
history_count: int | None = None
|
||||
|
||||
|
||||
class TerminalActionCollection(ActionCollection):
|
||||
"""MCP service for terminal command execution with safety controls.
|
||||
|
||||
Provides secure terminal command execution capabilities including:
|
||||
- Cross-platform command execution
|
||||
- Configurable timeout controls
|
||||
- Command history tracking
|
||||
- Safety checks for dangerous operations
|
||||
- LLM-friendly result formatting
|
||||
- Error handling and logging
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
|
||||
# Initialize command history
|
||||
self.command_history: list[dict] = []
|
||||
self.max_history_size = 50
|
||||
|
||||
# Define dangerous commands for safety
|
||||
self.dangerous_commands = [
|
||||
"rm -rf /",
|
||||
"mkfs",
|
||||
"dd if=",
|
||||
":(){ :|:& };:", # Unix
|
||||
"del /f /s /q",
|
||||
"format",
|
||||
"diskpart", # Windows
|
||||
"sudo rm",
|
||||
"sudo dd",
|
||||
"sudo mkfs", # Sudo variants
|
||||
]
|
||||
|
||||
# Get current platform info
|
||||
self.platform_info = {
|
||||
"system": platform.system(),
|
||||
"platform": platform.platform(),
|
||||
"architecture": platform.architecture()[0],
|
||||
}
|
||||
|
||||
self._color_log("Terminal service initialized", Color.green, "debug")
|
||||
self._color_log(f"Platform: {self.platform_info['system']}", Color.blue, "debug")
|
||||
|
||||
def _check_command_safety(self, command: str) -> tuple[bool, str | None]:
|
||||
"""Check if command is safe to execute.
|
||||
|
||||
Args:
|
||||
command: Command string to check
|
||||
|
||||
Returns:
|
||||
Tuple of (is_safe, reason_if_unsafe)
|
||||
"""
|
||||
command_lower = command.lower().strip()
|
||||
|
||||
for dangerous_cmd in self.dangerous_commands:
|
||||
if dangerous_cmd.lower() in command_lower:
|
||||
return False, f"Command contains dangerous pattern: {dangerous_cmd}"
|
||||
|
||||
return True, None
|
||||
|
||||
def _format_command_output(self, result: CommandResult, output_format: str = "markdown") -> str:
|
||||
"""Format command execution results for LLM consumption.
|
||||
|
||||
Args:
|
||||
result: Command execution result
|
||||
output_format: Format type ('markdown', 'json', 'text')
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if output_format == "json":
|
||||
return json.dumps(result.model_dump(), indent=2)
|
||||
|
||||
elif output_format == "text":
|
||||
output_parts = [
|
||||
f"Command: {result.command}",
|
||||
f"Status: {'SUCCESS' if result.success else 'FAILED'}",
|
||||
f"Duration: {result.duration}",
|
||||
f"Return Code: {result.return_code}",
|
||||
]
|
||||
|
||||
if result.stdout:
|
||||
output_parts.extend(["\nOutput:", result.stdout])
|
||||
|
||||
if result.stderr:
|
||||
output_parts.extend(["\nErrors/Warnings:", result.stderr])
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
else: # markdown (default)
|
||||
status_emoji = "✅" if result.success else "❌"
|
||||
|
||||
output_parts = [
|
||||
f"# Terminal Command Execution {status_emoji}",
|
||||
f"**Command:** `{result.command}`",
|
||||
f"**Status:** {'SUCCESS' if result.success else 'FAILED'}",
|
||||
f"**Duration:** {result.duration}",
|
||||
f"**Return Code:** {result.return_code}",
|
||||
f"**Timestamp:** {result.timestamp}",
|
||||
]
|
||||
|
||||
if result.stdout:
|
||||
output_parts.extend(["\n## Output", "```", result.stdout.strip(), "```"])
|
||||
|
||||
if result.stderr:
|
||||
output_parts.extend(["\n## Errors/Warnings", "```", result.stderr.strip(), "```"])
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
async def _execute_command_async(self, command: str, timeout: int) -> CommandResult:
|
||||
"""Execute command asynchronously with timeout.
|
||||
|
||||
Args:
|
||||
command: Command to execute
|
||||
timeout: Timeout in seconds
|
||||
|
||||
Returns:
|
||||
CommandResult with execution details
|
||||
"""
|
||||
start_time = datetime.now()
|
||||
|
||||
try:
|
||||
# Create appropriate subprocess for platform
|
||||
if self.platform_info["system"] == "Windows":
|
||||
process = await asyncio.create_subprocess_shell(
|
||||
command, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, shell=True
|
||||
)
|
||||
else:
|
||||
process = await asyncio.create_subprocess_shell(
|
||||
command,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
shell=True,
|
||||
executable="/bin/bash",
|
||||
)
|
||||
|
||||
try:
|
||||
stdout, stderr = await asyncio.wait_for(process.communicate(), timeout)
|
||||
stdout = stdout.decode("utf-8", errors="replace")
|
||||
stderr = stderr.decode("utf-8", errors="replace")
|
||||
return_code = process.returncode
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
try:
|
||||
process.kill()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
duration = str(datetime.now() - start_time)
|
||||
return CommandResult(
|
||||
command=command,
|
||||
success=False,
|
||||
stdout="",
|
||||
stderr=f"Command timed out after {timeout} seconds",
|
||||
return_code=-1,
|
||||
duration=duration,
|
||||
timestamp=start_time.isoformat(),
|
||||
)
|
||||
|
||||
duration = str(datetime.now() - start_time)
|
||||
result = CommandResult(
|
||||
command=command,
|
||||
success=return_code == 0,
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
return_code=return_code,
|
||||
duration=duration,
|
||||
timestamp=start_time.isoformat(),
|
||||
)
|
||||
|
||||
# Add to history
|
||||
self.command_history.append(
|
||||
{
|
||||
"timestamp": start_time.isoformat(),
|
||||
"command": command,
|
||||
"success": return_code == 0,
|
||||
"duration": duration,
|
||||
}
|
||||
)
|
||||
|
||||
# Maintain history size limit
|
||||
if len(self.command_history) > self.max_history_size:
|
||||
self.command_history.pop(0)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
duration = str(datetime.now() - start_time)
|
||||
return CommandResult(
|
||||
command=command,
|
||||
success=False,
|
||||
stdout="",
|
||||
stderr=f"Error executing command: {str(e)}",
|
||||
return_code=-1,
|
||||
duration=duration,
|
||||
timestamp=start_time.isoformat(),
|
||||
)
|
||||
|
||||
async def mcp_execute_command(
|
||||
self,
|
||||
command: str = Field(description="Terminal command to execute"),
|
||||
timeout: int = Field(default=30, description="Command timeout in seconds (default: 30)"),
|
||||
output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'"),
|
||||
) -> ActionResponse:
|
||||
"""Execute a terminal command with safety checks and timeout controls.
|
||||
|
||||
This tool provides secure command execution with:
|
||||
- Cross-platform compatibility (Windows, macOS, Linux)
|
||||
- Configurable timeout controls
|
||||
- Safety checks for dangerous commands
|
||||
- LLM-optimized result formatting
|
||||
- Command history tracking
|
||||
|
||||
Specialized Feature:
|
||||
- Execute Python code and output the result to stdout
|
||||
- Example (Directly execute simple Python code): `python -c "nums = [1, 2, 3, 4]\nsum_of_nums = sum(nums)\nprint(f'{sum_of_nums=}')"`
|
||||
- Example (Execute code from a file): `python my_script.py`
|
||||
|
||||
Args:
|
||||
command: The terminal command to execute
|
||||
timeout: Maximum execution time in seconds
|
||||
output_format: Format for the response output
|
||||
|
||||
Returns:
|
||||
ActionResponse with command execution results and metadata
|
||||
"""
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(command, FieldInfo):
|
||||
command = command.default
|
||||
if isinstance(timeout, FieldInfo):
|
||||
timeout = timeout.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
try:
|
||||
# Safety check
|
||||
is_safe, safety_reason = self._check_command_safety(command)
|
||||
if not is_safe:
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Command rejected for security reasons: {safety_reason}",
|
||||
metadata=TerminalMetadata(
|
||||
command=command,
|
||||
platform=self.platform_info["system"],
|
||||
working_directory=str(self.workspace),
|
||||
timeout_seconds=timeout,
|
||||
safety_check_passed=False,
|
||||
error_type="security_violation",
|
||||
).model_dump(),
|
||||
)
|
||||
|
||||
self._color_log(f"🔧 Executing command: {command}", Color.cyan)
|
||||
|
||||
# Execute command
|
||||
start_time = time.time()
|
||||
result = await self._execute_command_async(command, timeout)
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
# Format output
|
||||
formatted_output = self._format_command_output(result, output_format)
|
||||
|
||||
# Create metadata
|
||||
metadata = TerminalMetadata(
|
||||
command=command,
|
||||
platform=self.platform_info["system"],
|
||||
working_directory=str(self.workspace),
|
||||
timeout_seconds=timeout,
|
||||
execution_time=execution_time,
|
||||
return_code=result.return_code,
|
||||
safety_check_passed=True,
|
||||
)
|
||||
|
||||
if result.success:
|
||||
self._color_log("✅ Command completed successfully", Color.green)
|
||||
else:
|
||||
self._color_log(f"❌ Command failed with return code {result.return_code}", Color.red)
|
||||
metadata.error_type = "execution_failure"
|
||||
|
||||
return ActionResponse(success=result.success, message=formatted_output, metadata=metadata.model_dump())
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to execute command: {str(e)}"
|
||||
self.logger.error(f"Command execution error: {traceback.format_exc()}")
|
||||
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=error_msg,
|
||||
metadata=TerminalMetadata(
|
||||
command=command,
|
||||
platform=self.platform_info["system"],
|
||||
working_directory=str(self.workspace),
|
||||
timeout_seconds=timeout,
|
||||
safety_check_passed=True,
|
||||
error_type="internal_error",
|
||||
).model_dump(),
|
||||
)
|
||||
|
||||
def mcp_get_command_history(
|
||||
self,
|
||||
count: int = Field(default=10, description="Number of recent commands to return (default: 10)"),
|
||||
output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'"),
|
||||
) -> ActionResponse:
|
||||
"""Retrieve recent command execution history.
|
||||
|
||||
Args:
|
||||
count: Number of recent commands to return
|
||||
output_format: Format for the response output
|
||||
|
||||
Returns:
|
||||
ActionResponse with command history and metadata
|
||||
"""
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(count, FieldInfo):
|
||||
count = count.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
try:
|
||||
# Get recent history
|
||||
recent_history = self.command_history[-count:] if self.command_history else []
|
||||
|
||||
if not recent_history:
|
||||
message = "No command history available."
|
||||
else:
|
||||
if output_format == "json":
|
||||
message = json.dumps(recent_history, indent=2)
|
||||
elif output_format == "text":
|
||||
history_lines = []
|
||||
for i, entry in enumerate(recent_history, 1):
|
||||
status = "SUCCESS" if entry["success"] else "FAILED"
|
||||
history_lines.append(
|
||||
f"{i}. [{entry['timestamp']}] {entry['command']} - {status}"
|
||||
f" ({entry.get('duration', 'N/A')})"
|
||||
)
|
||||
message = "\n".join(history_lines)
|
||||
else: # markdown
|
||||
history_lines = ["# Command History", f"Showing {len(recent_history)} recent commands:\n"]
|
||||
|
||||
for i, entry in enumerate(recent_history, 1):
|
||||
status_emoji = "✅" if entry["success"] else "❌"
|
||||
history_lines.extend(
|
||||
[
|
||||
f"## {i}. {status_emoji} `{entry['command']}`",
|
||||
f"- **Timestamp:** {entry['timestamp']}",
|
||||
f"- **Duration:** {entry.get('duration', 'N/A')}",
|
||||
"",
|
||||
]
|
||||
)
|
||||
|
||||
message = "\n".join(history_lines)
|
||||
|
||||
metadata = TerminalMetadata(
|
||||
command="get_command_history",
|
||||
platform=self.platform_info["system"],
|
||||
working_directory=str(self.workspace),
|
||||
timeout_seconds=0,
|
||||
history_count=len(recent_history),
|
||||
)
|
||||
|
||||
return ActionResponse(success=True, message=message, metadata=metadata.model_dump())
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to retrieve command history: {str(e)}"
|
||||
self.logger.error(f"History retrieval error: {traceback.format_exc()}")
|
||||
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=error_msg,
|
||||
metadata=TerminalMetadata(
|
||||
command="get_command_history",
|
||||
platform=self.platform_info["system"],
|
||||
working_directory=str(self.workspace),
|
||||
timeout_seconds=0,
|
||||
error_type="internal_error",
|
||||
).model_dump(),
|
||||
)
|
||||
|
||||
def mcp_get_terminal_capabilities(self) -> ActionResponse:
|
||||
"""Get information about terminal service capabilities and configuration.
|
||||
|
||||
Returns:
|
||||
ActionResponse with terminal service capabilities and current configuration
|
||||
"""
|
||||
capabilities = {
|
||||
"platform_info": self.platform_info,
|
||||
"supported_features": [
|
||||
"Cross-platform command execution",
|
||||
"Configurable timeout controls",
|
||||
"Command history tracking",
|
||||
"Safety checks for dangerous commands",
|
||||
"Multiple output formats (markdown, json, text)",
|
||||
"LLM-optimized result formatting",
|
||||
"Async command execution",
|
||||
],
|
||||
"supported_formats": ["markdown", "json", "text"],
|
||||
"configuration": {
|
||||
"max_history_size": self.max_history_size,
|
||||
"current_history_count": len(self.command_history),
|
||||
"working_directory": str(self.workspace),
|
||||
"dangerous_commands_count": len(self.dangerous_commands),
|
||||
},
|
||||
"safety_features": [
|
||||
"Dangerous command detection",
|
||||
"Timeout controls",
|
||||
"Error handling and logging",
|
||||
"Command validation",
|
||||
],
|
||||
}
|
||||
|
||||
formatted_info = f"""# Terminal Service Capabilities
|
||||
|
||||
## Platform Information
|
||||
- **System:** {self.platform_info["system"]}
|
||||
- **Platform:** {self.platform_info["platform"]}
|
||||
- **Architecture:** {self.platform_info["architecture"]}
|
||||
|
||||
## Features
|
||||
{chr(10).join(f"- {feature}" for feature in capabilities["supported_features"])}
|
||||
|
||||
## Supported Output Formats
|
||||
{chr(10).join(f"- {fmt}" for fmt in capabilities["supported_formats"])}
|
||||
|
||||
## Current Configuration
|
||||
- **Max History Size:** {capabilities["configuration"]["max_history_size"]}
|
||||
- **Current History Count:** {capabilities["configuration"]["current_history_count"]}
|
||||
- **Working Directory:** {capabilities["configuration"]["working_directory"]}
|
||||
- **Dangerous Commands Monitored:** {capabilities["configuration"]["dangerous_commands_count"]}
|
||||
|
||||
## Safety Features
|
||||
{chr(10).join(f"- {feature}" for feature in capabilities["safety_features"])}
|
||||
"""
|
||||
|
||||
return ActionResponse(success=True, message=formatted_info, metadata=capabilities)
|
||||
|
||||
|
||||
# Default arguments for testing
|
||||
if __name__ == "__main__":
|
||||
load_dotenv()
|
||||
|
||||
arguments = ActionArguments(
|
||||
name="terminal",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
|
||||
)
|
||||
try:
|
||||
service = TerminalActionCollection(arguments)
|
||||
service.run()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}: {traceback.format_exc()}")
|
||||
@@ -0,0 +1,541 @@
|
||||
"""
|
||||
Wayback Machine MCP Server
|
||||
|
||||
This module provides MCP server functionality for interacting with the Wayback Machine.
|
||||
It supports listing archived versions, fetching archived content, and saving pages to the archive.
|
||||
|
||||
Key features:
|
||||
- List available archived versions of URLs with date filtering
|
||||
- Fetch content from specific archived page versions
|
||||
- Save current pages to the Wayback Machine
|
||||
- LLM-optimized output formatting with text extraction
|
||||
- Comprehensive error handling and logging
|
||||
|
||||
Main functions:
|
||||
- mcp_list_archived_versions: List available snapshots for a URL
|
||||
- mcp_get_archived_content: Fetch content from a specific archived version
|
||||
- mcp_get_wayback_capabilities: Get service capabilities information
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
|
||||
import requests
|
||||
from bs4 import BeautifulSoup
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic.fields import FieldInfo
|
||||
from waybackpy import WaybackMachineCDXServerAPI
|
||||
|
||||
from aworld.logs.util import Color
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
|
||||
|
||||
class ArchivedVersion(BaseModel):
|
||||
"""Individual archived version with structured data."""
|
||||
|
||||
timestamp: str
|
||||
url: str
|
||||
status_code: str
|
||||
digest: str
|
||||
length: str
|
||||
mime_type: str
|
||||
|
||||
|
||||
class WaybackMetadata(BaseModel):
|
||||
"""Metadata for Wayback Machine operation results."""
|
||||
|
||||
url: str
|
||||
operation: str # 'list_versions', 'get_content', 'save_page'
|
||||
timestamp: str | None = None
|
||||
total_versions: int | None = None
|
||||
date_range: dict[str, str | None] | None = None
|
||||
content_length: int | None = None
|
||||
text_extracted: bool = False
|
||||
truncated: bool = False
|
||||
execution_time: float | None = None
|
||||
error_type: str | None = None
|
||||
user_agent: str = "AWorld/1.0 (https://github.com/inclusionAI/AWorld; qintong.wqt@antgroup.com)"
|
||||
|
||||
|
||||
class WaybackActionCollection(ActionCollection):
|
||||
"""MCP service for Wayback Machine operations.
|
||||
|
||||
Provides comprehensive Wayback Machine functionality including:
|
||||
- Listing archived versions of URLs with flexible filtering
|
||||
- Fetching content from specific archived snapshots
|
||||
- LLM-friendly content formatting and text extraction
|
||||
- Robust error handling and detailed logging
|
||||
- Multiple output formats (markdown, JSON, text)
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
|
||||
# Configuration
|
||||
self.user_agent = "AWorld/1.0 (https://github.com/inclusionAI/AWorld; qintong.wqt@antgroup.com)"
|
||||
self.default_timeout = 30
|
||||
self.max_content_length = 8192
|
||||
|
||||
self._color_log("Wayback Machine service initialized", Color.green, "debug")
|
||||
self._color_log(f"User Agent: {self.user_agent}", Color.blue, "debug")
|
||||
|
||||
def _format_versions_for_llm(self, versions: list[ArchivedVersion], query_info: dict) -> str:
|
||||
"""Format archived versions list for LLM consumption.
|
||||
|
||||
Args:
|
||||
versions: List of archived versions
|
||||
query_info: Query information including URL and filters
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if not versions:
|
||||
return f"No archived versions found for URL: {query_info.get('url', 'Unknown')}"
|
||||
|
||||
output_parts = [
|
||||
f"# Wayback Machine Archives for {query_info.get('url', 'Unknown')}",
|
||||
f"\nFound **{len(versions)}** archived versions",
|
||||
]
|
||||
|
||||
if query_info.get("from_date") or query_info.get("to_date"):
|
||||
date_filter = []
|
||||
if query_info.get("from_date"):
|
||||
date_filter.append(f"From: {query_info['from_date']}")
|
||||
if query_info.get("to_date"):
|
||||
date_filter.append(f"To: {query_info['to_date']}")
|
||||
output_parts.append(f"\n**Date Filter:** {' | '.join(date_filter)}")
|
||||
|
||||
output_parts.append("\n## Available Versions:")
|
||||
|
||||
for i, version in enumerate(versions[:10], 1): # Show first 10
|
||||
timestamp_formatted = self._format_timestamp(version.timestamp)
|
||||
output_parts.append(
|
||||
f"\n{i}. **{timestamp_formatted}**\n"
|
||||
f" - Archive URL: {version.url}\n"
|
||||
f" - Status: {version.status_code} | Size: {version.length} bytes\n"
|
||||
f" - Type: {version.mime_type}"
|
||||
)
|
||||
|
||||
if len(versions) > 10:
|
||||
output_parts.append(f"\n... and {len(versions) - 10} more versions")
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
def _format_content_for_llm(self, content_data: dict, output_format: str = "markdown") -> str:
|
||||
"""Format archived content for LLM consumption.
|
||||
|
||||
Args:
|
||||
content_data: Content data dictionary
|
||||
output_format: Format type ('markdown', 'json', 'text')
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if output_format == "json":
|
||||
return json.dumps(content_data, indent=2)
|
||||
|
||||
elif output_format == "text":
|
||||
return content_data.get("content", "")
|
||||
|
||||
else: # markdown (default)
|
||||
output_parts = [
|
||||
f"# Archived Content from {content_data.get('url', 'Unknown')}",
|
||||
f"\n**Requested Timestamp:** {content_data.get('timestamp', 'Unknown')}",
|
||||
f"**Actual Timestamp:** {self._format_timestamp(content_data.get('fetched_timestamp', ''))}",
|
||||
f"**Content Length:** {content_data.get('original_content_length', 0):,} characters",
|
||||
]
|
||||
|
||||
if content_data.get("truncated"):
|
||||
output_parts.append(f"**Note:** Content truncated to {self.max_content_length:,} characters")
|
||||
|
||||
if content_data.get("extract_text_only"):
|
||||
output_parts.append("**Note:** Text-only extraction applied")
|
||||
|
||||
output_parts.extend(["\n## Content:", "\n---\n", content_data.get("content", ""), "\n---"])
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
def _format_timestamp(self, timestamp: str) -> str:
|
||||
"""Format Wayback Machine timestamp to human-readable format.
|
||||
|
||||
Args:
|
||||
timestamp: Wayback timestamp (YYYYMMDDhhmmss)
|
||||
|
||||
Returns:
|
||||
Human-readable timestamp
|
||||
"""
|
||||
try:
|
||||
if len(timestamp) >= 14:
|
||||
dt = datetime.strptime(timestamp[:14], "%Y%m%d%H%M%S")
|
||||
return dt.strftime("%Y-%m-%d %H:%M:%S UTC")
|
||||
return timestamp
|
||||
except (ValueError, TypeError):
|
||||
return timestamp or "Unknown"
|
||||
|
||||
def _validate_wayback_parameters(self, url: str, timestamp: str = None) -> tuple[str, str | None]:
|
||||
"""Validate and normalize Wayback Machine parameters.
|
||||
|
||||
Args:
|
||||
url: URL to validate
|
||||
timestamp: Optional timestamp to validate
|
||||
|
||||
Returns:
|
||||
Tuple of (validated_url, validated_timestamp)
|
||||
|
||||
Raises:
|
||||
ValueError: If parameters are invalid
|
||||
"""
|
||||
if not url or not url.strip():
|
||||
raise ValueError("URL cannot be empty")
|
||||
|
||||
url = url.strip()
|
||||
if not url.startswith(("http://", "https://")):
|
||||
url = "https://" + url
|
||||
|
||||
validated_timestamp = None
|
||||
if timestamp:
|
||||
timestamp = timestamp.strip()
|
||||
if len(timestamp) < 8:
|
||||
raise ValueError("Timestamp must be at least 8 characters (YYYYMMDD)")
|
||||
validated_timestamp = timestamp
|
||||
|
||||
return url, validated_timestamp
|
||||
|
||||
async def mcp_list_archived_versions(
|
||||
self,
|
||||
url: str = Field(description="The URL of the website to check for archived versions"),
|
||||
limit: int = Field(default=10, description="Maximum number of versions to return (0 for all)"),
|
||||
from_date: str | None = Field(default=None, description="Start date filter (YYYYMMDDhhmmss)"),
|
||||
to_date: str | None = Field(default=None, description="End date filter (YYYYMMDDhhmmss)"),
|
||||
output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'"),
|
||||
) -> ActionResponse:
|
||||
"""List available archived versions of a URL from the Wayback Machine.
|
||||
|
||||
This function queries the Wayback Machine CDX API to retrieve all available
|
||||
archived snapshots for a given URL, with optional date range filtering.
|
||||
|
||||
Args:
|
||||
url: The URL to search for archived versions
|
||||
limit: Maximum number of versions to return (0 for all, default: 10)
|
||||
from_date: Start date for filtering versions (YYYYMMDDhhmmss format)
|
||||
to_date: End date for filtering versions (YYYYMMDDhhmmss format)
|
||||
output_format: Format for the response ('markdown', 'json', or 'text')
|
||||
|
||||
Returns:
|
||||
ActionResponse with archived versions list and metadata
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(url, FieldInfo):
|
||||
url = url.default
|
||||
if isinstance(limit, FieldInfo):
|
||||
limit = limit.default
|
||||
if isinstance(from_date, FieldInfo):
|
||||
from_date = from_date.default
|
||||
if isinstance(to_date, FieldInfo):
|
||||
to_date = to_date.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
try:
|
||||
# Validate parameters
|
||||
url, _ = self._validate_wayback_parameters(url)
|
||||
|
||||
self._color_log(f"Listing archived versions for: {url}", Color.blue)
|
||||
|
||||
# Query Wayback Machine CDX API
|
||||
cdx_api = WaybackMachineCDXServerAPI(url, user_agent=self.user_agent)
|
||||
all_snapshots = list(cdx_api.snapshots())
|
||||
|
||||
# Apply date filtering
|
||||
if from_date or to_date:
|
||||
snapshots = [
|
||||
s
|
||||
for s in all_snapshots
|
||||
if (not from_date or s.timestamp >= from_date) and (not to_date or s.timestamp <= to_date)
|
||||
]
|
||||
else:
|
||||
snapshots = all_snapshots
|
||||
|
||||
if not snapshots:
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message="No archived versions found for the specified URL and date range.",
|
||||
metadata=WaybackMetadata(
|
||||
url=url,
|
||||
operation="list_versions",
|
||||
total_versions=0,
|
||||
date_range={"from_date": from_date, "to_date": to_date},
|
||||
execution_time=time.time() - start_time,
|
||||
error_type="no_results",
|
||||
).model_dump(),
|
||||
)
|
||||
|
||||
# Convert to structured format
|
||||
versions = [
|
||||
ArchivedVersion(
|
||||
timestamp=snapshot.timestamp,
|
||||
url=snapshot.archive_url,
|
||||
status_code=snapshot.statuscode,
|
||||
digest=snapshot.digest,
|
||||
length=snapshot.length,
|
||||
mime_type=snapshot.mimetype,
|
||||
)
|
||||
for snapshot in snapshots
|
||||
]
|
||||
|
||||
# Apply limit
|
||||
if limit > 0 and len(versions) > limit:
|
||||
versions = versions[:limit]
|
||||
|
||||
# Format output
|
||||
query_info = {"url": url, "from_date": from_date, "to_date": to_date, "total_found": len(snapshots)}
|
||||
|
||||
if output_format == "json":
|
||||
message = [version.model_dump() for version in versions]
|
||||
else:
|
||||
message = self._format_versions_for_llm(versions, query_info)
|
||||
|
||||
execution_time = time.time() - start_time
|
||||
self._color_log(f"Found {len(versions)} archived versions in {execution_time:.2f}s", Color.green)
|
||||
|
||||
return ActionResponse(
|
||||
success=True,
|
||||
message=message,
|
||||
metadata=WaybackMetadata(
|
||||
url=url,
|
||||
operation="list_versions",
|
||||
total_versions=len(snapshots),
|
||||
date_range={"from_date": from_date, "to_date": to_date},
|
||||
execution_time=execution_time,
|
||||
).model_dump(),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to list archived versions: {str(e)}"
|
||||
self._color_log(error_msg, Color.red)
|
||||
self.logger.error(f"Error in mcp_list_archived_versions: {traceback.format_exc()}")
|
||||
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=error_msg,
|
||||
metadata=WaybackMetadata(
|
||||
url=url or "unknown",
|
||||
operation="list_versions",
|
||||
execution_time=time.time() - start_time,
|
||||
error_type=type(e).__name__,
|
||||
).model_dump(),
|
||||
)
|
||||
|
||||
async def mcp_get_archived_content(
|
||||
self,
|
||||
url: str = Field(description="The URL of the website to fetch archived content from"),
|
||||
timestamp: str = Field(description="The timestamp of the desired version (YYYYMMDDhhmmss)"),
|
||||
extract_text_only: bool = Field(default=True, description="Extract only text content, removing HTML tags"),
|
||||
truncate_content: bool = Field(default=False, description="Truncate content to manageable length for LLMs"),
|
||||
output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'"),
|
||||
) -> ActionResponse:
|
||||
"""Fetch content from a specific archived page version.
|
||||
|
||||
This function retrieves the content of a specific archived snapshot from the
|
||||
Wayback Machine, with options for text extraction and content truncation.
|
||||
|
||||
Args:
|
||||
url: The URL of the website to fetch
|
||||
timestamp: The timestamp of the desired version (YYYYMMDDhhmmss)
|
||||
extract_text_only: Whether to extract only text content (default: True)
|
||||
truncate_content: Whether to truncate content for LLM consumption (default: False)
|
||||
output_format: Format for the response ('markdown', 'json', or 'text')
|
||||
|
||||
Returns:
|
||||
ActionResponse with archived content and metadata
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(url, FieldInfo):
|
||||
url = url.default
|
||||
if isinstance(timestamp, FieldInfo):
|
||||
timestamp = timestamp.default
|
||||
if isinstance(extract_text_only, FieldInfo):
|
||||
extract_text_only = extract_text_only.default
|
||||
if isinstance(truncate_content, FieldInfo):
|
||||
truncate_content = truncate_content.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
try:
|
||||
# Validate parameters
|
||||
url, timestamp = self._validate_wayback_parameters(url, timestamp)
|
||||
|
||||
self._color_log(f"Fetching archived content: {url} at {timestamp}", Color.blue)
|
||||
|
||||
# Query Wayback Machine for closest snapshot
|
||||
cdx_api = WaybackMachineCDXServerAPI(url, user_agent=self.user_agent)
|
||||
snapshot = cdx_api.near(wayback_machine_timestamp=timestamp)
|
||||
|
||||
if not snapshot or not snapshot.archive_url:
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"No archived version found for {url} at timestamp {timestamp}",
|
||||
metadata=WaybackMetadata(
|
||||
url=url,
|
||||
operation="get_content",
|
||||
timestamp=timestamp,
|
||||
execution_time=time.time() - start_time,
|
||||
error_type="no_snapshot",
|
||||
).model_dump(),
|
||||
)
|
||||
|
||||
# Fetch content
|
||||
response = requests.get(snapshot.archive_url, timeout=self.default_timeout)
|
||||
response.raise_for_status()
|
||||
content = response.text
|
||||
original_length = len(content)
|
||||
|
||||
# Extract text if requested
|
||||
if extract_text_only:
|
||||
soup = BeautifulSoup(content, "html.parser")
|
||||
content = soup.get_text(separator=" ", strip=True)
|
||||
|
||||
# Truncate if requested
|
||||
truncated = False
|
||||
if truncate_content and len(content) > self.max_content_length:
|
||||
content = content[: self.max_content_length] + "..."
|
||||
truncated = True
|
||||
|
||||
# Prepare content data
|
||||
content_data = {
|
||||
"url": url,
|
||||
"timestamp": timestamp,
|
||||
"fetched_timestamp": snapshot.timestamp,
|
||||
"content": content,
|
||||
"original_content_length": original_length,
|
||||
"truncated": truncated,
|
||||
"extract_text_only": extract_text_only,
|
||||
}
|
||||
|
||||
# Format output
|
||||
if output_format == "json":
|
||||
message = content_data
|
||||
elif output_format == "text":
|
||||
message = content
|
||||
else: # markdown
|
||||
message = self._format_content_for_llm(content_data, output_format)
|
||||
|
||||
execution_time = time.time() - start_time
|
||||
self._color_log(f"Retrieved {len(content):,} characters in {execution_time:.2f}s", Color.green)
|
||||
|
||||
return ActionResponse(
|
||||
success=True,
|
||||
message=message,
|
||||
metadata=WaybackMetadata(
|
||||
url=url,
|
||||
operation="get_content",
|
||||
timestamp=snapshot.timestamp,
|
||||
content_length=len(content),
|
||||
text_extracted=extract_text_only,
|
||||
truncated=truncated,
|
||||
execution_time=execution_time,
|
||||
).model_dump(),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to fetch archived content: {str(e)}"
|
||||
self._color_log(error_msg, Color.red)
|
||||
self.logger.error(f"Error in mcp_get_archived_content: {traceback.format_exc()}")
|
||||
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=error_msg,
|
||||
metadata=WaybackMetadata(
|
||||
url=url or "unknown",
|
||||
operation="get_content",
|
||||
timestamp=timestamp or "unknown",
|
||||
execution_time=time.time() - start_time,
|
||||
error_type=type(e).__name__,
|
||||
).model_dump(),
|
||||
)
|
||||
|
||||
def mcp_get_wayback_capabilities(self) -> ActionResponse:
|
||||
"""Get Wayback Machine service capabilities and configuration.
|
||||
|
||||
Returns:
|
||||
ActionResponse with service capabilities information
|
||||
"""
|
||||
capabilities = {
|
||||
"service": "Wayback Machine MCP Server",
|
||||
"version": "1.0.0",
|
||||
"description": "Interact with the Internet Archive's Wayback Machine",
|
||||
"operations": {
|
||||
"list_versions": "List archived versions of URLs with date filtering",
|
||||
"get_content": "Fetch content from specific archived snapshots",
|
||||
},
|
||||
"features": {
|
||||
"date_filtering": True,
|
||||
"text_extraction": True,
|
||||
"content_truncation": True,
|
||||
"multiple_formats": ["markdown", "json", "text"],
|
||||
"error_handling": True,
|
||||
"logging": True,
|
||||
},
|
||||
"configuration": {
|
||||
"user_agent": self.user_agent,
|
||||
"default_timeout": self.default_timeout,
|
||||
"max_content_length": self.max_content_length,
|
||||
},
|
||||
"limits": {"max_content_length": self.max_content_length, "request_timeout": self.default_timeout},
|
||||
}
|
||||
|
||||
message = f"""# Wayback Machine Service Capabilities
|
||||
|
||||
## Service Information
|
||||
- **Service:** {capabilities["service"]}
|
||||
- **Version:** {capabilities["version"]}
|
||||
- **Description:** {capabilities["description"]}
|
||||
|
||||
## Available Operations
|
||||
- **List Versions:** {capabilities["operations"]["list_versions"]}
|
||||
- **Get Content:** {capabilities["operations"]["get_content"]}
|
||||
- **Save Page:** {capabilities["operations"]["save_page"]}
|
||||
|
||||
## Features
|
||||
- **Date Filtering:** {capabilities["features"]["date_filtering"]}
|
||||
- **Text Extraction:** {capabilities["features"]["text_extraction"]}
|
||||
- **Content Truncation:** {capabilities["features"]["content_truncation"]}
|
||||
- **Output Formats:** {", ".join(capabilities["features"]["multiple_formats"])}
|
||||
- **Error Handling:** {capabilities["features"]["error_handling"]}
|
||||
- **Logging:** {capabilities["features"]["logging"]}
|
||||
|
||||
## Configuration
|
||||
- **User Agent:** {capabilities["configuration"]["user_agent"]}
|
||||
- **Default Timeout:** {capabilities["configuration"]["default_timeout"]} seconds
|
||||
- **Max Content Length:** {capabilities["configuration"]["max_content_length"]:,} characters
|
||||
|
||||
## Limits
|
||||
- **Max Content Length:** {capabilities["limits"]["max_content_length"]:,} characters
|
||||
- **Request Timeout:** {capabilities["limits"]["request_timeout"]} seconds
|
||||
"""
|
||||
|
||||
return ActionResponse(success=True, message=message, metadata=capabilities)
|
||||
|
||||
|
||||
# Default arguments for testing
|
||||
if __name__ == "__main__":
|
||||
load_dotenv()
|
||||
|
||||
arguments = ActionArguments(
|
||||
name="wayback-machine-server",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
|
||||
)
|
||||
try:
|
||||
service = WaybackActionCollection(arguments)
|
||||
service.run()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}")
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,621 @@
|
||||
"""
|
||||
Yahoo Finance MCP Action Collection
|
||||
|
||||
This module provides Yahoo Finance data access through the ActionCollection framework.
|
||||
It supports stock quotes, historical data, company information, financial statements,
|
||||
news search, and market summaries with LLM-optimized output formatting.
|
||||
|
||||
Key features:
|
||||
- Real-time stock quotes and market data
|
||||
- Historical price data with configurable intervals
|
||||
- Company information and financial statements
|
||||
- Financial news search
|
||||
- Market indices summaries
|
||||
- LLM-friendly data formatting
|
||||
- Comprehensive error handling
|
||||
|
||||
Main functions:
|
||||
- mcp_get_stock_quote: Get current stock quote information
|
||||
- mcp_get_historical_data: Retrieve historical OHLCV data
|
||||
- mcp_get_company_info: Fetch company details and business information
|
||||
- mcp_get_financial_statements: Access income statements, balance sheets, cash flow
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
import yfinance as yf
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from aworld.logs.util import Color
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
|
||||
|
||||
class YFinanceMetadata(BaseModel):
|
||||
"""Metadata for Yahoo Finance operation results."""
|
||||
|
||||
symbol: str
|
||||
operation: str
|
||||
execution_time: float | None = None
|
||||
data_points: int | None = None
|
||||
error_type: str | None = None
|
||||
timestamp: str | None = None
|
||||
|
||||
|
||||
class YahooFinanceActionCollection(ActionCollection):
|
||||
"""Yahoo Finance MCP service for financial data access.
|
||||
|
||||
Provides comprehensive financial data capabilities including:
|
||||
- Real-time stock quotes and market data
|
||||
- Historical price data with flexible time ranges
|
||||
- Company information and business details
|
||||
- Financial statements (income, balance sheet, cash flow)
|
||||
- Financial news search and aggregation
|
||||
- Market indices summaries and overviews
|
||||
- LLM-optimized data formatting
|
||||
- Error handling and validation
|
||||
"""
|
||||
|
||||
def _format_financial_data(self, data: Any, data_type: str) -> str:
|
||||
"""Format financial data for LLM consumption.
|
||||
|
||||
Args:
|
||||
data: Raw financial data
|
||||
data_type: Type of data for context
|
||||
|
||||
Returns:
|
||||
LLM-friendly formatted string
|
||||
"""
|
||||
if isinstance(data, dict):
|
||||
if data_type == "quote":
|
||||
return self._format_quote_data(data)
|
||||
elif data_type == "company":
|
||||
return self._format_company_data(data)
|
||||
elif isinstance(data, list):
|
||||
if data_type == "historical":
|
||||
return self._format_historical_data(data)
|
||||
elif data_type == "market_summary":
|
||||
return self._format_market_summary_data(data)
|
||||
elif data_type == "news":
|
||||
return self._format_news_list_data(data)
|
||||
|
||||
return str(data)
|
||||
|
||||
def _format_quote_data(self, quote: dict[str, Any]) -> str:
|
||||
"""Format stock quote data for LLM."""
|
||||
lines = [f"# Stock Quote: {quote.get('symbol', 'N/A')}"]
|
||||
|
||||
if quote.get("companyName"):
|
||||
lines.append(f"**Company:** {quote['companyName']}")
|
||||
|
||||
if quote.get("currentPrice"):
|
||||
lines.append(f"**Current Price:** ${quote['currentPrice']:.2f} {quote.get('currency', '')}")
|
||||
|
||||
if quote.get("previousClose"):
|
||||
change = quote.get("currentPrice", 0) - quote.get("previousClose", 0)
|
||||
change_pct = (change / quote["previousClose"]) * 100 if quote.get("previousClose") else 0
|
||||
direction = "📈" if change >= 0 else "📉"
|
||||
lines.append(f"**Change:** {direction} ${change:.2f} ({change_pct:.2f}%)")
|
||||
|
||||
if quote.get("dayHigh") and quote.get("dayLow"):
|
||||
lines.append(f"**Day Range:** ${quote['dayLow']:.2f} - ${quote['dayHigh']:.2f}")
|
||||
|
||||
if quote.get("volume"):
|
||||
lines.append(f"**Volume:** {quote['volume']:,}")
|
||||
|
||||
if quote.get("marketCap"):
|
||||
lines.append(f"**Market Cap:** ${quote['marketCap']:,}")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
def _format_company_data(self, company: dict[str, Any]) -> str:
|
||||
"""Format company information for LLM."""
|
||||
lines = [f"# Company Information: {company.get('symbol', 'N/A')}"]
|
||||
|
||||
if company.get("longName"):
|
||||
lines.append(f"**Company Name:** {company['longName']}")
|
||||
|
||||
if company.get("sector"):
|
||||
lines.append(f"**Sector:** {company['sector']}")
|
||||
|
||||
if company.get("industry"):
|
||||
lines.append(f"**Industry:** {company['industry']}")
|
||||
|
||||
if company.get("fullTimeEmployees"):
|
||||
lines.append(f"**Employees:** {company['fullTimeEmployees']:,}")
|
||||
|
||||
if company.get("city") and company.get("country"):
|
||||
location = f"{company['city']}, {company['country']}"
|
||||
if company.get("state"):
|
||||
location = f"{company['city']}, {company['state']}, {company['country']}"
|
||||
lines.append(f"**Location:** {location}")
|
||||
|
||||
if company.get("website"):
|
||||
lines.append(f"**Website:** {company['website']}")
|
||||
|
||||
if company.get("longBusinessSummary"):
|
||||
summary = (
|
||||
company["longBusinessSummary"][:500] + "..."
|
||||
if len(company["longBusinessSummary"]) > 500
|
||||
else company["longBusinessSummary"]
|
||||
)
|
||||
lines.extend(["\n**Business Summary:**", summary])
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
def _format_historical_data(self, data: list[dict[str, Any]]) -> str:
|
||||
"""Format historical data for LLM."""
|
||||
if not data:
|
||||
return "No historical data available."
|
||||
|
||||
lines = [f"# Historical Data ({len(data)} records)"]
|
||||
|
||||
# Show first few and last few records
|
||||
preview_count = min(3, len(data))
|
||||
|
||||
lines.append("\n**Recent Data:**")
|
||||
for record in data[-preview_count:]:
|
||||
date = record.get("Date", record.get("Datetime", "N/A"))
|
||||
close = record.get("Close", 0)
|
||||
volume = record.get("Volume", 0)
|
||||
lines.append(f"- {date}: Close ${close:.2f}, Volume {volume:,}")
|
||||
|
||||
if len(data) > preview_count * 2:
|
||||
lines.append(f"\n... {len(data) - preview_count * 2} more records ...")
|
||||
|
||||
if len(data) > preview_count:
|
||||
lines.append("\n**Earliest Data:**")
|
||||
for record in data[:preview_count]:
|
||||
date = record.get("Date", record.get("Datetime", "N/A"))
|
||||
close = record.get("Close", 0)
|
||||
volume = record.get("Volume", 0)
|
||||
lines.append(f"- {date}: Close ${close:.2f}, Volume {volume:,}")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
def _format_news_list_data(self, news_list: list[dict[str, Any]]) -> str:
|
||||
"""Format news list for LLM."""
|
||||
if not news_list:
|
||||
return "No news articles found."
|
||||
|
||||
lines = [f"# Financial News ({len(news_list)} articles)"]
|
||||
|
||||
for i, article in enumerate(news_list, 1):
|
||||
lines.append(f"\n## {i}. {article.get('title', 'No Title')}")
|
||||
if article.get("publisher"):
|
||||
lines.append(f"**Publisher:** {article['publisher']}")
|
||||
if article.get("providerPublishTime"):
|
||||
lines.append(f"**Published:** {article['providerPublishTime']}")
|
||||
if article.get("link"):
|
||||
lines.append(f"**Link:** {article['link']}")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
def _format_market_summary_data(self, summaries: list[dict[str, Any]]) -> str:
|
||||
"""Format market summary for LLM."""
|
||||
if not summaries:
|
||||
return "No market data available."
|
||||
|
||||
lines = ["# Market Summary"]
|
||||
|
||||
for summary in summaries:
|
||||
symbol = summary.get("symbol", "N/A")
|
||||
name = summary.get("name", symbol)
|
||||
price = summary.get("currentPrice", 0)
|
||||
change = summary.get("change", 0)
|
||||
change_pct = summary.get("percentChange", 0)
|
||||
|
||||
direction = "📈" if change >= 0 else "📉"
|
||||
lines.append(f"\n**{name} ({symbol})**")
|
||||
lines.append(f"Price: ${price:.2f} {direction} {change:+.2f} ({change_pct:+.2f}%)")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
async def mcp_get_stock_quote(
|
||||
self,
|
||||
symbol: str = Field(description="Stock ticker symbol (e.g., AAPL, MSFT)"),
|
||||
) -> ActionResponse:
|
||||
"""Get current stock quote information.
|
||||
|
||||
Fetches real-time stock quote data including current price, daily changes,
|
||||
volume, market cap, and other key metrics for the specified ticker symbol.
|
||||
|
||||
Args:
|
||||
symbol: The stock ticker symbol to fetch quote for
|
||||
|
||||
Returns:
|
||||
ActionResponse with formatted quote data and metadata
|
||||
"""
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(symbol, FieldInfo):
|
||||
symbol = symbol.default
|
||||
|
||||
try:
|
||||
start_time = time.time()
|
||||
self._color_log(f"📊 Fetching stock quote for: {symbol}", Color.cyan)
|
||||
|
||||
ticker = yf.Ticker(symbol)
|
||||
info = ticker.info
|
||||
|
||||
if not info or (info.get("regularMarketPrice") is None and info.get("currentPrice") is None):
|
||||
# Try to get basic history to validate symbol
|
||||
hist = ticker.history(period="1d")
|
||||
if hist.empty:
|
||||
raise ValueError(f"No data found for symbol: {symbol}. It might be invalid or delisted.")
|
||||
raise ValueError(f"Could not retrieve detailed quote for symbol: {symbol}. Limited data available.")
|
||||
|
||||
# Extract key quote information
|
||||
quote_data = {
|
||||
"symbol": symbol.upper(),
|
||||
"companyName": info.get("shortName", info.get("longName")),
|
||||
"currentPrice": info.get("regularMarketPrice", info.get("currentPrice")),
|
||||
"previousClose": info.get("previousClose"),
|
||||
"open": info.get("regularMarketOpen", info.get("open")),
|
||||
"dayHigh": info.get("regularMarketDayHigh", info.get("dayHigh")),
|
||||
"dayLow": info.get("regularMarketDayLow", info.get("dayLow")),
|
||||
"volume": info.get("regularMarketVolume", info.get("volume")),
|
||||
"averageVolume": info.get("averageVolume"),
|
||||
"marketCap": info.get("marketCap"),
|
||||
"fiftyTwoWeekHigh": info.get("fiftyTwoWeekHigh"),
|
||||
"fiftyTwoWeekLow": info.get("fiftyTwoWeekLow"),
|
||||
"currency": info.get("currency"),
|
||||
"exchange": info.get("exchange"),
|
||||
}
|
||||
|
||||
# Filter out None values
|
||||
quote_data = {k: v for k, v in quote_data.items() if v is not None}
|
||||
|
||||
execution_time = time.time() - start_time
|
||||
formatted_message = self._format_financial_data(quote_data, "quote")
|
||||
|
||||
metadata = YFinanceMetadata(
|
||||
symbol=symbol.upper(),
|
||||
operation="get_stock_quote",
|
||||
execution_time=execution_time,
|
||||
data_points=len(quote_data),
|
||||
yfinance_available=True,
|
||||
timestamp=datetime.now().isoformat(),
|
||||
)
|
||||
|
||||
self._color_log("✅ Stock quote retrieved successfully", Color.green)
|
||||
|
||||
return ActionResponse(success=True, message=formatted_message, metadata=metadata.model_dump())
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to fetch stock quote for {symbol}: {str(e)}"
|
||||
self.logger.error(f"Stock quote error: {traceback.format_exc()}")
|
||||
|
||||
metadata = YFinanceMetadata(
|
||||
symbol=symbol.upper(),
|
||||
operation="get_stock_quote",
|
||||
error_type=type(e).__name__,
|
||||
timestamp=datetime.now().isoformat(),
|
||||
)
|
||||
|
||||
return ActionResponse(success=False, message=error_msg, metadata=metadata.model_dump())
|
||||
|
||||
async def mcp_get_historical_data(
|
||||
self,
|
||||
symbol: str = Field(description="Stock ticker symbol (e.g., AAPL, MSFT)"),
|
||||
start: str = Field(description="Start date (YYYY-MM-DD)"),
|
||||
end: str = Field(description="End date (YYYY-MM-DD)"),
|
||||
interval: str = Field(default="1d", description="Data interval (1d, 1wk, 1mo, etc.)"),
|
||||
max_rows_preview: int = Field(default=10, description="Max rows for preview (0 for all data)"),
|
||||
) -> ActionResponse:
|
||||
"""Retrieve historical stock data.
|
||||
|
||||
Fetches historical OHLCV (Open, High, Low, Close, Volume) data for the
|
||||
specified ticker symbol within the given date range and interval.
|
||||
|
||||
Args:
|
||||
symbol: Stock ticker symbol
|
||||
start: Start date in YYYY-MM-DD format
|
||||
end: End date in YYYY-MM-DD format
|
||||
interval: Data interval (1d, 1wk, 1mo, etc.)
|
||||
max_rows_preview: Maximum rows to show in preview
|
||||
|
||||
Returns:
|
||||
ActionResponse with historical data and metadata
|
||||
"""
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(symbol, FieldInfo):
|
||||
symbol = symbol.default
|
||||
if isinstance(start, FieldInfo):
|
||||
start = start.default
|
||||
if isinstance(end, FieldInfo):
|
||||
end = end.default
|
||||
if isinstance(interval, FieldInfo):
|
||||
interval = interval.default
|
||||
if isinstance(max_rows_preview, FieldInfo):
|
||||
max_rows_preview = max_rows_preview.default
|
||||
|
||||
try:
|
||||
start_time = time.time()
|
||||
self._color_log(f"📈 Fetching historical data for: {symbol} ({start} to {end})", Color.cyan)
|
||||
|
||||
ticker = yf.Ticker(symbol)
|
||||
hist_df = ticker.history(start=start, end=end, interval=interval)
|
||||
|
||||
if hist_df.empty:
|
||||
raise ValueError(
|
||||
f"No historical data found for {symbol} with start={start}, end={end}, interval={interval}"
|
||||
)
|
||||
|
||||
# Convert DataFrame to list of dictionaries
|
||||
hist_df.reset_index(inplace=True)
|
||||
|
||||
# Ensure date columns are strings for JSON serialization
|
||||
if "Date" in hist_df.columns:
|
||||
hist_df["Date"] = hist_df["Date"].astype(str)
|
||||
if "Datetime" in hist_df.columns:
|
||||
hist_df["Datetime"] = hist_df["Datetime"].astype(str)
|
||||
|
||||
# Clean column names
|
||||
hist_df.columns = hist_df.columns.str.replace(" ", "")
|
||||
|
||||
historical_data = hist_df.to_dict(orient="records")
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
# Format message based on data size
|
||||
if max_rows_preview > 0 and len(historical_data) > max_rows_preview:
|
||||
preview_count = max_rows_preview // 2
|
||||
preview_count = max(1, preview_count)
|
||||
|
||||
preview_data = historical_data[:preview_count] + historical_data[-preview_count:]
|
||||
formatted_message = self._format_financial_data(preview_data, "historical")
|
||||
formatted_message += (
|
||||
f"\n\n*Note: Showing preview of {len(preview_data)} out of {len(historical_data)} total records*"
|
||||
)
|
||||
else:
|
||||
formatted_message = self._format_financial_data(historical_data, "historical")
|
||||
|
||||
metadata = YFinanceMetadata(
|
||||
symbol=symbol.upper(),
|
||||
operation="get_historical_data",
|
||||
execution_time=execution_time,
|
||||
data_points=len(historical_data),
|
||||
yfinance_available=True,
|
||||
timestamp=datetime.now().isoformat(),
|
||||
)
|
||||
|
||||
self._color_log(f"✅ Historical data retrieved: {len(historical_data)} records", Color.green)
|
||||
|
||||
return ActionResponse(success=True, message=formatted_message, metadata=metadata.model_dump())
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to fetch historical data for {symbol}: {str(e)}"
|
||||
self.logger.error(f"Historical data error: {traceback.format_exc()}")
|
||||
|
||||
metadata = YFinanceMetadata(
|
||||
symbol=symbol.upper(),
|
||||
operation="get_historical_data",
|
||||
error_type=type(e).__name__,
|
||||
timestamp=datetime.now().isoformat(),
|
||||
)
|
||||
|
||||
return ActionResponse(success=False, message=error_msg, metadata=metadata.model_dump())
|
||||
|
||||
async def mcp_get_company_info(
|
||||
self,
|
||||
symbol: str = Field(description="Stock ticker symbol (e.g., AAPL, MSFT)"),
|
||||
) -> ActionResponse:
|
||||
"""Get company information and business details.
|
||||
|
||||
Fetches comprehensive company information including sector, industry,
|
||||
employee count, business summary, location, and other key details.
|
||||
|
||||
Args:
|
||||
symbol: Stock ticker symbol
|
||||
|
||||
Returns:
|
||||
ActionResponse with company information and metadata
|
||||
"""
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(symbol, FieldInfo):
|
||||
symbol = symbol.default
|
||||
|
||||
try:
|
||||
start_time = time.time()
|
||||
self._color_log(f"🏢 Fetching company info for: {symbol}", Color.cyan)
|
||||
|
||||
ticker = yf.Ticker(symbol)
|
||||
info = ticker.info
|
||||
|
||||
if not info or not info.get("symbol"):
|
||||
raise ValueError(f"No company information found for symbol: {symbol}. It might be invalid.")
|
||||
|
||||
# Extract key company information
|
||||
company_data = {
|
||||
"symbol": info.get("symbol"),
|
||||
"shortName": info.get("shortName"),
|
||||
"longName": info.get("longName"),
|
||||
"sector": info.get("sector"),
|
||||
"industry": info.get("industry"),
|
||||
"fullTimeEmployees": info.get("fullTimeEmployees"),
|
||||
"longBusinessSummary": info.get("longBusinessSummary"),
|
||||
"city": info.get("city"),
|
||||
"state": info.get("state"),
|
||||
"country": info.get("country"),
|
||||
"website": info.get("website"),
|
||||
"exchange": info.get("exchange"),
|
||||
"currency": info.get("currency"),
|
||||
"marketCap": info.get("marketCap"),
|
||||
}
|
||||
|
||||
# Filter out None values
|
||||
company_data = {k: v for k, v in company_data.items() if v is not None}
|
||||
|
||||
execution_time = time.time() - start_time
|
||||
formatted_message = self._format_financial_data(company_data, "company")
|
||||
|
||||
metadata = YFinanceMetadata(
|
||||
symbol=symbol.upper(),
|
||||
operation="get_company_info",
|
||||
execution_time=execution_time,
|
||||
data_points=len(company_data),
|
||||
yfinance_available=True,
|
||||
timestamp=datetime.now().isoformat(),
|
||||
)
|
||||
|
||||
self._color_log("✅ Company information retrieved successfully", Color.green)
|
||||
|
||||
return ActionResponse(success=True, message=formatted_message, metadata=metadata.model_dump())
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to fetch company info for {symbol}: {str(e)}"
|
||||
self.logger.error(f"Company info error: {traceback.format_exc()}")
|
||||
|
||||
metadata = YFinanceMetadata(
|
||||
symbol=symbol.upper(),
|
||||
operation="get_company_info",
|
||||
error_type=type(e).__name__,
|
||||
timestamp=datetime.now().isoformat(),
|
||||
)
|
||||
|
||||
return ActionResponse(success=False, message=error_msg, metadata=metadata.model_dump())
|
||||
|
||||
async def mcp_get_financial_statements(
|
||||
self,
|
||||
symbol: str = Field(description="Stock ticker symbol (e.g., AAPL, MSFT)"),
|
||||
statement_type: str = Field(description="Statement type: income_statement, balance_sheet, or cash_flow"),
|
||||
period_type: str = Field(default="annual", description="Period type: annual or quarterly"),
|
||||
max_columns_preview: int = Field(default=4, description="Max periods to show (0 for all)"),
|
||||
) -> ActionResponse:
|
||||
"""Get financial statements for a company.
|
||||
|
||||
Fetches financial statements including income statement, balance sheet,
|
||||
or cash flow statement for the specified company and period.
|
||||
|
||||
Args:
|
||||
symbol: Stock ticker symbol
|
||||
statement_type: Type of statement (income_statement, balance_sheet, cash_flow)
|
||||
period_type: Period type (annual or quarterly)
|
||||
max_columns_preview: Maximum periods to show in preview
|
||||
|
||||
Returns:
|
||||
ActionResponse with financial statement data and metadata
|
||||
"""
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(symbol, FieldInfo):
|
||||
symbol = symbol.default
|
||||
if isinstance(statement_type, FieldInfo):
|
||||
statement_type = statement_type.default
|
||||
if isinstance(period_type, FieldInfo):
|
||||
period_type = period_type.default
|
||||
if isinstance(max_columns_preview, FieldInfo):
|
||||
max_columns_preview = max_columns_preview.default
|
||||
|
||||
try:
|
||||
start_time = time.time()
|
||||
self._color_log(f"📋 Fetching {statement_type} for: {symbol} ({period_type})", Color.cyan)
|
||||
|
||||
ticker = yf.Ticker(symbol)
|
||||
statement_df = None
|
||||
|
||||
# Get appropriate statement
|
||||
if statement_type == "income_statement":
|
||||
statement_df = ticker.income_stmt if period_type == "annual" else ticker.quarterly_income_stmt
|
||||
elif statement_type == "balance_sheet":
|
||||
statement_df = ticker.balance_sheet if period_type == "annual" else ticker.quarterly_balance_sheet
|
||||
elif statement_type == "cash_flow":
|
||||
statement_df = ticker.cashflow if period_type == "annual" else ticker.quarterly_cashflow
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid statement_type: {statement_type}. "
|
||||
"Must be one of: income_statement, balance_sheet, cash_flow"
|
||||
)
|
||||
|
||||
if statement_df is None or statement_df.empty:
|
||||
raise ValueError(f"No {period_type} {statement_type} data found for symbol {symbol}")
|
||||
|
||||
# Process DataFrame
|
||||
statement_df.reset_index(inplace=True)
|
||||
statement_df.rename(columns={"index": "Item"}, inplace=True)
|
||||
|
||||
# Convert date columns to strings
|
||||
for col in statement_df.columns:
|
||||
if col != "Item":
|
||||
try:
|
||||
if hasattr(col, "strftime"):
|
||||
statement_df.rename(columns={col: col.strftime("%Y-%m-%d")}, inplace=True)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
statement_data = statement_df.to_dict(orient="records")
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
# Format message
|
||||
if max_columns_preview > 0 and len(statement_df.columns) > (max_columns_preview + 1):
|
||||
columns_to_keep = ["Item"] + list(statement_df.columns[1 : max_columns_preview + 1])
|
||||
preview_df = statement_df[columns_to_keep]
|
||||
preview_data = preview_df.to_dict(orient="records")
|
||||
|
||||
formatted_message = f"# {statement_type.replace('_', ' ').title()} ({period_type.title()})\n\n"
|
||||
formatted_message += (
|
||||
f"Showing preview of most recent {max_columns_preview} periods "
|
||||
f"out of {len(statement_df.columns) - 1} available.\n\n"
|
||||
)
|
||||
|
||||
# Show key financial items
|
||||
for item in preview_data[:10]: # Show first 10 items
|
||||
item_name = item.get("Item", "N/A")
|
||||
formatted_message += f"**{item_name}:**\n"
|
||||
for col, value in item.items():
|
||||
if col != "Item" and value is not None:
|
||||
formatted_message += (
|
||||
f" - {col}: {value:,}\n"
|
||||
if isinstance(value, (int, float))
|
||||
else f" - {col}: {value}\n"
|
||||
)
|
||||
formatted_message += "\n"
|
||||
else:
|
||||
formatted_message = f"# {statement_type.replace('_', ' ').title()} ({period_type.title()})\n\n"
|
||||
formatted_message += f"Complete financial statement with {len(statement_data)} line items.\n"
|
||||
|
||||
metadata = YFinanceMetadata(
|
||||
symbol=symbol.upper(),
|
||||
operation="get_financial_statements",
|
||||
execution_time=execution_time,
|
||||
data_points=len(statement_data),
|
||||
yfinance_available=True,
|
||||
timestamp=datetime.now().isoformat(),
|
||||
)
|
||||
|
||||
self._color_log(f"✅ Financial statements retrieved: {len(statement_data)} items", Color.green)
|
||||
|
||||
return ActionResponse(success=True, message=formatted_message, metadata=metadata.model_dump())
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to fetch {statement_type} for {symbol}: {str(e)}"
|
||||
self.logger.error(f"Financial statements error: {traceback.format_exc()}")
|
||||
|
||||
metadata = YFinanceMetadata(
|
||||
symbol=symbol.upper(),
|
||||
operation="get_financial_statements",
|
||||
error_type=type(e).__name__,
|
||||
timestamp=datetime.now().isoformat(),
|
||||
)
|
||||
|
||||
return ActionResponse(success=False, message=error_msg, metadata=metadata.model_dump())
|
||||
|
||||
|
||||
# Default arguments for testing
|
||||
if __name__ == "__main__":
|
||||
load_dotenv()
|
||||
|
||||
arguments = ActionArguments(
|
||||
name="yahoo-finance",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
|
||||
)
|
||||
|
||||
try:
|
||||
service = YahooFinanceActionCollection(arguments)
|
||||
service.run()
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
@@ -0,0 +1,549 @@
|
||||
"""
|
||||
YouTube MCP Service
|
||||
|
||||
This module provides MCP service functionality for YouTube operations including:
|
||||
- Downloading videos from YouTube URLs
|
||||
- Extracting transcripts from YouTube videos
|
||||
|
||||
It handles various scenarios with proper validation, error handling,
|
||||
and progress tracking while providing LLM-friendly formatted results.
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import BaseModel, Field
|
||||
from selenium import webdriver
|
||||
from selenium.webdriver.chrome.service import Service
|
||||
from selenium.webdriver.common.by import By
|
||||
from youtube_transcript_api import FetchedTranscript, YouTubeTranscriptApi
|
||||
|
||||
from aworld.logs.util import Color
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
|
||||
# Default driver path for Chrome WebDriver
|
||||
_DEFAULT_DRIVER_PATH = os.environ.get(
|
||||
"CHROME_DRIVER_PATH", str(Path("~/Downloads/chromedriver-mac-arm64/chromedriver").expanduser())
|
||||
)
|
||||
|
||||
|
||||
class YoutubeDownloadResults(BaseModel):
|
||||
"""Download result model with file information"""
|
||||
|
||||
file_path: str
|
||||
file_name: str
|
||||
file_size: int
|
||||
content_type: str | None = None
|
||||
success: bool
|
||||
error: str | None = None
|
||||
|
||||
|
||||
class TranscriptResult(BaseModel):
|
||||
"""Transcript result model with transcript information"""
|
||||
|
||||
video_id: str
|
||||
transcript: FetchedTranscript
|
||||
success: bool
|
||||
error: str | None = None
|
||||
|
||||
|
||||
class YouTubeMetadata(BaseModel):
|
||||
"""Metadata for YouTube operation results"""
|
||||
|
||||
operation: str
|
||||
url: str | None = None
|
||||
video_id: str | None = None
|
||||
file_path: str | None = None
|
||||
file_name: str | None = None
|
||||
file_size: int | None = None
|
||||
content_type: str | None = None
|
||||
language_code: str | None = None
|
||||
translate_to_language: str | None = None
|
||||
execution_time: float | None = None
|
||||
error_type: str | None = None
|
||||
|
||||
|
||||
class YouTubeActionCollection(ActionCollection):
|
||||
"""MCP service for YouTube operations.
|
||||
|
||||
Provides YouTube capabilities including:
|
||||
- Video downloading with Selenium automation
|
||||
- Transcript extraction and translation
|
||||
- LLM-friendly result formatting
|
||||
- Error handling and logging
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
|
||||
# Initialize supported file extensions
|
||||
self.supported_extensions = {".mp4", ".webm", ".mkv"}
|
||||
|
||||
self._color_log("YouTube service initialized", Color.green, "debug")
|
||||
|
||||
def _format_transcript_output(self, result: TranscriptResult, format_type: str = "markdown") -> str:
|
||||
"""Format transcript results for LLM consumption.
|
||||
|
||||
Args:
|
||||
result: Transcript extraction result
|
||||
format_type: Output format ('markdown', 'json', 'text')
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if result is None or not result.success:
|
||||
return f"Failed to extract transcript: {result.error}"
|
||||
|
||||
if format_type == "json":
|
||||
return result.model_dump()
|
||||
elif format_type == "text":
|
||||
output = [f"Transcript for video ID: {result.video_id}\n"]
|
||||
|
||||
# Access snippets from FetchedTranscript
|
||||
for entry in result.transcript.snippets:
|
||||
start_time = entry["start"]
|
||||
text = entry["text"]
|
||||
|
||||
minutes, seconds = divmod(int(start_time), 60)
|
||||
timestamp = f"{minutes:02d}:{seconds:02d}"
|
||||
|
||||
output.append(f"[{timestamp}] {text}")
|
||||
|
||||
return "\n".join(output)
|
||||
else: # markdown (default)
|
||||
output = [f"# Transcript for YouTube Video: {result.video_id}\n"]
|
||||
output.append("| Timestamp | Text |")
|
||||
output.append("| --- | --- |")
|
||||
|
||||
# Access snippets from FetchedTranscript
|
||||
for entry in result.transcript.snippets:
|
||||
start_time = entry["start"]
|
||||
text: str = entry["text"]
|
||||
|
||||
minutes, seconds = divmod(int(start_time), 60)
|
||||
timestamp = f"{minutes:02d}:{seconds:02d}"
|
||||
|
||||
# Escape pipe characters in markdown table
|
||||
safe_text = text.replace("|", "\\|")
|
||||
output.append(f"| {timestamp} | {safe_text} |")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
def _format_download_output(self, result: YoutubeDownloadResults, format_type: str = "markdown") -> str:
|
||||
"""Format download results for LLM consumption.
|
||||
|
||||
Args:
|
||||
result: Download result
|
||||
format_type: Output format ('markdown', 'json', 'text')
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if not result.success:
|
||||
return f"Failed to download video: {result.error}"
|
||||
|
||||
if format_type == "json":
|
||||
return result.model_dump()
|
||||
elif format_type == "text":
|
||||
output_parts = [
|
||||
"Download completed successfully",
|
||||
f"File: {result.file_name}",
|
||||
f"Path: {result.file_path}",
|
||||
f"Size: {result.file_size} bytes",
|
||||
]
|
||||
if result.content_type:
|
||||
output_parts.append(f"Content Type: {result.content_type}")
|
||||
|
||||
return "\n".join(output_parts)
|
||||
else: # markdown (default)
|
||||
output_parts = [
|
||||
"# YouTube Download Results ✅",
|
||||
"",
|
||||
"## File Information",
|
||||
f"**Filename:** `{result.file_name}`",
|
||||
f"**Path:** `{result.file_path}`",
|
||||
f"**Size:** {result.file_size} bytes",
|
||||
]
|
||||
if result.content_type:
|
||||
output_parts.append(f"**Content Type:** {result.content_type}")
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
def _get_youtube_content(self, url: str, output_dir: str, timeout: int) -> None:
|
||||
"""Use Selenium to download YouTube content via cobalt.tools
|
||||
|
||||
Args:
|
||||
url: YouTube video URL
|
||||
output_dir: Directory to save downloaded content
|
||||
timeout: Maximum time to wait for download in seconds
|
||||
"""
|
||||
driver = None
|
||||
try:
|
||||
options = webdriver.ChromeOptions()
|
||||
options.add_argument("--disable-blink-features=AutomationControlled")
|
||||
# Set download file default path
|
||||
prefs = {
|
||||
"download.default_directory": output_dir,
|
||||
"download.prompt_for_download": False,
|
||||
"download.directory_upgrade": True,
|
||||
"safebrowsing.enabled": True,
|
||||
}
|
||||
options.add_experimental_option("prefs", prefs)
|
||||
# Create WebDriver object and launch Chrome browser
|
||||
service = Service(executable_path=_DEFAULT_DRIVER_PATH)
|
||||
driver = webdriver.Chrome(service=service, options=options)
|
||||
|
||||
self._color_log(f"Opening cobalt.tools to download from {url}", Color.blue)
|
||||
# Open target webpage
|
||||
driver.get("https://cobalt.tools/")
|
||||
# Wait for page to load
|
||||
time.sleep(5)
|
||||
# Find input field and enter YouTube link
|
||||
input_field = driver.find_element(By.ID, "link-area")
|
||||
input_field.send_keys(url)
|
||||
time.sleep(5)
|
||||
# Find download button and click
|
||||
download_button = driver.find_element(By.ID, "download-button")
|
||||
download_button.click()
|
||||
time.sleep(5)
|
||||
|
||||
try:
|
||||
# Handle bot detection popup
|
||||
driver.find_element(
|
||||
By.CLASS_NAME,
|
||||
"button.elevated.popup-button.undefined.svelte-nnawom.active",
|
||||
).click()
|
||||
except Exception as e:
|
||||
self._color_log(f"Bot detection handling: {str(e)}", Color.yellow)
|
||||
|
||||
# try:
|
||||
# t = 0
|
||||
# while t < timeout:
|
||||
# if (
|
||||
# "downloading" not in driver.find_element(By.CLASS_NAME, "status-text.svelte-dmosdd").text
|
||||
# and "starting" not in driver.find_element(By.CLASS_NAME, "status-text.svelte-dmosdd").text
|
||||
# ):
|
||||
# driver.find_element(By.CLASS_NAME, "button.action-button.svelte-dmosdd").click()
|
||||
# break
|
||||
# t += 3
|
||||
# time.sleep(3)
|
||||
# except Exception as e:
|
||||
# self._color_log(f"Bot detection handling: {str(e)}", Color.yellow)
|
||||
|
||||
# Wait for download to complete
|
||||
cnt = 0
|
||||
while len(os.listdir(output_dir)) == 0 or os.listdir(output_dir)[0].split(".")[-1] == "crdownload":
|
||||
time.sleep(3)
|
||||
cnt += 3
|
||||
if cnt >= timeout:
|
||||
self._color_log(f"Download timeout after {timeout} seconds", Color.yellow)
|
||||
break
|
||||
|
||||
self._color_log("Download process completed", Color.green)
|
||||
|
||||
except Exception as e:
|
||||
self._color_log(f"Error during YouTube content download: {str(e)}", Color.red)
|
||||
raise
|
||||
finally:
|
||||
# Close browser
|
||||
if driver:
|
||||
driver.quit()
|
||||
|
||||
def _find_existing_video(self, search_dir: str, video_id: str) -> str | None:
|
||||
"""Recursively search for an existing video file with the given ID.
|
||||
|
||||
Args:
|
||||
search_dir: Directory to search in
|
||||
video_id: YouTube video ID to look for
|
||||
|
||||
Returns:
|
||||
Path to existing file if found, None otherwise
|
||||
"""
|
||||
if not video_id:
|
||||
return None
|
||||
|
||||
search_path = Path(search_dir)
|
||||
if not search_path.exists():
|
||||
return None
|
||||
|
||||
for item in search_path.iterdir():
|
||||
if item.is_file() and video_id in item.name:
|
||||
return str(item)
|
||||
elif item.is_dir():
|
||||
found = self._find_existing_video(str(item), video_id)
|
||||
if found:
|
||||
return found
|
||||
|
||||
return None
|
||||
|
||||
async def mcp_download_youtube_video(
|
||||
self,
|
||||
url: str = Field(description="The URL of YouTube video to download."),
|
||||
timeout: int = Field(180, description="Download timeout in seconds (default: 180)."),
|
||||
output_format: str = Field(
|
||||
"markdown", description="Output format: 'markdown', 'json', or 'text' (default: markdown)."
|
||||
),
|
||||
) -> ActionResponse:
|
||||
"""Download a YouTube video from URL and save it to the local filesystem.
|
||||
|
||||
This tool provides YouTube video downloading with:
|
||||
- Selenium-based automation via cobalt.tools
|
||||
- Configurable timeout controls
|
||||
- Existing file detection to avoid redundant downloads
|
||||
- LLM-optimized result formatting
|
||||
|
||||
Args:
|
||||
url: The URL of YouTube video to download
|
||||
timeout: Maximum download time in seconds
|
||||
output_format: Format for the response output
|
||||
|
||||
Returns:
|
||||
ActionResponse with download results and metadata
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# Validate URL
|
||||
if not url.startswith(("http://", "https://")):
|
||||
raise ValueError("Invalid URL format. URL must start with http:// or https://")
|
||||
|
||||
if not ("youtube.com" in url or "youtu.be" in url):
|
||||
raise ValueError("URL must be a valid YouTube URL")
|
||||
|
||||
# Create output directory if it doesn't exist
|
||||
output_path = self.workspace / "youtube_downloads"
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Generate filename based on timestamp
|
||||
filename = f"youtube_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
|
||||
file_path = output_path / filename
|
||||
file_path.mkdir(parents=True, exist_ok=True)
|
||||
self._color_log(f"Output path: {file_path}", Color.blue)
|
||||
|
||||
# Extract video ID for existing file check
|
||||
video_id = url.split("?v=")[-1].split("&")[0] if "?v=" in url else ""
|
||||
if "youtu.be/" in url and not video_id:
|
||||
video_id = url.split("youtu.be/")[-1].split("?")[0]
|
||||
|
||||
# Check if video already exists
|
||||
base_path = self.workspace
|
||||
existing_file = self._find_existing_video(str(base_path), video_id)
|
||||
|
||||
if existing_file:
|
||||
existing_path = Path(existing_file)
|
||||
result = YoutubeDownloadResults(
|
||||
file_path=str(existing_path),
|
||||
file_name=existing_path.name,
|
||||
file_size=existing_path.stat().st_size,
|
||||
content_type="mp4",
|
||||
success=True,
|
||||
error=None,
|
||||
)
|
||||
self._color_log(f"Found {video_id} already downloaded in: {existing_file}", Color.green)
|
||||
|
||||
# Format output for LLM
|
||||
message = self._format_download_output(result, output_format)
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
# Create metadata
|
||||
metadata = YouTubeMetadata(
|
||||
operation="download",
|
||||
url=url,
|
||||
video_id=video_id,
|
||||
file_path=str(existing_path),
|
||||
file_name=existing_path.name,
|
||||
file_size=existing_path.stat().st_size,
|
||||
content_type="mp4",
|
||||
execution_time=execution_time,
|
||||
).model_dump()
|
||||
|
||||
return ActionResponse(success=True, message=message, metadata=metadata)
|
||||
|
||||
# Download the video
|
||||
self._color_log(f"Downloading video from {url} to {file_path}", Color.blue)
|
||||
self._get_youtube_content(url, str(file_path), timeout)
|
||||
|
||||
# Check if download was successful
|
||||
downloaded_files = list(file_path.iterdir())
|
||||
if not downloaded_files:
|
||||
raise FileNotFoundError("No files were downloaded")
|
||||
|
||||
download_file = downloaded_files[0]
|
||||
file_size = download_file.stat().st_size
|
||||
|
||||
self._color_log(f"File downloaded successfully to {download_file}", Color.green)
|
||||
|
||||
# Create result
|
||||
result = YoutubeDownloadResults(
|
||||
file_path=str(download_file),
|
||||
file_name=download_file.name,
|
||||
file_size=file_size,
|
||||
content_type="mp4",
|
||||
success=True,
|
||||
error=None,
|
||||
)
|
||||
|
||||
# Format output for LLM
|
||||
message = self._format_download_output(result, output_format)
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
# Create metadata
|
||||
metadata = YouTubeMetadata(
|
||||
operation="download",
|
||||
url=url,
|
||||
video_id=video_id,
|
||||
file_path=str(download_file),
|
||||
file_name=download_file.name,
|
||||
file_size=file_size,
|
||||
content_type="mp4",
|
||||
execution_time=execution_time,
|
||||
).model_dump()
|
||||
|
||||
return ActionResponse(success=True, message=message, metadata=metadata)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
self._color_log(f"Download error: {traceback.format_exc()}", Color.red)
|
||||
|
||||
# Format error for LLM
|
||||
message = f"Failed to download YouTube video: {error_msg}"
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
# Create metadata
|
||||
metadata = YouTubeMetadata(
|
||||
operation="download",
|
||||
url=url,
|
||||
error_type="download_failure",
|
||||
execution_time=execution_time,
|
||||
).model_dump()
|
||||
|
||||
return ActionResponse(success=False, message=message, metadata=metadata)
|
||||
|
||||
async def mcp_extract_youtube_transcript(
|
||||
self,
|
||||
video_id: str = Field(description="The YouTube video ID or URL to extract transcript from."),
|
||||
language_code: str = Field("en", description="Language code for the transcript (default: en)."),
|
||||
translate_to_language: str | None = Field(
|
||||
None, description="Translate transcript to this language code if provided."
|
||||
),
|
||||
output_format: str = Field(
|
||||
"markdown", description="Output format: 'markdown', 'json', or 'text' (default: markdown)."
|
||||
),
|
||||
) -> ActionResponse:
|
||||
"""Extract transcript from a YouTube video given its video ID or URL.
|
||||
|
||||
This tool provides transcript extraction with:
|
||||
- Support for multiple languages
|
||||
- Translation capabilities
|
||||
- URL or video ID input handling
|
||||
- LLM-optimized result formatting
|
||||
|
||||
Args:
|
||||
video_id: The YouTube video ID or URL to extract transcript from
|
||||
language_code: Language code for the transcript
|
||||
translate_to_language: Translate transcript to this language code if provided
|
||||
output_format: Format for the response output
|
||||
|
||||
Returns:
|
||||
ActionResponse with transcript data and metadata
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# Clean video_id if full URL was provided
|
||||
if "youtube.com" in video_id or "youtu.be" in video_id:
|
||||
if "?v=" in video_id:
|
||||
video_id = video_id.split("?v=")[-1].split("&")[0]
|
||||
elif "youtu.be/" in video_id:
|
||||
video_id = video_id.split("youtu.be/")[-1].split("?")[0]
|
||||
|
||||
self._color_log(f"Extracting transcript for video ID: {video_id}", Color.blue)
|
||||
|
||||
# Get transcript in specified language
|
||||
if translate_to_language:
|
||||
# Get transcript and translate it
|
||||
y_api = YouTubeTranscriptApi()
|
||||
transcript_list = y_api.list(video_id)
|
||||
transcript = None
|
||||
|
||||
try:
|
||||
# Try to get transcript in specified language
|
||||
transcript = transcript_list.find_transcript([language_code])
|
||||
except Exception:
|
||||
# If specified language not found, get any available transcript
|
||||
transcript = transcript_list.find_generated_transcript(["en"])
|
||||
|
||||
# Translate to target language
|
||||
transcript_data = transcript.translate(translate_to_language).fetch()
|
||||
|
||||
else:
|
||||
try:
|
||||
# Get transcript without translation
|
||||
transcript_data: FetchedTranscript = (
|
||||
YouTubeTranscriptApi()
|
||||
.list(video_id)
|
||||
.find_transcript((language_code,))
|
||||
.fetch(preserve_formatting=False)
|
||||
)
|
||||
except Exception:
|
||||
transcript_data = None
|
||||
|
||||
result = TranscriptResult(video_id=video_id, transcript=transcript_data, success=True, error=None)
|
||||
|
||||
self._color_log(f"Successfully extracted transcript for video ID: {video_id}", Color.green)
|
||||
|
||||
# Format output for LLM
|
||||
message = self._format_transcript_output(result, output_format)
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
# Create metadata
|
||||
metadata = YouTubeMetadata(
|
||||
operation="transcript",
|
||||
video_id=video_id,
|
||||
language_code=language_code,
|
||||
translate_to_language=translate_to_language,
|
||||
execution_time=execution_time,
|
||||
).model_dump()
|
||||
|
||||
return ActionResponse(success=True, message=message, metadata=metadata)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
self._color_log(f"Transcript extraction error: {traceback.format_exc()}", Color.red)
|
||||
|
||||
# Format error for LLM
|
||||
message = f"Failed to extract transcript: {error_msg}"
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
# Create metadata
|
||||
metadata = YouTubeMetadata(
|
||||
operation="transcript",
|
||||
video_id=video_id,
|
||||
language_code=language_code,
|
||||
translate_to_language=translate_to_language,
|
||||
error_type="transcript_failure",
|
||||
execution_time=execution_time,
|
||||
).model_dump()
|
||||
|
||||
return ActionResponse(success=False, message=message, metadata=metadata)
|
||||
|
||||
|
||||
# Default arguments for testing
|
||||
if __name__ == "__main__":
|
||||
load_dotenv()
|
||||
|
||||
arguments = ActionArguments(
|
||||
name="youtube_service",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
|
||||
)
|
||||
|
||||
try:
|
||||
youtube_service = YouTubeActionCollection(arguments)
|
||||
youtube_service.run()
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
@@ -0,0 +1,169 @@
|
||||
import os
|
||||
import tempfile
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import magic
|
||||
import requests
|
||||
|
||||
|
||||
def is_url(path_or_url: str) -> bool:
|
||||
"""
|
||||
Check if the given string is a URL.
|
||||
|
||||
Args:
|
||||
path_or_url: String to check
|
||||
|
||||
Returns:
|
||||
bool: True if the string is a URL, False otherwise
|
||||
"""
|
||||
parsed = urlparse(path_or_url)
|
||||
return bool(parsed.scheme and parsed.netloc)
|
||||
|
||||
|
||||
def get_mime_type(file_path: str, default_mime: str | None = None) -> str:
|
||||
"""
|
||||
Detect MIME type of a file using python-magic if available,
|
||||
otherwise fallback to extension-based detection.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
default_mime: Default MIME type to return if detection fails
|
||||
|
||||
Returns:
|
||||
str: Detected MIME type
|
||||
"""
|
||||
# Try using python-magic for accurate MIME type detection
|
||||
try:
|
||||
mime = magic.Magic(mime=True)
|
||||
return mime.from_file(file_path)
|
||||
except (AttributeError, IOError):
|
||||
# Fallback to extension-based detection
|
||||
extension_mime_map = {
|
||||
# Audio formats
|
||||
".mp3": "audio/mpeg",
|
||||
".wav": "audio/wav",
|
||||
".ogg": "audio/ogg",
|
||||
".m4a": "audio/mp4",
|
||||
".flac": "audio/flac",
|
||||
# Image formats
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".png": "image/png",
|
||||
".gif": "image/gif",
|
||||
".webp": "image/webp",
|
||||
".bmp": "image/bmp",
|
||||
".tiff": "image/tiff",
|
||||
# Video formats
|
||||
".mp4": "video/mp4",
|
||||
".avi": "video/x-msvideo",
|
||||
".mov": "video/quicktime",
|
||||
".mkv": "video/x-matroska",
|
||||
".webm": "video/webm",
|
||||
}
|
||||
|
||||
ext = Path(file_path).suffix.lower()
|
||||
return extension_mime_map.get(ext, default_mime or "application/octet-stream")
|
||||
|
||||
|
||||
def get_file_from_source(
|
||||
source: str,
|
||||
max_size_mb: float = 100.0,
|
||||
timeout: int = 60,
|
||||
) -> tuple[str, str, bytes]:
|
||||
"""
|
||||
Unified function to get file content from a URL or local path with validation.
|
||||
|
||||
Args:
|
||||
source: URL or local file path
|
||||
max_size_mb: Maximum allowed file size in MB
|
||||
timeout: Timeout for URL requests in seconds
|
||||
|
||||
Returns:
|
||||
Tuple[str, str, bytes]: (file_path, mime_type, file_content)
|
||||
- For URLs, file_path will be a temporary file path
|
||||
- For local files, file_path will be the original path
|
||||
|
||||
Raises:
|
||||
ValueError: When file doesn't exist, exceeds size limit, or has invalid MIME type
|
||||
IOError: When file cannot be read
|
||||
requests.RequestException: When URL request fails
|
||||
"""
|
||||
max_size_bytes = max_size_mb * 1024 * 1024
|
||||
|
||||
if is_url(source):
|
||||
# Handle URL source
|
||||
try:
|
||||
# Make a HEAD request first to check content length
|
||||
head_response = requests.head(source, timeout=timeout, allow_redirects=True)
|
||||
head_response.raise_for_status()
|
||||
|
||||
# Check content length if available
|
||||
content_length = head_response.headers.get("content-length")
|
||||
if content_length and int(content_length) > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({int(content_length) / (1024 * 1024):.2f} MB) "
|
||||
f"exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Download the file
|
||||
response = requests.get(source, timeout=timeout, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
# Read content with size checking
|
||||
content = b""
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
if len(content) + len(chunk) > max_size_bytes:
|
||||
raise ValueError(f"File size exceeds maximum allowed size ({max_size_mb} MB)")
|
||||
content += chunk
|
||||
|
||||
# Create temporary file
|
||||
parsed_url = urlparse(source)
|
||||
filename = os.path.basename(parsed_url.path) or "downloaded_file"
|
||||
|
||||
# Create temporary file with proper extension
|
||||
suffix = Path(filename).suffix or ".tmp"
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
|
||||
temp_file.write(content)
|
||||
temp_path = temp_file.name
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(temp_path)
|
||||
|
||||
return temp_path, mime_type, content
|
||||
|
||||
except requests.RequestException as e:
|
||||
raise requests.RequestException(f"Failed to download file from URL: {e}: {traceback.format_exc()}")
|
||||
except Exception as e:
|
||||
raise IOError(f"Error processing URL: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
else:
|
||||
# Handle local file path
|
||||
file_path = Path(source)
|
||||
|
||||
# Check if file exists
|
||||
if not file_path.exists():
|
||||
raise ValueError(f"File does not exist: {source}")
|
||||
|
||||
if not file_path.is_file():
|
||||
raise ValueError(f"Path is not a file: {source}")
|
||||
|
||||
# Check file size
|
||||
file_size = file_path.stat().st_size
|
||||
if file_size > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({file_size / (1024 * 1024):.2f} MB) exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Read file content
|
||||
try:
|
||||
with open(file_path, "rb") as f:
|
||||
content = f.read()
|
||||
except Exception as e:
|
||||
raise IOError(f"Cannot read file {source}: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(str(file_path))
|
||||
|
||||
return str(file_path), mime_type, content
|
||||
Reference in New Issue
Block a user