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,155 @@
|
||||
# BFCL Sample Synthesis using AWorld Runtime
|
||||
|
||||
This example demonstrates how to use AWorld to construct a runtime environment and synthesize function call samples for model training. The BFCL (Basic Function Call Learning) example shows how to create a virtual file system with MCP (Model Context Protocol) tools and generate training data from agent interactions.
|
||||
|
||||
## 📋 Overview
|
||||
|
||||
The BFCL example consists of:
|
||||
- **GorillaFileSystem**: A virtual file system with MCP tools
|
||||
- **Agent Runtime**: AWorld agent that interacts with the file system
|
||||
- **Function Call Synthesis**: Generation of training samples from agent trajectories
|
||||
|
||||
## 🏗️ Architecture
|
||||
|
||||
```
|
||||
┌─────────────────┐ ┌──────────────────┐ ┌─────────────────┐
|
||||
│ AWorld Agent │───▶│ GorillaFileSystem│───▶│ MCP Tools │
|
||||
│ │ │ (Virtual FS) │ │ (pwd, ls, cd, │
|
||||
│ - LLM Provider │ │ │ │ touch, echo, │
|
||||
│ - MCP Client │ │ - File/Directory │ │ cat, etc.) │
|
||||
│ - Trajectory │ │ - State Management│ │ │
|
||||
└─────────────────┘ └──────────────────┘ └─────────────────┘
|
||||
│ │ │
|
||||
▼ ▼ ▼
|
||||
┌─────────────────┐ ┌──────────────────┐ ┌─────────────────┐
|
||||
│ Trajectory │ │ File System │ │ Function Call │
|
||||
│ Collection │ │ Operations │ │ Samples │
|
||||
│ │ │ │ │ │
|
||||
│ - Agent Actions │ │ - Create/Read/ │ │ - Tool Calls │
|
||||
│ - Tool Calls │ │ Write Files │ │ - Parameters │
|
||||
│ - Results │ │ - Directory │ │ - Results │
|
||||
│ │ │ Navigation │ │ - Context │
|
||||
└─────────────────┘ └──────────────────┘ └─────────────────┘
|
||||
```
|
||||
|
||||
## 🚀 Quick Start
|
||||
|
||||
### 1. Environment Setup
|
||||
|
||||
```bash
|
||||
# Set your OpenRouter API key
|
||||
export OPENROUTER_API_KEY="your-api-key-here"
|
||||
```
|
||||
|
||||
### 2. Run the Example
|
||||
|
||||
```bash
|
||||
# Navigate to the BFCL example directory
|
||||
cd examples/BFCL
|
||||
|
||||
# Run the BFCL agent example
|
||||
python run.py
|
||||
```
|
||||
|
||||
### 3. Expected Output
|
||||
|
||||
The agent will:
|
||||
1. Connect to the GorillaFileSystem MCP server
|
||||
2. Perform file operations (create, read, write files)
|
||||
3. Generate trajectory data with function calls
|
||||
4. Display the results
|
||||
|
||||
## 📁 File Structure
|
||||
|
||||
```
|
||||
examples/BFCL/
|
||||
├── README.md # This file
|
||||
├── run.py # Main agent runner
|
||||
├── mcp_tools/
|
||||
│ ├── __init__.py # Package initialization
|
||||
│ ├── gorilla_file_system.py # Virtual file system
|
||||
│ └── test_server.py # Function testing
|
||||
└── requirements.txt # Dependencies
|
||||
```
|
||||
|
||||
## 🔧 Core Components
|
||||
|
||||
### 1. Agent Configuration (`run.py`)
|
||||
|
||||
```python
|
||||
# Environment-based API key configuration
|
||||
api_key = os.getenv('OPENROUTER_API_KEY')
|
||||
|
||||
agent_config = AgentConfig(
|
||||
llm_provider="openai",
|
||||
llm_model_name="openai/gpt-4o",
|
||||
llm_api_key=api_key,
|
||||
llm_base_url="https://openrouter.ai/api/v1"
|
||||
)
|
||||
```
|
||||
|
||||
### 2. MCP Server Configuration
|
||||
|
||||
```python
|
||||
mcp_config = {
|
||||
"mcpServers": {
|
||||
"GorillaFileSystem": {
|
||||
"type": "stdio",
|
||||
"command": "python",
|
||||
"args": ["mcp_tools/gorilla_file_system.py"],
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### 3. Agent Creation
|
||||
|
||||
```python
|
||||
file_sys_prompt = "You are a helpful agent to use the standard file system..."
|
||||
file_sys = Agent(
|
||||
conf=agent_config,
|
||||
name="file_sys_agent",
|
||||
system_prompt=file_sys_prompt,
|
||||
mcp_servers=mcp_config.get("mcpServers", []).keys(),
|
||||
mcp_config=mcp_config,
|
||||
)
|
||||
```
|
||||
|
||||
### 4. Trajectory Collection
|
||||
|
||||
```python
|
||||
result = Runners.sync_run(
|
||||
input="use mcp tools to perform file operations...",
|
||||
agent=file_sys,
|
||||
)
|
||||
|
||||
print("=" * 100)
|
||||
print(f"result.answer: {result.answer}")
|
||||
print("=" * 100)
|
||||
print(f"result.trajectory: {json.dumps(result.trajectory[0], indent=4)}")
|
||||
```
|
||||
|
||||
## 🛠️ MCP Tools (GorillaFileSystem)
|
||||
|
||||
The virtual file system provides the following MCP tools:
|
||||
|
||||
### File Operations
|
||||
- `mcp_touch(file_name)`: Create a new file
|
||||
- `mcp_echo(content, file_name)`: Write content to file
|
||||
- `mcp_cat(file_name)`: Read file content
|
||||
- `mcp_rm(file_name)`: Remove file
|
||||
|
||||
### Directory Operations
|
||||
- `mcp_pwd()`: Get current directory
|
||||
- `mcp_ls(a=False)`: List directory contents
|
||||
- `mcp_cd(folder)`: Change directory
|
||||
- `mcp_mkdir(dir_name)`: Create directory
|
||||
- `mcp_rmdir(dir_name)`: Remove directory
|
||||
|
||||
### Advanced Operations
|
||||
- `mcp_find(path, name)`: Search for files
|
||||
- `mcp_wc(file_name, mode)`: Word count
|
||||
- `mcp_sort(file_name)`: Sort file content
|
||||
- `mcp_grep(file_name, pattern)`: Search in file
|
||||
- `mcp_mv(source, destination)`: Move/rename
|
||||
- `mcp_cp(source, destination)`: Copy files
|
||||
@@ -0,0 +1,2 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
@@ -0,0 +1,14 @@
|
||||
"""
|
||||
MCP Tools Package
|
||||
|
||||
This package contains MCP (Model Context Protocol) tools and servers
|
||||
for various functionalities including file system operations.
|
||||
"""
|
||||
|
||||
from .gorilla_file_system import GorillaFileSystem
|
||||
|
||||
__all__ = [
|
||||
"GorillaFileSystem"
|
||||
]
|
||||
|
||||
__version__ = "1.0.0"
|
||||
@@ -0,0 +1,74 @@
|
||||
{
|
||||
"root": {
|
||||
"workspace": {
|
||||
"type": "directory",
|
||||
"contents": {
|
||||
"src": {
|
||||
"type": "directory",
|
||||
"contents": {
|
||||
"main.py": {
|
||||
"type": "file",
|
||||
"content": "#!/usr/bin/env python3\n\ndef main():\n print('Hello from GorillaFileSystem!')\n return 0\n\nif __name__ == '__main__':\n exit(main())"
|
||||
},
|
||||
"utils.py": {
|
||||
"type": "file",
|
||||
"content": "# Utility functions\n\ndef format_bytes(bytes_value):\n \"\"\"Format bytes into human readable format.\"\"\"\n for unit in ['B', 'KB', 'MB', 'GB']:\n if bytes_value < 1024.0:\n return f\"{bytes_value:.1f} {unit}\"\n bytes_value /= 1024.0\n return f\"{bytes_value:.1f} TB\"\n\ndef count_lines(text):\n \"\"\"Count lines in text.\"\"\"\n return len(text.splitlines())"
|
||||
},
|
||||
"config": {
|
||||
"type": "directory",
|
||||
"contents": {
|
||||
"settings.json": {
|
||||
"type": "file",
|
||||
"content": "{\n \"debug\": true,\n \"log_level\": \"INFO\",\n \"max_file_size\": 1048576,\n \"allowed_extensions\": [\".py\", \".txt\", \".md\", \".json\"]\n}"
|
||||
},
|
||||
"database.conf": {
|
||||
"type": "file",
|
||||
"content": "# Database configuration\nhost=localhost\nport=5432\ndatabase=testdb\nusername=user\npassword=secret"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"docs": {
|
||||
"type": "directory",
|
||||
"contents": {
|
||||
"README.md": {
|
||||
"type": "file",
|
||||
"content": "# Project Documentation\n\nThis is a sample project demonstrating file system operations.\n\n## Features\n\n- File creation and manipulation\n- Directory navigation\n- Text processing\n- Configuration management\n\n## Usage\n\nUse the various file system tools to explore and modify the project structure."
|
||||
},
|
||||
"API.md": {
|
||||
"type": "file",
|
||||
"content": "# API Documentation\n\n## File Operations\n\n### Creating Files\n- `touch filename.ext` - Create empty file\n- `echo \"content\" filename.ext` - Create file with content\n\n### Reading Files\n- `cat filename.ext` - Display file contents\n- `tail filename.ext` - Show last 10 lines\n\n### File Information\n- `wc filename.ext` - Count lines, words, characters\n- `du` - Show disk usage"
|
||||
}
|
||||
}
|
||||
},
|
||||
"tests": {
|
||||
"type": "directory",
|
||||
"contents": {
|
||||
"test_main.py": {
|
||||
"type": "file",
|
||||
"content": "import unittest\nfrom src.main import main\nfrom src.utils import format_bytes, count_lines\n\nclass TestMain(unittest.TestCase):\n def test_main_returns_zero(self):\n self.assertEqual(main(), 0)\n \n def test_format_bytes(self):\n self.assertEqual(format_bytes(1024), \"1.0 KB\")\n self.assertEqual(format_bytes(1048576), \"1.0 MB\")\n \n def test_count_lines(self):\n self.assertEqual(count_lines(\"line1\\nline2\\nline3\"), 3)\n\nif __name__ == '__main__':\n unittest.main()"
|
||||
}
|
||||
}
|
||||
},
|
||||
"data": {
|
||||
"type": "directory",
|
||||
"contents": {
|
||||
"sample.csv": {
|
||||
"type": "file",
|
||||
"content": "id,name,email,age\n1,Alice Johnson,alice@example.com,28\n2,Bob Smith,bob@example.com,34\n3,Carol Davis,carol@example.com,29\n4,David Wilson,david@example.com,42\n5,Eve Brown,eve@example.com,31"
|
||||
},
|
||||
"log.txt": {
|
||||
"type": "file",
|
||||
"content": "2024-01-01 10:00:00 INFO Application started\n2024-01-01 10:00:01 INFO Loading configuration\n2024-01-01 10:00:02 INFO Database connection established\n2024-01-01 10:00:03 INFO Server listening on port 4242\n2024-01-01 10:00:04 ERROR Failed to connect to external API\n2024-01-01 10:00:05 WARN Retrying API connection\n2024-01-01 10:00:06 INFO API connection restored\n2024-01-01 10:00:07 INFO Application ready"
|
||||
}
|
||||
}
|
||||
},
|
||||
"temp": {
|
||||
"type": "directory",
|
||||
"contents": {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,952 @@
|
||||
import datetime
|
||||
from copy import deepcopy
|
||||
from typing import Dict, List, Optional, Union
|
||||
from pydantic import Field
|
||||
from aworld.logs.util import Color
|
||||
import traceback
|
||||
|
||||
from examples.BFCL.mcp_tools.tool_base import ActionArguments, ActionCollection
|
||||
|
||||
FILE_CONTENT_EXTENSION = "The company's financials for the year reflect a period of steady growth and consistent revenue generation, with both top-line and bottom-line figures showing improvement compared to the previous year. Total revenue increased at a modest pace, driven primarily by strong performance in the company’s core markets. Despite some fluctuations in demand, the business maintained healthy margins, with cost controls and efficiency measures helping to offset any increase in operational expenses. As a result, gross profit grew at a stable rate, keeping in line with management’s expectations. The company’s operating income saw an uptick, indicating that the firm was able to manage its administrative and selling expenses effectively, while also benefiting from a more streamlined supply chain. This contributed to a higher operating margin, suggesting that the company’s core operations were becoming more efficient and profitable. Net income also rose, bolstered by favorable tax conditions and reduced interest expenses due to a restructuring of long-term debt. The company managed to reduce its financial leverage, leading to an improvement in its interest coverage ratio. On the balance sheet, the company maintained a solid financial position, with total assets increasing year over year. The growth in assets was largely due to strategic investments in new technology and facilities, aimed at expanding production capacity and improving operational efficiency. Cash reserves remained robust, supported by positive cash flow from operations. The company also reduced its short-term liabilities, improving its liquidity ratios, and signaling a stronger ability to meet near-term obligations.Shareholders’ equity grew as a result of retained earnings, reflecting the company’s profitability and its strategy of reinvesting profits back into the business rather than paying out large dividends. The company maintained a conservative approach to debt, with its debt-to-equity ratio remaining within industry norms, which reassured investors about the company’s long-term solvency and risk management practices. The cash flow statement highlighted the company’s ability to generate cash from its core operations, which remained a strong indicator of the business's health. Cash from operating activities was sufficient to cover both investing and financing needs, allowing the company to continue its capital expenditure plans without increasing its reliance on external financing. The company’s investment activities included expanding its production facilities and acquiring new technology to improve future productivity and efficiency. Meanwhile, the company’s financing activities reflected a balanced approach, with some debt repayments and a modest issuance of new equity, allowing for flexible capital management.Overall, the company's financials indicate a well-managed business with a clear focus on sustainable growth. Profitability remains strong, operational efficiency is improving, and the company’s balance sheet reflects a stable, low-risk financial structure. The management’s strategy of cautious expansion, combined with a disciplined approach to debt and investment, has positioned the company well for future growth and profitability."
|
||||
FILES_TAIL_USED = ['log.txt', 'report.txt', 'report.csv', 'DataSet1.csv', 'file1.txt', 'finance_report.txt', 'config.py', 'Q4_summary.doc', 'file3.txt']
|
||||
POPULATE_FILE_EXTENSION = ['image_344822349461074042.jpg', 'image_8219547643081662353.jpg', 'image_5421509146842474663.jpg', 'image_185391401034246046.jpg', 'image_6824007961180780019.jpg', 'image_2994974694593273051.jpg', 'image_2537728455072851196.jpg', 'image_2164918946836800275.jpg', 'image_1745133864906284051.jpg', 'image_7707563551789432679.jpg', 'image_8190489168166590809.jpg', 'image_2385660725381355820.jpg', 'image_4771211633166048374.jpg', 'image_3443718094055823214.jpg', 'image_6838087561356843690.jpg','image_605952633285970710.jpg', 'image_6341510244180179744.jpg', 'image_4119241148692325954.jpg', 'image_5651066601163181955.jpg', 'image_3747091333751395055.jpg', 'image_4623743619379194431.jpg', 'image_5072742684386583099.jpg', 'image_1978458056362464778.jpg', 'image_3090346927968358019.jpg', 'image_7193806748674265039.jpg', 'image_7169516574395086720.jpg', 'image_8618240224293913315.jpg', 'image_5514683852355062444.jpg', 'image_8749630317332649147.jpg', 'image_1912245706439755759.jpg']
|
||||
|
||||
class File:
|
||||
|
||||
def __init__(self, name: str, content: str = "") -> None:
|
||||
"""
|
||||
Initialize a file with a name and optional content.
|
||||
|
||||
Args:
|
||||
name (str): The name of the file.
|
||||
content (str, optional): The initial content of the file. Defaults to an empty string.
|
||||
"""
|
||||
self.name: str = name
|
||||
self.content: str = content
|
||||
self._last_modified: datetime.datetime = datetime.datetime.now()
|
||||
|
||||
def _write(self, new_content: str) -> None:
|
||||
"""
|
||||
Write new content to the file and update the last modified time.
|
||||
|
||||
Args:
|
||||
new_content (str): The new content to write to the file.
|
||||
"""
|
||||
self.content = new_content
|
||||
self._last_modified = datetime.datetime.now()
|
||||
|
||||
def _read(self) -> str:
|
||||
"""
|
||||
Read the content of the file.
|
||||
|
||||
Returns:
|
||||
content (str): The current content of the file.
|
||||
"""
|
||||
return self.content
|
||||
|
||||
def _append(self, additional_content: str) -> None:
|
||||
"""
|
||||
Append content to the existing file content.
|
||||
|
||||
Args:
|
||||
additional_content (str): The content to append to the file.
|
||||
"""
|
||||
self.content += additional_content
|
||||
self._last_modified = datetime.datetime.now()
|
||||
|
||||
def __repr__(self):
|
||||
return f"<<File: {self.name}, Content: {self.content}>>"
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if not isinstance(other, File):
|
||||
return False
|
||||
return self.name == other.name and self.content == other.content
|
||||
|
||||
|
||||
class Directory:
|
||||
|
||||
def __init__(self, name: str, parent: Optional["Directory"] = None) -> None:
|
||||
"""
|
||||
Initialize a directory with a name.
|
||||
|
||||
Args:
|
||||
name (str): The name of the directory.
|
||||
"""
|
||||
self.name: str = name
|
||||
self.parent: Optional["Directory"] = parent
|
||||
self.contents: Dict[str, Union["File", "Directory"]] = {}
|
||||
|
||||
def _add_file(self, file_name: str, content: str = "") -> None:
|
||||
"""
|
||||
Add a new file to the directory.
|
||||
|
||||
Args:
|
||||
file_name (str): The name of the file.
|
||||
content (str, optional): The content of the new file. Defaults to an empty string.
|
||||
"""
|
||||
if file_name in self.contents:
|
||||
raise ValueError(
|
||||
f"File '{file_name}' already exists in directory '{self.name}'."
|
||||
)
|
||||
new_file = File(file_name, content)
|
||||
self.contents[file_name] = new_file
|
||||
|
||||
def _add_directory(self, dir_name: str) -> None:
|
||||
"""
|
||||
Add a new subdirectory to the directory.
|
||||
|
||||
Args:
|
||||
dir_name (str): The name of the subdirectory.
|
||||
"""
|
||||
if dir_name in self.contents:
|
||||
raise ValueError(
|
||||
f"Directory '{dir_name}' already exists in directory '{self.name}'."
|
||||
)
|
||||
new_dir = Directory(dir_name, self)
|
||||
self.contents[dir_name] = new_dir
|
||||
|
||||
def _get_item(self, item_name: str) -> Union["File", "Directory", None]:
|
||||
"""
|
||||
Get an item (file or subdirectory) from the directory.
|
||||
|
||||
Args:
|
||||
item_name (str): The name of the item to retrieve.
|
||||
|
||||
Returns:
|
||||
item (any): The retrieved item or None if it does not exist.
|
||||
"""
|
||||
if item_name == ".":
|
||||
return self
|
||||
return self.contents.get(item_name)
|
||||
|
||||
def _list_contents(self) -> List[str]:
|
||||
"""
|
||||
List the names of all contents in the directory.
|
||||
|
||||
Returns:
|
||||
contents (List[str]): A list of names of the files and subdirectories in the directory.
|
||||
"""
|
||||
return list(self.contents.keys())
|
||||
|
||||
def __repr__(self):
|
||||
return f"<Directory: {self.name}, Parent: {self.parent.name if self.parent else None}, Contents: {self.contents}>"
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if not isinstance(other, Directory):
|
||||
return False
|
||||
return self.name == other.name and self.contents == other.contents
|
||||
|
||||
DEFAULT_STATE = {"root": Directory("/", None)}
|
||||
|
||||
class GorillaFileSystem(ActionCollection):
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
"""
|
||||
Initialize the Gorilla file system with a root directory
|
||||
"""
|
||||
super().__init__(arguments)
|
||||
self._color_log("GorillaFileSystem service initialized", Color.green, "debug")
|
||||
|
||||
self.root: Directory
|
||||
self._current_dir: Directory
|
||||
self._api_description = "This tool belongs to the Gorilla file system. It is a simple file system that allows users to perform basic file operations such as navigating directories, creating files and directories, reading and writing to files, etc."
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if not isinstance(other, GorillaFileSystem):
|
||||
return False
|
||||
return self.root == other.root
|
||||
|
||||
def _load_scenario(self, scenario: dict, long_context: bool = False) -> None:
|
||||
"""
|
||||
Load a scenario into the file system.
|
||||
|
||||
Args:
|
||||
scenario (dict): The scenario to load.
|
||||
|
||||
The scenario always starts with a root directory. Each directory can contain files or subdirectories.
|
||||
The key is the name of the file or directory, and the value is a dictionary with the following keys
|
||||
An example scenario:
|
||||
Here John is the root directory and it contains a home directory with a user directory inside it.
|
||||
The user directory contains a file named file1.txt and a directory named directory1.
|
||||
Root is not a part of the scenario and it's just easy for parsing. During generation, you should have at most 2 layers.
|
||||
{
|
||||
"root": {
|
||||
"john": {
|
||||
"type": "directory",
|
||||
"contents": {
|
||||
"home": {
|
||||
"type": "directory",
|
||||
"contents": {
|
||||
"user": {
|
||||
"type": "directory",
|
||||
"contents": {
|
||||
"file1.txt": {
|
||||
"type": "file",
|
||||
"content": "Hello, world!"
|
||||
},
|
||||
"directory1": {
|
||||
"type": "directory",
|
||||
"contents": {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
"""
|
||||
DEFAULT_STATE_COPY = deepcopy(DEFAULT_STATE)
|
||||
self.long_context = long_context
|
||||
self.root = DEFAULT_STATE_COPY["root"]
|
||||
if "root" in scenario:
|
||||
root_dir = Directory(list(scenario["root"].keys())[0], None)
|
||||
self.root = self._load_directory(
|
||||
scenario["root"][list(scenario["root"].keys())[0]]["contents"], root_dir
|
||||
)
|
||||
self._current_dir = self.root
|
||||
|
||||
def _load_directory(
|
||||
self, current: dict, parent: Optional[Directory] = None
|
||||
) -> Directory:
|
||||
"""
|
||||
Load a directory and its contents from a dictionary.
|
||||
|
||||
Args:
|
||||
data (dict): The dictionary representing the directory.
|
||||
parent (Directory, optional): The parent directory. Defaults to None.
|
||||
|
||||
Returns:
|
||||
Directory: The loaded directory.
|
||||
"""
|
||||
is_bottommost = True
|
||||
for dir_name, dir_data in current.items():
|
||||
|
||||
if dir_data["type"] == "directory":
|
||||
is_bottommost = False
|
||||
new_dir = Directory(dir_name, parent)
|
||||
new_dir = self._load_directory(dir_data["contents"], new_dir)
|
||||
parent.contents[dir_name] = new_dir
|
||||
|
||||
elif dir_data["type"] == "file":
|
||||
content = dir_data["content"]
|
||||
if self.long_context and dir_name not in FILES_TAIL_USED:
|
||||
content += FILE_CONTENT_EXTENSION
|
||||
new_file = File(dir_name, content)
|
||||
parent.contents[dir_name] = new_file
|
||||
|
||||
if is_bottommost and self.long_context:
|
||||
self._populate_directory(parent)
|
||||
|
||||
return parent
|
||||
|
||||
def _populate_directory(
|
||||
self, directory: Directory
|
||||
) -> None: # Used only for long context
|
||||
"""
|
||||
Populate an innermost directory with multiple empty files.
|
||||
|
||||
Args:
|
||||
directory (Directory): The innermost directory to populate.
|
||||
"""
|
||||
for i in range(len(POPULATE_FILE_EXTENSION)):
|
||||
name = POPULATE_FILE_EXTENSION[i]
|
||||
file_name = f"{name}"
|
||||
directory._add_file(file_name)
|
||||
|
||||
def mcp_pwd(self) -> Dict[str, str]:
|
||||
"""
|
||||
This tool belongs to the Gorilla file system. It is a simple file system that allows users to perform basic file operations such as navigating directories, creating files and directories, reading and writing to files, etc. Tool description: Return the current working directory path.
|
||||
|
||||
Args:
|
||||
None: This function takes no parameters.
|
||||
|
||||
Returns:
|
||||
Dict[str, str]: A dictionary containing:
|
||||
- current_working_directory (str): The current working directory path.
|
||||
|
||||
"""
|
||||
path = []
|
||||
dir = self._current_dir
|
||||
while dir is not None and dir.name != self.root:
|
||||
path.append(dir.name)
|
||||
dir = dir.parent
|
||||
return {"current_working_directory": "/" + "/".join(reversed(path))}
|
||||
|
||||
def mcp_ls(self, a: bool = Field(False, description="Show hidden files and directories. Defaults to False.")) -> Dict[str, List[str]]:
|
||||
"""
|
||||
This tool belongs to the Gorilla file system. It is a simple file system that allows users to perform basic file operations such as navigating directories, creating files and directories, reading and writing to files, etc. Tool description: List the contents of the current directory.
|
||||
|
||||
Args:
|
||||
a (bool): Show hidden files and directories. If True, includes files and directories that start with '.'. Defaults to False.
|
||||
|
||||
Returns:
|
||||
Dict[str, List[str]]: A dictionary containing:
|
||||
- current_directory_content (List[str]): A list of names of files and directories in the current directory.
|
||||
"""
|
||||
contents = self._current_dir._list_contents()
|
||||
if not a:
|
||||
contents = [item for item in contents if not item.startswith(".")]
|
||||
return {"current_directory_content": contents}
|
||||
|
||||
def mcp_cd(self, folder: str = Field(..., description="The folder of the directory to change to. You can only change one folder at a time.")) -> Union[None, Dict[str, str]]:
|
||||
"""
|
||||
This tool belongs to the Gorilla file system. It is a simple file system that allows users to perform basic file operations such as navigating directories, creating files and directories, reading and writing to files, etc. Tool description: Change the current working directory to the specified folder.
|
||||
|
||||
Args:
|
||||
folder (str): The folder of the directory to change to. You can only change one folder at a time. Use ".." to go to parent directory.
|
||||
|
||||
Returns:
|
||||
Union[None, Dict[str, str]]:
|
||||
- None: If operation successful (when using "..")
|
||||
- Dict[str, str]: Contains either:
|
||||
- current_working_directory (str): The new current working directory name on success
|
||||
- error (str): Error message if directory does not exist or operation fails
|
||||
"""
|
||||
# Handle navigating to the parent directory with "cd .."
|
||||
if folder == "..":
|
||||
if self._current_dir.parent:
|
||||
self._current_dir = self._current_dir.parent
|
||||
elif self.root == self._current_dir:
|
||||
return {"error": "Cuurent directory is already the root. Cannot go back."}
|
||||
else:
|
||||
return {"error": "cd: ..: No such directory"}
|
||||
return {}
|
||||
|
||||
# Handle absolute or relative paths
|
||||
target_dir = self._navigate_to_directory(folder)
|
||||
if isinstance(target_dir, dict): # Error condition check
|
||||
return {
|
||||
"error": f"cd: {folder}: No such directory. You cannot use path to change directory."
|
||||
}
|
||||
self._current_dir = target_dir
|
||||
return {"current_working_directory": target_dir.name}
|
||||
|
||||
def _validate_file_or_directory_name(self, dir_name: str) -> bool:
|
||||
if any(c in dir_name for c in '|/\\?%*:"><'):
|
||||
return False
|
||||
return True
|
||||
|
||||
def mcp_mkdir(self, dir_name: str = Field(..., description="The name of the new directory at current directory. You can only create directory at current directory.")) -> Union[None, Dict[str, str]]:
|
||||
"""
|
||||
This tool belongs to the Gorilla file system. It is a simple file system that allows users to perform basic file operations such as navigating directories, creating files and directories, reading and writing to files, etc. Tool description: Create a new directory in the current directory.
|
||||
|
||||
Args:
|
||||
dir_name (str): The name of the new directory to create in the current directory. Directory name cannot contain special characters like |/\\?%*:"><.
|
||||
|
||||
Returns:
|
||||
Union[None, Dict[str, str]]:
|
||||
- None: If directory creation is successful
|
||||
- Dict[str, str]: Contains error message if operation fails:
|
||||
- error (str): Description of why the directory creation failed
|
||||
"""
|
||||
if not self._validate_file_or_directory_name(dir_name):
|
||||
return {
|
||||
"error": f"mkdir: cannot create directory '{dir_name}': Invalid character"
|
||||
}
|
||||
if dir_name in self._current_dir.contents:
|
||||
return {"error": f"mkdir: cannot create directory '{dir_name}': File exists"}
|
||||
|
||||
self._current_dir._add_directory(dir_name)
|
||||
return None
|
||||
|
||||
def mcp_touch(self, file_name: str = Field(..., description="The name of the new file in the current directory. file_name is local to the current directory and does not allow path.")) -> Union[None, Dict[str, str]]:
|
||||
"""
|
||||
This tool belongs to the Gorilla file system. It is a simple file system that allows users to perform basic file operations such as navigating directories, creating files and directories, reading and writing to files, etc. Tool description: Create a new file of any extension in the current directory.
|
||||
|
||||
Args:
|
||||
file_name (str): The name of the new file to create in the current directory. File name is local to current directory and does not allow path. Cannot contain special characters.
|
||||
|
||||
Returns:
|
||||
Union[None, Dict[str, str]]:
|
||||
- None: If file creation is successful
|
||||
- Dict[str, str]: Contains error message if operation fails:
|
||||
- error (str): Description of why the file creation failed
|
||||
"""
|
||||
if not self._validate_file_or_directory_name(file_name):
|
||||
return {"error": f"touch: cannot touch '{file_name}': Invalid character"}
|
||||
|
||||
if file_name in self._current_dir.contents:
|
||||
return {"error": f"touch: cannot touch '{file_name}': File exists"}
|
||||
|
||||
self._current_dir._add_file(file_name)
|
||||
return None
|
||||
|
||||
def mcp_echo(
|
||||
self,
|
||||
content: str = Field(..., description="The content to write or display."),
|
||||
file_name: Optional[str] = Field(None, description="The name of the file at current directory to write the content to. Defaults to None.")
|
||||
) -> Union[Dict[str, str], None]:
|
||||
"""
|
||||
This tool belongs to the Gorilla file system. It is a simple file system that allows users to perform basic file operations such as navigating directories, creating files and directories, reading and writing to files, etc. Tool description: Write content to a file at current directory or display it in the terminal.
|
||||
|
||||
Args:
|
||||
content (str): The content to write to a file or display in terminal.
|
||||
file_name (Optional[str]): The name of the file at current directory to write content to. If None, content is displayed in terminal. Defaults to None.
|
||||
|
||||
Returns:
|
||||
Union[Dict[str, str], None]:
|
||||
- Dict[str, str]: When file_name is None, returns:
|
||||
- terminal_output (str): The content displayed in terminal
|
||||
- Dict[str, str]: When operation fails, returns:
|
||||
- error (str): Error message describing the failure
|
||||
- None: When content is successfully written to file
|
||||
"""
|
||||
if file_name is None:
|
||||
return {"terminal_output": content}
|
||||
if not self._validate_file_or_directory_name(file_name):
|
||||
return {"error": f"echo: cannot touch '{file_name}': Invalid character"}
|
||||
|
||||
if file_name:
|
||||
if file_name in self._current_dir.contents:
|
||||
self._current_dir._get_item(file_name)._write(content)
|
||||
else:
|
||||
self._current_dir._add_file(file_name, content)
|
||||
else:
|
||||
return {"terminal_output": content}
|
||||
|
||||
def mcp_cat(self, file_name: str = Field(..., description="The name of the file from current directory to display. No path is allowed.")) -> Dict[str, str]:
|
||||
"""
|
||||
This tool belongs to the Gorilla file system. It is a simple file system that allows users to perform basic file operations such as navigating directories, creating files and directories, reading and writing to files, etc. Tool description: Display the contents of a file of any extension from currrent directory.
|
||||
|
||||
Args:
|
||||
file_name (str): The name of the file from current directory to display contents. No path is allowed, only local file names.
|
||||
|
||||
Returns:
|
||||
Dict[str, str]: A dictionary containing either:
|
||||
- file_content (str): The complete content of the file
|
||||
- error (str): Error message if file doesn't exist or is a directory
|
||||
"""
|
||||
if not self._validate_file_or_directory_name(file_name):
|
||||
return {"error": f"cat: '{file_name}': Invalid character"}
|
||||
|
||||
if file_name in self._current_dir.contents:
|
||||
item = self._current_dir._get_item(file_name)
|
||||
if isinstance(item, File):
|
||||
return {"file_content": item._read()}
|
||||
else:
|
||||
return {"error": f"cat: {file_name}: Is a directory"}
|
||||
else:
|
||||
return {"error": f"cat: {file_name}: No such file or directory"}
|
||||
|
||||
def mcp_find(
|
||||
self,
|
||||
path: str = Field("", description="The directory path to start the search. Defaults to the current directory (\".\")."),
|
||||
name: Optional[str] = Field(None, description="The name of the file or directory to search for. If None, all items are returned.")
|
||||
) -> Dict[str, List[str]]:
|
||||
"""
|
||||
This tool belongs to the Gorilla file system. It is a simple file system that allows users to perform basic file operations such as navigating directories, creating files and directories, reading and writing to files, etc. Tool description: Find any file or directories under specific path that contain name in its file name.
|
||||
|
||||
This method searches for files of any extension and directories within a specified path that match
|
||||
the given name. If no name is provided, it returns all files and directories
|
||||
in the specified path and its subdirectories.
|
||||
Note: This method performs a recursive search through all subdirectories of the given path.
|
||||
|
||||
Args:
|
||||
path (str): The directory path to start the search from. Defaults to current directory ("."). Searches recursively through subdirectories.
|
||||
name (Optional[str]): The name pattern to search for in file and directory names. If None, returns all items found. Performs substring matching.
|
||||
|
||||
Returns:
|
||||
Dict[str, List[str]]: A dictionary containing:
|
||||
- matches (List[str]): List of matching file and directory paths relative to the given path, including subdirectory contents
|
||||
"""
|
||||
matches = []
|
||||
target_dir = self._current_dir
|
||||
|
||||
def recursive_search(directory: Directory, base_path: str) -> None:
|
||||
for item_name, item in directory.contents.items():
|
||||
item_path = f"{base_path}/{item_name}"
|
||||
if name is None or name in item_name:
|
||||
matches.append(item_path)
|
||||
if isinstance(item, Directory):
|
||||
recursive_search(item, item_path)
|
||||
|
||||
recursive_search(target_dir, path.rstrip("/"))
|
||||
return {"matches": matches}
|
||||
|
||||
def mcp_wc(
|
||||
self,
|
||||
file_name: str = Field(..., description="Name of the file of current directory to perform wc operation on."),
|
||||
mode: str = Field("l", description="Mode of operation ('l' for lines, 'w' for words, 'c' for characters).")
|
||||
) -> Dict[str, Union[int, str]]:
|
||||
"""
|
||||
This tool belongs to the Gorilla file system. It is a simple file system that allows users to perform basic file operations such as navigating directories, creating files and directories, reading and writing to files, etc. Tool description: Count the number of lines, words, and characters in a file of any extension from current directory.
|
||||
|
||||
Args:
|
||||
file_name (str): Name of the file in current directory to perform word count operation on.
|
||||
mode (str): Mode of operation - 'l' for counting lines, 'w' for counting words, 'c' for counting characters. Defaults to 'l'.
|
||||
|
||||
Returns:
|
||||
Dict[str, Union[int, str]]: A dictionary containing either:
|
||||
- count (int): The count of lines, words, or characters in the file
|
||||
- type (str): The type of unit being counted - "lines", "words", or "characters"
|
||||
- error (str): Error message if file doesn't exist or invalid mode specified
|
||||
"""
|
||||
if mode not in ["l", "w", "c"]:
|
||||
return {"error": f"wc: invalid mode '{mode}'"}
|
||||
|
||||
if file_name in self._current_dir.contents:
|
||||
file = self._current_dir._get_item(file_name)
|
||||
if isinstance(file, File):
|
||||
content = file._read()
|
||||
|
||||
if mode == "l":
|
||||
line_count = len(content.splitlines())
|
||||
return {"count": line_count, "type": "lines"}
|
||||
|
||||
elif mode == "w":
|
||||
word_count = len(content.split())
|
||||
return {"count": word_count, "type": "words"}
|
||||
|
||||
elif mode == "c":
|
||||
char_count = len(content)
|
||||
return {"count": char_count, "type": "characters"}
|
||||
|
||||
return {"error": f"wc: {file_name}: No such file or directory"}
|
||||
|
||||
def mcp_sort(self, file_name: str = Field(..., description="The name of the file appeared at current directory to sort.")) -> Dict[str, str]:
|
||||
"""
|
||||
This tool belongs to the Gorilla file system. It is a simple file system that allows users to perform basic file operations such as navigating directories, creating files and directories, reading and writing to files, etc. Tool description: Sort the contents of a file line by line.
|
||||
|
||||
Args:
|
||||
file_name (str): The name of the file in current directory to sort. Sorts content line by line alphabetically.
|
||||
|
||||
Returns:
|
||||
Dict[str, str]: A dictionary containing either:
|
||||
- sorted_content (str): The file content with lines sorted alphabetically
|
||||
- error (str): Error message if file doesn't exist in current directory
|
||||
"""
|
||||
if file_name in self._current_dir.contents:
|
||||
file = self._current_dir._get_item(file_name)
|
||||
if isinstance(file, File):
|
||||
content = file._read()
|
||||
|
||||
sorted_content = "\n".join(sorted(content.splitlines()))
|
||||
|
||||
return {"sorted_content": sorted_content}
|
||||
|
||||
return {"error": f"sort: {file_name}: No such file or directory"}
|
||||
|
||||
def mcp_grep(
|
||||
self,
|
||||
file_name: str = Field(..., description="The name of the file to search. No path is allowed and you can only perform on file at local directory."),
|
||||
pattern: str = Field(..., description="The pattern to search for.")
|
||||
) -> Dict[str, List[str]]:
|
||||
"""
|
||||
This tool belongs to the Gorilla file system. It is a simple file system that allows users to perform basic file operations such as navigating directories, creating files and directories, reading and writing to files, etc. Tool description: Search for lines in a file of any extension at current directory that contain the specified pattern.
|
||||
|
||||
Args:
|
||||
file_name (str): The name of the file to search in current directory. No path allowed, only local file names.
|
||||
pattern (str): The text pattern to search for within the file lines. Performs substring matching.
|
||||
|
||||
Returns:
|
||||
Dict[str, List[str]]: A dictionary containing either:
|
||||
- matching_lines (List[str]): List of lines from the file that contain the specified pattern
|
||||
- error (str): Error message if file doesn't exist in current directory
|
||||
"""
|
||||
if file_name in self._current_dir.contents:
|
||||
file = self._current_dir._get_item(file_name)
|
||||
if isinstance(file, File):
|
||||
content = file._read()
|
||||
|
||||
matching_lines = [line for line in content.splitlines() if pattern in line]
|
||||
|
||||
return {"matching_lines": matching_lines}
|
||||
|
||||
return {"error": f"grep: {file_name}: No such file or directory"}
|
||||
|
||||
def mcp_du(self, human_readable: bool = Field(False, description="If True, returns the size in human-readable format (e.g., KB, MB).")) -> Dict[str, str]:
|
||||
"""
|
||||
This tool belongs to the Gorilla file system. It is a simple file system that allows users to perform basic file operations such as navigating directories, creating files and directories, reading and writing to files, etc. Tool description: Estimate the disk usage of a directory and its contents.
|
||||
|
||||
Args:
|
||||
human_readable (bool): If True, returns size in human-readable format (B, KB, MB, GB, TB, PB). If False, returns size in bytes. Defaults to False.
|
||||
|
||||
Returns:
|
||||
Dict[str, str]: A dictionary containing:
|
||||
- disk_usage (str): The estimated disk usage of current directory and all its contents, either in bytes or human-readable format
|
||||
"""
|
||||
|
||||
def get_size(item: Union[File, Directory]) -> int:
|
||||
if isinstance(item, File):
|
||||
return len(item._read().encode("utf-8"))
|
||||
elif isinstance(item, Directory):
|
||||
return sum(get_size(child) for child in item.contents.values())
|
||||
return 0
|
||||
|
||||
target_dir = self._navigate_to_directory(None)
|
||||
if isinstance(target_dir, dict): # Error condition check
|
||||
return target_dir
|
||||
|
||||
total_size = get_size(target_dir)
|
||||
|
||||
if human_readable:
|
||||
for unit in ["B", "KB", "MB", "GB", "TB"]:
|
||||
if total_size < 1024:
|
||||
size_str = f"{total_size:.2f} {unit}"
|
||||
break
|
||||
total_size /= 1024
|
||||
else:
|
||||
size_str = f"{total_size:.2f} PB"
|
||||
else:
|
||||
size_str = f"{total_size} bytes"
|
||||
|
||||
return {"disk_usage": size_str}
|
||||
|
||||
def mcp_tail(
|
||||
self,
|
||||
file_name: str = Field(..., description="The name of the file to display. No path is allowed and you can only perform on file at local directory."),
|
||||
lines: int = Field(10, description="The number of lines to display from the end of the file. Defaults to 10.")
|
||||
) -> Dict[str, str]:
|
||||
"""
|
||||
This tool belongs to the Gorilla file system. It is a simple file system that allows users to perform basic file operations such as navigating directories, creating files and directories, reading and writing to files, etc. Tool description: Display the last part of a file of any extension.
|
||||
|
||||
Args:
|
||||
file_name (str): The name of the file in current directory to display. No path allowed, only local file names.
|
||||
lines (int): The number of lines to display from the end of the file. If file has fewer lines, shows all available lines. Defaults to 10.
|
||||
|
||||
Returns:
|
||||
Dict[str, str]: A dictionary containing either:
|
||||
- last_lines (str): The last N lines of the file content joined with newlines
|
||||
- error (str): Error message if file doesn't exist in current directory
|
||||
"""
|
||||
if file_name in self._current_dir.contents:
|
||||
file = self._current_dir._get_item(file_name)
|
||||
if isinstance(file, File):
|
||||
content = file._read().splitlines()
|
||||
|
||||
if lines > len(content):
|
||||
lines = len(content)
|
||||
|
||||
last_lines = content[-lines:]
|
||||
return {"last_lines": "\n".join(last_lines)}
|
||||
|
||||
return {"error": f"tail: {file_name}: No such file or directory"}
|
||||
|
||||
def mcp_diff(
|
||||
self,
|
||||
file_name1: str = Field(..., description="The name of the first file in current directory."),
|
||||
file_name2: str = Field(..., description="The name of the second file in current directorry.")
|
||||
) -> Dict[str, str]:
|
||||
"""
|
||||
This tool belongs to the Gorilla file system. It is a simple file system that allows users to perform basic file operations such as navigating directories, creating files and directories, reading and writing to files, etc. Tool description: Compare two files of any extension line by line at the current directory.
|
||||
|
||||
Args:
|
||||
file_name1 (str): The name of the first file in current directory to compare.
|
||||
file_name2 (str): The name of the second file in current directory to compare.
|
||||
|
||||
Returns:
|
||||
Dict[str, str]: A dictionary containing either:
|
||||
- diff_lines (str): Line-by-line differences between the two files, showing removed (-) and added (+) lines
|
||||
- error (str): Error message if either file doesn't exist in current directory
|
||||
"""
|
||||
if (
|
||||
file_name1 in self._current_dir.contents
|
||||
and file_name2 in self._current_dir.contents
|
||||
):
|
||||
file1 = self._current_dir._get_item(file_name1)
|
||||
file2 = self._current_dir._get_item(file_name2)
|
||||
|
||||
if isinstance(file1, File) and isinstance(file2, File):
|
||||
content1 = file1._read().splitlines()
|
||||
content2 = file2._read().splitlines()
|
||||
|
||||
diff_lines = [
|
||||
f"- {line1}\n+ {line2}"
|
||||
for line1, line2 in zip(content1, content2)
|
||||
if line1 != line2
|
||||
]
|
||||
|
||||
return {"diff_lines": "\n".join(diff_lines)}
|
||||
|
||||
return {"error": f"diff: {file_name1} or {file_name2}: No such file or directory"}
|
||||
|
||||
def mcp_mv(
|
||||
self,
|
||||
source: str = Field(..., description="Source name of the file or directory to move. Source must be local to the current directory."),
|
||||
destination: str = Field(..., description="The destination name to move the file or directory to. Destination must be local to the current directory and cannot be a path. If destination is not an existing directory like when renaming something, destination is the new file name.")
|
||||
) -> Dict[str, str]:
|
||||
"""
|
||||
This tool belongs to the Gorilla file system. It is a simple file system that allows users to perform basic file operations such as navigating directories, creating files and directories, reading and writing to files, etc. Tool description: Move a file or directory from one location to another. so
|
||||
|
||||
Args:
|
||||
source (str): Source name of the file or directory to move. Must be local to current directory, no paths allowed.
|
||||
destination (str): The destination name to move to. Must be local to current directory, no paths allowed. If destination is existing directory, source moves into it. Otherwise, source is renamed to destination.
|
||||
|
||||
Returns:
|
||||
Dict[str, str]: A dictionary containing either:
|
||||
- result (str): Success message describing the move operation
|
||||
- error (str): Error message if source doesn't exist, destination conflicts, or operation fails
|
||||
"""
|
||||
if source not in self._current_dir.contents:
|
||||
return {"error": f"mv: cannot move '{source}': No such file or directory"}
|
||||
|
||||
item = self._current_dir._get_item(source)
|
||||
|
||||
if not isinstance(item, (File, Directory)):
|
||||
return {"error": f"mv: cannot move '{source}': Not a file or directory"}
|
||||
|
||||
if "/" in destination:
|
||||
return {
|
||||
"error": f"mv: no path allowed in destination. Only file name and folder name is supported for this operation."
|
||||
}
|
||||
|
||||
# Check if the destination is an existing directory
|
||||
if destination in self._current_dir.contents:
|
||||
dest_item = self._current_dir._get_item(destination)
|
||||
if isinstance(dest_item, Directory):
|
||||
# Move source into the destination directory
|
||||
new_destination = f"{source}"
|
||||
if new_destination in dest_item.contents:
|
||||
return {
|
||||
"error": f"mv: cannot move '{source}' to '{destination}/{source}': File exists"
|
||||
}
|
||||
else:
|
||||
self._current_dir.contents.pop(source)
|
||||
if isinstance(item, File):
|
||||
dest_item._add_file(source, item.content)
|
||||
else:
|
||||
dest_item._add_directory(source)
|
||||
dest_item.contents[source].contents = item.contents
|
||||
return {"result": f"'{source}' moved to '{destination}/{source}'"}
|
||||
else:
|
||||
return {
|
||||
"error": f"mv: cannot move '{source}' to '{destination}': Not a directory"
|
||||
}
|
||||
else:
|
||||
# Destination is not an existing directory, move/rename the item
|
||||
self._current_dir.contents.pop(source)
|
||||
if isinstance(item, File):
|
||||
self._current_dir._add_file(destination, item.content)
|
||||
else:
|
||||
self._current_dir._add_directory(destination)
|
||||
self._current_dir.contents[destination].contents = item.contents
|
||||
return {"result": f"'{source}' moved to '{destination}'"}
|
||||
|
||||
def mcp_rm(self, file_name: str = Field(..., description="The name of the file or directory to remove.")) -> Dict[str, str]:
|
||||
"""
|
||||
This tool belongs to the Gorilla file system. It is a simple file system that allows users to perform basic file operations such as navigating directories, creating files and directories, reading and writing to files, etc. Tool description: Remove a file or directory.
|
||||
|
||||
Args:
|
||||
file_name (str): The name of the file or directory to remove from current directory. Can remove both files and directories (including non-empty ones).
|
||||
|
||||
Returns:
|
||||
Dict[str, str]: A dictionary containing either:
|
||||
- result (str): Success message confirming the file or directory was removed
|
||||
- error (str): Error message if file/directory doesn't exist or is not a valid file system item
|
||||
"""
|
||||
if file_name in self._current_dir.contents:
|
||||
item = self._current_dir._get_item(file_name)
|
||||
if isinstance(item, File) or isinstance(item, Directory):
|
||||
self._current_dir.contents.pop(file_name)
|
||||
return {"result": f"'{file_name}' removed"}
|
||||
else:
|
||||
return {
|
||||
"error": f"rm: cannot remove '{file_name}': Not a file or directory"
|
||||
}
|
||||
else:
|
||||
return {"error": f"rm: cannot remove '{file_name}': No such file or directory"}
|
||||
|
||||
def mcp_rmdir(self, dir_name: str = Field(..., description="The name of the directory to remove. Directory must be local to the current directory.")) -> Dict[str, str]:
|
||||
"""
|
||||
This tool belongs to the Gorilla file system. It is a simple file system that allows users to perform basic file operations such as navigating directories, creating files and directories, reading and writing to files, etc. Tool description: Remove a directory at current directory.
|
||||
|
||||
Args:
|
||||
dir_name (str): The name of the directory to remove from current directory. Directory must be empty to be removed.
|
||||
|
||||
Returns:
|
||||
Dict[str, str]: A dictionary containing either:
|
||||
- result (str): Success message confirming the directory was removed
|
||||
- error (str): Error message if directory doesn't exist, is not empty, or is not a directory
|
||||
"""
|
||||
if dir_name in self._current_dir.contents:
|
||||
item = self._current_dir._get_item(dir_name)
|
||||
if isinstance(item, Directory):
|
||||
if item.contents: # Check if directory is not empty
|
||||
return {
|
||||
"error": f"rmdir: failed to remove '{dir_name}': Directory not empty"
|
||||
}
|
||||
else:
|
||||
self._current_dir.contents.pop(dir_name)
|
||||
return {"result": f"'{dir_name}' removed"}
|
||||
else:
|
||||
return {"error": f"rmdir: cannot remove '{dir_name}': Not a directory"}
|
||||
else:
|
||||
return {
|
||||
"error": f"rmdir: cannot remove '{dir_name}': No such file or directory"
|
||||
}
|
||||
|
||||
def mcp_cp(
|
||||
self,
|
||||
source: str = Field(..., description="The name of the file or directory to copy."),
|
||||
destination: str = Field(..., description="The destination name to copy the file or directory to. If the destination is a directory, the source will be copied into this directory. No file paths allowed.")
|
||||
) -> Dict[str, str]:
|
||||
"""
|
||||
This tool belongs to the Gorilla file system. It is a simple file system that allows users to perform basic file operations such as navigating directories, creating files and directories, reading and writing to files, etc. Tool description: Copy a file or directory from one location to another.
|
||||
|
||||
If the destination is a directory, the source file or directory will be copied
|
||||
into the destination directory.
|
||||
|
||||
Both source and destination must be local to the current directory.
|
||||
|
||||
Args:
|
||||
source (str): The name of the file or directory to copy from current directory.
|
||||
destination (str): The destination name to copy to in current directory. If destination is existing directory, source is copied into it. Otherwise, source is copied with new name. No file paths allowed.
|
||||
|
||||
Returns:
|
||||
Dict[str, str]: A dictionary containing either:
|
||||
- result (str): Success message describing the copy operation performed
|
||||
- error (str): Error message if source doesn't exist, destination conflicts, or operation fails
|
||||
"""
|
||||
if source not in self._current_dir.contents:
|
||||
return {"error": f"cp: cannot copy '{source}': No such file or directory"}
|
||||
|
||||
item = self._current_dir._get_item(source)
|
||||
|
||||
if not isinstance(item, (File, Directory)):
|
||||
return {"error": f"cp: cannot copy '{source}': Not a file or directory"}
|
||||
|
||||
if "/" in destination:
|
||||
return {
|
||||
"error": f"cp: don't allow path in destination. Only file name and folder name is supported for this operation."
|
||||
}
|
||||
# Check if the destination is an existing directory
|
||||
if destination in self._current_dir.contents:
|
||||
dest_item = self._current_dir._get_item(destination)
|
||||
if isinstance(dest_item, Directory):
|
||||
# Copy source into the destination directory
|
||||
new_destination = f"{destination}/{source}"
|
||||
if new_destination in dest_item.contents:
|
||||
return {
|
||||
"error": f"cp: cannot copy '{source}' to '{destination}/{source}': File exists"
|
||||
}
|
||||
else:
|
||||
if isinstance(item, File):
|
||||
dest_item._add_file(source, item.content)
|
||||
else:
|
||||
dest_item._add_directory(source)
|
||||
dest_item.contents[source].contents = item.contents.copy()
|
||||
return {"result": f"'{source}' copied to '{destination}/{source}'"}
|
||||
else:
|
||||
return {
|
||||
"error": f"cp: cannot copy '{source}' to '{destination}': Not a directory"
|
||||
}
|
||||
else:
|
||||
# Destination is not an existing directory, perform the copy
|
||||
if isinstance(item, File):
|
||||
self._current_dir._add_file(destination, item.content)
|
||||
else:
|
||||
self._current_dir._add_directory(destination)
|
||||
self._current_dir.contents[destination].contents = item.contents.copy()
|
||||
return {"result": f"'{source}' copied to '{destination}'"}
|
||||
|
||||
def _navigate_to_directory(
|
||||
self, path: Optional[str]
|
||||
) -> Union[Directory, Dict[str, str]]:
|
||||
"""
|
||||
Navigate to a specified directory path from the current directory.
|
||||
|
||||
Args:
|
||||
path (str): [Optional] The path to navigate to. Defaults to None (current directory).
|
||||
|
||||
Returns:
|
||||
target_directory (Directory or dict): The target directory object or error message.
|
||||
"""
|
||||
if path is None or path == ".":
|
||||
return self._current_dir
|
||||
elif path == "/":
|
||||
return self.root
|
||||
|
||||
dirs = path.strip("/").split("/")
|
||||
temp_dir = self._current_dir if not path.startswith("/") else self.root
|
||||
|
||||
for dir_name in dirs:
|
||||
next_dir = temp_dir._get_item(dir_name)
|
||||
if isinstance(next_dir, Directory):
|
||||
temp_dir = next_dir
|
||||
else:
|
||||
return {"error": f"cd: '{path}': No such file or directory"}
|
||||
|
||||
return temp_dir
|
||||
|
||||
def _parse_positions(self, positions: str) -> List[int]:
|
||||
"""
|
||||
Helper function to parse position strings, e.g., '1,3,5', '1-5', '-3', or '3-'.
|
||||
|
||||
Args:
|
||||
positions (str): The position string to parse.
|
||||
|
||||
Returns:
|
||||
list (List[int]): A list of integers representing the positions.
|
||||
"""
|
||||
result = []
|
||||
if "," in positions:
|
||||
for part in positions.split(","):
|
||||
result.extend(self._parse_positions(part))
|
||||
elif "-" in positions:
|
||||
start, end = positions.split("-")
|
||||
start = int(start) if start else 1
|
||||
end = int(end) if end else float("inf")
|
||||
result.extend(range(start, end + 1))
|
||||
else:
|
||||
result.append(int(positions))
|
||||
return result
|
||||
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Create a global file system instance for the server
|
||||
default_scenario = {
|
||||
"root": {
|
||||
"workspace": {
|
||||
"type": "directory",
|
||||
"contents": {
|
||||
"projects": {
|
||||
"type": "directory",
|
||||
"contents": {
|
||||
"web_app": {
|
||||
"type": "directory",
|
||||
"contents": {
|
||||
"src": {
|
||||
"type": "directory",
|
||||
"contents": {
|
||||
"main.py": {
|
||||
"type": "file",
|
||||
"content": "#!/usr/bin/env python3\nfrom flask import Flask\napp = Flask(__name__)\n\n@app.route('/')\ndef hello():\n return 'Hello World!'\n\nif __name__ == '__main__':\n app.run(debug=True)"
|
||||
},
|
||||
"utils.py": {
|
||||
"type": "file",
|
||||
"content": "# Utility functions\nimport os\nimport json\n\ndef load_config(path):\n with open(path, 'r') as f:\n return json.load(f)\n\ndef get_file_size(path):\n return os.path.getsize(path)"
|
||||
}
|
||||
}
|
||||
},
|
||||
"README.md": {
|
||||
"type": "file",
|
||||
"content": "# MyWebApp\n\nA simple Flask web application.\n\n## Features\n- Hello World endpoint\n- Configuration management\n\n## Setup\n1. Install dependencies: `pip install flask`\n2. Run the app: `python src/main.py`"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"temp": {
|
||||
"type": "directory",
|
||||
"contents": {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
args = ActionArguments(
|
||||
name="GorillaFileSystem",
|
||||
transport="stdio"
|
||||
)
|
||||
|
||||
try:
|
||||
file_system = GorillaFileSystem(args)
|
||||
file_system._load_scenario(default_scenario)
|
||||
print("GorillaFileSystem service initialized")
|
||||
print("GorillaFileSystem service has the following tools:")
|
||||
print("- pwd, ls, cd, mkdir, touch, echo, cat, find, wc, sort, grep, du, tail, diff, mv, rm, rmdir, cp")
|
||||
file_system.run()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}: {traceback.format_exc()}")
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,318 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Test script for GorillaFileSystem MCP Functions
|
||||
|
||||
This script tests all mcp_ functions in the GorillaFileSystem class.
|
||||
Run this to verify all file system operations work correctly.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import traceback
|
||||
|
||||
# Add the current directory to the path to import gorilla_file_system
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
from gorilla_file_system import GorillaFileSystem, ActionArguments
|
||||
|
||||
|
||||
def test_mcp_functions():
|
||||
"""Test all mcp_ functions in the GorillaFileSystem class."""
|
||||
print("🧪 Testing GorillaFileSystem MCP Functions")
|
||||
print("=" * 60)
|
||||
|
||||
# Initialize the file system
|
||||
args = ActionArguments(
|
||||
name="GorillaFileSystem",
|
||||
transport="sse"
|
||||
)
|
||||
|
||||
file_system = GorillaFileSystem(args)
|
||||
|
||||
# Load default scenario
|
||||
default_scenario = {
|
||||
"root": {
|
||||
"workspace": {
|
||||
"type": "directory",
|
||||
"contents": {
|
||||
"projects": {
|
||||
"type": "directory",
|
||||
"contents": {
|
||||
"web_app": {
|
||||
"type": "directory",
|
||||
"contents": {
|
||||
"src": {
|
||||
"type": "directory",
|
||||
"contents": {
|
||||
"main.py": {
|
||||
"type": "file",
|
||||
"content": "#!/usr/bin/env python3\nfrom flask import Flask\napp = Flask(__name__)\n\n@app.route('/')\ndef hello():\n return 'Hello World!'\n\nif __name__ == '__main__':\n app.run(debug=True)"
|
||||
},
|
||||
"utils.py": {
|
||||
"type": "file",
|
||||
"content": "# Utility functions\nimport os\nimport json\n\ndef load_config(path):\n with open(path, 'r') as f:\n return json.load(f)\n\ndef get_file_size(path):\n return os.path.getsize(path)"
|
||||
}
|
||||
}
|
||||
},
|
||||
"README.md": {
|
||||
"type": "file",
|
||||
"content": "# MyWebApp\n\nA simple Flask web application.\n\n## Features\n- Hello World endpoint\n- Configuration management\n\n## Setup\n1. Install dependencies: `pip install flask`\n2. Run the app: `python src/main.py`"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"temp": {
|
||||
"type": "directory",
|
||||
"contents": {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
file_system._load_scenario(default_scenario)
|
||||
|
||||
# Test results tracking
|
||||
tests_passed = 0
|
||||
tests_failed = 0
|
||||
|
||||
def run_test(test_name, test_func):
|
||||
"""Run a test and track results."""
|
||||
nonlocal tests_passed, tests_failed
|
||||
try:
|
||||
print(f"\n🔍 Testing {test_name}...")
|
||||
result = test_func()
|
||||
print(f"✅ {test_name}: PASSED")
|
||||
print(f" Result: {result}")
|
||||
tests_passed += 1
|
||||
return result
|
||||
except Exception as e:
|
||||
print(f"❌ {test_name}: FAILED")
|
||||
print(f" Error: {e}")
|
||||
tests_failed += 1
|
||||
return None
|
||||
|
||||
# Test 1: mcp_pwd
|
||||
def test_pwd():
|
||||
return file_system.mcp_pwd()
|
||||
|
||||
run_test("mcp_pwd", test_pwd)
|
||||
|
||||
# Test 2: mcp_ls
|
||||
def test_ls():
|
||||
return file_system.mcp_ls()
|
||||
|
||||
run_test("mcp_ls", test_ls)
|
||||
|
||||
# Test 3: mcp_ls with hidden files
|
||||
def test_ls_hidden():
|
||||
return file_system.mcp_ls(a=True)
|
||||
|
||||
run_test("mcp_ls (with hidden files)", test_ls_hidden)
|
||||
|
||||
# Test 4: mcp_cd
|
||||
def test_cd():
|
||||
return file_system.mcp_cd("projects")
|
||||
|
||||
run_test("mcp_cd", test_cd)
|
||||
|
||||
# Test 5: mcp_mkdir
|
||||
def test_mkdir():
|
||||
return file_system.mcp_mkdir("test_dir")
|
||||
|
||||
run_test("mcp_mkdir", test_mkdir)
|
||||
|
||||
# Test 6: mcp_touch
|
||||
def test_touch():
|
||||
return file_system.mcp_touch("test_file.txt")
|
||||
|
||||
run_test("mcp_touch", test_touch)
|
||||
|
||||
# Test 7: mcp_echo (to file)
|
||||
def test_echo_to_file():
|
||||
return file_system.mcp_echo("Hello, FastMCP!", "test_file.txt")
|
||||
|
||||
run_test("mcp_echo (to file)", test_echo_to_file)
|
||||
|
||||
# Test 8: mcp_echo (to terminal)
|
||||
def test_echo_to_terminal():
|
||||
return file_system.mcp_echo("Hello, Terminal!")
|
||||
|
||||
run_test("mcp_echo (to terminal)", test_echo_to_terminal)
|
||||
|
||||
# Test 9: mcp_cat
|
||||
def test_cat():
|
||||
return file_system.mcp_cat("test_file.txt")
|
||||
|
||||
run_test("mcp_cat", test_cat)
|
||||
|
||||
# Test 10: mcp_find
|
||||
def test_find():
|
||||
return file_system.mcp_find(name="test")
|
||||
|
||||
run_test("mcp_find", test_find)
|
||||
|
||||
# Test 11: mcp_wc (words)
|
||||
def test_wc_words():
|
||||
return file_system.mcp_wc("test_file.txt", "w")
|
||||
|
||||
run_test("mcp_wc (words)", test_wc_words)
|
||||
|
||||
# Test 12: mcp_wc (lines)
|
||||
def test_wc_lines():
|
||||
return file_system.mcp_wc("test_file.txt", "l")
|
||||
|
||||
run_test("mcp_wc (lines)", test_wc_lines)
|
||||
|
||||
# Test 13: mcp_wc (characters)
|
||||
def test_wc_chars():
|
||||
return file_system.mcp_wc("test_file.txt", "c")
|
||||
|
||||
run_test("mcp_wc (characters)", test_wc_chars)
|
||||
|
||||
# Test 14: mcp_sort
|
||||
def test_sort():
|
||||
# Create a file with unsorted content
|
||||
file_system.mcp_echo("zebra\napple\nbanana\ncat", "unsorted.txt")
|
||||
return file_system.mcp_sort("unsorted.txt")
|
||||
|
||||
run_test("mcp_sort", test_sort)
|
||||
|
||||
# Test 15: mcp_grep
|
||||
def test_grep():
|
||||
return file_system.mcp_grep("test_file.txt", "Hello")
|
||||
|
||||
run_test("mcp_grep", test_grep)
|
||||
|
||||
# Test 16: mcp_du
|
||||
def test_du():
|
||||
return file_system.mcp_du()
|
||||
|
||||
run_test("mcp_du", test_du)
|
||||
|
||||
# Test 17: mcp_du (human readable)
|
||||
def test_du_human():
|
||||
return file_system.mcp_du(human_readable=True)
|
||||
|
||||
run_test("mcp_du (human readable)", test_du_human)
|
||||
|
||||
# Test 18: mcp_tail
|
||||
def test_tail():
|
||||
# Create a file with multiple lines
|
||||
file_system.mcp_echo("line1\nline2\nline3\nline4\nline5", "multi_line.txt")
|
||||
return file_system.mcp_tail("multi_line.txt", 3)
|
||||
|
||||
run_test("mcp_tail", test_tail)
|
||||
|
||||
# Test 19: mcp_diff
|
||||
def test_diff():
|
||||
# Create two different files
|
||||
file_system.mcp_echo("apple\nbanana\ncherry", "file1.txt")
|
||||
file_system.mcp_echo("apple\norange\ncherry", "file2.txt")
|
||||
return file_system.mcp_diff("file1.txt", "file2.txt")
|
||||
|
||||
run_test("mcp_diff", test_diff)
|
||||
|
||||
# Test 20: mcp_mv
|
||||
def test_mv():
|
||||
return file_system.mcp_mv("test_file.txt", "renamed_file.txt")
|
||||
|
||||
run_test("mcp_mv", test_mv)
|
||||
|
||||
# Test 21: mcp_cp
|
||||
def test_cp():
|
||||
return file_system.mcp_cp("renamed_file.txt", "copied_file.txt")
|
||||
|
||||
run_test("mcp_cp", test_cp)
|
||||
|
||||
# Test 22: mcp_rm
|
||||
def test_rm():
|
||||
return file_system.mcp_rm("copied_file.txt")
|
||||
|
||||
run_test("mcp_rm", test_rm)
|
||||
|
||||
# Test 23: mcp_rmdir
|
||||
def test_rmdir():
|
||||
# Create an empty directory first
|
||||
file_system.mcp_mkdir("empty_dir")
|
||||
return file_system.mcp_rmdir("empty_dir")
|
||||
|
||||
run_test("mcp_rmdir", test_rmdir)
|
||||
|
||||
# Test 24: mcp_cd (back to parent)
|
||||
def test_cd_parent():
|
||||
return file_system.mcp_cd("..")
|
||||
|
||||
run_test("mcp_cd (parent)", test_cd_parent)
|
||||
|
||||
# Test 25: mcp_cd (absolute path)
|
||||
def test_cd_absolute():
|
||||
return file_system.mcp_cd("/workspace")
|
||||
|
||||
run_test("mcp_cd (absolute)", test_cd_absolute)
|
||||
|
||||
# Test 26: Error handling - mcp_cat non-existent file
|
||||
def test_cat_error():
|
||||
return file_system.mcp_cat("non_existent_file.txt")
|
||||
|
||||
run_test("mcp_cat (error case)", test_cat_error)
|
||||
|
||||
# Test 27: Error handling - mcp_mkdir existing directory
|
||||
def test_mkdir_error():
|
||||
return file_system.mcp_mkdir("test_dir") # Should fail as it already exists
|
||||
|
||||
run_test("mcp_mkdir (error case)", test_mkdir_error)
|
||||
|
||||
# Test 28: Error handling - mcp_touch existing file
|
||||
def test_touch_error():
|
||||
return file_system.mcp_touch("renamed_file.txt") # Should fail as it already exists
|
||||
|
||||
run_test("mcp_touch (error case)", test_touch_error)
|
||||
|
||||
# Test 29: Error handling - mcp_rmdir non-empty directory
|
||||
def test_rmdir_error():
|
||||
return file_system.mcp_rmdir("test_dir") # Should fail as it's not empty
|
||||
|
||||
run_test("mcp_rmdir (error case)", test_rmdir_error)
|
||||
|
||||
# Test 30: Error handling - mcp_cd non-existent directory
|
||||
def test_cd_error():
|
||||
return file_system.mcp_cd("non_existent_dir")
|
||||
|
||||
run_test("mcp_cd (error case)", test_cd_error)
|
||||
|
||||
# Print final results
|
||||
print("\n" + "=" * 60)
|
||||
print("📊 TEST RESULTS SUMMARY")
|
||||
print("=" * 60)
|
||||
print(f"✅ Tests Passed: {tests_passed}")
|
||||
print(f"❌ Tests Failed: {tests_failed}")
|
||||
print(f"📈 Total Tests: {tests_passed + tests_failed}")
|
||||
|
||||
if tests_failed == 0:
|
||||
print("\n🎉 All tests passed! GorillaFileSystem is working correctly.")
|
||||
else:
|
||||
print(f"\n⚠️ {tests_failed} test(s) failed. Please check the errors above.")
|
||||
|
||||
return tests_failed == 0
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function to run the tests."""
|
||||
print("🔧 GorillaFileSystem MCP Functions Test Suite")
|
||||
print("Testing all mcp_ functions in the GorillaFileSystem class")
|
||||
|
||||
try:
|
||||
success = test_mcp_functions()
|
||||
if success:
|
||||
print("\n🚀 All tests completed successfully!")
|
||||
else:
|
||||
print("\n⚠️ Some tests failed. Please review the output above.")
|
||||
except Exception as e:
|
||||
print(f"\n💥 Test suite failed with error: {e}")
|
||||
print(f"Traceback: {traceback.format_exc()}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,52 @@
|
||||
import logging
|
||||
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")
|
||||
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.logger: logging.Logger = setup_logger(self.__class__.__name__, output_folder_path='./logs')
|
||||
|
||||
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)
|
||||
@@ -0,0 +1,5 @@
|
||||
openai>=1.0.0
|
||||
pydantic>=2.0.0
|
||||
dataclasses-json>=0.6.0
|
||||
fastmcp>=0.1.0
|
||||
tabulate>=0.9.0
|
||||
@@ -0,0 +1,54 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
|
||||
import os
|
||||
from aworld.config.conf import AgentConfig
|
||||
from aworld.agents.llm_agent import Agent
|
||||
from aworld.runner import Runners
|
||||
import json
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
agent_config = AgentConfig(
|
||||
llm_provider=os.getenv("LLM_PROVIDER", "openai"),
|
||||
llm_model_name=os.getenv("LLM_MODEL_NAME"),
|
||||
llm_base_url=os.getenv("LLM_BASE_URL"),
|
||||
llm_api_key=os.getenv("LLM_API_KEY"),
|
||||
llm_temperature=os.getenv("LLM_TEMPERATURE", 0.0)
|
||||
)
|
||||
|
||||
# Register the MCP tool here, or create a separate configuration file.
|
||||
mcp_config = {
|
||||
"mcpServers": {
|
||||
"GorillaFileSystem": {
|
||||
"type": "stdio",
|
||||
"command": "python",
|
||||
"args": ["mcp_tools/gorilla_file_system.py"],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
file_sys_prompt = "You are a helpful agent to use the standard file system to perform file operations."
|
||||
file_sys = Agent(
|
||||
conf=agent_config,
|
||||
name="file_sys_agent",
|
||||
system_prompt=file_sys_prompt,
|
||||
mcp_servers=mcp_config.get("mcpServers", []).keys(),
|
||||
mcp_config=mcp_config,
|
||||
)
|
||||
|
||||
# run
|
||||
result = Runners.sync_run(
|
||||
input=(
|
||||
"use mcp tools in the GorillaFileSystem server to perform file operations: "
|
||||
"write the content 'AWorld' into the hello_world.py file with a new line "
|
||||
"and keep the original content of the file. Make sure the new and old "
|
||||
"content are all in the file; and display the content of the file"
|
||||
),
|
||||
agent=file_sys,
|
||||
)
|
||||
|
||||
print("=" * 100)
|
||||
print(f"result.answer: {result.answer}")
|
||||
print("=" * 100)
|
||||
print(f"result.trajectory: {json.dumps(result.trajectory[0], indent=4)}")
|
||||
@@ -0,0 +1,2 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
@@ -0,0 +1,34 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
import os
|
||||
|
||||
from aworld.agents.llm_agent import Agent
|
||||
from aworld.config.conf import AgentConfig
|
||||
from aworld.core.agent.swarm import Swarm
|
||||
from aworld.logs.util import color_log, Color
|
||||
from aworld.runner import Runners
|
||||
from aworld.tools.human.human import HUMAN
|
||||
import examples
|
||||
|
||||
if __name__ == '__main__':
|
||||
conf = AgentConfig(
|
||||
llm_provider=os.getenv("LLM_PROVIDER", "openai"),
|
||||
llm_model_name=os.getenv("LLM_MODEL_NAME"),
|
||||
llm_base_url=os.getenv("LLM_BASE_URL"),
|
||||
llm_api_key=os.getenv("LLM_API_KEY"),
|
||||
)
|
||||
# human in the loop
|
||||
agent = Agent(
|
||||
conf=conf,
|
||||
name='human_test',
|
||||
system_prompt="You are a helpful assistant.",
|
||||
tool_names=[HUMAN]
|
||||
)
|
||||
|
||||
swarm = Swarm(agent, max_steps=1)
|
||||
result = Runners.sync_run(
|
||||
input="use human tool to ask a question, e.g. what is the weather in beijing?" \
|
||||
"please use HUMAN tool only once",
|
||||
swarm=swarm
|
||||
)
|
||||
color_log(f"agent result:{result.answer}", color=Color.pink)
|
||||
@@ -0,0 +1,59 @@
|
||||
# AWorld Examples
|
||||
|
||||
This directory contains a wide range of examples demonstrating the capabilities of the AWorld framework.
|
||||
The examples cover single-agent and multi-agent scenarios. Each subdirectory focuses on a specific paradigm or
|
||||
application area, making it easy for developers to explore and extend.
|
||||
|
||||
## Directory Overview
|
||||
|
||||
- **common/**
|
||||
Shared tools, utilities, and components used by other examples.
|
||||
|
||||
- **multi_agents/**
|
||||
Multi-agent system examples demonstrating three core paradigms:
|
||||
- **collaborative/**: Agents working together (e.g., debate, travel planning)
|
||||
- **coordination/**: Orchestrated agent teams (e.g., master-worker, deep research)
|
||||
- **workflow/**: Multi-agent workflow automation (e.g., search and summary)
|
||||
See `multi_agents/README.md` for details.
|
||||
|
||||
- **web/**
|
||||
Aworld web for visual interaction.
|
||||
|
||||
**Run agent in build-in WebUI**
|
||||
|
||||
- **Configure Environment**: Navigate to `examples/web/agent_deploy/` and you'll find 3 demo agents: `single_agent`, `team_agent`, and `deep_research`. Copy `.env.template` to `.env` in your chosen agent directory, then update the configuration values with your own settings.
|
||||
- **Launch WebUI**: Start the web server by running: `cd examples/web/ && aworld web`
|
||||
|
||||
## Application Overview
|
||||
|
||||
- **browser_use/**
|
||||
Agents specialized in web browser, capable of browsing, interacting with, and extracting information from web pages.
|
||||
|
||||
- **BFCL/**
|
||||
Demonstrates Basic Function Call Learning using a virtual file system and MCP tools. Useful for generating training data and testing function call synthesis.
|
||||
|
||||
- **gaia/**
|
||||
Advanced agent runner and server examples, including integration with MCP collections and OpenWebUI.
|
||||
|
||||
- **gym_demo/**
|
||||
Example of using an agent to interact with OpenAI Gym environments, such as CartPole, to showcase reinforcement learning and environment control.
|
||||
|
||||
- **phone_use/**
|
||||
Examples of agents for Android device, including app operation, UI analysis, and task execution.
|
||||
|
||||
- **text_to_audio/**
|
||||
Example of text-to-audio conversion using MCP servers and agents.
|
||||
|
||||
|
||||
## Usage
|
||||
|
||||
Create .env file in the examples' dir, the file content is the environment variables required for runtime,
|
||||
such as LLM_MODEL_NAME, LLM_API_KEY, LLM_BASE_URL, LLM_TEMPERATURE = 0.0 etc.
|
||||
|
||||
- Each subdirectory contains its own entry point (usually `run.py`) and may include additional configuration or requirements files.
|
||||
- Before running any example, ensure you have installed all required dependencies and set the necessary environment variables (e.g., LLM provider credentials, API keys).
|
||||
- For detailed instructions, refer to the README or comments within each subdirectory.
|
||||
|
||||
---
|
||||
|
||||
If you need more detailed usage instructions or want to add new examples, refer to the documentation and code samples in each subdirectory.
|
||||
@@ -0,0 +1,5 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
@@ -0,0 +1,4 @@
|
||||
# Browser use
|
||||
|
||||
Agents specialized in web browser automation.
|
||||
The implementation of browser agent version is now derived from [browser use](https://github.com/browser-use/browser-use), which we have made a lot of modifications to integrate into our own framework
|
||||
@@ -0,0 +1,2 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
@@ -0,0 +1,565 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
|
||||
import re
|
||||
import time
|
||||
import traceback
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Dict, Any, Optional, List, Union, Tuple
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from langchain_core.messages import HumanMessage, BaseMessage, AIMessage, ToolMessage
|
||||
from pydantic import ValidationError
|
||||
|
||||
from aworld.core.agent.base import AgentFactory, AgentResult
|
||||
from aworld.agents.llm_agent import Agent
|
||||
from examples.browser_use.prompts import SystemPrompt
|
||||
from examples.browser_use.utils import convert_input_messages, extract_json_from_model_output, estimate_messages_tokens
|
||||
from examples.browser_use.common import AgentState, AgentStepInfo, AgentHistory, PolicyMetadata, AgentBrain
|
||||
from aworld.config.conf import AgentConfig, ConfigDict
|
||||
from aworld.core.common import Observation, ActionModel, ToolActionInfo, ActionResult
|
||||
from aworld.logs.util import logger
|
||||
from examples.browser_use.prompts import AgentMessagePrompt
|
||||
from examples.common.tools.common import Tools
|
||||
from examples.common.tools.tool_action import BrowserAction
|
||||
|
||||
|
||||
@dataclass
|
||||
class Trajectory:
|
||||
"""A class to store agent history records, including all observations, info and AgentResult"""
|
||||
history: List[tuple[List[BaseMessage], Observation, Dict[str, Any], AIMessage, AgentResult]] = field(
|
||||
default_factory=list)
|
||||
|
||||
def add_step(self, input_messages: List[BaseMessage], observation: Observation, info: Dict[str, Any],
|
||||
output_message: AIMessage, agent_result: AgentResult):
|
||||
"""Add a step to the history"""
|
||||
self.history.append((input_messages, observation, info, output_message, agent_result))
|
||||
|
||||
def get_history(self) -> List[tuple[List[BaseMessage], Observation, Dict[str, Any], AIMessage, AgentResult]]:
|
||||
"""Get the complete history"""
|
||||
return self.history
|
||||
|
||||
def save_history(self, file_path: str):
|
||||
his_li = []
|
||||
for input_messages, observation, info, output_message, agent_result in self.get_history():
|
||||
llm_input = [{"type": input_message.type, "content": input_message.content} for input_message in
|
||||
input_messages]
|
||||
llm_output = output_message.content
|
||||
his_li.append({"llm_input": llm_input, "llm_output": llm_output})
|
||||
with open(file_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(his_li, f, ensure_ascii=False, indent=4)
|
||||
|
||||
|
||||
class BrowserAgent(Agent):
|
||||
def __init__(self, conf: Union[Dict[str, Any], ConfigDict, AgentConfig], name: str, **kwargs):
|
||||
super(BrowserAgent, self).__init__(conf=conf, name=name, **kwargs)
|
||||
self.state = AgentState()
|
||||
self.settings = self.conf
|
||||
provider = self.conf.llm_config.llm_provider
|
||||
if self.conf.llm_config.llm_provider:
|
||||
self.conf.llm_config.llm_provider = "chat" + provider
|
||||
else:
|
||||
raise Exception("no llm provider")
|
||||
|
||||
self.save_file_path = self.conf.save_file_path
|
||||
self.available_actions = self._build_action_prompt()
|
||||
# Note: Removed _message_manager initialization as it's no longer used
|
||||
# Initialize trajectory
|
||||
self.trajectory = Trajectory()
|
||||
self._init = False
|
||||
|
||||
def reset(self, options: Dict[str, Any] = None):
|
||||
super(BrowserAgent, self).reset(options)
|
||||
|
||||
# Reset trajectory
|
||||
self.trajectory = Trajectory()
|
||||
|
||||
# Note: Removed _message_manager initialization as it's no longer used
|
||||
# _estimate_tokens_for_messages method now directly uses functions from utils.py
|
||||
|
||||
self._init = True
|
||||
|
||||
def _build_action_prompt(self) -> str:
|
||||
def _prompt(info: ToolActionInfo) -> str:
|
||||
s = f'{info.desc}: \n'
|
||||
s += '{' + str(info.name) + ': '
|
||||
if info.input_params:
|
||||
s += str({k: {"title": k, "type": v.type} for k, v in info.input_params.items()})
|
||||
s += '}'
|
||||
return s
|
||||
|
||||
val = "\n".join([_prompt(v.value) for k, v in BrowserAction.__members__.items()])
|
||||
return val
|
||||
|
||||
def _log_message_sequence(self, input_messages: List[BaseMessage]) -> None:
|
||||
"""Log the sequence of messages for debugging purposes"""
|
||||
logger.info(f"[agent] 🔍 Invoking LLM with {len(input_messages)} messages")
|
||||
logger.info("[agent] 📝 Messages sequence:")
|
||||
for i, msg in enumerate(input_messages):
|
||||
prefix = msg.type
|
||||
logger.info(f"[agent] Message {i + 1}: {prefix} ===================================")
|
||||
if isinstance(msg.content, list):
|
||||
for item in msg.content:
|
||||
if item.get('type') == 'text':
|
||||
logger.info(f"[agent] Text content: {item.get('text')}")
|
||||
elif item.get('type') == 'image_url':
|
||||
# Only print the first 30 characters of image URL to avoid printing entire base64
|
||||
image_url = item.get('image_url', {}).get('url', '')
|
||||
if image_url.startswith('data:image'):
|
||||
logger.info(f"[agent] Image: [Base64 image data]")
|
||||
else:
|
||||
logger.info(f"[agent] Image URL: {image_url[:30]}...")
|
||||
else:
|
||||
content = str(msg.content)
|
||||
chunk_size = 500
|
||||
for j in range(0, len(content), chunk_size):
|
||||
chunk = content[j:j + chunk_size]
|
||||
if j == 0:
|
||||
logger.info(f"[agent] Content: {chunk}")
|
||||
else:
|
||||
logger.info(f"[agent] Content (continued): {chunk}")
|
||||
|
||||
if isinstance(msg, AIMessage) and hasattr(msg, 'tool_calls') and msg.tool_calls:
|
||||
for tool_call in msg.tool_calls:
|
||||
logger.info(f"[agent] Tool call: {tool_call.get('name')} - ID: {tool_call.get('id')}")
|
||||
args = str(tool_call.get('args', {}))[:1000]
|
||||
logger.info(f"[agent] Tool args: {args}...")
|
||||
|
||||
def save_process(self, file_path: str):
|
||||
self.trajectory.save_history(file_path)
|
||||
|
||||
def policy(self,
|
||||
observation: Observation,
|
||||
info: Dict[str, Any] = None, **kwargs) -> Union[List[ActionModel], None]:
|
||||
start_time = time.time()
|
||||
|
||||
if self._init is False:
|
||||
self.reset({"task": observation.content})
|
||||
|
||||
self._finished = False
|
||||
# Save current observation to state for message construction
|
||||
self.state.last_result = observation.action_result
|
||||
|
||||
if self.conf.max_steps <= self.state.n_steps:
|
||||
logger.info('Last step finishing up')
|
||||
|
||||
logger.info(f'[agent] step {self.state.n_steps}')
|
||||
|
||||
# Use the new method to build messages, passing the current observation
|
||||
input_messages = self.build_messages_from_trajectory_and_observation(observation=observation)
|
||||
|
||||
# Note: Special message addition has been moved to build_messages_from_trajectory_and_observation
|
||||
|
||||
# Estimate token count
|
||||
tokens = self._estimate_tokens_for_messages(input_messages)
|
||||
|
||||
llm_result = None
|
||||
output_message = None
|
||||
try:
|
||||
# Log the message sequence
|
||||
self._log_message_sequence(input_messages)
|
||||
|
||||
output_message, llm_result = self._do_policy(input_messages)
|
||||
|
||||
if not llm_result:
|
||||
logger.error("[agent] ❌ Failed to parse LLM response")
|
||||
return [ActionModel(tool_name=Tools.BROWSER.value, action_name="stop")]
|
||||
|
||||
self.state.n_steps += 1
|
||||
|
||||
# No longer need to remove the last state message
|
||||
# self._message_manager._remove_last_state_message()
|
||||
|
||||
if self.state.stopped or self.state.paused:
|
||||
logger.info('Browser gent paused after getting state')
|
||||
return [ActionModel(tool_name=Tools.BROWSER.value, action_name="stop")]
|
||||
|
||||
tool_action = llm_result.actions
|
||||
|
||||
# Add the current step to the trajectory
|
||||
self.trajectory.add_step(input_messages, observation, info, output_message, llm_result)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(traceback.format_exc())
|
||||
# No longer need to remove the last state message
|
||||
# self._message_manager._remove_last_state_message()
|
||||
logger.error(f"[agent] ❌ Error parsing LLM response: {str(e)}")
|
||||
|
||||
# Create an AgentResult object with an empty actions list
|
||||
error_result = AgentResult(
|
||||
current_state=AgentBrain(
|
||||
evaluation_previous_goal="Failed due to error",
|
||||
memory=f"Error occurred: {str(e)}",
|
||||
thought="Recover from error",
|
||||
next_goal="Recover from error"
|
||||
),
|
||||
actions=[] # Empty actions list
|
||||
)
|
||||
|
||||
# Add the error state to the trajectory
|
||||
self.trajectory.add_step(input_messages, observation, info, output_message, error_result)
|
||||
|
||||
raise RuntimeError("Browser agent encountered exception while making the policy.", e)
|
||||
finally:
|
||||
if llm_result:
|
||||
# Only keep the history_item creation part
|
||||
metadata = PolicyMetadata(
|
||||
number=self.state.n_steps,
|
||||
start_time=start_time,
|
||||
end_time=time.time(),
|
||||
input_tokens=tokens,
|
||||
)
|
||||
self._make_history_item(llm_result, observation, observation.action_result, metadata)
|
||||
else:
|
||||
logger.warning("no result to record!")
|
||||
|
||||
return tool_action
|
||||
|
||||
def _do_policy(self, input_messages: list[BaseMessage]) -> Tuple[AIMessage, AgentResult]:
|
||||
THINK_TAGS = re.compile(r'<think>.*?</think>', re.DOTALL)
|
||||
|
||||
def _remove_think_tags(text: str) -> str:
|
||||
"""Remove think tags from text"""
|
||||
return re.sub(THINK_TAGS, '', text)
|
||||
|
||||
input_messages = self._convert_input_messages(input_messages)
|
||||
output_message = None
|
||||
try:
|
||||
|
||||
output_message = self.llm.invoke(input_messages)
|
||||
|
||||
if not output_message or not output_message.content:
|
||||
logger.warning("[agent] LLM returned empty response")
|
||||
return output_message, AgentResult(
|
||||
current_state=AgentBrain(evaluation_previous_goal="", memory="", thought="", next_goal=""),
|
||||
actions=[ActionModel(agent_name=self.id(), tool_name='browser', action_name="stop")])
|
||||
except:
|
||||
logger.error(f"[agent] Response content: {output_message}")
|
||||
raise RuntimeError('call llm fail, please check llm conf and network.')
|
||||
|
||||
if self.model_name == 'deepseek-reasoner':
|
||||
output_message.content = _remove_think_tags(output_message.content)
|
||||
try:
|
||||
# Get max retries from config
|
||||
max_retries = self.settings.get('max_llm_json_retries', 3)
|
||||
retry_count = 0
|
||||
json_parse_error = None
|
||||
|
||||
while retry_count < max_retries:
|
||||
try:
|
||||
parsed_json = extract_json_from_model_output(output_message.content)
|
||||
# If parsing succeeds, break out of the retry loop
|
||||
json_parse_error = None
|
||||
break
|
||||
except ValueError as e:
|
||||
# Store the error and retry
|
||||
json_parse_error = e
|
||||
retry_count += 1
|
||||
logger.warning(f"[agent] Failed to parse JSON (attempt {retry_count}/{max_retries}): {str(e)}")
|
||||
|
||||
if retry_count < max_retries:
|
||||
# Add a reminder message about JSON format with specific structure guidance
|
||||
format_reminder = HumanMessage(
|
||||
content="Your responses must be always JSON with the specified format. Make sure your response includes a 'current_state' object with 'evaluation_previous_goal', 'memory', and 'next_goal' fields, and an 'action' array with the actions to perform. Do not include any explanatory text, only return the raw JSON.")
|
||||
retry_messages = input_messages.copy()
|
||||
retry_messages.append(format_reminder)
|
||||
|
||||
# Retry with the updated messages
|
||||
logger.info(
|
||||
f"[agent] Retrying LLM invocation ({retry_count}/{max_retries}) with format reminder")
|
||||
output_message = self.llm.invoke(retry_messages)
|
||||
|
||||
# Check for empty response during retry
|
||||
if not output_message or not output_message.content:
|
||||
logger.warning(
|
||||
f"[agent] LLM returned empty response on retry attempt {retry_count}/{max_retries}")
|
||||
# Continue to next retry instead of immediately returning
|
||||
continue
|
||||
|
||||
if self.model_name == 'deepseek-reasoner':
|
||||
output_message.content = _remove_think_tags(output_message.content)
|
||||
|
||||
# If all retries failed, raise the last error
|
||||
if json_parse_error:
|
||||
logger.error(f"[agent] ❌ All {max_retries} attempts to parse JSON failed")
|
||||
raise json_parse_error
|
||||
|
||||
logger.info((f"llm response: {parsed_json}"))
|
||||
try:
|
||||
agent_brain = AgentBrain(**parsed_json['current_state'])
|
||||
except:
|
||||
agent_brain = None
|
||||
actions = parsed_json.get('action')
|
||||
result = []
|
||||
if not actions:
|
||||
actions = parsed_json.get("actions")
|
||||
if not actions:
|
||||
logger.warning("agent not policy an action.")
|
||||
self._finished = True
|
||||
return output_message, AgentResult(current_state=agent_brain,
|
||||
actions=[ActionModel(tool_name='browser',
|
||||
agent_name=self.id(),
|
||||
action_name="done")])
|
||||
|
||||
for action in actions:
|
||||
if "action_name" in action:
|
||||
action_name = action['action_name']
|
||||
browser_action = BrowserAction.get_value_by_name(action_name)
|
||||
if not browser_action:
|
||||
logger.warning(f"Unsupported action: {action_name}")
|
||||
if action_name == "done":
|
||||
self._finished = True
|
||||
action_model = ActionModel(agent_name=self.id(),
|
||||
tool_name='browser',
|
||||
action_name=action_name,
|
||||
params=action.get('params', {}))
|
||||
result.append(action_model)
|
||||
else:
|
||||
for k, v in action.items():
|
||||
browser_action = BrowserAction.get_value_by_name(k)
|
||||
if not browser_action:
|
||||
logger.warning(f"Unsupported action: {k}")
|
||||
|
||||
action_model = ActionModel(agent_name=self.id(), tool_name='browser', action_name=k, params=v)
|
||||
result.append(action_model)
|
||||
if k == "done":
|
||||
self._finished = True
|
||||
return output_message, AgentResult(current_state=agent_brain, actions=result)
|
||||
except (ValueError, ValidationError) as e:
|
||||
logger.warning(f'Failed to parse model output: {output_message} {str(e)}')
|
||||
raise ValueError('Could not parse response.')
|
||||
|
||||
def _convert_input_messages(self, input_messages: list[BaseMessage]) -> list[BaseMessage]:
|
||||
"""Convert input messages to the correct format"""
|
||||
if self.model_name == 'deepseek-reasoner' or self.model_name.startswith('deepseek-r1'):
|
||||
return convert_input_messages(input_messages, self.model_name)
|
||||
else:
|
||||
return input_messages
|
||||
|
||||
def _make_history_item(self,
|
||||
model_output: AgentResult | None,
|
||||
state: Observation,
|
||||
result: list[ActionResult],
|
||||
metadata: Optional[PolicyMetadata] = None) -> None:
|
||||
content = ""
|
||||
if hasattr(state, 'dom_tree') and state.dom_tree is not None:
|
||||
if hasattr(state.dom_tree, 'element_tree'):
|
||||
content = state.dom_tree.element_tree.__repr__()
|
||||
else:
|
||||
content = str(state.dom_tree)
|
||||
|
||||
history_item = AgentHistory(model_output=model_output,
|
||||
result=state.action_result,
|
||||
metadata=metadata,
|
||||
content=content,
|
||||
base64_img=state.image if hasattr(state, 'image') else None)
|
||||
|
||||
self.state.history.history.append(history_item)
|
||||
|
||||
def _process_action_result(self, action_result, messages, tool_call=None):
|
||||
"""Helper method to process an action result and add appropriate messages"""
|
||||
if action_result.content is not None:
|
||||
messages.append(HumanMessage(content='Action result: ' + action_result.content))
|
||||
elif action_result.error is not None:
|
||||
# Assemble error message when error information exists
|
||||
messages.append(HumanMessage(content='Action result: ' + action_result.error))
|
||||
if tool_call is not None:
|
||||
logger.warning(f"Action {tool_call} failed: {action_result.error}")
|
||||
else:
|
||||
logger.warning(f"Action failed: {action_result.error}")
|
||||
# If there is an error but success is true, log the error and terminate the program as the result is invalid
|
||||
if action_result.success is True:
|
||||
error_msg = f"Invalid result: success=True but error message exists: {action_result.error}"
|
||||
logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
return action_result.error is not None
|
||||
|
||||
def build_messages_from_trajectory_and_observation(self, observation: Optional[Observation] = None) -> List[
|
||||
BaseMessage]:
|
||||
"""
|
||||
Build complete message history from trajectory and current observation
|
||||
|
||||
Args:
|
||||
observation: Current observation object, if None current observation won't be added
|
||||
"""
|
||||
messages = []
|
||||
# Add system message
|
||||
system_message = SystemPrompt(
|
||||
max_actions_per_step=self.settings.get('max_actions_per_step')
|
||||
).get_system_message()
|
||||
if isinstance(system_message, tuple):
|
||||
system_message = system_message[0]
|
||||
messages.append(system_message)
|
||||
|
||||
tool_calling_method = self.settings.get("tool_calling_method")
|
||||
llm_provider = self.conf.llm_config.llm_provider
|
||||
|
||||
if tool_calling_method == 'raw' or (tool_calling_method == 'auto' and (
|
||||
llm_provider == 'deepseek-reasoner' or llm_provider.startswith('deepseek-r1'))):
|
||||
message_context = f'\n\nAvailable actions: {self.available_actions}'
|
||||
else:
|
||||
message_context = None
|
||||
|
||||
# Add task context (if any)
|
||||
if message_context:
|
||||
context_message = HumanMessage(content='Context for the task' + message_context)
|
||||
messages.append(context_message)
|
||||
|
||||
# Add task message
|
||||
task_message = HumanMessage(
|
||||
content=f'Your ultimate task is: """{self.task}""". If you achieved your ultimate task, stop everything and use the done action in the next step to complete the task. If not, continue as usual.'
|
||||
)
|
||||
messages.append(task_message)
|
||||
|
||||
# Add example output
|
||||
placeholder_message = HumanMessage(content='Example output:')
|
||||
messages.append(placeholder_message)
|
||||
|
||||
# Add example tool call
|
||||
tool_calls = [
|
||||
{
|
||||
'name': 'AgentOutput',
|
||||
'args': {
|
||||
'current_state': {
|
||||
'evaluation_previous_goal': 'Success - I opend the first page',
|
||||
'memory': 'Starting with the new task. I have completed 1/10 steps',
|
||||
'thought': 'From the current page I can get information about all the companies.',
|
||||
'next_goal': 'Click on company a',
|
||||
},
|
||||
'action': [{'click_element': {'index': 0}}],
|
||||
},
|
||||
'id': '1',
|
||||
'type': 'tool_call',
|
||||
}
|
||||
]
|
||||
example_tool_call = AIMessage(
|
||||
content='',
|
||||
tool_calls=tool_calls,
|
||||
)
|
||||
messages.append(example_tool_call)
|
||||
|
||||
# Add first tool message with "Browser started" content
|
||||
messages.append(ToolMessage(content='Browser started', tool_call_id='1'))
|
||||
|
||||
# Add task history marker
|
||||
messages.append(HumanMessage(content='[Your task history memory starts here]'))
|
||||
|
||||
# Add available file paths (if any)
|
||||
if self.settings.get('available_file_paths'):
|
||||
filepaths_msg = HumanMessage(
|
||||
content=f'Here are file paths you can use: {self.settings.get("available_file_paths")}')
|
||||
messages.append(filepaths_msg)
|
||||
previous_action_entries = []
|
||||
# Add messages from the history trajectory
|
||||
for input_msgs, obs, info, output_msg, llm_result in self.trajectory.get_history():
|
||||
# Check the previous step's actionResult
|
||||
has_error = False
|
||||
if obs.action_result is not None:
|
||||
# The previous action entries should match with action results
|
||||
if len(previous_action_entries) == 0:
|
||||
# if previous_action_entries is empty,process action_result directly
|
||||
logger.info(
|
||||
f"History item with action_result count ({len(obs.action_result)}) with empty previous actions - skipping count check")
|
||||
elif len(previous_action_entries) == len(obs.action_result):
|
||||
for i, one_action_result in enumerate(obs.action_result):
|
||||
has_error = self._process_action_result(one_action_result, messages,
|
||||
previous_action_entries[i]) or has_error
|
||||
else:
|
||||
# If sizes don't match, this is a critical error
|
||||
error_msg = f"Action results count ({len(obs.action_result)}) doesn't match action entries count ({len(previous_action_entries)})"
|
||||
logger.error(error_msg)
|
||||
has_error = True
|
||||
# raise ValueError(error_msg)
|
||||
|
||||
# Add agent response
|
||||
if llm_result:
|
||||
# Create AI message
|
||||
output_data = llm_result.model_dump(mode='json', exclude_unset=True)
|
||||
action_entries = [{action.action_name: action.params} for action in llm_result.actions]
|
||||
output_data["action"] = action_entries
|
||||
if "actions" in output_data:
|
||||
del output_data["actions"]
|
||||
|
||||
# Calculate tool_id based on trajectory history. If no actions yet, start with ID 1
|
||||
tool_id = 1 if len(self.trajectory.get_history()) == 0 else len(self.trajectory.get_history()) + 1
|
||||
tool_calls = [
|
||||
{
|
||||
'name': 'AgentOutput',
|
||||
'args': output_data,
|
||||
'id': str(tool_id),
|
||||
'type': 'tool_call',
|
||||
}
|
||||
]
|
||||
previous_action_entries = action_entries
|
||||
ai_message = AIMessage(
|
||||
content='',
|
||||
tool_calls=tool_calls,
|
||||
)
|
||||
messages.append(ai_message)
|
||||
|
||||
# Add empty tool message after each AIMessage
|
||||
messages.append(ToolMessage(content='', tool_call_id=str(tool_id)))
|
||||
|
||||
# Add current observation - using the passed observation parameter instead of self.state.current_observation
|
||||
if observation:
|
||||
# Check if the current observation has an action_result with error
|
||||
has_error = False
|
||||
if hasattr(observation, 'action_result') and observation.action_result is not None:
|
||||
# Match action results with previous actions
|
||||
if len(previous_action_entries) == 0:
|
||||
# if previous_action_entries is empty,process action_result directly
|
||||
logger.info(
|
||||
f"Current observation with action_result count ({len(observation.action_result)}) with empty previous actions - skipping count check")
|
||||
elif len(previous_action_entries) == len(observation.action_result):
|
||||
for i, one_action_result in enumerate(observation.action_result):
|
||||
has_error = self._process_action_result(one_action_result, messages,
|
||||
previous_action_entries[i]) or has_error
|
||||
else:
|
||||
# If sizes don't match, this is a critical error
|
||||
error_msg = f"Action results count ({len(observation.action_result)}) doesn't match action entries count ({len(previous_action_entries)})"
|
||||
logger.error(error_msg)
|
||||
has_error = True
|
||||
|
||||
# If there's an error, append observation content outside the loop
|
||||
if has_error and observation.content:
|
||||
messages.append(HumanMessage(content=observation.content))
|
||||
# If no error, process the observation normally
|
||||
elif not has_error:
|
||||
step_info = AgentStepInfo(number=self.state.n_steps, max_steps=self.conf.max_steps)
|
||||
if hasattr(observation, 'dom_tree') and observation.dom_tree:
|
||||
state_message = AgentMessagePrompt(
|
||||
observation,
|
||||
self.state.last_result,
|
||||
include_attributes=self.settings.get('include_attributes',
|
||||
["title", "type", "name", "role", "aria-label",
|
||||
"placeholder", "value", "alt", "aria-expanded",
|
||||
"data-date-format"]),
|
||||
step_info=step_info,
|
||||
).get_user_message(self.settings.get('use_vision'))
|
||||
messages.append(state_message)
|
||||
elif observation.content:
|
||||
messages.append(HumanMessage(content=observation.content))
|
||||
|
||||
# Add special message for the last step
|
||||
# Note: Moved here from policy method to centralize all message building logic
|
||||
if self.conf.max_steps <= self.state.n_steps:
|
||||
last_step_message = f"""
|
||||
Now comes your last step. Use only the "done" action now. No other actions - so here your action sequence must have length 1.
|
||||
\nIf the task is not yet fully finished as requested by the user, set success in "done" to false! E.g. if not all steps are fully completed.
|
||||
\nIf the task is fully finished, set success in "done" to true.
|
||||
\nInclude everything you found out for the ultimate task in the done text.
|
||||
"""
|
||||
messages.append(HumanMessage(content=[{'type': 'text', 'text': last_step_message}]))
|
||||
|
||||
return messages
|
||||
|
||||
def _estimate_tokens_for_messages(self, messages: List[BaseMessage]) -> int:
|
||||
"""Roughly estimate token count for message list"""
|
||||
# Note: Using estimate_messages_tokens function from utils.py instead of calling _message_manager
|
||||
# This decouples the dependency on MessageManager
|
||||
return estimate_messages_tokens(
|
||||
messages,
|
||||
image_tokens=self.settings.get('image_tokens', 800),
|
||||
estimated_characters_per_token=self.settings.get('estimated_characters_per_token', 3)
|
||||
)
|
||||
@@ -0,0 +1,130 @@
|
||||
# coding: utf-8
|
||||
|
||||
import json
|
||||
import traceback
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional, Dict, List
|
||||
|
||||
from openai import RateLimitError
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from aworld.core.common import ActionResult
|
||||
|
||||
|
||||
class PolicyMetadata(BaseModel):
|
||||
"""Metadata for a single step including timing information"""
|
||||
start_time: float
|
||||
end_time: float
|
||||
number: int
|
||||
input_tokens: int
|
||||
|
||||
@property
|
||||
def duration_seconds(self) -> float:
|
||||
"""Calculate step duration in seconds"""
|
||||
return self.end_time - self.start_time
|
||||
|
||||
|
||||
class AgentBrain(BaseModel):
|
||||
"""Current state of the agent"""
|
||||
evaluation_previous_goal: str = None
|
||||
memory: str = None
|
||||
thought: str = None
|
||||
next_goal: str = None
|
||||
|
||||
|
||||
class AgentHistory(BaseModel):
|
||||
"""History item for agent actions"""
|
||||
model_output: Optional[BaseModel] = None
|
||||
result: List[ActionResult]
|
||||
metadata: Optional[PolicyMetadata] = None
|
||||
content: Optional[str] = None
|
||||
base64_img: Optional[str] = None
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
def model_dump(self, **kwargs) -> Dict[str, Any]:
|
||||
"""Custom serialization handling"""
|
||||
return {
|
||||
'model_output': self.model_output.model_dump() if self.model_output else None,
|
||||
'result': [r.model_dump(exclude_none=True) for r in self.result],
|
||||
'metadata': self.metadata.model_dump() if self.metadata else None,
|
||||
'content': self.xml_content,
|
||||
'base64_img': self.base64_img
|
||||
}
|
||||
|
||||
|
||||
class AgentHistoryList(BaseModel):
|
||||
"""List of agent history items"""
|
||||
history: List[AgentHistory]
|
||||
|
||||
def total_duration_seconds(self) -> float:
|
||||
"""Get total duration of all steps in seconds"""
|
||||
total = 0.0
|
||||
for h in self.history:
|
||||
if h.metadata:
|
||||
total += h.metadata.duration_seconds
|
||||
return total
|
||||
|
||||
def save_to_file(self, filepath: str | Path) -> None:
|
||||
"""Save history to JSON file with proper serialization"""
|
||||
try:
|
||||
Path(filepath).parent.mkdir(parents=True, exist_ok=True)
|
||||
data = self.model_dump()
|
||||
with open(filepath, 'w', encoding='utf-8') as f:
|
||||
json.dump(data, f, indent=2)
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
def model_dump(self, **kwargs) -> Dict[str, Any]:
|
||||
"""Custom serialization that properly uses AgentHistory's model_dump"""
|
||||
return {
|
||||
'history': [h.model_dump(**kwargs) for h in self.history],
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def load_from_file(cls, filepath: str | Path) -> 'AgentHistoryList':
|
||||
"""Load history from JSON file"""
|
||||
with open(filepath, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
return cls.model_validate(data)
|
||||
|
||||
|
||||
class AgentError:
|
||||
"""Container for agent error handling"""
|
||||
VALIDATION_ERROR = 'Invalid model output format. Please follow the correct schema.'
|
||||
RATE_LIMIT_ERROR = 'Rate limit reached. Waiting before retry.'
|
||||
NO_VALID_ACTION = 'No valid action found'
|
||||
|
||||
@staticmethod
|
||||
def format_error(error: Exception, include_trace: bool = False) -> str:
|
||||
"""Format error message based on error type and optionally include trace"""
|
||||
if isinstance(error, RateLimitError):
|
||||
return AgentError.RATE_LIMIT_ERROR
|
||||
if include_trace:
|
||||
return f'{str(error)}\nStacktrace:\n{traceback.format_exc()}'
|
||||
return f'{str(error)}'
|
||||
|
||||
|
||||
class AgentState(BaseModel):
|
||||
"""Holds all state information for an Agent"""
|
||||
|
||||
agent_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
|
||||
n_steps: int = 1
|
||||
consecutive_failures: int = 0
|
||||
last_result: Optional[List['ActionResult']] = None
|
||||
history: AgentHistoryList = Field(default_factory=lambda: AgentHistoryList(history=[]))
|
||||
last_plan: Optional[str] = None
|
||||
paused: bool = False
|
||||
stopped: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentStepInfo:
|
||||
number: int
|
||||
max_steps: int
|
||||
|
||||
def is_last_step(self) -> bool:
|
||||
"""Check if this is the last step"""
|
||||
return self.number >= self.max_steps - 1
|
||||
@@ -0,0 +1,25 @@
|
||||
# coding: utf-8
|
||||
|
||||
from typing import Optional
|
||||
from aworld.config.conf import AgentConfig
|
||||
from typing import Literal
|
||||
|
||||
ToolCallingMethod = Literal['function_calling', 'json_mode', 'raw', 'auto']
|
||||
|
||||
|
||||
class BrowserAgentConfig(AgentConfig):
|
||||
use_vision: bool = True
|
||||
use_vision_for_planner: bool = False
|
||||
save_conversation_path: Optional[str] = None
|
||||
save_conversation_path_encoding: Optional[str] = 'utf-8'
|
||||
max_failures: int = 3
|
||||
retry_delay: int = 10
|
||||
validate_output: bool = False
|
||||
message_context: Optional[str] = None
|
||||
generate_gif: bool | str = False
|
||||
available_file_paths: Optional[list[str]] = None
|
||||
override_system_message: Optional[str] = None
|
||||
extend_system_message: Optional[str] = None
|
||||
tool_calling_method: Optional[ToolCallingMethod] = 'auto'
|
||||
max_llm_json_retries: int = 3
|
||||
save_file_path: str = "browser_agent_history.json"
|
||||
@@ -0,0 +1,212 @@
|
||||
# coding: utf-8
|
||||
|
||||
from datetime import datetime
|
||||
from typing import List, Optional
|
||||
|
||||
from langchain_core.messages import HumanMessage, SystemMessage
|
||||
|
||||
from examples.browser_use.common import AgentStepInfo
|
||||
from aworld.core.common import Observation, ActionResult
|
||||
|
||||
PROMPT_TEMPLATE = """
|
||||
You are an AI agent designed to automate browser tasks. Your goal is to accomplish the ultimate task following the rules.
|
||||
|
||||
# Input Format
|
||||
Task
|
||||
Previous steps
|
||||
Current URL
|
||||
Open Tabs
|
||||
Interactive Elements
|
||||
[index]<type>text</type>
|
||||
- index: Numeric identifier for interaction
|
||||
- type: HTML element type (button, input, etc.)
|
||||
- text: Element description
|
||||
Example:
|
||||
[33]<button>Submit Form</button>
|
||||
|
||||
- Only elements with numeric indexes in [] are interactive
|
||||
- elements without [] provide only context
|
||||
|
||||
# Response Rules
|
||||
1. RESPONSE FORMAT: You must ALWAYS respond with valid JSON in this exact format:
|
||||
{{"current_state": {{"evaluation_previous_goal": "Success|Failed|Unknown - Analyze the current elements and the image to check if the previous goals/actions are successful like intended by the task. Mention if something unexpected happened. Shortly state why/why not",
|
||||
"memory": "Description of what has been done and what you need to remember. Be very specific. Count here ALWAYS how many times you have done something and how many remain. E.g. 0 out of 10 websites analyzed. Continue with abc and xyz",
|
||||
"thought": "Your thought or reasoning based on the ultimate task and current observations",
|
||||
"next_goal": "What needs to be done with the next immediate action"}},
|
||||
"action":[{{"one_action_name": {{// action-specific parameter}}}}, // ... more actions in sequence]}}
|
||||
|
||||
2. ACTIONS: You can specify multiple actions in the list to be executed in sequence. But always specify only one action name per item. Use maximum {max_actions} actions per sequence.
|
||||
Common action sequences:
|
||||
- Form filling: [{{"input_text": {{"index": 1, "text": "username"}}}}, {{"input_text": {{"index": 2, "text": "password"}}}}, {{"click_element": {{"index": 3}}}}]
|
||||
- Navigation and extraction: [{{"go_to_url": {{"url": "https://example.com"}}}}, {{"extract_content": {{"goal": "extract the names"}}}}]
|
||||
- Actions are executed in the given order
|
||||
- If the page changes after an action, the sequence is interrupted and you get the new state.
|
||||
- Only provide the action sequence until an action which changes the page state significantly.
|
||||
- Try to be efficient, e.g. fill forms at once, or chain actions where nothing changes on the page
|
||||
- only use multiple actions if it makes sense.
|
||||
|
||||
3. ELEMENT INTERACTION:
|
||||
- Only use indexes of the interactive elements
|
||||
- Elements marked with "[]Non-interactive text" are non-interactive
|
||||
|
||||
4. NAVIGATION & ERROR HANDLING:
|
||||
- If no suitable elements exist, use other functions to complete the task
|
||||
- If stuck, try alternative approaches - like going back to a previous page, new search, new tab etc.
|
||||
- Handle popups/cookies by accepting or closing them
|
||||
- Use scroll to find elements you are looking for
|
||||
- If you want to research something, open a new tab instead of using the current tab
|
||||
- If captcha pops up, try to solve it - else try a different approach
|
||||
- If the page is not fully loaded, use wait action
|
||||
|
||||
5. TASK COMPLETION:
|
||||
- Use the done action as the last action as soon as the ultimate task is complete
|
||||
- Dont use "done" before you are done with everything the user asked you, except you reach the last step of max_steps.
|
||||
- If you reach your last step, use the done action even if the task is not fully finished. Provide all the information you have gathered so far. If the ultimate task is completly finished set success to true. If not everything the user asked for is completed set success in done to false!
|
||||
- If you have to do something repeatedly for example the task says for "each", or "for all", or "x times", count always inside "memory" how many times you have done it and how many remain. Don't stop until you have completed like the task asked you. Only call done after the last step.
|
||||
- Don't hallucinate actions
|
||||
- Make sure you include everything you found out for the ultimate task in the done text parameter. Do not just say you are done, but include the requested information of the task.
|
||||
|
||||
6. VISUAL CONTEXT:
|
||||
- When an image is provided, use it to understand the page layout
|
||||
- Bounding boxes with labels on their top right corner correspond to element indexes
|
||||
|
||||
7. Form filling:
|
||||
- If you fill an input field and your action sequence is interrupted, most often something changed e.g. suggestions popped up under the field.
|
||||
|
||||
8. Long tasks:
|
||||
- Keep track of the status and subresults in the memory.
|
||||
|
||||
9. Extraction:
|
||||
- If your task is to find information - call extract_content on the specific pages to get and store the information.
|
||||
Your responses must be always JSON with the specified format.
|
||||
"""
|
||||
|
||||
|
||||
class SystemPrompt:
|
||||
def __init__(self,
|
||||
max_actions_per_step: int = 10,
|
||||
override_system_message: Optional[str] = None,
|
||||
extend_system_message: Optional[str] = None):
|
||||
self.max_actions_per_step = max_actions_per_step
|
||||
if override_system_message:
|
||||
prompt = override_system_message
|
||||
else:
|
||||
prompt = PROMPT_TEMPLATE.format(max_actions=self.max_actions_per_step)
|
||||
|
||||
if extend_system_message:
|
||||
prompt += f'\n{extend_system_message}'
|
||||
|
||||
self.system_message = SystemMessage(content=prompt)
|
||||
|
||||
def get_system_message(self) -> SystemMessage:
|
||||
"""
|
||||
Get the system prompt for the agent.
|
||||
|
||||
Returns:
|
||||
SystemMessage: Formatted system prompt
|
||||
"""
|
||||
return self.system_message
|
||||
|
||||
|
||||
class AgentMessagePrompt:
|
||||
def __init__(
|
||||
self,
|
||||
state: Observation,
|
||||
result: Optional[List[ActionResult]] = None,
|
||||
include_attributes: list[str] = [],
|
||||
step_info: Optional[AgentStepInfo] = None,
|
||||
):
|
||||
self.state = state
|
||||
self.result = result
|
||||
self.include_attributes = include_attributes
|
||||
self.step_info = step_info
|
||||
|
||||
def get_user_message(self, use_vision: bool = True) -> HumanMessage:
|
||||
elements_text = self.state.dom_tree.element_tree.clickable_elements_to_string(
|
||||
include_attributes=self.include_attributes)
|
||||
|
||||
pixels_above = self.state.info.get('pixels_above', 0)
|
||||
pixels_below = self.state.info.get('pixels_below', 0)
|
||||
|
||||
if elements_text != '':
|
||||
if pixels_above > 0:
|
||||
elements_text = (
|
||||
f'... {pixels_above} pixels above - scroll or extract content to see more ...\n{elements_text}'
|
||||
)
|
||||
else:
|
||||
elements_text = f'[Start of page]\n{elements_text}'
|
||||
if pixels_below > 0:
|
||||
elements_text = (
|
||||
f'{elements_text}\n... {pixels_below} pixels below - scroll or extract content to see more ...'
|
||||
)
|
||||
else:
|
||||
elements_text = f'{elements_text}\n[End of page]'
|
||||
else:
|
||||
elements_text = 'empty page'
|
||||
|
||||
if self.step_info:
|
||||
step_info_description = f'Current step: {self.step_info.number}/{self.step_info.max_steps}'
|
||||
else:
|
||||
step_info_description = ''
|
||||
time_str = datetime.now().strftime('%Y-%m-%d %H:%M')
|
||||
step_info_description += f'Current date and time: {time_str}'
|
||||
|
||||
state_description = f"""
|
||||
[Task history memory ends]
|
||||
[Current state starts here]
|
||||
The following is one-time information - if you need to remember it write it to memory:
|
||||
Current url: {self.state.info.get("url")}
|
||||
Interactive elements from top layer of the current page inside the viewport:
|
||||
{elements_text}
|
||||
{step_info_description}
|
||||
"""
|
||||
|
||||
if self.result:
|
||||
for i, result in enumerate(self.result):
|
||||
if result.content:
|
||||
state_description += f'\nAction result {i + 1}/{len(self.result)}: {result.content}'
|
||||
if result.error:
|
||||
# only use last line of error
|
||||
error = result.error.split('\n')[-1]
|
||||
state_description += f'\nAction error {i + 1}/{len(self.result)}: ...{error}'
|
||||
|
||||
if self.state.image and use_vision == True:
|
||||
# Format message for vision model
|
||||
return HumanMessage(
|
||||
content=[
|
||||
{'type': 'text', 'text': state_description},
|
||||
{
|
||||
'type': 'image_url',
|
||||
'image_url': {'url': f'data:image/png;base64,{self.state.image}'}, # , 'detail': 'low'
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
return HumanMessage(content=state_description)
|
||||
|
||||
|
||||
class PlannerPrompt(SystemPrompt):
|
||||
def get_system_message(self) -> SystemMessage:
|
||||
return SystemMessage(
|
||||
content="""You are a planning agent that helps break down tasks into smaller steps and reason about the current state.
|
||||
Your role is to:
|
||||
1. Analyze the current state and history
|
||||
2. Evaluate progress towards the ultimate goal
|
||||
3. Identify potential challenges or roadblocks
|
||||
4. Suggest the next high-level steps to take
|
||||
|
||||
Inside your messages, there will be AI messages from different agents with different formats.
|
||||
|
||||
Your output format should be always a JSON object with the following fields:
|
||||
{
|
||||
"state_analysis": "Brief analysis of the current state and what has been done so far",
|
||||
"progress_evaluation": "Evaluation of progress towards the ultimate goal (as percentage and description)",
|
||||
"challenges": "List any potential challenges or roadblocks",
|
||||
"next_steps": "List 2-3 concrete next steps to take",
|
||||
"reasoning": "Explain your reasoning for the suggested next steps"
|
||||
}
|
||||
|
||||
Ignore the other AI messages output structures.
|
||||
don't forget the index param for input_text action.
|
||||
Keep your responses concise and focused on actionable insights."""
|
||||
)
|
||||
@@ -0,0 +1,6 @@
|
||||
langchain~=0.3.20
|
||||
langchain-openai~=0.3.8
|
||||
langchain-ollama~=0.2.3
|
||||
langchain-anthropic~=0.3.9
|
||||
langchain-mistralai~=0.2.7
|
||||
langchain-google-genai~=2.1.0
|
||||
@@ -0,0 +1,56 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
import os
|
||||
|
||||
from aworld.config.conf import ModelConfig
|
||||
from aworld.core.task import Task
|
||||
from aworld.runner import Runners
|
||||
from examples.browser_use.agent import BrowserAgent
|
||||
from examples.browser_use.config import BrowserAgentConfig
|
||||
from examples.common.tools.common import Agents, Tools
|
||||
from examples.common.tools.conf import BrowserToolConfig
|
||||
|
||||
# os.environ["LLM_MODEL_NAME"] = "YOUR_LLM_MODEL_NAME"
|
||||
# os.environ["LLM_BASE_URL"] = "YOUR_LLM_BASE_URL"
|
||||
# os.environ["LLM_API_KEY"] = "YOUR_LLM_API_KEY"
|
||||
|
||||
if __name__ == '__main__':
|
||||
llm_config = ModelConfig(
|
||||
llm_provider=os.getenv("LLM_PROVIDER", "openai"),
|
||||
llm_model_name=os.getenv("LLM_MODEL_NAME"),
|
||||
llm_base_url=os.getenv("LLM_BASE_URL"),
|
||||
llm_api_key=os.getenv("LLM_API_KEY"),
|
||||
llm_temperature=os.getenv("LLM_TEMPERATURE", 0.0)
|
||||
)
|
||||
browser_tool_config = BrowserToolConfig(width=1280,
|
||||
height=720,
|
||||
headless=False,
|
||||
keep_browser_open=True,
|
||||
use_async=True,
|
||||
custom_executor=True,
|
||||
llm_config=llm_config)
|
||||
agent_config = BrowserAgentConfig(
|
||||
name=Agents.BROWSER.value,
|
||||
tool_calling_method="raw",
|
||||
llm_config=llm_config,
|
||||
max_actions_per_step=10,
|
||||
max_input_tokens=128000,
|
||||
working_dir="",
|
||||
# llm model not supported vision, need to set `False`
|
||||
# use_vision=False
|
||||
)
|
||||
|
||||
task_config = {
|
||||
'max_steps': 100,
|
||||
'max_actions_per_step': 100
|
||||
}
|
||||
|
||||
task = Task(
|
||||
input="""step1: first go to https://www.dangdang.com/ and search for 'the little prince' and rank by sales from high to low, get the first 5 results and put the products info in memory.
|
||||
step 2: write each product's title, price, discount, and publisher information to a fully structured HTML document with write_to_file, ensuring that the data is presented in a table with visible grid lines.
|
||||
step3: open the html file in browser by go_to_url""",
|
||||
agent=BrowserAgent(conf=agent_config, name=Agents.BROWSER.value, tool_names=[Tools.BROWSER.name]),
|
||||
tools_conf={Tools.BROWSER.value: browser_tool_config},
|
||||
conf=task_config
|
||||
)
|
||||
Runners.sync_run_task(task)
|
||||
@@ -0,0 +1,199 @@
|
||||
# coding: utf-8
|
||||
import requests
|
||||
import json
|
||||
from io import BytesIO
|
||||
import os
|
||||
from typing import Any, Optional, Type
|
||||
import base64
|
||||
|
||||
from langchain_core.messages import (
|
||||
AIMessage,
|
||||
BaseMessage,
|
||||
HumanMessage,
|
||||
SystemMessage,
|
||||
ToolMessage,
|
||||
)
|
||||
|
||||
from aworld.logs.util import logger
|
||||
|
||||
|
||||
def extract_json_from_model_output(content: str) -> dict:
|
||||
"""Extract JSON from model output, handling both plain JSON and code-block-wrapped JSON."""
|
||||
try:
|
||||
# If content is wrapped in code blocks, extract just the JSON part
|
||||
if '```' in content:
|
||||
# Find the JSON content between code blocks
|
||||
content = content.split('```')[1]
|
||||
# Remove language identifier if present (e.g., 'json\n')
|
||||
if '\n' in content:
|
||||
content = content.split('\n', 1)[1]
|
||||
# Parse the cleaned content
|
||||
return json.loads(content)
|
||||
except json.JSONDecodeError as e:
|
||||
logger.warning(f'Failed to parse model output: {content} {str(e)}')
|
||||
raise ValueError('Could not parse response.')
|
||||
|
||||
|
||||
def convert_input_messages(input_messages: list[BaseMessage], model_name: Optional[str]) -> list[BaseMessage]:
|
||||
"""Convert input messages to a format that is compatible with the planner model"""
|
||||
if model_name is None:
|
||||
return input_messages
|
||||
if model_name == 'deepseek-reasoner' or model_name.startswith('deepseek-r1'):
|
||||
converted_input_messages = _convert_messages_for_non_function_calling_models(input_messages)
|
||||
merged_input_messages = _merge_successive_messages(converted_input_messages, HumanMessage)
|
||||
merged_input_messages = _merge_successive_messages(merged_input_messages, AIMessage)
|
||||
return merged_input_messages
|
||||
return input_messages
|
||||
|
||||
|
||||
def _convert_messages_for_non_function_calling_models(input_messages: list[BaseMessage]) -> list[BaseMessage]:
|
||||
"""Convert messages for non-function-calling models"""
|
||||
output_messages = []
|
||||
for message in input_messages:
|
||||
if isinstance(message, HumanMessage):
|
||||
output_messages.append(message)
|
||||
elif isinstance(message, SystemMessage):
|
||||
output_messages.append(message)
|
||||
elif isinstance(message, ToolMessage):
|
||||
output_messages.append(HumanMessage(content=message.content))
|
||||
elif isinstance(message, AIMessage):
|
||||
# check if tool_calls is a valid JSON object
|
||||
if message.tool_calls:
|
||||
tool_calls = json.dumps(message.tool_calls)
|
||||
output_messages.append(AIMessage(content=tool_calls))
|
||||
else:
|
||||
output_messages.append(message)
|
||||
else:
|
||||
raise ValueError(f'Unknown message type: {type(message)}')
|
||||
return output_messages
|
||||
|
||||
|
||||
def _merge_successive_messages(messages: list[BaseMessage], class_to_merge: Type[BaseMessage]) -> list[BaseMessage]:
|
||||
"""Some models like deepseek-reasoner dont allow multiple human messages in a row. This function merges them into one."""
|
||||
merged_messages = []
|
||||
streak = 0
|
||||
for message in messages:
|
||||
if isinstance(message, class_to_merge):
|
||||
streak += 1
|
||||
if streak > 1:
|
||||
if isinstance(message.content, list):
|
||||
merged_messages[-1].content += message.content[0]['text'] # type:ignore
|
||||
else:
|
||||
merged_messages[-1].content += message.content
|
||||
else:
|
||||
merged_messages.append(message)
|
||||
else:
|
||||
merged_messages.append(message)
|
||||
streak = 0
|
||||
return merged_messages
|
||||
|
||||
|
||||
def save_conversation(input_messages: list[BaseMessage], response: Any, target: str,
|
||||
encoding: Optional[str] = None) -> None:
|
||||
"""Save conversation history to file."""
|
||||
|
||||
# create folders if not exists
|
||||
os.makedirs(os.path.dirname(target), exist_ok=True)
|
||||
|
||||
with open(
|
||||
target,
|
||||
'w',
|
||||
encoding=encoding,
|
||||
) as f:
|
||||
_write_messages_to_file(f, input_messages)
|
||||
_write_response_to_file(f, response)
|
||||
|
||||
|
||||
def _write_messages_to_file(f: Any, messages: list[BaseMessage]) -> None:
|
||||
"""Write messages to conversation file"""
|
||||
for message in messages:
|
||||
f.write(f' {message.__class__.__name__} \n')
|
||||
|
||||
if isinstance(message.content, list):
|
||||
for item in message.content:
|
||||
if isinstance(item, dict) and item.get('type') == 'text':
|
||||
f.write(item['text'].strip() + '\n')
|
||||
elif isinstance(message.content, str):
|
||||
try:
|
||||
content = json.loads(message.content)
|
||||
f.write(json.dumps(content, indent=2) + '\n')
|
||||
except json.JSONDecodeError:
|
||||
f.write(message.content.strip() + '\n')
|
||||
|
||||
f.write('\n')
|
||||
|
||||
|
||||
def _write_response_to_file(f: Any, response: Any) -> None:
|
||||
"""Write model response to conversation file"""
|
||||
f.write(' RESPONSE\n')
|
||||
f.write(json.dumps(json.loads(response.model_dump_json(exclude_unset=True)), indent=2))
|
||||
|
||||
|
||||
# Add token counting related functions
|
||||
# Note: These functions have been moved from memory.py and agent.py to utils.py, removing the dependency on MessageManager class
|
||||
|
||||
def estimate_text_tokens(text: str, estimated_characters_per_token: int = 3) -> int:
|
||||
"""Roughly estimate token count in text
|
||||
|
||||
Args:
|
||||
text: The text to estimate tokens for
|
||||
estimated_characters_per_token: Estimated characters per token, default is 3
|
||||
|
||||
Returns:
|
||||
Estimated token count
|
||||
"""
|
||||
if not text:
|
||||
return 0
|
||||
# Use character count divided by average characters per token to estimate tokens
|
||||
return len(text) // estimated_characters_per_token
|
||||
|
||||
|
||||
def estimate_message_tokens(message: BaseMessage, image_tokens: int = 800,
|
||||
estimated_characters_per_token: int = 3) -> int:
|
||||
"""Roughly estimate token count for a single message
|
||||
|
||||
Args:
|
||||
message: The message to estimate tokens for
|
||||
image_tokens: Estimated tokens per image, default is 800
|
||||
estimated_characters_per_token: Estimated characters per token, default is 3
|
||||
|
||||
Returns:
|
||||
Estimated token count
|
||||
"""
|
||||
tokens = 0
|
||||
# Handle tuple case
|
||||
if isinstance(message, tuple):
|
||||
# Convert to string and estimate tokens
|
||||
message_str = str(message)
|
||||
return estimate_text_tokens(message_str, estimated_characters_per_token)
|
||||
|
||||
if isinstance(message.content, list):
|
||||
for item in message.content:
|
||||
if 'image_url' in item:
|
||||
tokens += image_tokens
|
||||
elif isinstance(item, dict) and 'text' in item:
|
||||
tokens += estimate_text_tokens(item['text'], estimated_characters_per_token)
|
||||
else:
|
||||
msg = message.content
|
||||
if hasattr(message, 'tool_calls'):
|
||||
msg += str(message.tool_calls) # type: ignore
|
||||
tokens += estimate_text_tokens(msg, estimated_characters_per_token)
|
||||
return tokens
|
||||
|
||||
|
||||
def estimate_messages_tokens(messages: list[BaseMessage], image_tokens: int = 800,
|
||||
estimated_characters_per_token: int = 3) -> int:
|
||||
"""Roughly estimate total token count for a list of messages
|
||||
|
||||
Args:
|
||||
messages: The list of messages to estimate tokens for
|
||||
image_tokens: Estimated tokens per image, default is 800
|
||||
estimated_characters_per_token: Estimated characters per token, default is 3
|
||||
|
||||
Returns:
|
||||
Estimated total token count
|
||||
"""
|
||||
total_tokens = 0
|
||||
for msg in messages:
|
||||
total_tokens += estimate_message_tokens(msg, image_tokens, estimated_characters_per_token)
|
||||
return total_tokens
|
||||
@@ -0,0 +1,3 @@
|
||||
# Common module
|
||||
|
||||
Tools commonly used in examples.
|
||||
@@ -0,0 +1,9 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
|
||||
from aworld.core.tool.base import Tool, AsyncTool
|
||||
from aworld.core.tool.action import ExecutableAction
|
||||
from aworld.utils.common import scan_packages
|
||||
|
||||
scan_packages("examples.common.tools", [Tool, AsyncTool, ExecutableAction])
|
||||
from examples.common.tools.browsers.action.actions import *
|
||||
@@ -0,0 +1,79 @@
|
||||
## Android Environment Setup Guide
|
||||
|
||||
This guide will help you set up a local Android environment for AgentWorld.
|
||||
|
||||
### Installation Steps
|
||||
|
||||
1. **Download and Install Android Studio**
|
||||
- Visit [https://developer.android.com/studio](https://developer.android.com/studio)
|
||||
- Download and install the latest version for your operating system
|
||||
|
||||
2. **Install ADB and Android Emulator**
|
||||
- Open Android Studio
|
||||
- Click on the top menu: Tools → SDK Manager
|
||||
<img src="../../../readme_assets/android_step1.png" width="70%" alt="SDK Manager">
|
||||
<!-- {:style="width:200px; height:auto;"} -->
|
||||
- Check the following components:
|
||||
- Android SDK Build-Tools
|
||||
- Android SDK Command-line Tools
|
||||
- Android Emulator
|
||||
- Android SDK Platform-Tools
|
||||
- Click "Apply" to install these components
|
||||
<img src="../../../readme_assets/android_step2.png" width="70%" alt="Check components">
|
||||
- **Important**: Copy the installation directory path (you'll need it later for configuration)
|
||||
|
||||
3. **Create a Virtual Device**
|
||||
- From the main menu, select: View → Tool Windows → Device Manager
|
||||
<img src="../../../readme_assets/android_step3.png" width="70%" alt="Device Manager">
|
||||
- Click the "+" button, then "Create Virtual Device"
|
||||
<img src="../../../readme_assets/android_step4.png" width="70%" alt="button">
|
||||
- Select a device (e.g., Medium Phone), then click "Next"
|
||||
<img src="../../../readme_assets/android_step5.png" width="70%" alt="next">
|
||||
- Select a image (e.g., VanillalceCream), then click "Next"
|
||||
<img src="../../../readme_assets/android_step6.png" width="70%" alt="next">
|
||||
- Configure device settings as needed, then click "Finish"
|
||||
- **Important**: Note down the AVD ID (device name) for later use
|
||||
<img src="../../../readme_assets/android_step7.png" width="70%" alt="avd id">
|
||||
|
||||
4. **Configure in Your Code**
|
||||
- Method 1: Default Acquisition of Emulator and ADB Installation Paths
|
||||
- Only set the AVD_ID copied during the earlier installation process.
|
||||
- Method 2: Manually Specify Emulator and ADB Installation Paths.Provide the following:
|
||||
- AVD_ID: The name of the virtual device you created
|
||||
- ADB path: Your SDK directory + "/platform-tools/adb"
|
||||
- Emulator path: Your SDK directory + "/emulator/emulator"
|
||||
### Example Code
|
||||
#### Method 1
|
||||
|
||||
```python
|
||||
from examples.common.tools.android.action.adb_controller import ADBController
|
||||
|
||||
# Initialize the Android controller
|
||||
android_controller = ADBController(avd_name="Medium_Phone_API_35")
|
||||
```
|
||||
#### Method 2
|
||||
|
||||
```python
|
||||
from examples.common.tools.android.action.adb_controller import ADBController
|
||||
|
||||
# Initialize the Android controller
|
||||
android_controller = ADBController(
|
||||
avd_name="Medium_Phone_API_35",
|
||||
adb_path="/Users/username/Library/Android/sdk/platform-tools/adb",
|
||||
emulator_path="/Users/username/Library/Android/sdk/emulator/emulator"
|
||||
)
|
||||
|
||||
# Now you can use this controller with your agent
|
||||
```
|
||||
|
||||
### Troubleshooting
|
||||
|
||||
- If the emulator fails to start, try increasing the memory allocation in the AVD settings
|
||||
- Make sure your paths are correct for your operating system:
|
||||
- Windows: Use backslashes or raw strings (r"C:\path\to\sdk")
|
||||
- macOS/Linux: Use forward slashes as shown in the example
|
||||
|
||||
### Additional Resources
|
||||
|
||||
- [Android SDK Official Documentation](https://developer.android.com/studio/intro)
|
||||
- [Android Emulator Documentation](https://developer.android.com/studio/run/emulator)
|
||||
@@ -0,0 +1,2 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
@@ -0,0 +1,77 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
|
||||
import json
|
||||
|
||||
from examples.common.tools.tool_action import AndroidAction
|
||||
from aworld.core.tool.action_factory import ActionFactory
|
||||
from aworld.core.common import ActionModel, ActionResult
|
||||
from examples.common.tools.android.action.adb_controller import ADBController
|
||||
from examples.common.tools.android.config.android_action_space import AndroidActionParamEnum
|
||||
from aworld.core.tool.action import ExecutableAction
|
||||
|
||||
|
||||
@ActionFactory.register(name=AndroidAction.TAP.value.name,
|
||||
desc=AndroidAction.TAP.value.desc,
|
||||
tool_name="android")
|
||||
class Tap(ExecutableAction):
|
||||
def act(self, action: ActionModel, **kwargs) -> ActionResult:
|
||||
controller: ADBController = kwargs.get('controller')
|
||||
tap_index = action.params[AndroidActionParamEnum.TAP_INDEX.value]
|
||||
if tap_index is None:
|
||||
raise Exception(f'Invalid action: {action}')
|
||||
controller.tap(tap_index)
|
||||
return ActionResult(content="", keep=True)
|
||||
|
||||
|
||||
@ActionFactory.register(name=AndroidAction.INPUT_TEXT.value.name,
|
||||
desc=AndroidAction.INPUT_TEXT.value.desc,
|
||||
tool_name="android")
|
||||
class InputText(ExecutableAction):
|
||||
def act(self, action: ActionModel, **kwargs) -> ActionResult:
|
||||
controller: ADBController = kwargs.get('controller')
|
||||
input_text = action.params[AndroidActionParamEnum.INPUT_TEXT.value]
|
||||
if input_text is None:
|
||||
raise Exception(f'Invalid action: {action}')
|
||||
controller.text(input_text)
|
||||
return ActionResult(content="", keep=True)
|
||||
|
||||
|
||||
@ActionFactory.register(name=AndroidAction.LONG_PRESS.value.name,
|
||||
desc=AndroidAction.LONG_PRESS.value.desc,
|
||||
tool_name="android")
|
||||
class LongPress(ExecutableAction):
|
||||
def act(self, action: ActionModel, **kwargs) -> ActionResult:
|
||||
controller: ADBController = kwargs.get('controller')
|
||||
long_press_index = action.params[AndroidActionParamEnum.LONG_PRESS_INDEX.value]
|
||||
if long_press_index is None:
|
||||
raise Exception(f'Invalid action: {action}')
|
||||
controller.long_press(long_press_index)
|
||||
return ActionResult(content="", keep=True)
|
||||
|
||||
|
||||
@ActionFactory.register(name=AndroidAction.SWIPE.value.name,
|
||||
desc=AndroidAction.SWIPE.value.desc,
|
||||
tool_name="android")
|
||||
class Swipe(ExecutableAction):
|
||||
def act(self, action: ActionModel, **kwargs) -> ActionResult:
|
||||
controller: ADBController = kwargs.get('controller')
|
||||
swipe_start_index = action.params[AndroidActionParamEnum.SWIPE_START_INDEX.value]
|
||||
direction = action.params[AndroidActionParamEnum.DIRECTION.value]
|
||||
dist = action.params.get(AndroidActionParamEnum.DIST.value, None)
|
||||
if swipe_start_index is None or direction is None:
|
||||
raise Exception(f'Invalid action: {action}')
|
||||
if dist:
|
||||
controller.swipe(swipe_start_index, direction, dist)
|
||||
else:
|
||||
controller.swipe(swipe_start_index, direction)
|
||||
return ActionResult(content="", keep=True)
|
||||
|
||||
|
||||
@ActionFactory.register(name=AndroidAction.DONE.value.name,
|
||||
desc=AndroidAction.DONE.value.desc,
|
||||
tool_name="android")
|
||||
class Done(ExecutableAction):
|
||||
def act(self, action: ActionModel, **kwargs) -> ActionResult:
|
||||
output_dict = action.model_dump(exclude={'success'})
|
||||
return ActionResult(is_done=True, success=True, content=json.dumps(output_dict))
|
||||
+541
@@ -0,0 +1,541 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
|
||||
import subprocess
|
||||
import time
|
||||
import re
|
||||
import traceback
|
||||
from time import sleep
|
||||
from typing import Optional, Tuple, List
|
||||
import base64
|
||||
|
||||
import xml.etree.ElementTree as ET
|
||||
import os
|
||||
|
||||
from aworld.logs.util import logger, color_log, Color
|
||||
from aworld.utils import import_package
|
||||
|
||||
configs = {"MIN_DIST": 30}
|
||||
|
||||
|
||||
class AndroidElement:
|
||||
def __init__(self, uid, bbox, attrib):
|
||||
self.uid = uid
|
||||
self.bbox = bbox
|
||||
self.attrib = attrib
|
||||
import_package('cv2', install_name='opencv-python')
|
||||
import_package('pyshine')
|
||||
|
||||
def get_id_from_element(elem):
|
||||
bounds = elem.attrib["bounds"][1:-1].split("][")
|
||||
x1, y1 = map(int, bounds[0].split(","))
|
||||
x2, y2 = map(int, bounds[1].split(","))
|
||||
elem_w, elem_h = x2 - x1, y2 - y1
|
||||
if "resource-id" in elem.attrib and elem.attrib["resource-id"]:
|
||||
elem_id = elem.attrib["resource-id"].replace(":", ".").replace("/", "_")
|
||||
else:
|
||||
elem_id = f"{elem.attrib['class']}_{elem_w}_{elem_h}"
|
||||
if "content-desc" in elem.attrib and elem.attrib["content-desc"] and len(elem.attrib["content-desc"]) < 20:
|
||||
content_desc = elem.attrib['content-desc'].replace("/", "_").replace(" ", "").replace(":", "_")
|
||||
elem_id += f"_{content_desc}"
|
||||
return elem_id
|
||||
|
||||
|
||||
def traverse_tree(xml_path, elem_list, attrib, add_index=False):
|
||||
path = []
|
||||
for event, elem in ET.iterparse(xml_path, ['start', 'end']):
|
||||
if event == 'start':
|
||||
path.append(elem)
|
||||
if attrib in elem.attrib and elem.attrib[attrib] == "true":
|
||||
parent_prefix = ""
|
||||
if len(path) > 1:
|
||||
parent_elem = path[-2]
|
||||
# Checks if the parent element has the required attributes
|
||||
has_bounds = "bounds" in parent_elem.attrib
|
||||
has_rid_or_class = "resource-id" in parent_elem.attrib or "class" in parent_elem.attrib
|
||||
if has_bounds and has_rid_or_class:
|
||||
parent_prefix = get_id_from_element(parent_elem)
|
||||
bounds = elem.attrib["bounds"][1:-1].split("][")
|
||||
x1, y1 = map(int, bounds[0].split(","))
|
||||
x2, y2 = map(int, bounds[1].split(","))
|
||||
center = (x1 + x2) // 2, (y1 + y2) // 2
|
||||
elem_id = get_id_from_element(elem)
|
||||
if parent_prefix:
|
||||
elem_id = parent_prefix + "_" + elem_id
|
||||
if add_index:
|
||||
elem_id += f"_{elem.attrib['index']}"
|
||||
close = False
|
||||
for e in elem_list:
|
||||
bbox = e.bbox
|
||||
center_ = (bbox[0][0] + bbox[1][0]) // 2, (bbox[0][1] + bbox[1][1]) // 2
|
||||
dist = (abs(center[0] - center_[0]) ** 2 + abs(center[1] - center_[1]) ** 2) ** 0.5
|
||||
if dist <= configs["MIN_DIST"]:
|
||||
close = True
|
||||
break
|
||||
if not close:
|
||||
elem_list.append(AndroidElement(elem_id, ((x1, y1), (x2, y2)), attrib))
|
||||
|
||||
if event == 'end':
|
||||
path.pop()
|
||||
|
||||
|
||||
def create_directory_for_file(file_path):
|
||||
# Extract the directory from the file path
|
||||
directory = os.path.dirname(file_path)
|
||||
|
||||
# Check if the directory exists
|
||||
if not os.path.exists(directory):
|
||||
# Create the directory
|
||||
os.makedirs(directory)
|
||||
# Print the absolute path of the directory
|
||||
absolute_directory_path = os.path.abspath(directory)
|
||||
logger.info(f"Directory absolute path: {absolute_directory_path}")
|
||||
|
||||
|
||||
def draw_bbox_multi(img_path, output_path, elem_list):
|
||||
import cv2
|
||||
import pyshine as ps
|
||||
|
||||
imgcv = cv2.imread(img_path)
|
||||
count = 1
|
||||
for elem in elem_list:
|
||||
try:
|
||||
top_left = elem.bbox[0]
|
||||
bottom_right = elem.bbox[1]
|
||||
left, top = top_left[0], top_left[1]
|
||||
right, bottom = bottom_right[0], bottom_right[1]
|
||||
|
||||
# draw rectangle
|
||||
cv2.rectangle(imgcv,
|
||||
(left, top),
|
||||
(right, bottom),
|
||||
(0, 0, 221),
|
||||
3)
|
||||
|
||||
label = str(count)
|
||||
imgcv = ps.putBText(imgcv, label, text_offset_x=(left + right) // 2 + 10,
|
||||
text_offset_y=(top + bottom) // 2 + 10,
|
||||
vspace=10, hspace=10, font_scale=1, thickness=2, background_RGB=(221, 0, 0),
|
||||
text_RGB=(255, 255, 255), alpha=0.0)
|
||||
|
||||
except Exception as e:
|
||||
color_log(f"ERROR: An exception occurs while labeling the image\n{e}", Color.red)
|
||||
logger.info(traceback.print_exc())
|
||||
count += 1
|
||||
cv2.imwrite(output_path, imgcv)
|
||||
return imgcv
|
||||
|
||||
|
||||
def draw_grid(img_path, output_path):
|
||||
import cv2
|
||||
|
||||
def get_unit_len(n):
|
||||
for i in range(1, n + 1):
|
||||
if n % i == 0 and 120 <= i <= 180:
|
||||
return i
|
||||
return -1
|
||||
|
||||
image = cv2.imread(img_path)
|
||||
height, width, _ = image.shape
|
||||
color = (255, 116, 113)
|
||||
unit_height = get_unit_len(height)
|
||||
if unit_height < 0:
|
||||
unit_height = 120
|
||||
unit_width = get_unit_len(width)
|
||||
if unit_width < 0:
|
||||
unit_width = 120
|
||||
thick = int(unit_width // 50)
|
||||
rows = height // unit_height
|
||||
cols = width // unit_width
|
||||
for i in range(rows):
|
||||
for j in range(cols):
|
||||
label = i * cols + j + 1
|
||||
left = int(j * unit_width)
|
||||
top = int(i * unit_height)
|
||||
right = int((j + 1) * unit_width)
|
||||
bottom = int((i + 1) * unit_height)
|
||||
cv2.rectangle(image, (left, top), (right, bottom), color, thick // 2)
|
||||
cv2.putText(image, str(label), (left + int(unit_width * 0.05) + 3, top + int(unit_height * 0.3) + 3), 0,
|
||||
int(0.01 * unit_width), (0, 0, 0), thick)
|
||||
cv2.putText(image, str(label), (left + int(unit_width * 0.05), top + int(unit_height * 0.3)), 0,
|
||||
int(0.01 * unit_width), color, thick)
|
||||
cv2.imwrite(output_path, image)
|
||||
return rows, cols
|
||||
|
||||
|
||||
def encode_image(image_path):
|
||||
with open(image_path, "rb") as image_file:
|
||||
return base64.b64encode(image_file.read()).decode('utf-8')
|
||||
|
||||
|
||||
class ADBController:
|
||||
def __init__(self, avd_name: str = None,
|
||||
adb_path: str = os.path.expanduser('~') + "/Library/Android/sdk/platform-tools/adb",
|
||||
emulator_path: str = os.path.expanduser('~') + "/Library/Android/sdk/emulator/emulator",
|
||||
timeout: int = 30):
|
||||
self.avd_name = avd_name
|
||||
self.adb_path = adb_path
|
||||
self.emulator_path = emulator_path
|
||||
self.timeout = timeout
|
||||
self.emulator_process = None
|
||||
self.device_serial = "emulator-5554" # default
|
||||
self.current_elem_list = []
|
||||
self.width, self.height = 0, 0
|
||||
|
||||
def start_emulator(self, avd_name: str = None, headless: bool = False,
|
||||
max_retry: int = 2) -> bool:
|
||||
avd = avd_name or self.avd_name
|
||||
if not avd:
|
||||
raise ValueError("AVD name must be specified")
|
||||
|
||||
for attempt in range(max_retry + 1):
|
||||
if self._start_emulator_process(avd, headless):
|
||||
if self._wait_for_device():
|
||||
logger.info(f"start success,attempt count:{attempt + 1}")
|
||||
self.width, self.height = self.get_screen_size()
|
||||
return True
|
||||
self.stop_emulator()
|
||||
return False
|
||||
|
||||
def _start_emulator_process(self, avd: str, headless: bool) -> bool:
|
||||
try:
|
||||
cmd = [
|
||||
self.emulator_path,
|
||||
f"@{avd}",
|
||||
"-no-snapshot",
|
||||
"-no-audio",
|
||||
"-gpu", "swiftshader",
|
||||
"-wipe-data"
|
||||
]
|
||||
if headless:
|
||||
cmd.append("-no-window")
|
||||
|
||||
self.emulator_process = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.STDOUT
|
||||
)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.warning(f"adb start fail: {str(e)}")
|
||||
return False
|
||||
|
||||
def stop_emulator(self) -> bool:
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[self.adb_path, "-s", self.device_serial, "emu", "kill"],
|
||||
timeout=self.timeout,
|
||||
capture_output=True,
|
||||
text=True
|
||||
)
|
||||
return "OK" in result.stdout
|
||||
except subprocess.TimeoutExpired:
|
||||
return False
|
||||
finally:
|
||||
if self.emulator_process:
|
||||
self.emulator_process.terminate()
|
||||
|
||||
def execute_adb(self, command: list, device_serial: str = None) -> Tuple[bool, str]:
|
||||
"""execute adb command"""
|
||||
device = device_serial or self.device_serial
|
||||
full_cmd = [self.adb_path, "-s", device] + command
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
full_cmd,
|
||||
timeout=self.timeout,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True
|
||||
)
|
||||
return True, result.stdout.strip()
|
||||
except subprocess.CalledProcessError as e:
|
||||
return False, f"Command failed: {e.stderr}"
|
||||
except Exception as e:
|
||||
return False, str(e)
|
||||
|
||||
def execute_adb_with_stdout(self, command: List[str]) -> Tuple[bool, Optional[str]]:
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["adb", "-s", self.device_serial] + command,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
timeout=10
|
||||
)
|
||||
if result.returncode == 0:
|
||||
return True, result.stdout.strip()
|
||||
else:
|
||||
return False, None
|
||||
except subprocess.TimeoutExpired:
|
||||
return False, None
|
||||
except Exception as e:
|
||||
return False, None
|
||||
|
||||
# ---------- device operate ----------
|
||||
|
||||
def screenshot(self, save_path: str) -> bool:
|
||||
timestamp = int(time.time())
|
||||
remote_path = f"/sdcard/screenshot_{timestamp}.png"
|
||||
|
||||
success, _ = self.execute_adb(["shell", "screencap", "-p", remote_path])
|
||||
if not success:
|
||||
return False
|
||||
|
||||
return self._pull_file(remote_path, save_path)
|
||||
|
||||
def dump_ui_xml(self, save_path: str) -> Optional[str]:
|
||||
remote_path = "/sdcard/ui_dump.xml"
|
||||
success, _ = self.execute_adb(["shell", "uiautomator", "dump", remote_path])
|
||||
if not success:
|
||||
logger.info("dump ui xml fail")
|
||||
return None
|
||||
success = self._pull_file(remote_path, save_path)
|
||||
if not success:
|
||||
logger.info("pull ui xml fail")
|
||||
return None
|
||||
|
||||
with open(save_path, 'r', encoding='utf-8') as f:
|
||||
xml_content = f.read()
|
||||
return xml_content
|
||||
|
||||
def tap(self, element: int):
|
||||
x, y = self.__get_element_center(element)
|
||||
self.__tap_coordinate(x, y)
|
||||
|
||||
def text(self, text: str):
|
||||
"""
|
||||
Input text, automatically replacing spaces with %s for proper ADB text input.
|
||||
|
||||
Parameters:
|
||||
text: The text to input
|
||||
"""
|
||||
# Replace spaces with %s for proper handling in ADB
|
||||
formatted_text = text.replace(" ", "%s")
|
||||
success, _ = self.execute_adb(["shell", "input", "text", formatted_text])
|
||||
return success
|
||||
|
||||
def long_press(self, element: int):
|
||||
x, y = self.__get_element_center(element)
|
||||
self.__swipe_coordinate(x, y, x, y, 2000)
|
||||
|
||||
def swipe(self, element: int, direction: str, dist: str = "medium"):
|
||||
"""
|
||||
Perform swipe operations based on screen element labels
|
||||
|
||||
Parameters:
|
||||
element_tag: digital label displayed on the interface (1-based)
|
||||
direction: swipe direction ["up", "down", "left", "right"]
|
||||
dist: swipe distance ["short", "medium", "long"]
|
||||
"""
|
||||
|
||||
# 获取元素坐标
|
||||
x, y = self.__get_element_center(element)
|
||||
|
||||
unit_dist = int(self.width / 10)
|
||||
if dist == "long":
|
||||
unit_dist *= 3
|
||||
elif dist == "medium":
|
||||
unit_dist *= 2
|
||||
if direction == "up":
|
||||
offset = 0, -2 * unit_dist
|
||||
elif direction == "down":
|
||||
offset = 0, 2 * unit_dist
|
||||
elif direction == "left":
|
||||
offset = -1 * unit_dist, 0
|
||||
elif direction == "right":
|
||||
offset = unit_dist, 0
|
||||
else:
|
||||
return False
|
||||
|
||||
self.__swipe_coordinate(x, y, x + offset[0], y + offset[1])
|
||||
|
||||
def screenshot_and_annotate(self, name_prefix=None, return_base64=True):
|
||||
import cv2
|
||||
|
||||
"""Collect screen information and mark interactive elements, and return data containing Base64 images"""
|
||||
sleep(3)
|
||||
if name_prefix is None:
|
||||
name_prefix = str(time.time())
|
||||
tmp_files_dir = os.path.join(os.path.dirname(__file__), "tmp_files")
|
||||
os.makedirs(tmp_files_dir, exist_ok=True)
|
||||
screenshot_path = os.path.join(tmp_files_dir, f"{name_prefix}_origin.png")
|
||||
screenshot_res = self.screenshot(screenshot_path)
|
||||
xml_path = os.path.join(tmp_files_dir, f"{name_prefix}.xml")
|
||||
xml_res = self.dump_ui_xml(xml_path)
|
||||
if screenshot_res == "ERROR" or xml_res is None:
|
||||
logger.warning(f"Failed to take screenshot or read XML")
|
||||
return None, None
|
||||
|
||||
# Parsing interactive elements
|
||||
clickable_list = []
|
||||
focusable_list = []
|
||||
traverse_tree(xml_path, clickable_list, "clickable", True)
|
||||
traverse_tree(xml_path, focusable_list, "focusable", True)
|
||||
|
||||
# Merge a list of duplicate elements
|
||||
elem_list = clickable_list.copy()
|
||||
for elem in focusable_list:
|
||||
bbox = elem.bbox
|
||||
center = (bbox[0][0] + bbox[1][0]) // 2, (bbox[0][1] + bbox[1][1]) // 2
|
||||
if not any(
|
||||
((center[0] - ((e.bbox[0][0] + e.bbox[1][0]) // 2)) ** 2 +
|
||||
(center[1] - ((e.bbox[0][1] + e.bbox[1][1]) // 2)) ** 2) ** 0.5 <= configs["MIN_DIST"]
|
||||
for e in clickable_list
|
||||
):
|
||||
elem_list.append(elem)
|
||||
|
||||
# Generate annotated images
|
||||
labeled_path = os.path.join(tmp_files_dir, f"{name_prefix}_labeled.png")
|
||||
labeled_img = draw_bbox_multi(screenshot_path, labeled_path, elem_list)
|
||||
|
||||
# Show Image Window
|
||||
# cv2.imshow("image", labeled_img)
|
||||
# cv2.waitKey(0)
|
||||
# cv2.destroyAllWindows()
|
||||
|
||||
# Base64 encoding
|
||||
base64_str = None
|
||||
if return_base64:
|
||||
# Convert color space BGR->RGB
|
||||
rgb_image = cv2.cvtColor(labeled_img, cv2.COLOR_BGR2RGB)
|
||||
# Compress to JPEG format (with adjustable quality parameters)
|
||||
success, buffer = cv2.imencode(".jpg", rgb_image, [int(cv2.IMWRITE_JPEG_QUALITY), 85])
|
||||
if success:
|
||||
base64_str = base64.b64encode(buffer).decode("utf-8")
|
||||
|
||||
self.current_elem_list = elem_list.copy()
|
||||
logger.info(f"Current elem size{len(self.current_elem_list)}")
|
||||
return xml_res, base64_str
|
||||
|
||||
def setup_connection(self) -> bool:
|
||||
"""Intelligent initialization device connection"""
|
||||
# Prioritize physical equipment testing
|
||||
if self.__connect_physical_device():
|
||||
return True
|
||||
|
||||
# Try connecting to the simulator
|
||||
if self.avd_name and self.start_emulator():
|
||||
return True
|
||||
|
||||
raise ConnectionError("No available device found, please connect your phone or configure the simulator")
|
||||
|
||||
# ---------- Helper Methods ----------
|
||||
def __connect_physical_device(self) -> bool:
|
||||
"""Connect an authorized USB device"""
|
||||
devices = self.__get_authorized_devices()
|
||||
if not devices:
|
||||
return False
|
||||
|
||||
self.device = devices[0]
|
||||
logger.info(f"Connected physical device: {self.device}")
|
||||
self.device_serial = self.device
|
||||
self.width, self.height = self.get_screen_size()
|
||||
return True
|
||||
|
||||
def __get_authorized_devices(self) -> list:
|
||||
"""Get a list of authorized devices"""
|
||||
success, output = self.execute_adb(["devices"])
|
||||
if not success:
|
||||
return []
|
||||
|
||||
return [
|
||||
line.split("\t")[0]
|
||||
for line in output.splitlines()
|
||||
if "\tdevice" in line and "emulator" not in line
|
||||
]
|
||||
|
||||
def __tap_coordinate(self, x: int, y: int) -> bool:
|
||||
"""Click screen coordinates"""
|
||||
success, _ = self.execute_adb(["shell", "input", "tap", str(x), str(y)])
|
||||
return success
|
||||
|
||||
def __get_element_center(self, elem_idx: int) -> tuple:
|
||||
"""Calculate the coordinates of the center of the element"""
|
||||
tl, br = self.current_elem_list[int(elem_idx) - 1].bbox
|
||||
return (tl[0] + br[0]) // 2, (tl[1] + br[1]) // 2
|
||||
|
||||
def __swipe_coordinate(self, x1: int, y1: int, x2: int, y2: int, duration: int = 300) -> bool:
|
||||
"""Slide Operation"""
|
||||
success, _ = self.execute_adb([
|
||||
"shell", "input", "swipe",
|
||||
str(x1), str(y1), str(x2), str(y2),
|
||||
str(duration)
|
||||
])
|
||||
return success
|
||||
|
||||
def _wait_for_device(self, timeout: int = 300) -> bool:
|
||||
"""Three-level waiting detection strategy"""
|
||||
start_time = time.time()
|
||||
stages = {
|
||||
"adb_connected": False,
|
||||
"boot_completed": False,
|
||||
"services_ready": False
|
||||
}
|
||||
|
||||
while time.time() - start_time < timeout:
|
||||
# Step 1: Detect adb connection
|
||||
if not stages["adb_connected"]:
|
||||
_, devices = self.execute_adb(["devices"])
|
||||
if self.device_serial in devices:
|
||||
stages["adb_connected"] = True
|
||||
|
||||
# Step 2: Detection system boot completed
|
||||
if stages["adb_connected"] and not stages["boot_completed"]:
|
||||
_, output = self.execute_adb([
|
||||
"shell", "getprop", "sys.boot_completed"
|
||||
])
|
||||
if output.strip() == "1":
|
||||
stages["boot_completed"] = True
|
||||
|
||||
# Step 3: Detecting Graphics Service Readiness
|
||||
if stages["boot_completed"] and not stages["services_ready"]:
|
||||
_, output = self.execute_adb([
|
||||
"shell", "service check SurfaceFlinger"
|
||||
])
|
||||
if "found" in output.lower():
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def _pull_file(self, remote: str, local: str) -> bool:
|
||||
"""Pull device files to local"""
|
||||
create_directory_for_file(local)
|
||||
success, _ = self.execute_adb(["pull", remote, local])
|
||||
if success:
|
||||
self.execute_adb(["shell", "rm", remote]) # 清理临时文件
|
||||
return success
|
||||
|
||||
def get_screen_size(self) -> Optional[Tuple[int, int]]:
|
||||
"""Get screen resolution"""
|
||||
success, output = self.execute_adb(["shell", "wm", "size"])
|
||||
if not success:
|
||||
return None
|
||||
|
||||
match = re.search(r"(\d+)x(\d+)", output)
|
||||
if match:
|
||||
return int(match.group(1)), int(match.group(2))
|
||||
return None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Examples
|
||||
controller = ADBController(avd_name="Medium_Phone_API_35")
|
||||
|
||||
# controller.stop_emulator()
|
||||
if controller.setup_connection():
|
||||
logger.info("Simulator started successfully")
|
||||
width, height = controller.get_screen_size()
|
||||
logger.info(f"Get the screen size{width},{height}")
|
||||
|
||||
# Take screenshots and annotate them
|
||||
controller.screenshot_and_annotate()
|
||||
controller.swipe(6, "up")
|
||||
|
||||
# controller.screenshot_and_annotate()
|
||||
# controller.tap(6)
|
||||
xml_txt, base64_txt = controller.screenshot_and_annotate()
|
||||
logger.info(xml_txt)
|
||||
|
||||
# controller.stop_emulator()
|
||||
logger.info("Close the simulator")
|
||||
@@ -0,0 +1,42 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
|
||||
from typing import List
|
||||
from aworld.core.tool.action_factory import ActionFactory
|
||||
from aworld.core.common import ActionModel, ActionResult
|
||||
from aworld.logs.util import logger
|
||||
from examples.common.tools.android.action.adb_controller import ADBController
|
||||
from aworld.core.tool.base import ToolActionExecutor
|
||||
|
||||
|
||||
class AndroidToolActionExecutor(ToolActionExecutor):
|
||||
|
||||
def __init__(self, controller: ADBController):
|
||||
self.controller = controller
|
||||
|
||||
def execute_action(self, actions: List[ActionModel], **kwargs) -> list[ActionResult]:
|
||||
"""Execute the specified android action sequence by agent policy.
|
||||
|
||||
Args:
|
||||
actions: Tool action sequence.
|
||||
|
||||
Returns:
|
||||
Browser action result list.
|
||||
"""
|
||||
action_results = []
|
||||
for action in actions:
|
||||
action_result = self._exec(action, **kwargs)
|
||||
action_results.append(action_result)
|
||||
return action_results
|
||||
|
||||
def _exec(self, action_model: ActionModel, **kwargs):
|
||||
action_name = action_model.action_name
|
||||
if action_name not in ActionFactory:
|
||||
action_name = action_model.tool_name + action_model.action_name
|
||||
if action_name not in ActionFactory:
|
||||
raise ValueError(f'Action {action_name} not found')
|
||||
|
||||
action = ActionFactory(action_name)
|
||||
action_result = action.act(action_model, controller=self.controller, **kwargs)
|
||||
logger.info(f"{action_name} execute finished")
|
||||
return action_result
|
||||
@@ -0,0 +1,92 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from typing import Any, Tuple, List, Dict
|
||||
|
||||
from examples.common.tools.tool_action import AndroidAction
|
||||
from aworld.core.common import ActionModel, Observation, ActionResult
|
||||
from aworld.logs.util import logger
|
||||
from examples.common.tools.android.action.adb_controller import ADBController
|
||||
from examples.common.tools.android.action.executor import AndroidToolActionExecutor
|
||||
from examples.common.tools.conf import AndroidToolConfig
|
||||
from aworld.core.tool.base import ToolFactory, Tool
|
||||
from aworld.tools.utils import build_observation
|
||||
|
||||
ALL_UNICODE_CHARS = frozenset(chr(i) for i in range(0x10FFFF + 1))
|
||||
|
||||
|
||||
@ToolFactory.register(name="android",
|
||||
desc="android",
|
||||
supported_action=AndroidAction,
|
||||
conf_file_name=f'android_tool.yaml',
|
||||
dir=f"{Path(__file__).parent.absolute()}")
|
||||
class AndroidTool(Tool):
|
||||
|
||||
def __init__(self, conf: AndroidToolConfig, **kwargs):
|
||||
super(AndroidTool, self).__init__(conf, **kwargs)
|
||||
self.controller = ADBController(avd_name=self.conf.get('avd_name'),
|
||||
adb_path=self.conf.get('adb_path'),
|
||||
emulator_path=self.conf.get('emulator_path'))
|
||||
|
||||
if self.conf.get("custom_executor"):
|
||||
self.action_executor = AndroidToolActionExecutor(self.controller)
|
||||
|
||||
def reset(self, *, seed: int | None = None, options: Dict[str, str] | None = None) -> Tuple[
|
||||
Observation, Dict[str, Any]]:
|
||||
# self.controller.stop_emulator()
|
||||
# self.controller.start_emulator()
|
||||
self.controller.setup_connection()
|
||||
logger.info("start emulator successfully...")
|
||||
# snapshot screen and annotate
|
||||
xml, pic_base64 = self.get_observation()
|
||||
action_result_list = [ActionResult(content='start', keep=True)]
|
||||
return build_observation(observer=self.name(),
|
||||
ability='',
|
||||
dom_tree=xml,
|
||||
image=pic_base64,
|
||||
action_result=action_result_list), {}
|
||||
|
||||
def do_step(self, action_list: List[ActionModel], **kwargs) -> Tuple[
|
||||
Observation, float, bool, bool, Dict[str, Any]]:
|
||||
|
||||
exec_state = 0
|
||||
fail_error = ""
|
||||
action_result_list = None
|
||||
try:
|
||||
action_result_list = self.action_executor.execute_action(action_list, **kwargs)
|
||||
exec_state = 1
|
||||
except Exception as e:
|
||||
traceback.print_exc()
|
||||
fail_error = str(e)
|
||||
|
||||
terminated = kwargs.get("terminated", False)
|
||||
if action_result_list:
|
||||
for action_result in action_result_list:
|
||||
if action_result.is_done:
|
||||
terminated = action_result.is_done
|
||||
self._finish = True
|
||||
|
||||
info = {"exception": fail_error}
|
||||
info.update(kwargs)
|
||||
xml, pic_base64 = self.get_observation()
|
||||
|
||||
return (build_observation(observer=self.name(),
|
||||
ability=action_list[-1].action_name,
|
||||
dom_tree=xml,
|
||||
image=pic_base64,
|
||||
action_result=action_result_list),
|
||||
exec_state,
|
||||
terminated,
|
||||
kwargs.get("truncated", False),
|
||||
info)
|
||||
|
||||
def close(self):
|
||||
self.controller.stop_emulator()
|
||||
|
||||
def get_controller(self):
|
||||
return self.controller
|
||||
|
||||
def get_observation(self) -> Observation:
|
||||
return self.controller.screenshot_and_annotate()
|
||||
@@ -0,0 +1,8 @@
|
||||
avd_name:
|
||||
adb_path:
|
||||
emulator_path:
|
||||
headless: False
|
||||
custom_executor: True
|
||||
enable_recording: False
|
||||
working_dir:
|
||||
max_retry: 3
|
||||
+26
@@ -0,0 +1,26 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class AndroidActionParamEnum(Enum):
|
||||
TAP_INDEX = "index"
|
||||
LONG_PRESS_INDEX = "index"
|
||||
INPUT_TEXT = "text"
|
||||
SWIPE_START_INDEX = "index"
|
||||
DIRECTION = "direction"
|
||||
DIST = "dist"
|
||||
|
||||
|
||||
class DirectionParamEnum(Enum):
|
||||
UP = "up"
|
||||
DOWN = "down"
|
||||
LEFT = "left"
|
||||
RIGHT = "right"
|
||||
|
||||
|
||||
class DistParamEnum(Enum):
|
||||
SHORT = "short"
|
||||
MEDIUM = "medium"
|
||||
LONG = "long"
|
||||
@@ -0,0 +1,2 @@
|
||||
opencv-python~=4.11.0.86
|
||||
pyshine~=0.0.9
|
||||
@@ -0,0 +1,259 @@
|
||||
# coding: utf-8
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
import requests
|
||||
|
||||
from typing import Tuple, Any, List, Dict
|
||||
|
||||
from examples.common.tools.tool_action import SearchAction
|
||||
from aworld.core.tool.action_factory import ActionFactory
|
||||
from aworld.core.common import ActionModel, ActionResult
|
||||
from aworld.logs.util import logger
|
||||
from aworld.utils import import_package
|
||||
from aworld.core.tool.action import ExecutableAction
|
||||
|
||||
|
||||
# @ActionFactory.register(name=SearchAction.WIKI.value.name,
|
||||
# desc=SearchAction.WIKI.value.desc,
|
||||
# tool_name='search_api')
|
||||
class SearchWiki(ExecutableAction):
|
||||
def __init__(self):
|
||||
import_package("wikipedia")
|
||||
|
||||
def act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
import wikipedia
|
||||
|
||||
query = action.params.get("query")
|
||||
logger.info(f"Calling search_wiki api with query: {query}")
|
||||
|
||||
result: str = ''
|
||||
try:
|
||||
page = wikipedia.page(query)
|
||||
result_dict = {
|
||||
'url': page.url,
|
||||
'title': page.title,
|
||||
'content': page.content,
|
||||
}
|
||||
result = str(result_dict)
|
||||
except wikipedia.exceptions.DisambiguationError as e:
|
||||
result = wikipedia.summary(
|
||||
e.options[0], sentences=5, auto_suggest=False
|
||||
)
|
||||
except wikipedia.exceptions.PageError:
|
||||
result = (
|
||||
"There is no page in Wikipedia corresponding to entity "
|
||||
f"{query}, please specify another word to describe the"
|
||||
" entity to be searched."
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"An exception occurred during the search: {e}")
|
||||
result = f"An exception occurred during the search: {e}"
|
||||
logger.debug(f"wiki result: {result}")
|
||||
return ActionResult(content=result, keep=True, is_done=True), None
|
||||
|
||||
|
||||
# @ActionFactory.register(name=SearchAction.DUCK_GO.value.name,
|
||||
# desc=SearchAction.DUCK_GO.value.desc,
|
||||
# tool_name="search_api")
|
||||
class Duckduckgo(ExecutableAction):
|
||||
def __init__(self):
|
||||
import_package("duckduckgo_search")
|
||||
|
||||
def act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
r"""Use DuckDuckGo search engine to search information for
|
||||
the given query.
|
||||
|
||||
This function queries the DuckDuckGo API for related topics to
|
||||
the given search term. The results are formatted into a list of
|
||||
dictionaries, each representing a search result.
|
||||
|
||||
Args:
|
||||
query (str): The query to be searched.
|
||||
source (str): The type of information to query (e.g., "text",
|
||||
"images", "videos"). Defaults to "text".
|
||||
max_results (int): Max number of results, defaults to `5`.
|
||||
|
||||
Returns:
|
||||
List[Dict[str, Any]]: A list of dictionaries where each dictionary
|
||||
represents a search result.
|
||||
"""
|
||||
|
||||
from duckduckgo_search import DDGS
|
||||
|
||||
params = action.params
|
||||
query = params.get("query")
|
||||
max_results = params.get("max_results", 5)
|
||||
source = params.get("source", "text")
|
||||
|
||||
logger.debug(f"Calling search_duckduckgo function with query: {query}")
|
||||
|
||||
ddgs = DDGS()
|
||||
responses: List[Dict[str, Any]] = []
|
||||
|
||||
if source == "text":
|
||||
try:
|
||||
results = ddgs.text(keywords=query, max_results=max_results)
|
||||
except Exception as e:
|
||||
# Handle specific exceptions or general request exceptions
|
||||
responses.append({"error": f"duckduckgo search failed.{e}"})
|
||||
return ActionResult(content="duckduckgo search failed", keep=True), responses
|
||||
|
||||
for i, result in enumerate(results, start=1):
|
||||
# Creating a response object with a similar structure
|
||||
response = {
|
||||
"result_id": i,
|
||||
"title": result["title"],
|
||||
"description": result["body"],
|
||||
"url": result["href"],
|
||||
}
|
||||
responses.append(response)
|
||||
elif source == "images":
|
||||
try:
|
||||
results = ddgs.images(keywords=query, max_results=max_results)
|
||||
except Exception as e:
|
||||
# Handle specific exceptions or general request exceptions
|
||||
responses.append({"error": f"duckduckgo search failed.{e}"})
|
||||
return ActionResult(content="duckduckgo search failed", keep=True), responses
|
||||
|
||||
# Iterate over results found
|
||||
for i, result in enumerate(results, start=1):
|
||||
# Creating a response object with a similar structure
|
||||
response = {
|
||||
"result_id": i,
|
||||
"title": result["title"],
|
||||
"image": result["image"],
|
||||
"url": result["url"],
|
||||
"source": result["source"],
|
||||
}
|
||||
responses.append(response)
|
||||
elif source == "videos":
|
||||
try:
|
||||
results = ddgs.videos(keywords=query, max_results=max_results)
|
||||
except Exception as e:
|
||||
# Handle specific exceptions or general request exceptions
|
||||
responses.append({"error": f"duckduckgo search failed.{e}"})
|
||||
return ActionResult(content="duckduckgo search failed", keep=True), responses
|
||||
|
||||
# Iterate over results found
|
||||
for i, result in enumerate(results, start=1):
|
||||
# Creating a response object with a similar structure
|
||||
response = {
|
||||
"result_id": i,
|
||||
"title": result["title"],
|
||||
"description": result["description"],
|
||||
"embed_url": result["embed_url"],
|
||||
"publisher": result["publisher"],
|
||||
"duration": result["duration"],
|
||||
"published": result["published"],
|
||||
}
|
||||
responses.append(response)
|
||||
logger.debug(f"Search results: {responses}")
|
||||
return ActionResult(content=json.dumps(responses), keep=True, is_done=True), None
|
||||
|
||||
|
||||
# @ActionFactory.register(name=SearchAction.GOOGLE.value.name,
|
||||
# desc=SearchAction.GOOGLE.value.desc,
|
||||
# tool_name="search_api")
|
||||
class SearchGoogle(ExecutableAction):
|
||||
def act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
query = action.params.get("query")
|
||||
num_result_pages = action.params.get("num_result_pages", 6)
|
||||
# https://developers.google.com/custom-search/v1/overview
|
||||
api_key = action.params.get("api_key", os.environ.get("GOOGLE_API_KEY"))
|
||||
# https://cse.google.com/cse/all
|
||||
engine_id = action.params.get("engine_id", os.environ.get("GOOGLE_ENGINE_ID"))
|
||||
logger.debug(f"Calling search_google function with query: {query}")
|
||||
|
||||
# Using the first page
|
||||
start_page_idx = 1
|
||||
# Different language may get different result
|
||||
search_language = "en"
|
||||
# How many pages to return
|
||||
num_result_pages = num_result_pages
|
||||
# Constructing the URL
|
||||
# Doc: https://developers.google.com/custom-search/v1/using_rest
|
||||
url = f"https://www.googleapis.com/customsearch/v1?key={api_key}&cx={engine_id}&q={query}&start={start_page_idx}&lr={search_language}&num={num_result_pages}"
|
||||
responses = []
|
||||
try:
|
||||
result = requests.get(url)
|
||||
result.raise_for_status()
|
||||
data = result.json()
|
||||
|
||||
# Get the result items
|
||||
if "items" in data:
|
||||
search_items = data.get("items")
|
||||
|
||||
for i, search_item in enumerate(search_items, start=1):
|
||||
# Check metatags are present
|
||||
if "pagemap" not in search_item:
|
||||
continue
|
||||
if "metatags" not in search_item["pagemap"]:
|
||||
continue
|
||||
if "og:description" in search_item["pagemap"]["metatags"][0]:
|
||||
long_description = search_item["pagemap"]["metatags"][0]["og:description"]
|
||||
else:
|
||||
long_description = "N/A"
|
||||
# Get the page title
|
||||
title = search_item.get("title")
|
||||
# Page snippet
|
||||
snippet = search_item.get("snippet")
|
||||
|
||||
# Extract the page url
|
||||
link = search_item.get("link")
|
||||
response = {
|
||||
"result_id": i,
|
||||
"title": title,
|
||||
"description": snippet,
|
||||
"long_description": long_description,
|
||||
"url": link,
|
||||
}
|
||||
if "huggingface.co" in link:
|
||||
logger.warning(f"Filter out the link: {link}")
|
||||
continue
|
||||
responses.append(response)
|
||||
else:
|
||||
responses.append({"error": f"google search failed with response: {data}"})
|
||||
except Exception as e:
|
||||
logger.error(f"Google search failed with error: {e}")
|
||||
responses.append({"error": f"google search failed with error: {e}"})
|
||||
|
||||
if len(responses) == 0:
|
||||
responses.append(
|
||||
"No relevant webpages found. Please simplify your query and expand the search space as much as you can, then try again.")
|
||||
logger.debug(f"search result: {responses}")
|
||||
responses.append(
|
||||
"If the search result does not contain the information you want, please make reflection on your query: what went well, what didn't, then refine your search plan.")
|
||||
return ActionResult(content=json.dumps(responses), keep=True, is_done=True), None
|
||||
|
||||
|
||||
@ActionFactory.register(name=SearchAction.BAIDU.value.name,
|
||||
desc=SearchAction.BAIDU.value.desc,
|
||||
tool_name="search_api")
|
||||
class SearchBaidu(ExecutableAction):
|
||||
def __init__(self):
|
||||
import_package("baidusearch")
|
||||
|
||||
def act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
from baidusearch.baidusearch import search
|
||||
|
||||
query = action.params.get("query")
|
||||
num_results = action.params.get("num_results", 6)
|
||||
num_results = int(num_results)
|
||||
logger.debug(f"Calling search_baidu with query: {query}")
|
||||
|
||||
responses = []
|
||||
try:
|
||||
responses = search(query, num_results=num_results)
|
||||
except Exception as e:
|
||||
logger.error(f"Baidu search failed with error: {e}")
|
||||
responses.append({"error": f"baidu search failed with error: {e}"})
|
||||
|
||||
if len(responses) == 0:
|
||||
responses.append(
|
||||
"No relevant webpages found. Please simplify your query and expand the search space as much as you can, then try again.")
|
||||
logger.debug(f"search result: {responses}")
|
||||
responses.append(
|
||||
"If the search result does not contain the information you want, please make reflection on your query: what went well, what didn't, then refine your search plan.")
|
||||
return ActionResult(content=json.dumps(responses), keep=True, is_done=True), None
|
||||
@@ -0,0 +1,16 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
from pathlib import Path
|
||||
|
||||
from aworld.core.tool.base import ToolFactory
|
||||
from aworld.tools.template_tool import TemplateTool
|
||||
from examples.common.tools.tool_action import SearchAction
|
||||
|
||||
|
||||
@ToolFactory.register(name="search_api",
|
||||
desc="search tool",
|
||||
supported_action=SearchAction,
|
||||
conf_file_name=f'search_api_tool.yaml',
|
||||
dir=f"{Path(__file__).parent.absolute()}")
|
||||
class SearchTool(TemplateTool):
|
||||
"""Search Tool"""
|
||||
@@ -0,0 +1,5 @@
|
||||
custom_executor: False
|
||||
enable_recording: False
|
||||
working_dir:
|
||||
max_retry: 3
|
||||
|
||||
@@ -0,0 +1,824 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
import os
|
||||
import traceback
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from typing import Tuple, Any
|
||||
|
||||
from examples.common.tools.tool_action import BrowserAction
|
||||
from aworld.core.tool.action_factory import ActionFactory
|
||||
from aworld.core.common import ActionModel, ActionResult, Observation
|
||||
from examples.common.tools.browsers.util.dom import DOMElementNode
|
||||
from aworld.logs.util import logger
|
||||
from examples.common.tools.browsers.action.utils import DomUtil
|
||||
from aworld.core.tool.action import ExecutableAction
|
||||
from aworld.utils import import_packages
|
||||
from aworld.models.llm import get_llm_model, call_llm_model
|
||||
|
||||
|
||||
def get_page(**kwargs):
|
||||
tool = kwargs.get("tool")
|
||||
if tool is None:
|
||||
page = kwargs.get('page')
|
||||
else:
|
||||
page = tool.page
|
||||
return page
|
||||
|
||||
|
||||
def get_browser(**kwargs):
|
||||
tool = kwargs.get("tool")
|
||||
if tool is None:
|
||||
page = kwargs.get('browser')
|
||||
else:
|
||||
page = tool.context
|
||||
return page
|
||||
|
||||
|
||||
@ActionFactory.register(name=BrowserAction.GO_TO_URL.value.name,
|
||||
desc=BrowserAction.GO_TO_URL.value.desc,
|
||||
tool_name="browser")
|
||||
class GotoUrl(ExecutableAction):
|
||||
def act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.GO_TO_URL.value.name} action")
|
||||
page = get_page(**kwargs)
|
||||
if page is None:
|
||||
logger.warning(f"{BrowserAction.GO_TO_URL.name} page is none")
|
||||
return ActionResult(content="no page", keep=True), page
|
||||
|
||||
params = action.params
|
||||
url = params.get("url")
|
||||
if not url:
|
||||
logger.warning("empty url, go to nothing.")
|
||||
return ActionResult(content="empty url", keep=True), page
|
||||
items = url.split('://')
|
||||
if len(items) == 1:
|
||||
if items[0][0] != '/':
|
||||
url = "file://" + os.path.join(os.getcwd(), url)
|
||||
|
||||
page.goto(url)
|
||||
page.wait_for_load_state()
|
||||
msg = f'Navigated to {url}'
|
||||
logger.info(msg)
|
||||
return ActionResult(content=msg, keep=True), page
|
||||
|
||||
async def async_act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.GO_TO_URL.value.name} action")
|
||||
page = get_page(**kwargs)
|
||||
if page is None:
|
||||
logger.warning(f"{BrowserAction.GO_TO_URL.name} page is none")
|
||||
return ActionResult(content="no page", keep=True), page
|
||||
|
||||
url = action.params.get("url")
|
||||
if not url:
|
||||
logger.warning("empty url, go to nothing.")
|
||||
return ActionResult(content="empty url", keep=True), page
|
||||
|
||||
items = url.split('://')
|
||||
if len(items) == 1:
|
||||
if items[0][0] != '/':
|
||||
url = "file://" + os.path.join(os.getcwd(), url)
|
||||
|
||||
await page.goto(url)
|
||||
await page.wait_for_load_state()
|
||||
msg = f'Navigated to {url}'
|
||||
logger.info(msg)
|
||||
return ActionResult(content=msg, keep=True), page
|
||||
|
||||
|
||||
@ActionFactory.register(name=BrowserAction.INPUT_TEXT.value.name,
|
||||
desc=BrowserAction.INPUT_TEXT.value.desc,
|
||||
tool_name="browser")
|
||||
class InputText(ExecutableAction):
|
||||
def act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.INPUT_TEXT.value.name} action")
|
||||
page = get_page(**kwargs)
|
||||
if page is None:
|
||||
logger.warning(f"{BrowserAction.INPUT_TEXT.name} page is none")
|
||||
return ActionResult(content="input text no page", keep=True), page
|
||||
|
||||
params = action.params
|
||||
index = params.get("index", 0)
|
||||
# compatible with int and str datatype
|
||||
index = int(index)
|
||||
input = params.get("text", "")
|
||||
|
||||
ob: Observation = kwargs.get("observation")
|
||||
if not ob or index not in ob.dom_tree.element_map:
|
||||
raise RuntimeError(f'Element index {index} does not exist')
|
||||
if not input:
|
||||
raise ValueError(f'No input to the page')
|
||||
|
||||
element_node = ob.dom_tree.element_map[index]
|
||||
self.input_to_element(input, page, element_node)
|
||||
msg = f'Input {input} into index {index}'
|
||||
logger.info(f"action {msg}")
|
||||
logger.debug(f'Element xpath: {element_node.xpath}')
|
||||
return ActionResult(content=msg, keep=True), page
|
||||
|
||||
async def async_act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.INPUT_TEXT.value.name} action")
|
||||
page = get_page(**kwargs)
|
||||
if page is None:
|
||||
logger.warning(f"{BrowserAction.INPUT_TEXT.name} page is none")
|
||||
return ActionResult(content="input text no page", keep=True), page
|
||||
|
||||
params = action.params
|
||||
index = params.get("index")
|
||||
# compatible with int and str datatype
|
||||
index = int(index)
|
||||
input = params.get("text", "")
|
||||
|
||||
ob: Observation = kwargs.get("observation")
|
||||
if not ob or index not in ob.dom_tree.element_map:
|
||||
raise RuntimeError(f'Element index {index} does not exist')
|
||||
if not input:
|
||||
raise ValueError(f'No input to the page')
|
||||
|
||||
element_node = ob.dom_tree.element_map[index]
|
||||
await self.async_input_to_element(input, page, element_node)
|
||||
msg = f'Input {input} into index {index}'
|
||||
logger.info(f"action {msg}")
|
||||
logger.debug(f'Element xpath: {element_node.xpath}')
|
||||
return ActionResult(content=msg, keep=True), page
|
||||
|
||||
def input_to_element(self, input: str, page, element_node: DOMElementNode):
|
||||
try:
|
||||
# Highlight before typing
|
||||
# if element_node.highlight_index is not None:
|
||||
# await self._update_state(focus_element=element_node.highlight_index)
|
||||
|
||||
element_handle = DomUtil.get_locate_element(page, element_node)
|
||||
|
||||
if element_handle is None:
|
||||
raise RuntimeError(f'Element: {repr(element_node)} not found')
|
||||
|
||||
# Ensure element is ready for input
|
||||
try:
|
||||
element_handle.wait_for_element_state('stable', timeout=1000)
|
||||
element_handle.scroll_into_view_if_needed(timeout=1000)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Get element properties to determine input method
|
||||
is_contenteditable = element_handle.get_property('isContentEditable')
|
||||
|
||||
# Different handling for contenteditable vs input fields
|
||||
if is_contenteditable.json_value():
|
||||
element_handle.evaluate('el => el.textContent = ""')
|
||||
element_handle.type(input, delay=5)
|
||||
else:
|
||||
element_handle.fill(input)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f'Failed to input text into element: {repr(element_node)}. Error: {str(e)}')
|
||||
raise RuntimeError(f'Failed to input text into index {element_node.highlight_index}')
|
||||
|
||||
async def async_input_to_element(self, input: str, page, element_node: DOMElementNode):
|
||||
try:
|
||||
element_handle = await DomUtil.async_get_locate_element(page, element_node)
|
||||
|
||||
if element_handle is None:
|
||||
raise RuntimeError(f'Element: {repr(element_node)} not found')
|
||||
|
||||
# Ensure element is ready for input
|
||||
try:
|
||||
await element_handle.wait_for_element_state('stable', timeout=1000)
|
||||
await element_handle.scroll_into_view_if_needed(timeout=1000)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Get element properties to determine input method
|
||||
is_contenteditable = await element_handle.get_property('isContentEditable')
|
||||
|
||||
# Different handling for contenteditable vs input fields
|
||||
if await is_contenteditable.json_value():
|
||||
await element_handle.evaluate('el => el.textContent = ""')
|
||||
await element_handle.type(input, delay=5)
|
||||
else:
|
||||
await element_handle.fill(input)
|
||||
except Exception as e:
|
||||
logger.warning(f'Failed to input text into element: {repr(element_node)}. Error: {str(e)}')
|
||||
raise RuntimeError(f'Failed to input text into index {element_node.highlight_index}')
|
||||
|
||||
|
||||
@ActionFactory.register(name=BrowserAction.CLICK_ELEMENT.value.name,
|
||||
desc=BrowserAction.CLICK_ELEMENT.value.desc,
|
||||
tool_name="browser")
|
||||
class ClickElement(ExecutableAction):
|
||||
def __init__(self):
|
||||
import_packages(['playwright', 'markdownify'])
|
||||
|
||||
def act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
from playwright.sync_api import BrowserContext
|
||||
|
||||
logger.info(f"exec {BrowserAction.CLICK_ELEMENT.value.name} action")
|
||||
page = get_page(**kwargs)
|
||||
if page is None:
|
||||
logger.warning(f"{BrowserAction.CLICK_ELEMENT.name} page is none")
|
||||
return ActionResult(content="input text no page", keep=True), page
|
||||
|
||||
browser: BrowserContext = get_browser(**kwargs)
|
||||
if browser is None:
|
||||
logger.warning(f"{BrowserAction.CLICK_ELEMENT.name} browser context is none")
|
||||
return ActionResult(content="none browser context", keep=True), page
|
||||
|
||||
index = action.params.get("index")
|
||||
# compatible with int and str datatype
|
||||
index = int(index)
|
||||
ob: Observation = kwargs.get("observation")
|
||||
if not ob or index not in ob.dom_tree.element_map:
|
||||
raise RuntimeError(f'Element index {index} does not exist')
|
||||
if not input:
|
||||
raise ValueError(f'No input to the page')
|
||||
element_node = ob.dom_tree.element_map[index]
|
||||
|
||||
try:
|
||||
pages = len(browser.pages)
|
||||
msg = f'Clicked button with index {index}: {element_node.get_all_text_till_next_clickable_element(max_depth=2)}'
|
||||
logger.info(msg)
|
||||
|
||||
DomUtil.click_element(page, element_node, browser=browser)
|
||||
logger.debug(f'Element xpath: {element_node.xpath}')
|
||||
if len(browser.pages) > pages:
|
||||
new_tab_msg = 'Open the new tab'
|
||||
msg += f' - {new_tab_msg}'
|
||||
logger.info(new_tab_msg)
|
||||
page = browser.pages[-1]
|
||||
page.bring_to_front()
|
||||
page.wait_for_load_state(timeout=60000)
|
||||
return ActionResult(content=msg, keep=True), page
|
||||
except Exception as e:
|
||||
logger.warning(f'Element not clickable with index {index} - most likely the page changed')
|
||||
return ActionResult(error=str(e)), page
|
||||
|
||||
async def async_act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.CLICK_ELEMENT.value.name} action")
|
||||
page = get_page(**kwargs)
|
||||
if page is None:
|
||||
logger.warn(f"{BrowserAction.CLICK_ELEMENT.name} page is none")
|
||||
return ActionResult(content="input text no page", keep=True), page
|
||||
|
||||
browser = get_browser(**kwargs)
|
||||
if browser is None:
|
||||
logger.warning(f"{BrowserAction.CLICK_ELEMENT.name} browser context is none")
|
||||
return ActionResult(content="none browser context", keep=True), page
|
||||
|
||||
index = action.params.get("index")
|
||||
# compatible with int and str datatype
|
||||
index = int(index)
|
||||
ob: Observation = kwargs.get("observation")
|
||||
if not ob or index not in ob.dom_tree.element_map:
|
||||
raise RuntimeError(f'Element index {index} does not exist')
|
||||
if not input:
|
||||
raise ValueError(f'No input to the page')
|
||||
element_node = ob.dom_tree.element_map[index]
|
||||
pages = len(browser.pages)
|
||||
|
||||
try:
|
||||
await DomUtil.async_click_element(page, element_node, browser=browser)
|
||||
msg = f'Clicked button with index {index}: {element_node.get_all_text_till_next_clickable_element(max_depth=2)}'
|
||||
|
||||
logger.info(msg)
|
||||
logger.debug(f'Element xpath: {element_node.xpath}')
|
||||
if len(browser.pages) > pages:
|
||||
new_tab_msg = 'Open the new tab'
|
||||
msg += f' - {new_tab_msg}'
|
||||
logger.info(new_tab_msg)
|
||||
page = browser.pages[-1]
|
||||
await page.bring_to_front()
|
||||
await page.wait_for_load_state(timeout=60000)
|
||||
return ActionResult(content=msg, keep=True), page
|
||||
except Exception as e:
|
||||
logger.warning(f'Element not clickable with index {index} - most likely the page changed')
|
||||
return ActionResult(error=str(e)), page
|
||||
|
||||
|
||||
# SEARCH_ENGINE = {"": "https://www.google.com/search?udm=14&q=",
|
||||
# "google": "https://www.google.com/search?udm=14&q="}
|
||||
|
||||
SEARCH_ENGINE = {"": "https://www.bing.com/search?q=",
|
||||
"google": "https://www.bing.com/search?q="}
|
||||
|
||||
|
||||
@ActionFactory.register(name=BrowserAction.SEARCH.value.name,
|
||||
desc=BrowserAction.SEARCH.value.desc,
|
||||
tool_name="browser")
|
||||
class Search(ExecutableAction):
|
||||
def act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.SEARCH.value.name} action")
|
||||
page = get_page(**kwargs)
|
||||
if page is None:
|
||||
logger.warning(f"{BrowserAction.SEARCH.name} page is none")
|
||||
return ActionResult(content="search no page", keep=True), page
|
||||
|
||||
params = action.params if action.params else {}
|
||||
engine = params.get("engine", "")
|
||||
url = SEARCH_ENGINE.get(engine)
|
||||
query = params.get("query")
|
||||
page.goto(f'{url}{query}')
|
||||
page.wait_for_load_state()
|
||||
msg = f'Searched for "{query}" in {url}'
|
||||
logger.info(msg)
|
||||
return ActionResult(content=msg, keep=True), page
|
||||
|
||||
async def async_act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.SEARCH.value.name} action")
|
||||
page = get_page(**kwargs)
|
||||
if page is None:
|
||||
logger.warning(f"{BrowserAction.SEARCH.name} page is none")
|
||||
return ActionResult(content="search no page", keep=True), page
|
||||
|
||||
params = action.params if action.params else {}
|
||||
engine = params.get("engine", "")
|
||||
url = SEARCH_ENGINE.get(engine)
|
||||
query = params.get("query")
|
||||
await page.goto(f'{url}{query}')
|
||||
await page.wait_for_load_state()
|
||||
msg = f'Searched for "{query}" in {url}'
|
||||
logger.info(msg)
|
||||
return ActionResult(content=msg, keep=True), page
|
||||
|
||||
|
||||
@ActionFactory.register(name=BrowserAction.SEARCH_GOOGLE.value.name,
|
||||
desc=BrowserAction.SEARCH_GOOGLE.value.desc,
|
||||
tool_name="browser")
|
||||
class SearchGoogle(ExecutableAction):
|
||||
def act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.SEARCH_GOOGLE.value.name} action")
|
||||
page = get_page(**kwargs)
|
||||
if page is None:
|
||||
logger.warning(f"{BrowserAction.SEARCH_GOOGLE.name} page is none")
|
||||
return ActionResult(content="search no page", keep=True), page
|
||||
|
||||
query = action.params.get("query")
|
||||
page.goto(f'{SEARCH_ENGINE.get("")}{query}')
|
||||
page.wait_for_load_state()
|
||||
msg = f'Searched for "{query}" in Google'
|
||||
logger.info(msg)
|
||||
return ActionResult(content=msg, keep=True), page
|
||||
|
||||
async def async_act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.SEARCH_GOOGLE.value.name} action")
|
||||
page = get_page(**kwargs)
|
||||
if page is None:
|
||||
logger.warning(f"{BrowserAction.SEARCH_GOOGLE.name} page is none")
|
||||
return ActionResult(content="search no page", keep=True), page
|
||||
|
||||
query = action.params.get("query")
|
||||
await page.goto(f'{SEARCH_ENGINE.get("")}{query}')
|
||||
await page.wait_for_load_state()
|
||||
msg = f'Searched for "{query}" in Google'
|
||||
logger.info(msg)
|
||||
return ActionResult(content=msg, keep=True), page
|
||||
|
||||
|
||||
@ActionFactory.register(name=BrowserAction.NEW_TAB.value.name,
|
||||
desc=BrowserAction.NEW_TAB.value.desc,
|
||||
tool_name="browser")
|
||||
class NewTab(ExecutableAction):
|
||||
def act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.NEW_TAB.value.name} action")
|
||||
browser = get_browser(**kwargs)
|
||||
url = action.params.get("url")
|
||||
|
||||
new_page = browser.new_page()
|
||||
new_page.wait_for_load_state()
|
||||
|
||||
if url:
|
||||
new_page.goto(url)
|
||||
DomUtil.wait_for_stable_network(new_page)
|
||||
|
||||
msg = f'Opened new tab with {url}'
|
||||
logger.debug(msg)
|
||||
return ActionResult(content=msg, keep=True), new_page
|
||||
|
||||
async def async_act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.NEW_TAB.value.name} action")
|
||||
browser = get_browser(**kwargs)
|
||||
url = action.params.get("url")
|
||||
new_page = await browser.new_page()
|
||||
await new_page.wait_for_load_state()
|
||||
|
||||
if url:
|
||||
await new_page.goto(url)
|
||||
DomUtil.wait_for_stable_network(new_page)
|
||||
msg = f'Opened new tab with {url}'
|
||||
logger.debug(msg)
|
||||
return ActionResult(content=msg, keep=True), get_page(**kwargs)
|
||||
|
||||
|
||||
@ActionFactory.register(name=BrowserAction.GO_BACK.value.name,
|
||||
desc=BrowserAction.GO_BACK.value.desc,
|
||||
tool_name="browser")
|
||||
class GoBack(ExecutableAction):
|
||||
def act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.GO_BACK.value.name} action")
|
||||
page = get_page(**kwargs)
|
||||
if page is None:
|
||||
logger.warning(f"{BrowserAction.GO_BACK.name} page is none")
|
||||
return ActionResult(content="search no page", keep=True), page
|
||||
|
||||
page.go_back()
|
||||
msg = 'Navigated back'
|
||||
logger.info(msg)
|
||||
return ActionResult(content=msg, keep=True), page
|
||||
|
||||
async def async_act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.GO_BACK.value.name} action")
|
||||
page = get_page(**kwargs)
|
||||
if page is None:
|
||||
logger.warning(f"{BrowserAction.GO_BACK.name} page is none")
|
||||
return ActionResult(content="search no page", keep=True), page
|
||||
|
||||
await page.go_back()
|
||||
msg = 'Navigated back'
|
||||
logger.info(msg)
|
||||
return ActionResult(content=msg, keep=True), page
|
||||
|
||||
|
||||
@ActionFactory.register(name=BrowserAction.EXTRACT_CONTENT.value.name,
|
||||
desc=BrowserAction.EXTRACT_CONTENT.value.desc,
|
||||
tool_name="browser")
|
||||
class ExtractContent(ExecutableAction):
|
||||
def act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
import markdownify
|
||||
from langchain_core.prompts import PromptTemplate
|
||||
|
||||
logger.info(f"exec {BrowserAction.EXTRACT_CONTENT.value.name} action")
|
||||
page = get_page(**kwargs)
|
||||
if page is None:
|
||||
logger.warning(f"{BrowserAction.EXTRACT_CONTENT.name} page is none")
|
||||
return ActionResult(content="extract content no page", keep=True), page
|
||||
|
||||
goal = action.params.get("goal")
|
||||
llm_config = kwargs.get("llm_config")
|
||||
if llm_config and llm_config.llm_api_key:
|
||||
llm = get_llm_model(llm_config)
|
||||
max_extract_content_output_tokens = kwargs.get("max_extract_content_output_tokens")
|
||||
max_extract_content_input_tokens = kwargs.get("max_extract_content_input_tokens")
|
||||
|
||||
content = markdownify.markdownify(page.content())
|
||||
|
||||
# Truncate content if it exceeds max input tokens
|
||||
if max_extract_content_input_tokens and len(content) > max_extract_content_input_tokens:
|
||||
logger.warning(
|
||||
f"Content length ({len(content)}) exceeds max input tokens ({max_extract_content_input_tokens}). Truncating content.")
|
||||
content = content[:max_extract_content_input_tokens]
|
||||
|
||||
prompt = 'Your task is to extract the content of the page. You will be given a page and a goal and you should extract all relevant information around this goal from the page. If the goal is vague, summarize the page. Respond in json format. Extraction goal: {goal}, Page: {page}'
|
||||
prompt_with_outputlimit = 'Your task is to extract the content of the page. You will be given a page and a goal and you should extract all relevant information around this goal from the page. If the goal is vague, summarize the page. Respond in json format. Extraction goal: {goal}, Page: {page} \n\n#The length of the returned result must be less than {max_extract_content_output_tokens} characters.'
|
||||
template = PromptTemplate(input_variables=['goal', 'page'], template=prompt)
|
||||
|
||||
messages = [{'role': 'user', 'content': template.format(goal=goal, page=content)}]
|
||||
try:
|
||||
output = call_llm_model(llm,
|
||||
messages=messages,
|
||||
model=llm_config.llm_model_name,
|
||||
temperature=llm_config.llm_temperature)
|
||||
result_content = output.content
|
||||
|
||||
# Check if output exceeds the token limit and retry with length-limited prompt if needed
|
||||
if max_extract_content_output_tokens and len(result_content) > max_extract_content_output_tokens:
|
||||
logger.warning(
|
||||
f"Output exceeds maximum length ({len(result_content)} > {max_extract_content_output_tokens}). Retrying with limited prompt.")
|
||||
template_with_limit = PromptTemplate(
|
||||
input_variables=['goal', 'page', 'max_extract_content_output_tokens'],
|
||||
template=prompt_with_outputlimit
|
||||
)
|
||||
messages = [{'role': 'user', 'content': template_with_limit.format(
|
||||
goal=goal,
|
||||
page=content,
|
||||
max_extract_content_output_tokens=max_extract_content_output_tokens,
|
||||
max_tokens=max_extract_content_output_tokens
|
||||
)}]
|
||||
# extract content with length limit
|
||||
output = call_llm_model(llm,
|
||||
messages=messages,
|
||||
model=llm_config.llm_model_name,
|
||||
temperature=llm_config.llm_temperature)
|
||||
result_content = output.content
|
||||
|
||||
msg = f'Extracted from page\n: {result_content}\n'
|
||||
logger.info(msg)
|
||||
return ActionResult(content=msg, keep=True), page
|
||||
except Exception as e:
|
||||
logger.debug(f'Error extracting content: {e}')
|
||||
msg = f'Extracted from page\n: {content}\n'
|
||||
logger.info(msg)
|
||||
return ActionResult(content=msg), page
|
||||
|
||||
async def async_act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
import markdownify
|
||||
from langchain_core.prompts import PromptTemplate
|
||||
|
||||
logger.info(f"exec {BrowserAction.EXTRACT_CONTENT.value.name} action")
|
||||
page = get_page(**kwargs)
|
||||
if page is None:
|
||||
logger.warning(f"{BrowserAction.EXTRACT_CONTENT.name} page is none")
|
||||
return ActionResult(content="extract content no page", keep=True), page
|
||||
|
||||
goal = action.params.get("goal")
|
||||
llm_config = kwargs.get("llm_config")
|
||||
if llm_config and llm_config.llm_api_key:
|
||||
llm = get_llm_model(llm_config)
|
||||
content = markdownify.markdownify(await page.content())
|
||||
max_extract_content_output_tokens = kwargs.get("max_extract_content_output_tokens")
|
||||
max_extract_content_input_tokens = kwargs.get("max_extract_content_input_tokens")
|
||||
|
||||
# Truncate content if it exceeds max input tokens
|
||||
if max_extract_content_input_tokens and len(content) > max_extract_content_input_tokens:
|
||||
logger.warning(
|
||||
f"Content length ({len(content)}) exceeds max input tokens ({max_extract_content_input_tokens}). Truncating content.")
|
||||
content = content[:max_extract_content_input_tokens]
|
||||
|
||||
prompt = 'Your task is to extract the content of the page. You will be given a page and a goal and you should extract all relevant information around this goal from the page. If the goal is vague, summarize the page. Respond in json format. Extraction goal: {goal}, Page: {page}'
|
||||
prompt_with_outputlimit = 'Your task is to extract the content of the page. You will be given a page and a goal and you should extract all relevant information around this goal from the page. If the goal is vague, summarize the page. Respond in json format. Extraction goal: {goal}, Page: {page} \n\n#The length of the returned result must be less than {max_extract_content_output_tokens} characters.'
|
||||
template = PromptTemplate(input_variables=['goal', 'page'], template=prompt)
|
||||
|
||||
messages = [{'role': 'user', 'content': template.format(goal=goal, page=content)}]
|
||||
try:
|
||||
output = call_llm_model(llm,
|
||||
messages=messages,
|
||||
model=llm_config.llm_model_name,
|
||||
temperature=llm_config.llm_temperature)
|
||||
result_content = output.content
|
||||
|
||||
# Check if output exceeds the token limit and retry with length-limited prompt if needed
|
||||
if max_extract_content_output_tokens and len(result_content) > max_extract_content_output_tokens:
|
||||
logger.info(
|
||||
f"Output exceeds maximum length ({len(result_content)} > {max_extract_content_output_tokens}). Retrying with limited prompt.")
|
||||
template_with_limit = PromptTemplate(
|
||||
input_variables=['goal', 'page', 'max_extract_content_output_tokens'],
|
||||
template=prompt_with_outputlimit
|
||||
)
|
||||
messages = [{'role': 'user', 'content': template_with_limit.format(
|
||||
goal=goal,
|
||||
page=content,
|
||||
max_extract_content_output_tokens=max_extract_content_output_tokens,
|
||||
max_tokens=max_extract_content_output_tokens
|
||||
)}]
|
||||
# extract content with length limit
|
||||
output = call_llm_model(llm,
|
||||
messages=messages,
|
||||
model=llm_config.llm_model_name,
|
||||
temperature=llm_config.llm_temperature)
|
||||
result_content = output.content
|
||||
|
||||
msg = f'Extracted from page\n: {result_content}\n'
|
||||
logger.info(msg)
|
||||
return ActionResult(content=msg, keep=True), page
|
||||
except Exception as e:
|
||||
logger.debug(f'Error extracting content: {e}')
|
||||
msg = f'Extracted from page\n: {content}\n'
|
||||
logger.info(msg)
|
||||
return ActionResult(content=msg), page
|
||||
|
||||
|
||||
@ActionFactory.register(name=BrowserAction.SCROLL_DOWN.value.name,
|
||||
desc=BrowserAction.SCROLL_DOWN.value.desc,
|
||||
tool_name="browser")
|
||||
class ScrollDown(ExecutableAction):
|
||||
def act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.SCROLL_DOWN.value.name} action")
|
||||
page = get_page(**kwargs)
|
||||
if page is None:
|
||||
logger.warning(f"{BrowserAction.SCROLL_DOWN.name} page is none")
|
||||
return ActionResult(content="scroll no page", keep=True), page
|
||||
|
||||
amount = action.params.get("amount")
|
||||
if not amount:
|
||||
page.evaluate('window.scrollBy(0, window.innerHeight);')
|
||||
else:
|
||||
amount = int(amount)
|
||||
page.evaluate(f'window.scrollBy(0, {amount});')
|
||||
|
||||
amount = f'{amount} pixels' if amount else 'one page'
|
||||
msg = f'Scrolled down the page by {amount}'
|
||||
logger.info(msg)
|
||||
return ActionResult(content=msg, keep=True), page
|
||||
|
||||
async def async_act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.SCROLL_DOWN.value.name} action")
|
||||
page = get_page(**kwargs)
|
||||
if page is None:
|
||||
logger.warning(f"{BrowserAction.SCROLL_DOWN.name} page is none")
|
||||
return ActionResult(content="scroll no page", keep=True), page
|
||||
|
||||
amount = action.params.get("amount")
|
||||
if not amount:
|
||||
await page.evaluate('window.scrollBy(0, window.innerHeight);')
|
||||
else:
|
||||
amount = int(amount)
|
||||
await page.evaluate(f'window.scrollBy(0, {amount});')
|
||||
|
||||
amount = f'{amount} pixels' if amount else 'one page'
|
||||
msg = f'Scrolled down the page by {amount}'
|
||||
logger.info(msg)
|
||||
return ActionResult(content=msg, keep=True), page
|
||||
|
||||
|
||||
@ActionFactory.register(name=BrowserAction.SCROLL_UP.value.name,
|
||||
desc=BrowserAction.SCROLL_UP.value.desc,
|
||||
tool_name="browser")
|
||||
class ScrollUp(ExecutableAction):
|
||||
def act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.SCROLL_UP.value.name} action")
|
||||
page = get_page(**kwargs)
|
||||
if page is None:
|
||||
logger.warning(f"{BrowserAction.SCROLL_UP.name} page is none")
|
||||
return ActionResult(content="scroll no page", keep=True), page
|
||||
|
||||
amount = action.params.get("amount")
|
||||
if not amount:
|
||||
page.evaluate('window.scrollBy(0, -window.innerHeight);')
|
||||
else:
|
||||
amount = int(amount)
|
||||
page.evaluate(f'window.scrollBy(0, -{amount});')
|
||||
|
||||
amount = f'{amount} pixels' if amount else 'one page'
|
||||
msg = f'Scrolled down the page by {amount}'
|
||||
logger.info(msg)
|
||||
return ActionResult(content=msg, keep=True), page
|
||||
|
||||
async def async_act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.SCROLL_UP.value.name} action")
|
||||
page = get_page(**kwargs)
|
||||
if page is None:
|
||||
logger.warning(f"{BrowserAction.SCROLL_UP.name} page is none")
|
||||
return ActionResult(content="scroll no page", keep=True), page
|
||||
|
||||
amount = action.params.get("amount")
|
||||
if not amount:
|
||||
await page.evaluate('window.scrollBy(0, -window.innerHeight);')
|
||||
else:
|
||||
amount = int(amount)
|
||||
await page.evaluate(f'window.scrollBy(0, -{amount});')
|
||||
|
||||
amount = f'{amount} pixels' if amount else 'one page'
|
||||
msg = f'Scrolled down the page by {amount}'
|
||||
logger.info(msg)
|
||||
return ActionResult(content=msg, keep=True), page
|
||||
|
||||
|
||||
@ActionFactory.register(name=BrowserAction.WAIT.value.name,
|
||||
desc=BrowserAction.WAIT.value.desc,
|
||||
tool_name="browser")
|
||||
class Wait(ExecutableAction):
|
||||
def act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
seconds = action.params.get("seconds")
|
||||
if not seconds:
|
||||
seconds = action.params.get("duration", 0)
|
||||
seconds = int(seconds)
|
||||
msg = f'Waiting for {seconds} seconds'
|
||||
logger.info(msg)
|
||||
time.sleep(seconds)
|
||||
return ActionResult(content=msg, keep=True), kwargs.get('page')
|
||||
|
||||
async def async_act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
seconds = action.params.get("seconds")
|
||||
if not seconds:
|
||||
seconds = action.params.get("duration", 0)
|
||||
seconds = int(seconds)
|
||||
msg = f'Waiting for {seconds} seconds'
|
||||
logger.info(msg)
|
||||
await asyncio.sleep(seconds)
|
||||
return ActionResult(content=msg, keep=True), kwargs.get('page')
|
||||
|
||||
|
||||
@ActionFactory.register(name=BrowserAction.SWITCH_TAB.value.name,
|
||||
desc=BrowserAction.SWITCH_TAB.value.desc,
|
||||
tool_name="browser")
|
||||
class SwitchTab(ExecutableAction):
|
||||
def act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.SWITCH_TAB.value.name} action")
|
||||
browser = get_browser(**kwargs)
|
||||
if browser is None:
|
||||
logger.warning(f"{BrowserAction.SWITCH_TAB.name} browser context is none")
|
||||
return ActionResult(content="switch tab no browser context", keep=True), get_page(**kwargs)
|
||||
|
||||
page_id = action.params.get("page_id", 0)
|
||||
page_id = int(page_id)
|
||||
pages = browser.pages
|
||||
|
||||
if page_id >= len(pages):
|
||||
raise RuntimeError(f'No tab found with page_id: {page_id}')
|
||||
|
||||
page = pages[page_id]
|
||||
page.bring_to_front()
|
||||
page.wait_for_load_state()
|
||||
msg = f'Switched to tab {page_id}'
|
||||
logger.info(msg)
|
||||
return ActionResult(content=msg, keep=True), page
|
||||
|
||||
async def async_act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.SWITCH_TAB.value.name} action")
|
||||
browser = get_browser(**kwargs)
|
||||
if browser is None:
|
||||
logger.warning(f"{BrowserAction.SWITCH_TAB.name} browser context is none")
|
||||
return ActionResult(content="switch tab no browser context", keep=True), get_page(**kwargs)
|
||||
|
||||
page_id = action.params.get("page_id", 0)
|
||||
page_id = int(page_id)
|
||||
pages = browser.pages
|
||||
|
||||
if page_id >= len(pages):
|
||||
raise RuntimeError(f'No tab found with page_id: {page_id}')
|
||||
|
||||
page = pages[page_id]
|
||||
await page.bring_to_front()
|
||||
await page.wait_for_load_state()
|
||||
msg = f'Switched to tab {page_id}'
|
||||
logger.info(msg)
|
||||
return ActionResult(content=msg, keep=True), page
|
||||
|
||||
|
||||
@ActionFactory.register(name=BrowserAction.SEND_KEYS.value.name,
|
||||
desc=BrowserAction.SEND_KEYS.value.desc,
|
||||
tool_name="browser")
|
||||
class SendKeys(ExecutableAction):
|
||||
def act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.SEND_KEYS.value.name} action")
|
||||
page = get_page(**kwargs)
|
||||
if page is None:
|
||||
logger.warning(f"{BrowserAction.SEND_KEYS.name} page is none")
|
||||
return ActionResult(content="scroll no page", keep=True), page
|
||||
|
||||
keys = action.params.get("keys")
|
||||
if not keys:
|
||||
return ActionResult(success=False, content="no keys", keep=True), page
|
||||
|
||||
try:
|
||||
page.keyboard.press(keys)
|
||||
except Exception as e:
|
||||
logger.warning(f"{keys} press fail. \n{traceback.format_exc()}")
|
||||
raise e
|
||||
return ActionResult(content=f"Sent keys: {keys}", keep=True), page
|
||||
|
||||
async def async_act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.SEND_KEYS.value.name} action")
|
||||
page = get_page(**kwargs)
|
||||
if page is None:
|
||||
logger.warning(f"{BrowserAction.SEND_KEYS.name} page is none")
|
||||
return ActionResult(content="scroll no page", keep=True), page
|
||||
|
||||
keys = action.params.get("keys")
|
||||
if not keys:
|
||||
return ActionResult(success=False, content="no keys", keep=True), page
|
||||
|
||||
try:
|
||||
await page.keyboard.press(keys)
|
||||
except Exception as e:
|
||||
logger.warning(f"{keys} press fail. \n{traceback.format_exc()}")
|
||||
raise e
|
||||
|
||||
return ActionResult(content=f"Sent keys: {keys}", keep=True), page
|
||||
|
||||
|
||||
@ActionFactory.register(name=BrowserAction.WRITE_TO_FILE.value.name,
|
||||
desc=BrowserAction.WRITE_TO_FILE.value.desc,
|
||||
tool_name="browser")
|
||||
class WriteToFile(ExecutableAction):
|
||||
def act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
# 设置默认文件路径
|
||||
file_path = "tmp_result.md"
|
||||
# 检查参数中是否有file_path
|
||||
if "file_path" in action.params:
|
||||
file_path = action.params.get("file_path", "tmp_result.md")
|
||||
# 检查参数中是否有file_name
|
||||
elif "file_name" in action.params:
|
||||
file_path = action.params.get("file_name", "tmp_result.md")
|
||||
elif "filename" in action.params:
|
||||
file_path = action.params.get("filename", "tmp_result.md")
|
||||
content = action.params.get("content", "")
|
||||
mode = action.params.get("mode", "a") # Default to append mode
|
||||
# 获取文件的绝对路径
|
||||
abs_file_path = os.path.abspath(file_path)
|
||||
try:
|
||||
with open(file_path, mode, encoding='utf-8') as f:
|
||||
f.write(content + '\n')
|
||||
msg = f'Successfully wrote content to {abs_file_path}'
|
||||
logger.info(msg)
|
||||
return ActionResult(content=msg, keep=True), get_page(**kwargs)
|
||||
except Exception as e:
|
||||
error_msg = f'Failed to write to file {abs_file_path}: {str(e)}'
|
||||
logger.error(error_msg)
|
||||
return ActionResult(content=error_msg, keep=True, error=error_msg), get_page(**kwargs)
|
||||
|
||||
async def async_act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
# For file operations, we don't need to make this asynchronous
|
||||
return self.act(action, **kwargs)
|
||||
|
||||
|
||||
@ActionFactory.register(name=BrowserAction.DONE.value.name,
|
||||
desc=BrowserAction.DONE.value.desc,
|
||||
tool_name="browser")
|
||||
class Done(ExecutableAction):
|
||||
def act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.DONE.value.name} action")
|
||||
return ActionResult(is_done=True, success=True, content="done", keep=True), get_page(**kwargs)
|
||||
|
||||
async def async_act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info(f"exec {BrowserAction.DONE.value.name} action")
|
||||
return ActionResult(is_done=True, success=True, content="done", keep=True), get_page(**kwargs)
|
||||
@@ -0,0 +1,71 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
|
||||
from typing import Tuple, List, Any
|
||||
|
||||
from aworld.core.tool.action_factory import ActionFactory
|
||||
from aworld.core.common import ActionModel, ActionResult, Observation
|
||||
from aworld.logs.util import logger
|
||||
from aworld.core.tool.base import Tool, ToolActionExecutor
|
||||
|
||||
|
||||
class BrowserToolActionExecutor(ToolActionExecutor):
|
||||
def __init__(self, tool: Tool = None):
|
||||
super(BrowserToolActionExecutor, self).__init__(tool)
|
||||
|
||||
def execute_action(self, actions: List[ActionModel], **kwargs) -> Tuple[
|
||||
List[ActionResult], Any]:
|
||||
"""Execute the specified browser action sequence by agent policy.
|
||||
|
||||
Args:
|
||||
actions: Tool action sequence.
|
||||
|
||||
Returns:
|
||||
Browser page and action result list.
|
||||
"""
|
||||
action_results = []
|
||||
page = self.tool.page
|
||||
for action in actions:
|
||||
action_result, page = self._exec(action, **kwargs)
|
||||
action_results.append(action_result)
|
||||
return action_results, page
|
||||
|
||||
async def async_execute_action(self, actions: List[ActionModel], **kwargs) -> Tuple[
|
||||
List[ActionResult], Any]:
|
||||
"""Execute the specified browser action sequence by agent policy.
|
||||
|
||||
Args:
|
||||
actions: Tool action sequence.
|
||||
|
||||
Returns:
|
||||
Browser page and action result list.
|
||||
"""
|
||||
action_results = []
|
||||
page = self.tool.page
|
||||
for action in actions:
|
||||
action_result, page = await self._async_exec(action, **kwargs)
|
||||
action_results.append(action_result)
|
||||
return action_results, page
|
||||
|
||||
def _exec(self, action_model: ActionModel, **kwargs):
|
||||
action_name = action_model.action_name
|
||||
if action_name not in ActionFactory:
|
||||
raise ValueError(f'Action {action_name} not found')
|
||||
|
||||
action = ActionFactory(action_name)
|
||||
action_result, page = action.act(action_model, page=self.tool.page, browser=self.tool.browser_context, **kwargs)
|
||||
logger.info(f"{action_name} execute finished")
|
||||
return action_result, page
|
||||
|
||||
async def _async_exec(self, action_model: ActionModel, **kwargs):
|
||||
action_name = action_model.action_name
|
||||
if action_name not in ActionFactory:
|
||||
action_name = action_model.tool_name + action_model.action_name
|
||||
if action_name not in ActionFactory:
|
||||
raise ValueError(f'Action {action_name} not found')
|
||||
|
||||
action = ActionFactory(action_name)
|
||||
action_result, page = await action.async_act(action_model, page=self.tool.page,
|
||||
browser=self.tool.browser_context, **kwargs)
|
||||
logger.info(f"{action_name} execute finished")
|
||||
return action_result, page
|
||||
@@ -0,0 +1,507 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
|
||||
import re
|
||||
import time
|
||||
import traceback
|
||||
from typing import Optional
|
||||
|
||||
from examples.common.tools.browsers.util.dom import DOMElementNode
|
||||
from aworld.logs.util import logger
|
||||
from aworld.utils import import_package
|
||||
|
||||
|
||||
class DomUtil:
|
||||
def __init__(self):
|
||||
import_package("playwright")
|
||||
|
||||
@staticmethod
|
||||
async def async_click_element(page, element_node: DOMElementNode, **kwargs) -> Optional[str]:
|
||||
from playwright.async_api import ElementHandle as AElementHandle, BrowserContext as ABrowserContext
|
||||
|
||||
try:
|
||||
element_handle: AElementHandle = await DomUtil.async_get_locate_element(page, element_node)
|
||||
if element_handle is None:
|
||||
raise Exception(f'Element: {repr(element_node)} not found')
|
||||
|
||||
bound = await element_handle.bounding_box()
|
||||
try:
|
||||
# todo: iframe.
|
||||
center_x = bound['x'] + bound['width'] / 2
|
||||
center_y = bound['y'] + bound['height'] / 2
|
||||
|
||||
try:
|
||||
browser: ABrowserContext = kwargs.get('browser')
|
||||
async with browser.expect_page() as new_page_info:
|
||||
await page.mouse.click(center_x, center_y)
|
||||
await page.mouse.click(center_x, center_y)
|
||||
await page.wait_for_load_state()
|
||||
except:
|
||||
logger.warning(traceback.format_exc())
|
||||
except:
|
||||
logger.info(f"click {element_handle}!!")
|
||||
if await element_handle.text_content():
|
||||
browser: ABrowserContext = kwargs.get('browser')
|
||||
if browser:
|
||||
try:
|
||||
async with browser.expect_page() as new_page_info:
|
||||
await page.click(f"text={element_handle.text_content()}")
|
||||
page = await new_page_info.value
|
||||
await page.wait_for_load_state()
|
||||
except:
|
||||
logger.warning(traceback.format_exc())
|
||||
else:
|
||||
await element_handle.click()
|
||||
await page.wait_for_load_state()
|
||||
else:
|
||||
await element_handle.click()
|
||||
await page.wait_for_load_state()
|
||||
except Exception as e:
|
||||
logger.error(traceback.format_exc())
|
||||
raise Exception(f'Failed to click element: {repr(element_node)}. Error: {str(e)}')
|
||||
|
||||
@staticmethod
|
||||
def click_element(page, element_node: DOMElementNode, **kwargs) -> Optional[str]:
|
||||
from playwright.sync_api import ElementHandle, BrowserContext
|
||||
|
||||
try:
|
||||
element_handle: ElementHandle = DomUtil.get_locate_element(page, element_node)
|
||||
if element_handle is None:
|
||||
raise Exception(f'Element: {repr(element_node)} not found')
|
||||
|
||||
bound = element_handle.bounding_box()
|
||||
try:
|
||||
# todo: iframe.
|
||||
center_x = bound['x'] + bound['width'] / 2
|
||||
center_y = bound['y'] + bound['height'] / 2
|
||||
|
||||
try:
|
||||
browser: BrowserContext = kwargs.get('browser')
|
||||
with browser.expect_page() as new_page_info:
|
||||
page.mouse.click(center_x, center_y)
|
||||
page = new_page_info.value
|
||||
page.wait_for_load_state()
|
||||
except:
|
||||
logger.warning(traceback.format_exc())
|
||||
except:
|
||||
logger.info(f"click {element_handle}!!")
|
||||
if element_handle.text_content():
|
||||
browser: BrowserContext = kwargs.get('browser')
|
||||
if browser:
|
||||
try:
|
||||
with browser.expect_page() as new_page_info:
|
||||
page.click(f"text={element_handle.text_content()}")
|
||||
page = new_page_info.value
|
||||
page.wait_for_load_state()
|
||||
except:
|
||||
logger.warning(traceback.format_exc())
|
||||
else:
|
||||
element_handle.click()
|
||||
page.wait_for_load_state()
|
||||
else:
|
||||
element_handle.click()
|
||||
page.wait_for_load_state()
|
||||
except Exception as e:
|
||||
logger.error(traceback.format_exc())
|
||||
raise Exception(f'Failed to click element: {repr(element_node)}. Error: {str(e)}')
|
||||
|
||||
@staticmethod
|
||||
async def async_get_locate_element(current_frame, element: DOMElementNode):
|
||||
# Start with the target element and collect all parents, return Optional[AElementHandle]
|
||||
from playwright.async_api import FrameLocator as AFrameLocator
|
||||
|
||||
parents: list[DOMElementNode] = []
|
||||
current = element
|
||||
while current.parent is not None:
|
||||
parent = current.parent
|
||||
parents.append(parent)
|
||||
current = parent
|
||||
|
||||
# Reverse the parents list to process from top to bottom
|
||||
parents.reverse()
|
||||
|
||||
# Process all iframe parents in sequence
|
||||
iframes = [item for item in parents if item.tag_name == 'iframe']
|
||||
for parent in iframes:
|
||||
css_selector = DomUtil._enhanced_css_selector_for_element(
|
||||
parent,
|
||||
include_dynamic_attributes=True,
|
||||
)
|
||||
current_frame = current_frame.frame_locator(css_selector)
|
||||
|
||||
css_selector = DomUtil._enhanced_css_selector_for_element(
|
||||
element, include_dynamic_attributes=True
|
||||
)
|
||||
|
||||
try:
|
||||
if isinstance(current_frame, AFrameLocator):
|
||||
element_handle = await current_frame.locator(css_selector).element_handle()
|
||||
return element_handle
|
||||
else:
|
||||
# Try to scroll into view if hidden
|
||||
element_handle = await current_frame.query_selector(css_selector)
|
||||
if element_handle:
|
||||
await element_handle.scroll_into_view_if_needed()
|
||||
return element_handle
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error(f'Failed to locate element: {str(e)}')
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def get_locate_element(current_frame, element: DOMElementNode):
|
||||
# Start with the target element and collect all parents
|
||||
from playwright.sync_api import FrameLocator
|
||||
|
||||
parents: list[DOMElementNode] = []
|
||||
current = element
|
||||
while current.parent is not None:
|
||||
parent = current.parent
|
||||
parents.append(parent)
|
||||
current = parent
|
||||
|
||||
# Reverse the parents list to process from top to bottom
|
||||
parents.reverse()
|
||||
|
||||
# Process all iframe parents in sequence
|
||||
iframes = [item for item in parents if item.tag_name == 'iframe']
|
||||
for parent in iframes:
|
||||
css_selector = DomUtil._enhanced_css_selector_for_element(
|
||||
parent,
|
||||
include_dynamic_attributes=True,
|
||||
)
|
||||
current_frame = current_frame.frame_locator(css_selector)
|
||||
|
||||
css_selector = DomUtil._enhanced_css_selector_for_element(
|
||||
element, include_dynamic_attributes=True
|
||||
)
|
||||
|
||||
try:
|
||||
if isinstance(current_frame, FrameLocator):
|
||||
element_handle = current_frame.locator(css_selector).element_handle()
|
||||
return element_handle
|
||||
else:
|
||||
# Try to scroll into view if hidden
|
||||
element_handle = current_frame.query_selector(css_selector)
|
||||
if element_handle:
|
||||
element_handle.scroll_into_view_if_needed()
|
||||
return element_handle
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error(f'Failed to locate element: {str(e)}')
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def wait_for_stable_network(page, **kwargs):
|
||||
pending_requests = set()
|
||||
last_activity = time.time()
|
||||
|
||||
# Define relevant resource types and content types
|
||||
RELEVANT_RESOURCE_TYPES = {
|
||||
'document',
|
||||
'stylesheet',
|
||||
'image',
|
||||
'font',
|
||||
'script',
|
||||
'iframe',
|
||||
}
|
||||
|
||||
RELEVANT_CONTENT_TYPES = {
|
||||
'text/html',
|
||||
'text/css',
|
||||
'application/javascript',
|
||||
'image/',
|
||||
'font/',
|
||||
'application/json',
|
||||
}
|
||||
|
||||
# Additional patterns to filter out
|
||||
IGNORED_URL_PATTERNS = {
|
||||
# Analytics and tracking
|
||||
'analytics',
|
||||
'tracking',
|
||||
'telemetry',
|
||||
'beacon',
|
||||
'metrics',
|
||||
# Ad-related
|
||||
'doubleclick',
|
||||
'adsystem',
|
||||
'adserver',
|
||||
'advertising',
|
||||
# Social media widgets
|
||||
'facebook.com/plugins',
|
||||
'platform.twitter',
|
||||
'linkedin.com/embed',
|
||||
# Live chat and support
|
||||
'livechat',
|
||||
'zendesk',
|
||||
'intercom',
|
||||
'crisp.chat',
|
||||
'hotjar',
|
||||
# Push notifications
|
||||
'push-notifications',
|
||||
'onesignal',
|
||||
'pushwoosh',
|
||||
# Background sync/heartbeat
|
||||
'heartbeat',
|
||||
'ping',
|
||||
'alive',
|
||||
# WebRTC and streaming
|
||||
'webrtc',
|
||||
'rtmp://',
|
||||
'wss://',
|
||||
# Common CDNs for dynamic content
|
||||
'cloudfront.net',
|
||||
'fastly.net',
|
||||
}
|
||||
|
||||
def on_request(request):
|
||||
# Filter by resource type
|
||||
if request.resource_type not in RELEVANT_RESOURCE_TYPES:
|
||||
return
|
||||
|
||||
# Filter out streaming, websocket, and other real-time requests
|
||||
if request.resource_type in {
|
||||
'websocket',
|
||||
'media',
|
||||
'eventsource',
|
||||
'manifest',
|
||||
'other',
|
||||
}:
|
||||
return
|
||||
|
||||
# Filter out by URL patterns
|
||||
url = request.url.lower()
|
||||
if any(pattern in url for pattern in IGNORED_URL_PATTERNS):
|
||||
return
|
||||
|
||||
# Filter out data URLs and blob URLs
|
||||
if url.startswith(('data:', 'blob:')):
|
||||
return
|
||||
|
||||
# Filter out requests with certain headers
|
||||
headers = request.headers
|
||||
if headers.get('purpose') == 'prefetch' or headers.get('sec-fetch-dest') in [
|
||||
'video',
|
||||
'audio',
|
||||
]:
|
||||
return
|
||||
|
||||
nonlocal last_activity
|
||||
pending_requests.add(request)
|
||||
last_activity = time.time()
|
||||
|
||||
def on_response(response):
|
||||
request = response.request
|
||||
if request not in pending_requests:
|
||||
return
|
||||
|
||||
# Filter by content type if available
|
||||
content_type = response.headers.get('content-type', '').lower()
|
||||
|
||||
# Skip if content type indicates streaming or real-time data
|
||||
if any(t in content_type
|
||||
for t in [
|
||||
'streaming',
|
||||
'video',
|
||||
'audio',
|
||||
'webm',
|
||||
'mp4',
|
||||
'event-stream',
|
||||
'websocket',
|
||||
'protobuf']):
|
||||
pending_requests.remove(request)
|
||||
return
|
||||
|
||||
# Only process relevant content types
|
||||
if not any(ct in content_type for ct in RELEVANT_CONTENT_TYPES):
|
||||
pending_requests.remove(request)
|
||||
return
|
||||
|
||||
# Skip if response is too large (likely not essential for page load)
|
||||
content_length = response.headers.get('content-length')
|
||||
if content_length and int(content_length) > 5 * 1024 * 1024: # 5MB
|
||||
pending_requests.remove(request)
|
||||
return
|
||||
|
||||
nonlocal last_activity
|
||||
pending_requests.remove(request)
|
||||
last_activity = time.time()
|
||||
|
||||
# Attach event listeners
|
||||
page.on('request', on_request)
|
||||
page.on('response', on_response)
|
||||
|
||||
try:
|
||||
start_time = time.time()
|
||||
while True:
|
||||
time.sleep(0.1)
|
||||
now = time.time()
|
||||
|
||||
if len(pending_requests) == 0 and (now - last_activity) >= kwargs.get('idle_wait_time', 0.5):
|
||||
break
|
||||
if now - start_time > kwargs.get('max_wait_time', 5):
|
||||
logger.debug(
|
||||
f'Network timeout after {kwargs.get("max_wait_time", 5)}s with {len(pending_requests)} '
|
||||
f'pending requests: {[r.url for r in pending_requests]}'
|
||||
)
|
||||
break
|
||||
|
||||
finally:
|
||||
# Clean up event listeners
|
||||
page.remove_listener('request', on_request)
|
||||
page.remove_listener('response', on_response)
|
||||
logger.debug(f'Network stabilized for {kwargs.get("idle_wait_time", 0.5)} seconds')
|
||||
|
||||
@staticmethod
|
||||
def _enhanced_css_selector_for_element(element: DOMElementNode, include_dynamic_attributes: bool = True) -> str:
|
||||
"""Creates a CSS selector for a DOM element, handling various edge cases and special characters.
|
||||
|
||||
Args:
|
||||
element: The DOM element to create a selector for
|
||||
|
||||
Returns:
|
||||
A valid CSS selector string
|
||||
"""
|
||||
try:
|
||||
# Get base selector from XPath
|
||||
css_selector = DomUtil._convert_simple_xpath_to_css_selector(element.xpath)
|
||||
|
||||
# Handle class attributes
|
||||
if 'class' in element.attributes and element.attributes['class'] and include_dynamic_attributes:
|
||||
# Define a regex pattern for valid class names in CSS
|
||||
valid_class_name_pattern = re.compile(r'^[a-zA-Z_][a-zA-Z0-9_-]*$')
|
||||
|
||||
# Iterate through the class attribute values
|
||||
classes = element.attributes['class'].split()
|
||||
for class_name in classes:
|
||||
# Skip empty class names
|
||||
if not class_name.strip():
|
||||
continue
|
||||
|
||||
# Check if the class name is valid
|
||||
if valid_class_name_pattern.match(class_name):
|
||||
# Append the valid class name to the CSS selector
|
||||
css_selector += f'.{class_name}'
|
||||
else:
|
||||
# Skip invalid class names
|
||||
continue
|
||||
|
||||
# Expanded set of safe attributes that are stable and useful for selection
|
||||
SAFE_ATTRIBUTES = {
|
||||
# Data attributes (if they're stable in your application)
|
||||
'id',
|
||||
# Standard HTML attributes
|
||||
'name',
|
||||
'type',
|
||||
'placeholder',
|
||||
# Accessibility attributes
|
||||
'aria-label',
|
||||
'aria-labelledby',
|
||||
'aria-describedby',
|
||||
'role',
|
||||
# Common form attributes
|
||||
'for',
|
||||
'autocomplete',
|
||||
'required',
|
||||
'readonly',
|
||||
# Media attributes
|
||||
'alt',
|
||||
'title',
|
||||
'src',
|
||||
# Custom stable attributes (add any application-specific ones)
|
||||
'href',
|
||||
'target',
|
||||
}
|
||||
|
||||
if include_dynamic_attributes:
|
||||
dynamic_attributes = {
|
||||
'data-id',
|
||||
'data-qa',
|
||||
'data-cy',
|
||||
'data-testid',
|
||||
}
|
||||
SAFE_ATTRIBUTES.update(dynamic_attributes)
|
||||
|
||||
# Handle other attributes
|
||||
for attribute, value in element.attributes.items():
|
||||
if attribute == 'class':
|
||||
continue
|
||||
|
||||
# Skip invalid attribute names
|
||||
if not attribute.strip():
|
||||
continue
|
||||
|
||||
if attribute not in SAFE_ATTRIBUTES:
|
||||
continue
|
||||
|
||||
# Escape special characters in attribute names
|
||||
safe_attribute = attribute.replace(':', r'\:')
|
||||
|
||||
# Handle different value cases
|
||||
if value == '':
|
||||
css_selector += f'[{safe_attribute}]'
|
||||
elif any(char in value for char in '"\'<>`\n\r\t'):
|
||||
# Use contains for values with special characters
|
||||
# Regex-substitute *any* whitespace with a single space, then strip.
|
||||
collapsed_value = re.sub(r'\s+', ' ', value).strip()
|
||||
# Escape embedded double-quotes.
|
||||
safe_value = collapsed_value.replace('"', '\\"')
|
||||
css_selector += f'[{safe_attribute}*="{safe_value}"]'
|
||||
else:
|
||||
css_selector += f'[{safe_attribute}="{value}"]'
|
||||
|
||||
return css_selector
|
||||
|
||||
except Exception:
|
||||
# Fallback to a more basic selector if something goes wrong
|
||||
tag_name = element.tag_name or '*'
|
||||
return f"{tag_name}[highlight_index='{element.highlight_index}']"
|
||||
|
||||
@staticmethod
|
||||
def _convert_simple_xpath_to_css_selector(xpath: str) -> str:
|
||||
"""Converts simple XPath expressions to CSS selectors."""
|
||||
if not xpath:
|
||||
return ''
|
||||
|
||||
# Remove leading slash if present
|
||||
xpath = xpath.lstrip('/')
|
||||
|
||||
# Split into parts
|
||||
parts = xpath.split('/')
|
||||
css_parts = []
|
||||
|
||||
for part in parts:
|
||||
if not part:
|
||||
continue
|
||||
|
||||
# Handle index notation [n]
|
||||
if '[' in part:
|
||||
base_part = part[: part.find('[')]
|
||||
index_part = part[part.find('['):]
|
||||
|
||||
# Handle multiple indices
|
||||
indices = [i.strip('[]') for i in index_part.split(']')[:-1]]
|
||||
|
||||
for idx in indices:
|
||||
try:
|
||||
# Handle numeric indices
|
||||
if idx.isdigit():
|
||||
index = int(idx) - 1
|
||||
base_part += f':nth-of-type({index + 1})'
|
||||
# Handle last() function
|
||||
elif idx == 'last()':
|
||||
base_part += ':last-of-type'
|
||||
# Handle position() functions
|
||||
elif 'position()' in idx:
|
||||
if '>1' in idx:
|
||||
base_part += ':nth-of-type(n+2)'
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
css_parts.append(base_part)
|
||||
else:
|
||||
css_parts.append(part)
|
||||
|
||||
base_selector = ' > '.join(css_parts)
|
||||
return base_selector
|
||||
@@ -0,0 +1,361 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import traceback
|
||||
from importlib import resources
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Tuple, List
|
||||
|
||||
from examples.common.tools.common import package
|
||||
from examples.common.tools.tool_action import BrowserAction
|
||||
from aworld.core.common import Observation, ActionModel, ActionResult
|
||||
from aworld.logs.util import logger
|
||||
from aworld.core.tool.base import action_executor, ToolFactory, AsyncTool
|
||||
from aworld.utils.import_package import is_package_installed
|
||||
from examples.common.tools.browsers.action.executor import BrowserToolActionExecutor
|
||||
from examples.common.tools.browsers.util.dom import DomTree
|
||||
from examples.common.tools.conf import BrowserToolConfig
|
||||
from examples.common.tools.browsers.util.dom_build import async_build_dom_tree
|
||||
from aworld.utils import import_package
|
||||
from aworld.tools.utils import build_observation
|
||||
|
||||
URL_MAX_LENGTH = 4096
|
||||
UTF8 = "".join(chr(x) for x in range(0, 55290))
|
||||
ASCII = "".join(chr(x) for x in range(32, 128))
|
||||
|
||||
|
||||
@ToolFactory.register(name="browser",
|
||||
desc="browser",
|
||||
asyn=True,
|
||||
supported_action=BrowserAction,
|
||||
conf_file_name=f'browser_tool.yaml',
|
||||
dir=f"{Path(__file__).parent.absolute()}")
|
||||
class BrowserTool(AsyncTool):
|
||||
def __init__(self, conf: BrowserToolConfig, **kwargs) -> None:
|
||||
super(BrowserTool, self).__init__(conf)
|
||||
|
||||
self.initialized = False
|
||||
self._finish = False
|
||||
self.record_trace = self.conf.get("working_dir", False)
|
||||
self.sleep_after_init = self.conf.get("sleep_after_init", False)
|
||||
dom_js_path = self.conf.get('dom_js_path')
|
||||
if dom_js_path and os.path.exists(dom_js_path):
|
||||
with open(dom_js_path, 'r') as read:
|
||||
self.js_code = read.read()
|
||||
else:
|
||||
self.js_code = resources.read_text(f'{package}.browsers.script',
|
||||
'buildDomTree.js')
|
||||
self.cur_observation = None
|
||||
if not is_package_installed('playwright'):
|
||||
import_package("playwright")
|
||||
logger.info("playwright install...")
|
||||
try:
|
||||
subprocess.check_call('playwright install', shell=True, timeout=300)
|
||||
except Exception as e:
|
||||
logger.error(f"Fail to auto execute playwright install, you can install manually\n {e}")
|
||||
|
||||
async def init(self) -> None:
|
||||
from playwright.async_api import async_playwright
|
||||
|
||||
if self.initialized:
|
||||
return
|
||||
|
||||
self.context_manager = async_playwright()
|
||||
self.playwright = await self.context_manager.start()
|
||||
|
||||
self.browser = await self._create_browser()
|
||||
self.browser_context = await self._create_browser_context()
|
||||
|
||||
if self.record_trace:
|
||||
await self.browser_context.tracing.start(screenshots=True, snapshots=True)
|
||||
|
||||
self.page = await self.browser_context.new_page()
|
||||
if self.conf.get("custom_executor"):
|
||||
self.action_executor = BrowserToolActionExecutor(self)
|
||||
else:
|
||||
self.action_executor = action_executor
|
||||
self.initialized = True
|
||||
|
||||
async def _create_browser(self):
|
||||
browse_name = self.conf.get("browse_name", "chromium")
|
||||
browse = getattr(self.playwright, browse_name)
|
||||
cdp_url = self.conf.get("cdp_url")
|
||||
wss_url = self.conf.get("wss_url")
|
||||
if cdp_url:
|
||||
if browse_name != "chromium":
|
||||
logger.warning(f"{browse_name} unsupported CDP, will use chromium browser")
|
||||
browse = self.playwright.chromium
|
||||
logger.info(f"Connecting to remote browser via CDP {cdp_url}")
|
||||
browser = await browse.connect_over_cdp(cdp_url)
|
||||
elif wss_url:
|
||||
logger.info(f"Connecting to remote browser via wss {wss_url}")
|
||||
browser = await browse.connect(wss_url)
|
||||
else:
|
||||
headless = self.conf.get("headless", False)
|
||||
slow_mo = self.conf.get("slow_mo", 0)
|
||||
disable_security_args = []
|
||||
if self.conf.get('disable_security', False):
|
||||
disable_security_args = ['--disable-web-security',
|
||||
'--disable-site-isolation-trials',
|
||||
'--disable-features=IsolateOrigins,site-per-process']
|
||||
args = ['--no-sandbox',
|
||||
'--disable-crash-reporte',
|
||||
'--disable-blink-features=AutomationControlled',
|
||||
'--disable-infobars',
|
||||
'--disable-background-timer-throttling',
|
||||
'--disable-popup-blocking',
|
||||
'--disable-backgrounding-occluded-windows',
|
||||
'--disable-renderer-backgrounding',
|
||||
'--disable-window-activation',
|
||||
'--disable-focus-on-load',
|
||||
'--no-first-run',
|
||||
'--no-default-browser-check',
|
||||
'--no-startup-window',
|
||||
'--window-position=0,0',
|
||||
'--window-size=1280,720'] + disable_security_args
|
||||
browser = await browse.launch(
|
||||
headless=headless,
|
||||
slow_mo=slow_mo,
|
||||
args=args,
|
||||
proxy=self.conf.get('proxy'),
|
||||
)
|
||||
return browser
|
||||
|
||||
async def _create_browser_context(self):
|
||||
"""Creates a new browser context with anti-detection measures and loads cookies if available."""
|
||||
from playwright.async_api import ViewportSize
|
||||
|
||||
browser = self.browser
|
||||
if self.conf.get("cdp_url") and len(browser.contexts) > 0:
|
||||
context = browser.contexts[0]
|
||||
else:
|
||||
viewport_size = ViewportSize(width=self.conf.get("width", 1280),
|
||||
height=self.conf.get("height", 720))
|
||||
disable_security = self.conf.get('disable_security', False)
|
||||
|
||||
context = await browser.new_context(viewport=viewport_size,
|
||||
no_viewport=False,
|
||||
user_agent=self.conf.get('user_agent'),
|
||||
java_script_enabled=True,
|
||||
bypass_csp=disable_security,
|
||||
ignore_https_errors=disable_security,
|
||||
record_video_dir=self.conf.get('working_dir'),
|
||||
record_video_size=viewport_size,
|
||||
locale=self.conf.get('locale'),
|
||||
storage_state=self.conf.get("storage_state", None),
|
||||
geolocation=self.conf.get("geolocation", None),
|
||||
device_scale_factor=1)
|
||||
if "chromium" == self.conf.get("browse_name", "chromium"):
|
||||
await context.grant_permissions(['camera', 'microphone'])
|
||||
|
||||
if self.conf.get('trace_path'):
|
||||
await context.tracing.start(screenshots=True, snapshots=True, sources=True)
|
||||
|
||||
cookie_file = self.conf.get('cookies_file')
|
||||
if cookie_file and os.path.exists(cookie_file):
|
||||
with open(cookie_file, 'r') as read:
|
||||
cookies = json.loads(read.read())
|
||||
await context.add_cookies(cookies)
|
||||
logger.info(f'Cookies load from {cookie_file} finished')
|
||||
|
||||
if self.conf.get('private'):
|
||||
js = resources.read_text(f"{package}.browsers.script", "stealth.min.js")
|
||||
await context.add_init_script(js)
|
||||
|
||||
return context
|
||||
|
||||
async def get_cur_page(self):
|
||||
return self.page
|
||||
|
||||
async def screenshot(self, full_page: bool = False) -> str:
|
||||
"""Returns a base64 encoded screenshot of the current page.
|
||||
|
||||
Args:
|
||||
full_page: When true, takes a screenshot of the full scrollable page, instead of the currently visible viewport.
|
||||
|
||||
Returns:
|
||||
Base64 of the page screenshot
|
||||
"""
|
||||
page = await self.get_cur_page()
|
||||
|
||||
try:
|
||||
await page.bring_to_front()
|
||||
await page.wait_for_load_state(timeout=2000)
|
||||
except:
|
||||
logger.warning("bring to front load timeout")
|
||||
pass
|
||||
|
||||
screenshot = await page.screenshot(
|
||||
full_page=full_page,
|
||||
animations='disabled',
|
||||
timeout=600000
|
||||
)
|
||||
logger.info("page screenshot finished")
|
||||
screenshot_base64 = base64.b64encode(screenshot).decode('utf-8')
|
||||
return screenshot_base64
|
||||
|
||||
async def _get_observation(self, info: Dict[str, Any] = {}) -> Observation:
|
||||
fail_error = info.get('exception')
|
||||
if fail_error:
|
||||
return Observation(observer=self.name(), action_result=[ActionResult(error=fail_error)])
|
||||
|
||||
try:
|
||||
dom_tree = await self._parse_dom_tree()
|
||||
image = await self.screenshot()
|
||||
pixels_above, pixels_below = await self._scroll_info()
|
||||
info.update({"pixels_above": pixels_above,
|
||||
"pixels_below": pixels_below,
|
||||
"url": self.page.url})
|
||||
return Observation(observer=self.name(), dom_tree=dom_tree, image=image, info=info)
|
||||
except Exception as e:
|
||||
try:
|
||||
try:
|
||||
await self.page.go_back()
|
||||
except:
|
||||
logger.warning("current page abnormal, new page to use.")
|
||||
self.page = await self.browser_context.new_page()
|
||||
dom_tree = await self._parse_dom_tree()
|
||||
image = await self.screenshot()
|
||||
pixels_above, pixels_below = await self._scroll_info()
|
||||
info.update({"pixels_above": pixels_above,
|
||||
"pixels_below": pixels_below,
|
||||
"url": self.page.url})
|
||||
return Observation(observer=self.name(), dom_tree=dom_tree, image=image, info=info)
|
||||
except Exception as e:
|
||||
logger.warning(f"build observation fail, {traceback.format_exc()}")
|
||||
return Observation(observer=self.name(), action_result=[ActionResult(error=traceback.format_exc())])
|
||||
|
||||
async def _parse_dom_tree(self) -> DomTree:
|
||||
args = {
|
||||
'doHighlightElements': self.conf.get("do_highlight", True),
|
||||
'focusHighlightIndex': self.conf.get("focus_highlight", -1),
|
||||
'viewportExpansion': self.conf.get("viewport_expansion", 0),
|
||||
'debugMode': logger.getEffectiveLevel() == 10,
|
||||
}
|
||||
element_tree, element_map = await async_build_dom_tree(self.page, self.js_code, args)
|
||||
return DomTree(element_tree=element_tree, element_map=element_map)
|
||||
|
||||
async def _scroll_info(self) -> tuple[int, int]:
|
||||
"""Get scroll position information for the current page."""
|
||||
scroll_y = await self.page.evaluate('window.scrollY')
|
||||
viewport_height = await self.page.evaluate('window.innerHeight')
|
||||
total_height = await self.page.evaluate('document.documentElement.scrollHeight')
|
||||
pixels_above = scroll_y
|
||||
pixels_below = total_height - (scroll_y + viewport_height)
|
||||
return pixels_above, pixels_below
|
||||
|
||||
async def reset(self, *, seed: int | None = None, options: Dict[str, str] | None = None) -> Tuple[
|
||||
Observation, Dict[str, Any]]:
|
||||
await super().reset(seed=seed, options=options)
|
||||
if self.initialized:
|
||||
observation = await self._get_observation()
|
||||
observation.action_result = [ActionResult(content='start', keep=True)]
|
||||
self.cur_observation = observation
|
||||
return observation, {}
|
||||
|
||||
await self.close()
|
||||
await self.init()
|
||||
|
||||
if self.sleep_after_init > 0:
|
||||
await asyncio.sleep(self.sleep_after_init)
|
||||
|
||||
observation = await self._get_observation()
|
||||
observation.action_result = [ActionResult(content='start', keep=True)]
|
||||
observation.ability = ''
|
||||
self.cur_observation = observation
|
||||
return observation, {}
|
||||
|
||||
async def save_trace(self, trace_path: str | Path) -> None:
|
||||
if self.record_trace:
|
||||
await self.browser_context.tracing.stop(path=trace_path)
|
||||
|
||||
@property
|
||||
async def finished(self) -> bool:
|
||||
return self._finish
|
||||
|
||||
async def close(self) -> None:
|
||||
if hasattr(self, 'context') and self.browser_context:
|
||||
await self.browser_context.close()
|
||||
if hasattr(self, 'browser') and self.browser:
|
||||
await self.browser.close()
|
||||
if hasattr(self, 'playwright') and self.playwright:
|
||||
await self.playwright.stop()
|
||||
if self.initialized:
|
||||
await self.context_manager.__aexit__()
|
||||
|
||||
async def do_step(self, action: List[ActionModel], **kwargs) -> Tuple[
|
||||
Observation, float, bool, bool, Dict[str, Any]]:
|
||||
if not self.initialized:
|
||||
raise RuntimeError("Call init first before calling step.")
|
||||
|
||||
if not action:
|
||||
logger.warning(f"{self.name()} has no action")
|
||||
return build_observation(observer=self.name(), ability='', content='no action'), 0., False, False, {}
|
||||
|
||||
reward = 0
|
||||
fail_error = ""
|
||||
action_result = None
|
||||
|
||||
invalid_acts: List[int] = []
|
||||
for i, act in enumerate(action):
|
||||
if act.tool_name != 'browser':
|
||||
logger.warning(f"tool {act.tool_name} is not a browser!")
|
||||
invalid_acts.append(i)
|
||||
|
||||
if invalid_acts:
|
||||
for i in invalid_acts:
|
||||
action[i] = None
|
||||
|
||||
try:
|
||||
action_result, self.page = await self.action_executor.async_execute_action(action,
|
||||
observation=self.cur_observation,
|
||||
llm_config=self.conf.llm_config,
|
||||
**kwargs)
|
||||
reward = 1
|
||||
except Exception as e:
|
||||
fail_error = str(e)
|
||||
|
||||
info = {"exception": fail_error}
|
||||
terminated = kwargs.get("terminated", False)
|
||||
for res in action_result:
|
||||
if res.is_done:
|
||||
terminated = res.is_done
|
||||
info['done'] = True
|
||||
self._finish = True
|
||||
if res.error:
|
||||
fail_error += res.error
|
||||
|
||||
contains_write_to_file = any(act.action_name == BrowserAction.WRITE_TO_FILE.value.name for act in action if act)
|
||||
if contains_write_to_file:
|
||||
msg = ""
|
||||
for action_result_elem in action_result:
|
||||
msg = action_result_elem.content
|
||||
# write_to_file observation
|
||||
return (Observation(content=msg, action_result=action_result, info=info),
|
||||
reward,
|
||||
terminated,
|
||||
kwargs.get("truncated", False),
|
||||
info)
|
||||
elif fail_error:
|
||||
# failed error observation
|
||||
return (Observation(action_result=action_result, observer=self.name()),
|
||||
reward,
|
||||
terminated,
|
||||
kwargs.get("truncated", False),
|
||||
info)
|
||||
else:
|
||||
# normal observation
|
||||
observation = await self._get_observation(info)
|
||||
observation.action_result = action_result
|
||||
observation.ability = action[-1].action_name
|
||||
self.cur_observation = observation
|
||||
return (observation,
|
||||
reward,
|
||||
terminated,
|
||||
kwargs.get("truncated", False),
|
||||
info)
|
||||
@@ -0,0 +1,368 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
import traceback
|
||||
from importlib import resources
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Tuple, List, Union
|
||||
|
||||
from aworld.config import ConfigDict
|
||||
from examples.common.tools.common import package
|
||||
from examples.common.tools.tool_action import BrowserAction
|
||||
from aworld.core.common import Observation, ActionModel, ActionResult
|
||||
from aworld.logs.util import logger
|
||||
from aworld.core.tool.base import action_executor, ToolFactory
|
||||
from aworld.core.tool.base import Tool
|
||||
from aworld.utils.import_package import is_package_installed
|
||||
from examples.common.tools.browsers.action.executor import BrowserToolActionExecutor
|
||||
from examples.common.tools.browsers.util.dom import DomTree
|
||||
from examples.common.tools.conf import BrowserToolConfig
|
||||
from examples.common.tools.browsers.util.dom_build import build_dom_tree
|
||||
from aworld.utils import import_package
|
||||
from aworld.tools.utils import build_observation
|
||||
|
||||
URL_MAX_LENGTH = 4096
|
||||
UTF8 = "".join(chr(x) for x in range(0, 55290))
|
||||
ASCII = "".join(chr(x) for x in range(32, 128))
|
||||
|
||||
BROWSER = "browser"
|
||||
|
||||
|
||||
@ToolFactory.register(name=BROWSER,
|
||||
desc="browser",
|
||||
supported_action=BrowserAction,
|
||||
conf_file_name=f'browser_tool.yaml',
|
||||
dir=f"{Path(__file__).parent.absolute()}")
|
||||
class BrowserTool(Tool):
|
||||
def __init__(self, conf: Union[ConfigDict, BrowserToolConfig], **kwargs) -> None:
|
||||
super(BrowserTool, self).__init__(conf, **kwargs)
|
||||
|
||||
self.initialized = False
|
||||
self._finish = False
|
||||
self.record_trace = self.conf.get("enable_recording", False)
|
||||
self.sleep_after_init = self.conf.get("sleep_after_init", False)
|
||||
|
||||
dom_js_path = self.conf.get('dom_js_path')
|
||||
if dom_js_path and os.path.exists(dom_js_path):
|
||||
with open(dom_js_path, 'r') as read:
|
||||
self.js_code = read.read()
|
||||
else:
|
||||
self.js_code = resources.read_text(f'{package}.browsers.script',
|
||||
'buildDomTree.js')
|
||||
self.cur_observation = None
|
||||
if not is_package_installed('playwright'):
|
||||
import_package("playwright")
|
||||
logger.info("playwright install...")
|
||||
try:
|
||||
subprocess.check_call('playwright install', shell=True, timeout=300)
|
||||
except Exception as e:
|
||||
logger.error(f"Fail to auto execute playwright install, you can install manually\n {e}")
|
||||
|
||||
def init(self) -> None:
|
||||
from playwright.sync_api import sync_playwright
|
||||
|
||||
if self.initialized:
|
||||
return
|
||||
|
||||
self.context_manager = sync_playwright()
|
||||
self.playwright = self.context_manager.start()
|
||||
self.browser = self._create_browser()
|
||||
self.browser_context = self._create_browser_context()
|
||||
|
||||
if self.record_trace:
|
||||
self.browser_context.tracing.start(screenshots=True, snapshots=True)
|
||||
|
||||
self.page = self.browser_context.new_page()
|
||||
if self.conf.get("custom_executor"):
|
||||
self.action_executor = BrowserToolActionExecutor(self)
|
||||
else:
|
||||
self.action_executor = action_executor
|
||||
self.initialized = True
|
||||
|
||||
def _create_browser(self):
|
||||
browse_name = self.conf.get("browse_name", "chromium")
|
||||
browse = getattr(self.playwright, browse_name)
|
||||
cdp_url = self.conf.get("cdp_url")
|
||||
wss_url = self.conf.get("wss_url")
|
||||
if cdp_url:
|
||||
if browse_name != "chromium":
|
||||
logger.warning(f"{browse_name} unsupported CDP, will use chromium browser")
|
||||
browse = self.playwright.chromium
|
||||
logger.info(f"Connecting to remote browser via CDP {cdp_url}")
|
||||
browser = browse.connect_over_cdp(cdp_url)
|
||||
elif wss_url:
|
||||
logger.info(f"Connecting to remote browser via wss {wss_url}")
|
||||
browser = browse.connect(wss_url)
|
||||
else:
|
||||
headless = self.conf.get("headless", False)
|
||||
slow_mo = self.conf.get("slow_mo", 0)
|
||||
disable_security_args = []
|
||||
if self.conf.get('disable_security', False):
|
||||
disable_security_args = ['--disable-web-security',
|
||||
'--disable-site-isolation-trials',
|
||||
'--disable-features=IsolateOrigins,site-per-process']
|
||||
args = ['--no-sandbox',
|
||||
'--disable-crash-reporte',
|
||||
'--disable-blink-features=AutomationControlled',
|
||||
'--disable-infobars',
|
||||
'--disable-background-timer-throttling',
|
||||
'--disable-popup-blocking',
|
||||
'--disable-backgrounding-occluded-windows',
|
||||
'--disable-renderer-backgrounding',
|
||||
'--disable-window-activation',
|
||||
'--disable-focus-on-load',
|
||||
'--no-first-run',
|
||||
'--no-default-browser-check',
|
||||
'--no-startup-window',
|
||||
'--window-position=0,0',
|
||||
'--window-size=1280,720'] + disable_security_args
|
||||
browser = browse.launch(
|
||||
headless=headless,
|
||||
slow_mo=slow_mo,
|
||||
args=args,
|
||||
proxy=self.conf.get('proxy'),
|
||||
)
|
||||
return browser
|
||||
|
||||
def _create_browser_context(self):
|
||||
"""Creates a new browser context with anti-detection measures and loads cookies if available."""
|
||||
from playwright.sync_api import ViewportSize
|
||||
|
||||
browser = self.browser
|
||||
if self.conf.get("cdp_url") and len(browser.contexts) > 0:
|
||||
context = browser.contexts[0]
|
||||
else:
|
||||
viewport_size = ViewportSize(width=self.conf.get("width", 1280),
|
||||
height=self.conf.get("height", 720))
|
||||
disable_security = self.conf.get('disable_security', False)
|
||||
|
||||
context = browser.new_context(viewport=viewport_size,
|
||||
no_viewport=False,
|
||||
user_agent=self.conf.get('user_agent'),
|
||||
java_script_enabled=True,
|
||||
bypass_csp=disable_security,
|
||||
ignore_https_errors=disable_security,
|
||||
record_video_dir=self.conf.get('working_dir'),
|
||||
record_video_size=viewport_size,
|
||||
locale=self.conf.get('locale'),
|
||||
storage_state=self.conf.get("storage_state", None),
|
||||
geolocation=self.conf.get("geolocation", None),
|
||||
device_scale_factor=1)
|
||||
if "chromium" == self.conf.get("browse_name", "chromium"):
|
||||
context.grant_permissions(['camera', 'microphone'])
|
||||
|
||||
if self.conf.get('working_dir'):
|
||||
context.tracing.start(screenshots=True, snapshots=True, sources=True)
|
||||
|
||||
cookie_file = self.conf.get('cookies_file')
|
||||
if cookie_file and os.path.exists(cookie_file):
|
||||
with open(cookie_file, 'r') as read:
|
||||
cookies = json.loads(read.read())
|
||||
context.add_cookies(cookies)
|
||||
logger.info(f'Cookies load from {cookie_file} finished')
|
||||
|
||||
if self.conf.get('private'):
|
||||
js = resources.read_text(f"{package}.browsers.script", "stealth.min.js")
|
||||
context.add_init_script(js)
|
||||
|
||||
return context
|
||||
|
||||
def get_cur_page(self):
|
||||
return self.page
|
||||
|
||||
def screenshot(self, full_page: bool = False) -> str:
|
||||
"""Returns a base64 encoded screenshot of the current page.
|
||||
|
||||
Args:
|
||||
full_page: When true, takes a screenshot of the full scrollable page, instead of the currently visible viewport.
|
||||
|
||||
Returns:
|
||||
Base64 of the page screenshot
|
||||
"""
|
||||
page = self.get_cur_page()
|
||||
|
||||
try:
|
||||
page.bring_to_front()
|
||||
page.wait_for_load_state(timeout=2000)
|
||||
except:
|
||||
logger.warning("bring to front load timeout")
|
||||
pass
|
||||
|
||||
screenshot = page.screenshot(
|
||||
full_page=full_page,
|
||||
animations='disabled',
|
||||
timeout=600000
|
||||
)
|
||||
logger.info("page screenshot finished")
|
||||
screenshot_base64 = base64.b64encode(screenshot).decode('utf-8')
|
||||
return screenshot_base64
|
||||
|
||||
def _get_observation(self, info: Dict[str, Any] = {}) -> Observation:
|
||||
fail_error = info.get('exception')
|
||||
if fail_error:
|
||||
return Observation(observer=self.name(), action_result=[ActionResult(error=fail_error)])
|
||||
|
||||
try:
|
||||
dom_tree = self._parse_dom_tree()
|
||||
image = self.screenshot()
|
||||
pixels_above, pixels_below = self._scroll_info()
|
||||
info.update({"pixels_above": pixels_above,
|
||||
"pixels_below": pixels_below,
|
||||
"url": self.page.url})
|
||||
return Observation(observer=self.name(),
|
||||
dom_tree=dom_tree,
|
||||
image=image,
|
||||
info=info)
|
||||
except Exception as e:
|
||||
try:
|
||||
self.page.go_back()
|
||||
except:
|
||||
logger.warning("current page abnormal, new page to use.")
|
||||
self.page = self.browser_context.new_page()
|
||||
try:
|
||||
dom_tree = self._parse_dom_tree()
|
||||
image = self.screenshot()
|
||||
pixels_above, pixels_below = self._scroll_info()
|
||||
info.update({"pixels_above": pixels_above,
|
||||
"pixels_below": pixels_below,
|
||||
"url": self.page.url})
|
||||
return Observation(observer=self.name(), dom_tree=dom_tree, image=image, info=info)
|
||||
except Exception as e:
|
||||
logger.warning(f"build observation fail, {traceback.format_exc()}")
|
||||
return Observation(observer=self.name(), action_result=[ActionResult(error=traceback.format_exc())])
|
||||
|
||||
def _parse_dom_tree(self) -> DomTree:
|
||||
args = {
|
||||
'doHighlightElements': self.conf.get("do_highlight", True),
|
||||
'focusHighlightIndex': self.conf.get("focus_highlight", -1),
|
||||
'viewportExpansion': self.conf.get("viewport_expansion", 0),
|
||||
'debugMode': logger.getEffectiveLevel() == 10,
|
||||
}
|
||||
element_tree, element_map = build_dom_tree(self.page, self.js_code, args)
|
||||
return DomTree(element_tree=element_tree, element_map=element_map)
|
||||
|
||||
def _scroll_info(self) -> tuple[int, int]:
|
||||
"""Get scroll position information for the current page."""
|
||||
scroll_y = self.page.evaluate('window.scrollY')
|
||||
viewport_height = self.page.evaluate('window.innerHeight')
|
||||
total_height = self.page.evaluate('document.documentElement.scrollHeight')
|
||||
pixels_above = scroll_y
|
||||
pixels_below = total_height - (scroll_y + viewport_height)
|
||||
return pixels_above, pixels_below
|
||||
|
||||
def reset(self, *, seed: int | None = None, options: Dict[str, str] | None = None) -> Tuple[
|
||||
Observation, Dict[str, Any]]:
|
||||
super().reset(seed=seed, options=options)
|
||||
if self.initialized:
|
||||
observation = self._get_observation()
|
||||
observation.action_result = [ActionResult(content='start', keep=True)]
|
||||
self.cur_observation = observation
|
||||
return observation, {}
|
||||
|
||||
self.close()
|
||||
self.init()
|
||||
|
||||
if self.sleep_after_init > 0:
|
||||
time.sleep(self.sleep_after_init)
|
||||
|
||||
observation = self._get_observation()
|
||||
observation.action_result = [ActionResult(content='start', keep=True)]
|
||||
self.cur_observation = observation
|
||||
return observation, {}
|
||||
|
||||
@property
|
||||
def finished(self) -> bool:
|
||||
return self._finish
|
||||
|
||||
def save_trace(self, trace_path: str | Path) -> None:
|
||||
if self.record_trace:
|
||||
self.browser_context.tracing.stop(path=trace_path)
|
||||
|
||||
def close(self) -> None:
|
||||
if hasattr(self, 'context') and self.browser_context:
|
||||
self.browser_context.close()
|
||||
if hasattr(self, 'browser') and self.browser:
|
||||
self.browser.close()
|
||||
if hasattr(self, 'playwright') and self.playwright:
|
||||
self.playwright.stop()
|
||||
|
||||
if self.initialized:
|
||||
self.context_manager.__exit__()
|
||||
|
||||
def do_step(self, action: List[ActionModel], **kwargs) -> Tuple[
|
||||
Observation, float, bool, bool, Dict[str, Any]]:
|
||||
if not self.initialized:
|
||||
raise RuntimeError("Call init first before calling step.")
|
||||
|
||||
if not action:
|
||||
logger.warning(f"{self.name()} has no action")
|
||||
return build_observation(observer=self.name(), ability='', content='no action'), 0., False, False, {}
|
||||
|
||||
reward = 0
|
||||
fail_error = ""
|
||||
action_result = None
|
||||
|
||||
invalid_acts: List[int] = []
|
||||
for i, act in enumerate(action):
|
||||
if act.tool_name != BROWSER:
|
||||
logger.warning(f"tool {act.tool_name} is not a browser!")
|
||||
invalid_acts.append(i)
|
||||
|
||||
if invalid_acts:
|
||||
for i in invalid_acts:
|
||||
action[i] = None
|
||||
|
||||
try:
|
||||
action_result, self.page = self.action_executor.execute_action(action,
|
||||
observation=self.cur_observation,
|
||||
llm_config=self.conf.llm_config,
|
||||
**kwargs)
|
||||
reward = 1
|
||||
except Exception as e:
|
||||
fail_error = str(e)
|
||||
|
||||
info = {"exception": fail_error}
|
||||
terminated = kwargs.get("terminated", False)
|
||||
if action_result:
|
||||
for res in action_result:
|
||||
if res.is_done:
|
||||
terminated = res.is_done
|
||||
info['done'] = True
|
||||
self._finish = True
|
||||
if res.error:
|
||||
fail_error += res.error
|
||||
|
||||
contains_write_to_file = any(act.action_name == BrowserAction.WRITE_TO_FILE.value.name for act in action if act)
|
||||
if contains_write_to_file:
|
||||
msg = ""
|
||||
for action_result_elem in action_result:
|
||||
msg = action_result_elem.content
|
||||
# write_to_file observation
|
||||
return (Observation(content=msg, action_result=action_result, info=info),
|
||||
reward,
|
||||
terminated,
|
||||
kwargs.get("truncated", False),
|
||||
info)
|
||||
elif fail_error:
|
||||
# failed error observation
|
||||
return (Observation(action_result=action_result, observer=self.name()),
|
||||
reward,
|
||||
terminated,
|
||||
kwargs.get("truncated", False),
|
||||
info)
|
||||
else:
|
||||
# normal observation
|
||||
observation = self._get_observation(info)
|
||||
observation.ability = action[-1].action_name
|
||||
observation.action_result = action_result
|
||||
self.cur_observation = observation
|
||||
return (observation,
|
||||
reward,
|
||||
terminated,
|
||||
kwargs.get("truncated", False),
|
||||
info)
|
||||
@@ -0,0 +1,24 @@
|
||||
browse_name: chromium
|
||||
headless: False
|
||||
width: 1280
|
||||
height: 720
|
||||
slow_mo: 0
|
||||
disable_security: False
|
||||
custom_executor: False
|
||||
dom_js_path:
|
||||
private:
|
||||
locale:
|
||||
geolocation:
|
||||
storage_state:
|
||||
do_highlight: True
|
||||
focus_highlight: -1
|
||||
viewport_expansion: 0
|
||||
cdp_url:
|
||||
wss_url:
|
||||
proxy:
|
||||
cookies_file:
|
||||
working_dir:
|
||||
enable_recording: False
|
||||
sleep_after_init: 0
|
||||
max_retry: 3
|
||||
reuse: True
|
||||
@@ -0,0 +1,2 @@
|
||||
playwright
|
||||
markdownify
|
||||
+1056
File diff suppressed because it is too large
Load Diff
+7
File diff suppressed because one or more lines are too long
@@ -0,0 +1,210 @@
|
||||
# coding: utf-8
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Dict, List
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class Coordinates(BaseModel):
|
||||
x: int
|
||||
y: int
|
||||
|
||||
|
||||
class CoordinateSet(BaseModel):
|
||||
top_left: Coordinates
|
||||
top_right: Coordinates
|
||||
bottom_left: Coordinates
|
||||
bottom_right: Coordinates
|
||||
center: Coordinates
|
||||
width: int
|
||||
height: int
|
||||
|
||||
|
||||
class ViewportInfo(BaseModel):
|
||||
width: int
|
||||
height: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class HashedDomElement:
|
||||
"""
|
||||
Hash of the dom element to be used as a unique identifier
|
||||
"""
|
||||
|
||||
branch_path_hash: str
|
||||
attributes_hash: str
|
||||
xpath_hash: str
|
||||
|
||||
|
||||
@dataclass(frozen=False)
|
||||
class DOMBaseNode:
|
||||
is_visible: bool
|
||||
# Use None as default and set parent later to avoid circular reference issues
|
||||
parent: Optional['DOMElementNode']
|
||||
|
||||
|
||||
@dataclass(frozen=False)
|
||||
class DOMTextNode(DOMBaseNode):
|
||||
text: str
|
||||
type: str = 'TEXT_NODE'
|
||||
|
||||
def has_parent_with_highlight_index(self) -> bool:
|
||||
current = self.parent
|
||||
while current is not None:
|
||||
# stop if the element has a highlight index (will be handled separately)
|
||||
if current.highlight_index is not None:
|
||||
return True
|
||||
|
||||
current = current.parent
|
||||
return False
|
||||
|
||||
def is_parent_in_viewport(self) -> bool:
|
||||
if self.parent is None:
|
||||
return False
|
||||
return self.parent.is_in_viewport
|
||||
|
||||
def is_parent_top_element(self) -> bool:
|
||||
if self.parent is None:
|
||||
return False
|
||||
return self.parent.is_top_element
|
||||
|
||||
|
||||
@dataclass(frozen=False)
|
||||
class DOMElementNode(DOMBaseNode):
|
||||
"""
|
||||
xpath: the xpath of the element from the last root node (shadow root or iframe OR document if no shadow root or iframe).
|
||||
To properly reference the element we need to recursively switch the root node until we find the element (work you way up the tree with `.parent`)
|
||||
"""
|
||||
|
||||
tag_name: str
|
||||
xpath: str
|
||||
attributes: Dict[str, str]
|
||||
children: List[DOMBaseNode]
|
||||
is_interactive: bool = False
|
||||
is_top_element: bool = False
|
||||
is_in_viewport: bool = False
|
||||
shadow_root: bool = False
|
||||
highlight_index: Optional[int] = None
|
||||
viewport_coordinates: Optional[CoordinateSet] = None
|
||||
page_coordinates: Optional[CoordinateSet] = None
|
||||
viewport_info: Optional[ViewportInfo] = None
|
||||
|
||||
def __repr__(self) -> str:
|
||||
tag_str = f'<{self.tag_name}'
|
||||
|
||||
# Add attributes
|
||||
for key, value in self.attributes.items():
|
||||
tag_str += f' {key}="{value}"'
|
||||
tag_str += '>'
|
||||
|
||||
# Add extra info
|
||||
extras = []
|
||||
if self.is_interactive:
|
||||
extras.append('interactive')
|
||||
if self.is_top_element:
|
||||
extras.append('top')
|
||||
if self.shadow_root:
|
||||
extras.append('shadow-root')
|
||||
if self.highlight_index is not None:
|
||||
extras.append(f'highlight:{self.highlight_index}')
|
||||
if self.is_in_viewport:
|
||||
extras.append('in-viewport')
|
||||
|
||||
if extras:
|
||||
tag_str += f' [{", ".join(extras)}]'
|
||||
|
||||
return tag_str
|
||||
|
||||
def get_all_text_till_next_clickable_element(self, max_depth: int = -1) -> str:
|
||||
text_parts = []
|
||||
|
||||
def collect_text(node: DOMBaseNode, current_depth: int) -> None:
|
||||
if max_depth != -1 and current_depth > max_depth:
|
||||
return
|
||||
|
||||
# Skip this branch if we hit a highlighted element (except for the current node)
|
||||
if isinstance(node, DOMElementNode) and node != self and node.highlight_index is not None:
|
||||
return
|
||||
|
||||
if isinstance(node, DOMTextNode):
|
||||
text_parts.append(node.text)
|
||||
elif isinstance(node, DOMElementNode):
|
||||
for child in node.children:
|
||||
collect_text(child, current_depth + 1)
|
||||
|
||||
collect_text(self, 0)
|
||||
return '\n'.join(text_parts).strip()
|
||||
|
||||
def clickable_elements_to_string(self, include_attributes: list[str] | None = None) -> str:
|
||||
"""Convert the processed DOM content to HTML."""
|
||||
formatted_text = []
|
||||
|
||||
def process_node(node: DOMBaseNode, depth: int) -> None:
|
||||
if isinstance(node, DOMElementNode):
|
||||
# Add element with highlight_index
|
||||
if node.highlight_index is not None:
|
||||
attributes_str = ''
|
||||
text = node.get_all_text_till_next_clickable_element()
|
||||
if include_attributes:
|
||||
attributes = list(
|
||||
set(
|
||||
[
|
||||
str(value)
|
||||
for key, value in node.attributes.items()
|
||||
if key in include_attributes and value != node.tag_name
|
||||
]
|
||||
)
|
||||
)
|
||||
if text in attributes:
|
||||
attributes.remove(text)
|
||||
attributes_str = ';'.join(attributes)
|
||||
line = f'[{node.highlight_index}]<{node.tag_name} '
|
||||
if attributes_str:
|
||||
line += f'{attributes_str}'
|
||||
if text:
|
||||
if attributes_str:
|
||||
line += f'>{text}'
|
||||
else:
|
||||
line += f'{text}'
|
||||
line += '/>'
|
||||
formatted_text.append(line)
|
||||
|
||||
# Process children regardless
|
||||
for child in node.children:
|
||||
process_node(child, depth + 1)
|
||||
|
||||
elif isinstance(node, DOMTextNode):
|
||||
# Add text only if it doesn't have a highlighted parent
|
||||
if not node.has_parent_with_highlight_index() and node.is_visible: # and node.is_parent_top_element()
|
||||
formatted_text.append(f'{node.text}')
|
||||
|
||||
process_node(self, 0)
|
||||
return '\n'.join(formatted_text)
|
||||
|
||||
def get_file_upload_element(self, check_siblings: bool = True) -> Optional['DOMElementNode']:
|
||||
# Check if current element is a file input
|
||||
if self.tag_name == 'input' and self.attributes.get('type') == 'file':
|
||||
return self
|
||||
|
||||
# Check children
|
||||
for child in self.children:
|
||||
if isinstance(child, DOMElementNode):
|
||||
result = child.get_file_upload_element(check_siblings=False)
|
||||
if result:
|
||||
return result
|
||||
|
||||
# Check siblings only for the initial call
|
||||
if check_siblings and self.parent:
|
||||
for sibling in self.parent.children:
|
||||
if sibling is not self and isinstance(sibling, DOMElementNode):
|
||||
result = sibling.get_file_upload_element(check_siblings=False)
|
||||
if result:
|
||||
return result
|
||||
|
||||
return None
|
||||
|
||||
|
||||
class DomTree(BaseModel):
|
||||
element_tree: DOMElementNode
|
||||
element_map: Dict[int, DOMElementNode]
|
||||
@@ -0,0 +1,138 @@
|
||||
# coding: utf-8
|
||||
|
||||
# Derived from browser_use DomService, we use it as a utility method, and supports sync and async.
|
||||
|
||||
import gc
|
||||
import json
|
||||
|
||||
from typing import Dict, Any, Tuple, Optional
|
||||
|
||||
from aworld.utils.async_func import async_func
|
||||
from examples.common.tools.browsers.util.dom import DOMElementNode, DOMBaseNode, DOMTextNode, ViewportInfo
|
||||
from aworld.logs.util import logger
|
||||
|
||||
|
||||
async def async_build_dom_tree(page, js_code: str, args: Dict[str, Any]) -> Tuple[DOMElementNode, Dict[int, DOMElementNode]]:
|
||||
if await page.evaluate('1+1') != 2:
|
||||
raise ValueError('The page cannot evaluate javascript code properly')
|
||||
|
||||
# NOTE: We execute JS code in the browser to extract important DOM information.
|
||||
# The returned hash map contains information about the DOM tree and the
|
||||
# relationship between the DOM elements.
|
||||
try:
|
||||
eval_page = await page.evaluate(js_code, args)
|
||||
except Exception as e:
|
||||
logger.error('Error evaluating JavaScript: %s', e)
|
||||
raise
|
||||
|
||||
# Only log performance metrics in debug mode
|
||||
if args.get("debugMode") and 'perfMetrics' in eval_page:
|
||||
logger.debug('DOM Tree Building Performance Metrics:\n%s', json.dumps(eval_page['perfMetrics'], indent=2))
|
||||
|
||||
return await async_func(_construct_dom_tree)(eval_page)
|
||||
|
||||
|
||||
def build_dom_tree(page, js_code: str, args: Dict[str, Any]) -> Tuple[DOMElementNode, Dict[int, DOMElementNode]]:
|
||||
if page.evaluate('1+1') != 2:
|
||||
raise ValueError('The page cannot evaluate javascript code properly')
|
||||
|
||||
# NOTE: We execute JS code in the browser to extract important DOM information.
|
||||
# The returned hash map contains information about the DOM tree and the
|
||||
# relationship between the DOM elements.
|
||||
try:
|
||||
eval_page = page.evaluate(js_code, args)
|
||||
except Exception as e:
|
||||
logger.error('Error evaluating JavaScript: %s', e)
|
||||
raise
|
||||
|
||||
# Only log performance metrics in debug mode
|
||||
if args.get("debugMode") and 'perfMetrics' in eval_page:
|
||||
logger.debug('DOM Tree Building Performance Metrics:\n%s', json.dumps(eval_page['perfMetrics'], indent=2))
|
||||
|
||||
return _construct_dom_tree(eval_page)
|
||||
|
||||
|
||||
def _construct_dom_tree(eval_page: dict, ) -> tuple[DOMElementNode, Dict[int, DOMElementNode]]:
|
||||
js_node_map = eval_page['map']
|
||||
js_root_id = eval_page['rootId']
|
||||
|
||||
selector_map = {}
|
||||
node_map = {}
|
||||
|
||||
for id, node_data in js_node_map.items():
|
||||
node, children_ids = _parse_node(node_data)
|
||||
if node is None:
|
||||
continue
|
||||
|
||||
node_map[id] = node
|
||||
|
||||
if isinstance(node, DOMElementNode) and node.highlight_index is not None:
|
||||
selector_map[node.highlight_index] = node
|
||||
|
||||
# NOTE: We know that we are building the tree bottom up
|
||||
# and all children are already processed.
|
||||
if isinstance(node, DOMElementNode):
|
||||
for child_id in children_ids:
|
||||
if child_id not in node_map:
|
||||
continue
|
||||
|
||||
child_node = node_map[child_id]
|
||||
|
||||
child_node.parent = node
|
||||
node.children.append(child_node)
|
||||
|
||||
html_to_dict = node_map[str(js_root_id)]
|
||||
|
||||
del node_map
|
||||
del js_node_map
|
||||
del js_root_id
|
||||
|
||||
gc.collect()
|
||||
|
||||
if html_to_dict is None or not isinstance(html_to_dict, DOMElementNode):
|
||||
raise ValueError('Failed to parse HTML to dictionary')
|
||||
|
||||
return html_to_dict, selector_map
|
||||
|
||||
|
||||
def _parse_node(node_data: dict, ) -> Tuple[Optional[DOMBaseNode], list[int]]:
|
||||
if not node_data:
|
||||
return None, []
|
||||
|
||||
# Process text nodes immediately
|
||||
if node_data.get('type') == 'TEXT_NODE':
|
||||
text_node = DOMTextNode(
|
||||
text=node_data['text'],
|
||||
is_visible=node_data['isVisible'],
|
||||
parent=None,
|
||||
)
|
||||
return text_node, []
|
||||
|
||||
# Process coordinates if they exist for element nodes
|
||||
|
||||
viewport_info = None
|
||||
|
||||
if 'viewport' in node_data:
|
||||
viewport_info = ViewportInfo(
|
||||
width=node_data['viewport']['width'],
|
||||
height=node_data['viewport']['height'],
|
||||
)
|
||||
|
||||
element_node = DOMElementNode(
|
||||
tag_name=node_data['tagName'],
|
||||
xpath=node_data['xpath'],
|
||||
attributes=node_data.get('attributes', {}),
|
||||
children=[],
|
||||
is_visible=node_data.get('isVisible', False),
|
||||
is_interactive=node_data.get('isInteractive', False),
|
||||
is_top_element=node_data.get('isTopElement', False),
|
||||
is_in_viewport=node_data.get('isInViewport', False),
|
||||
highlight_index=node_data.get('highlightIndex'),
|
||||
shadow_root=node_data.get('shadowRoot', False),
|
||||
parent=None,
|
||||
viewport_info=viewport_info,
|
||||
)
|
||||
|
||||
children_ids = node_data.get('children', [])
|
||||
|
||||
return element_node, children_ids
|
||||
@@ -0,0 +1,37 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
from enum import Enum
|
||||
|
||||
package = 'examples.common.tools'
|
||||
|
||||
|
||||
class Tools(Enum):
|
||||
"""Tool list supported in the framework, pre-defined to avoid spelling errors."""
|
||||
BROWSER = "browser"
|
||||
ANDROID = "android"
|
||||
GYM = "openai_gym"
|
||||
SEARCH_API = "search_api"
|
||||
SHELL = "shell"
|
||||
PYTHON_EXECUTE = "python_execute"
|
||||
CODE_EXECUTE = "code_execute"
|
||||
FILE = "file"
|
||||
IMAGE_ANALYSIS = "image_analysis"
|
||||
DOCUMENT_ANALYSIS = "document_analysis"
|
||||
HTML = "html"
|
||||
MCP = "mcp"
|
||||
|
||||
|
||||
class Agents(Enum):
|
||||
"""Agent supported in the framework, pre-defined to avoid spelling errors."""
|
||||
BROWSER = "browser_agent"
|
||||
ANDROID = "android_agent"
|
||||
SEARCH = "search_agent"
|
||||
CODE_EXECUTE = "code_execute_agent"
|
||||
FILE = "file_agent"
|
||||
IMAGE_ANALYSIS = "image_analysis_agent"
|
||||
SHELL = "shell_agent"
|
||||
DOCUMENT = "document_agent"
|
||||
GYM = "gym_agent"
|
||||
PLAN = "plan_agent"
|
||||
EXECUTE = "execute_agent"
|
||||
SUMMARY = "summary_agent"
|
||||
@@ -0,0 +1,43 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
import os
|
||||
|
||||
from aworld.config.conf import ToolConfig, ModelConfig
|
||||
|
||||
|
||||
class BrowserToolConfig(ToolConfig):
|
||||
headless: bool = False
|
||||
keep_browser_open: bool = True
|
||||
private: bool = True
|
||||
browse_name: str = "chromium"
|
||||
custom_executor: bool = False
|
||||
width: int = 1280
|
||||
height: int = 720
|
||||
slow_mo: int = 0
|
||||
disable_security: bool = False
|
||||
dom_js_path: str = None
|
||||
locale: str = None
|
||||
geolocation: str = None
|
||||
storage_state: str = None
|
||||
do_highlight: bool = True
|
||||
focus_highlight: int = -1
|
||||
viewport_expansion: int = 0
|
||||
cdp_url: str = None
|
||||
wss_url: str = None
|
||||
proxy: str = None
|
||||
cookies_file: str = None
|
||||
working_dir: str = None
|
||||
enable_recording: bool = False
|
||||
sleep_after_init: float = 0
|
||||
max_retry: int = 3
|
||||
llm_config: ModelConfig = ModelConfig()
|
||||
max_extract_content_input_tokens: int = 64000
|
||||
max_extract_content_output_tokens: int = 5000
|
||||
reuse: bool = True
|
||||
|
||||
|
||||
class AndroidToolConfig(ToolConfig):
|
||||
avd_name: str | None = None
|
||||
adb_path: str | None = os.path.expanduser('~') + "/Library/Android/sdk/platform-tools/adb"
|
||||
emulator_path: str | None = os.path.expanduser('~') + "/Library/Android/sdk/emulator/emulator"
|
||||
headless: bool | None = None
|
||||
@@ -0,0 +1,2 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
@@ -0,0 +1,12 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
from examples.common.tools.tool_action import DocumentExecuteAction
|
||||
from aworld.core.tool.action_factory import ActionFactory
|
||||
from aworld.core.tool.action import ExecutableAction
|
||||
|
||||
|
||||
@ActionFactory.register(name=DocumentExecuteAction.DOCUMENT_ANALYSIS.value.name,
|
||||
desc=DocumentExecuteAction.DOCUMENT_ANALYSIS.value.desc,
|
||||
tool_name="document_analysis")
|
||||
class ExecuteAction(ExecutableAction):
|
||||
"""Only one action, define it, implemented can be omitted. Act in tool."""
|
||||
@@ -0,0 +1,529 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
|
||||
import json
|
||||
import os
|
||||
import base64
|
||||
import tempfile
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Tuple
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from aworld.config import ToolConfig
|
||||
from examples.common.tools.tool_action import DocumentExecuteAction
|
||||
from aworld.core.common import Observation, ActionModel, ActionResult
|
||||
from aworld.core.tool.base import ToolFactory, Tool
|
||||
from aworld.logs.util import logger
|
||||
from examples.common.tools.document.utils import encode_image_from_file, encode_image_from_url
|
||||
from aworld.utils import import_package, import_packages
|
||||
from aworld.tools.utils import build_observation
|
||||
|
||||
|
||||
class InputDocument(BaseModel):
|
||||
document_path: str | None = None
|
||||
|
||||
|
||||
@ToolFactory.register(name="document_analysis",
|
||||
desc="document analysis",
|
||||
supported_action=DocumentExecuteAction,
|
||||
conf_file_name=f'document_analysis_tool.yaml',
|
||||
dir=f"{Path(__file__).parent.absolute()}")
|
||||
class DocumentTool(Tool):
|
||||
def __init__(self, conf: ToolConfig, **kwargs) -> None:
|
||||
"""Init document tool."""
|
||||
import_package('cv2', install_name='opencv-python')
|
||||
import_packages(['xmltodict', 'pandas', 'docx2markdown', 'PyPDF2', 'numpy'])
|
||||
super(DocumentTool, self).__init__(conf, **kwargs)
|
||||
self.cur_observation = None
|
||||
self.content = None
|
||||
self.keyframes = []
|
||||
self.init()
|
||||
self.step_finished = True
|
||||
|
||||
def reset(self, *, seed: int | None = None, options: Dict[str, str] | None = None) -> Tuple[
|
||||
Observation, dict[str, Any]]:
|
||||
super().reset(seed=seed, options=options)
|
||||
|
||||
self.close()
|
||||
self.step_finished = True
|
||||
return build_observation(observer=self.name(),
|
||||
ability=DocumentExecuteAction.DOCUMENT_ANALYSIS.value.name), {}
|
||||
|
||||
def init(self) -> None:
|
||||
self.initialized = True
|
||||
|
||||
def close(self) -> None:
|
||||
pass
|
||||
|
||||
def finished(self) -> bool:
|
||||
return self.step_finished
|
||||
|
||||
def do_step(self, actions: list[ActionModel], **kwargs) -> Tuple[Observation, float, bool, bool, Dict[str, Any]]:
|
||||
self.step_finished = False
|
||||
reward = 0.
|
||||
fail_error = ""
|
||||
observation = build_observation(observer=self.name(),
|
||||
ability=DocumentExecuteAction.DOCUMENT_ANALYSIS.value.name)
|
||||
info = {}
|
||||
try:
|
||||
if not actions:
|
||||
raise ValueError("actions is empty")
|
||||
action = actions[0]
|
||||
document_path = action.params.get("document_path", "")
|
||||
if not document_path:
|
||||
raise ValueError("document path invalid")
|
||||
output, keyframes, error = self.document_analysis(document_path)
|
||||
observation.content = output
|
||||
observation.action_result.append(
|
||||
ActionResult(is_done=True,
|
||||
success=False if error else True,
|
||||
content=f"{output}",
|
||||
error=f"{error}",
|
||||
keep=False))
|
||||
info['key_frame'] = f"{keyframes}"
|
||||
reward = 1.
|
||||
except Exception as e:
|
||||
fail_error = str(e)
|
||||
finally:
|
||||
self.step_finished = True
|
||||
info["exception"] = fail_error
|
||||
info.update(kwargs)
|
||||
return (observation, reward, kwargs.get("terminated", False),
|
||||
kwargs.get("truncated", False), info)
|
||||
|
||||
def document_analysis(self, document_path):
|
||||
import xmltodict
|
||||
error = None
|
||||
# Initialize content to empty list to avoid None return
|
||||
self.content = []
|
||||
try:
|
||||
if any(document_path.endswith(ext) for ext in [".jpg", ".jpeg", ".png"]):
|
||||
parsed_url = urlparse(document_path)
|
||||
is_url = all([parsed_url.scheme, parsed_url.netloc])
|
||||
if not is_url:
|
||||
base64_image = encode_image_from_file(document_path)
|
||||
else:
|
||||
base64_image = encode_image_from_url(document_path)
|
||||
self.content = f"data:image/jpeg;base64,{base64_image}"
|
||||
|
||||
if any(document_path.endswith(ext) for ext in ["xls", "xlsx"]):
|
||||
try:
|
||||
try:
|
||||
import pandas as pd
|
||||
except ImportError:
|
||||
error = "pandas library not found. Please install pandas: pip install pandas"
|
||||
return self.content, self.keyframes, error
|
||||
|
||||
excel_data = {}
|
||||
|
||||
with pd.ExcelFile(document_path) as xls:
|
||||
sheet_names = xls.sheet_names
|
||||
for sheet_name in sheet_names:
|
||||
df = pd.read_excel(xls, sheet_name=sheet_name)
|
||||
sheet_data = df.to_dict(orient='records')
|
||||
excel_data[sheet_name] = sheet_data
|
||||
|
||||
self.content = json.dumps(excel_data, ensure_ascii=False)
|
||||
logger.info(f"Successfully processed Excel file: {document_path}")
|
||||
logger.info(f"Found {len(sheet_names)} sheets: {', '.join(sheet_names)}")
|
||||
|
||||
except Exception as excel_error:
|
||||
error = str(excel_error)
|
||||
|
||||
if any(document_path.endswith(ext) for ext in ["json", "jsonl", "jsonld"]):
|
||||
with open(document_path, "r", encoding="utf-8") as f:
|
||||
self.content = json.load(f)
|
||||
f.close()
|
||||
|
||||
if any(document_path.endswith(ext) for ext in ["xml"]):
|
||||
data = None
|
||||
with open(document_path, "r", encoding="utf-8") as f:
|
||||
data = f.read()
|
||||
f.close()
|
||||
|
||||
try:
|
||||
self.content = xmltodict.parse(data)
|
||||
logger.info(f"The extracted xml data is: {self.content}")
|
||||
|
||||
except Exception as e:
|
||||
logger.info(f"The raw xml data is: {data}")
|
||||
error = str(e)
|
||||
self.content = data
|
||||
|
||||
if any(document_path.endswith(ext) for ext in ["doc", "docx"]):
|
||||
from docx2markdown._docx_to_markdown import docx_to_markdown
|
||||
file_name = os.path.basename(document_path)
|
||||
md_file_path = f"{file_name}.md"
|
||||
docx_to_markdown(document_path, md_file_path)
|
||||
with open(md_file_path, "r") as f:
|
||||
self.content = f.read()
|
||||
f.close()
|
||||
|
||||
if any(document_path.endswith(ext) for ext in ["pdf"]):
|
||||
# try using pypdf to extract text from pdf
|
||||
try:
|
||||
from PyPDF2 import PdfReader
|
||||
|
||||
# Open file in binary mode for PdfReader
|
||||
f = open(document_path, "rb")
|
||||
reader = PdfReader(f)
|
||||
extracted_text = ""
|
||||
for page in reader.pages:
|
||||
extracted_text += page.extract_text()
|
||||
self.content = extracted_text
|
||||
f.close()
|
||||
except Exception as pdf_error:
|
||||
error = str(pdf_error)
|
||||
|
||||
# audio
|
||||
if any(document_path.endswith(ext.lower()) for ext in [".mp3", ".wav", ".wave"]):
|
||||
try:
|
||||
# audio-> base64
|
||||
with open(document_path, "rb") as audio_file:
|
||||
audio_bytes = audio_file.read()
|
||||
audio_base64 = base64.b64encode(audio_bytes).decode('utf-8')
|
||||
|
||||
# ext
|
||||
ext = os.path.splitext(document_path)[1].lower()
|
||||
mime_type = "audio/mpeg" if ext == ".mp3" else "audio/wav"
|
||||
|
||||
# data URI
|
||||
self.content = f"data:{mime_type};base64,{audio_base64}"
|
||||
except Exception as audio_error:
|
||||
error = str(audio_error)
|
||||
logger.error(f"Error processing audio file: {error}")
|
||||
|
||||
# video
|
||||
if any(document_path.endswith(ext.lower()) for ext in [".mp4", ".avi", ".mov", ".mkv", ".flv", ".wmv"]):
|
||||
try:
|
||||
try:
|
||||
import cv2
|
||||
import numpy as np
|
||||
except ImportError:
|
||||
error = "Required libraries not found. Please install opencv-python: pip install opencv-python"
|
||||
return None, None, error
|
||||
|
||||
# create temp dir
|
||||
temp_dir = tempfile.mkdtemp()
|
||||
|
||||
# 1.get audio -> base64
|
||||
audio_path = os.path.join(temp_dir, "extracted_audio.mp3")
|
||||
|
||||
# get audio by ffmpeg
|
||||
try:
|
||||
subprocess.run([
|
||||
"ffmpeg", "-i", document_path, "-q:a", "0",
|
||||
"-map", "a", audio_path, "-y"
|
||||
], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
||||
|
||||
# audio->base64
|
||||
with open(audio_path, "rb") as audio_file:
|
||||
audio_bytes = audio_file.read()
|
||||
audio_base64 = base64.b64encode(audio_bytes).decode('utf-8')
|
||||
|
||||
audio_data_uri = f"data:audio/mpeg;base64,{audio_base64}"
|
||||
except (subprocess.SubprocessError, FileNotFoundError) as e:
|
||||
logger.warning(f"Failed to extract audio: {str(e)}")
|
||||
audio_data_uri = None
|
||||
|
||||
# 2. get keyframes
|
||||
cap = cv2.VideoCapture(document_path)
|
||||
|
||||
if not cap.isOpened():
|
||||
raise ValueError(f"Could not open video file: {document_path}")
|
||||
|
||||
# get video message
|
||||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||||
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
duration = frame_count / fps if fps > 0 else 0
|
||||
|
||||
# keyframes policy- per duration/10s,max 10
|
||||
keyframes_count = min(10, int(frame_count))
|
||||
frames_interval = max(1, int(frame_count / keyframes_count))
|
||||
|
||||
self.keyframes = []
|
||||
frame_index = 0
|
||||
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
break
|
||||
|
||||
# per frames_interval save
|
||||
if frame_index % frames_interval == 0:
|
||||
# save JPEG -> base64
|
||||
_, buffer = cv2.imencode(".jpg", frame)
|
||||
img_base64 = base64.b64encode(buffer).decode('utf-8')
|
||||
time_position = frame_index / fps if fps > 0 else 0
|
||||
|
||||
self.keyframes.append(f"data:image/jpeg;base64,{img_base64}")
|
||||
|
||||
if len(self.keyframes) >= keyframes_count:
|
||||
break
|
||||
|
||||
frame_index += 1
|
||||
|
||||
cap.release()
|
||||
|
||||
self.content = audio_data_uri
|
||||
logger.info(f"Successfully processed video file: {document_path}")
|
||||
logger.info(f"Extracted {len(self.keyframes)} keyframes and audio track")
|
||||
# clean tmp files
|
||||
try:
|
||||
os.remove(audio_path)
|
||||
os.rmdir(temp_dir)
|
||||
except Exception as cleanup_error:
|
||||
logger.warning(f"Error cleaning up temp files: {str(cleanup_error)}")
|
||||
|
||||
except Exception as video_error:
|
||||
error = str(video_error)
|
||||
logger.error(f"Error processing video file: {error}")
|
||||
|
||||
if any(document_path.endswith(ext) for ext in ["pptx"]):
|
||||
try:
|
||||
# Initialize content list and empty keyframes
|
||||
self.content = []
|
||||
self.keyframes = []
|
||||
|
||||
# Check if file exists
|
||||
if not os.path.exists(document_path):
|
||||
error = f"File does not exist: {document_path}"
|
||||
return self.content, self.keyframes, error
|
||||
|
||||
# Check if file is readable
|
||||
if not os.access(document_path, os.R_OK):
|
||||
error = f"File is not readable: {document_path}"
|
||||
return self.content, self.keyframes, error
|
||||
|
||||
# Check file size
|
||||
try:
|
||||
file_size = os.path.getsize(document_path)
|
||||
|
||||
if file_size == 0:
|
||||
error = "File is empty"
|
||||
return self.content, self.keyframes, error
|
||||
except Exception as size_error:
|
||||
logger.warning(f"Cannot get file size: {str(size_error)}")
|
||||
|
||||
try:
|
||||
# Import required libraries
|
||||
from pptx import Presentation
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
import io
|
||||
except ImportError as import_error:
|
||||
error = f"Missing required libraries: {str(import_error)}. Please install: pip install python-pptx Pillow"
|
||||
return self.content, self.keyframes, error
|
||||
|
||||
# Create temporary directory for images
|
||||
try:
|
||||
temp_dir = tempfile.mkdtemp()
|
||||
except Exception as temp_dir_error:
|
||||
error = f"Failed to create temporary directory: {str(temp_dir_error)}"
|
||||
return self.content, self.keyframes, error
|
||||
|
||||
# Open presentation
|
||||
try:
|
||||
presentation = Presentation(document_path)
|
||||
# Get total slides count
|
||||
total_slides = len(presentation.slides)
|
||||
|
||||
if total_slides == 0:
|
||||
error = "PPTX file does not contain any slides"
|
||||
return self.content, self.keyframes, error
|
||||
|
||||
# Process each slide
|
||||
for i, slide in enumerate(presentation.slides):
|
||||
|
||||
# Generate temporary file path for current slide
|
||||
img_path = os.path.join(temp_dir, f"slide_{i + 1}.jpg")
|
||||
|
||||
# Get slide dimensions
|
||||
try:
|
||||
slide_width = presentation.slide_width
|
||||
slide_height = presentation.slide_height
|
||||
|
||||
# PPTX dimensions are in EMU (English Metric Unit)
|
||||
# 1 inch = 914400 EMU, 1 cm = 360000 EMU
|
||||
# Convert to pixels (assuming 96 DPI)
|
||||
slide_width_px = int(slide_width / 914400 * 96 * 10)
|
||||
slide_height_px = int(slide_height / 914400 * 96 * 10)
|
||||
|
||||
# Ensure dimensions are reasonable positive integers
|
||||
slide_width_px = max(1, min(slide_width_px, 4000)) # Limit max width to 4000px
|
||||
slide_height_px = max(1, min(slide_height_px, 3000)) # Limit max height to 3000px
|
||||
|
||||
except Exception as size_error:
|
||||
# Use default dimensions
|
||||
slide_width_px = 960 # Default width 960px
|
||||
slide_height_px = 720 # Default height 720px
|
||||
|
||||
# Create blank image
|
||||
try:
|
||||
# Log operation start
|
||||
|
||||
# Create blank image
|
||||
try:
|
||||
slide_img = Image.new('RGB', (slide_width_px, slide_height_px), 'white')
|
||||
draw = ImageDraw.Draw(slide_img)
|
||||
except Exception as img_create_error:
|
||||
logger.error(
|
||||
f"Slide {i + 1} blank image creation failed: {str(img_create_error) or 'Unknown error'}")
|
||||
raise
|
||||
|
||||
# Draw slide number
|
||||
try:
|
||||
font = ImageFont.load_default()
|
||||
draw.text((20, 20), f"Slide {i + 1}/{total_slides}", fill="black", font=font)
|
||||
except Exception as font_error:
|
||||
logger.warning(f"Failed to draw slide number: {str(font_error) or 'Unknown error'}")
|
||||
|
||||
# Record shape count
|
||||
try:
|
||||
shape_count = len(slide.shapes)
|
||||
except Exception as shape_count_error:
|
||||
logger.warning(
|
||||
f"Failed to get slide {i + 1} shape count: {str(shape_count_error) or 'Unknown error'}")
|
||||
shape_count = 0
|
||||
|
||||
# Try to render shapes on image
|
||||
shape_success_count = 0
|
||||
shape_fail_count = 0
|
||||
|
||||
try:
|
||||
for j, shape in enumerate(slide.shapes):
|
||||
try:
|
||||
shape_type = type(shape).__name__
|
||||
|
||||
# Process images
|
||||
if hasattr(shape, 'image') and shape.image:
|
||||
try:
|
||||
# Extract image from shape
|
||||
image_stream = io.BytesIO(shape.image.blob)
|
||||
img = Image.open(image_stream)
|
||||
|
||||
# Calculate position
|
||||
left = shape.left
|
||||
top = shape.top
|
||||
|
||||
# Paste image onto slide
|
||||
slide_img.paste(img, (left, top))
|
||||
shape_success_count += 1
|
||||
except Exception as img_error:
|
||||
logger.warning(
|
||||
f"Failed to process image {j + 1} in slide {i + 1}: {str(img_error) or 'Unknown error'}")
|
||||
if not str(img_error):
|
||||
import traceback
|
||||
logger.warning(
|
||||
f"Image processing stack: {traceback.format_exc()}")
|
||||
shape_fail_count += 1
|
||||
|
||||
# Process text
|
||||
elif hasattr(shape, 'text') and shape.text:
|
||||
try:
|
||||
text = shape.text[:30] + "..." if len(
|
||||
shape.text) > 30 else shape.text
|
||||
# Simple text rendering
|
||||
text_left = shape.left
|
||||
text_top = shape.top
|
||||
draw.text((text_left, text_top), shape.text, fill="black",
|
||||
font=font)
|
||||
shape_success_count += 1
|
||||
except Exception as text_error:
|
||||
logger.warning(
|
||||
f"Failed to process text {j + 1} in slide {i + 1}: {str(text_error) or 'Unknown error'}")
|
||||
if not str(text_error):
|
||||
import traceback
|
||||
logger.warning(
|
||||
f"Text processing stack: {traceback.format_exc()}")
|
||||
shape_fail_count += 1
|
||||
else:
|
||||
logger.info(
|
||||
f"Shape {j + 1} in slide {i + 1} is neither image nor text, skipping")
|
||||
except Exception as shape_error:
|
||||
if not str(shape_error):
|
||||
import traceback
|
||||
logger.warning(f"Shape processing stack: {traceback.format_exc()}")
|
||||
shape_fail_count += 1
|
||||
except Exception as shapes_iteration_error:
|
||||
logger.error(
|
||||
f"Failed while iterating through shapes in slide {i + 1}: {str(shapes_iteration_error) or 'Unknown error'}")
|
||||
if not str(shapes_iteration_error):
|
||||
import traceback
|
||||
logger.error(f"Shape iteration stack: {traceback.format_exc()}")
|
||||
|
||||
# Save slide image
|
||||
try:
|
||||
slide_img.save(img_path, 'JPEG')
|
||||
|
||||
# Check if image was saved successfully
|
||||
if not os.path.exists(img_path):
|
||||
raise ValueError(f"Saved image file does not exist: {img_path}")
|
||||
|
||||
file_size = os.path.getsize(img_path)
|
||||
if file_size == 0:
|
||||
raise ValueError(
|
||||
f"Saved image file is empty: {img_path}, size: {file_size} bytes")
|
||||
|
||||
# Convert to base64
|
||||
try:
|
||||
base64_image = encode_image_from_file(img_path)
|
||||
self.content.append(f"data:image/jpeg;base64,{base64_image}")
|
||||
except Exception as base64_error:
|
||||
error_msg = str(base64_error) or "Unknown base64 conversion error"
|
||||
if not str(base64_error):
|
||||
import traceback
|
||||
logger.error(f"Base64 conversion stack: {traceback.format_exc()}")
|
||||
raise ValueError(f"Base64 conversion error: {error_msg}")
|
||||
|
||||
except Exception as save_error:
|
||||
error_msg = str(save_error) or "Unknown save error"
|
||||
logger.error(f"Failed to save slide {i + 1} as image: {error_msg}")
|
||||
if not str(save_error):
|
||||
import traceback
|
||||
logger.error(f"Image save stack: {traceback.format_exc()}")
|
||||
raise ValueError(f"Image save error: {error_msg}")
|
||||
|
||||
except Exception as slide_render_error:
|
||||
error_msg = str(slide_render_error) or "Unknown rendering error"
|
||||
logger.error(f"Failed to render slide {i + 1}: {error_msg}")
|
||||
if not str(slide_render_error):
|
||||
import traceback
|
||||
logger.error(f"Slide rendering stack: {traceback.format_exc()}")
|
||||
# Continue processing next slide, don't interrupt the entire process
|
||||
continue
|
||||
|
||||
except Exception as pptx_error:
|
||||
error = f"Failed to process PPTX file: {str(pptx_error)}"
|
||||
import traceback
|
||||
|
||||
# Clean up temporary files
|
||||
try:
|
||||
for file in os.listdir(temp_dir):
|
||||
try:
|
||||
file_path = os.path.join(temp_dir, file)
|
||||
os.remove(file_path)
|
||||
except Exception as file_error:
|
||||
logger.warning(f"Failed to delete temporary file: {str(file_error)}")
|
||||
os.rmdir(temp_dir)
|
||||
except Exception as cleanup_error:
|
||||
logger.warning(f"Failed to clean up temporary files: {str(cleanup_error)}")
|
||||
|
||||
if len(self.content) > 0:
|
||||
logger.info(f"Extracted {len(self.content)} slides")
|
||||
else:
|
||||
error = error or "Could not extract any slides from PPTX file"
|
||||
logger.error(error)
|
||||
|
||||
except Exception as outer_error:
|
||||
error = f"Error occurred during PPTX file processing: {str(outer_error)}"
|
||||
import traceback
|
||||
|
||||
return self.content, self.keyframes, error
|
||||
|
||||
finally:
|
||||
pass
|
||||
|
||||
return self.content, self.keyframes, error
|
||||
+4
@@ -0,0 +1,4 @@
|
||||
custom_executor: False
|
||||
enable_recording: False
|
||||
working_dir:
|
||||
max_retry: 3
|
||||
@@ -0,0 +1,32 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
import base64
|
||||
from io import BytesIO
|
||||
|
||||
|
||||
def encode_image_from_url(image_url):
|
||||
from aworld.utils.import_package import import_package
|
||||
import_package("requests")
|
||||
|
||||
import requests
|
||||
from PIL import Image
|
||||
|
||||
response = requests.get(image_url)
|
||||
image = Image.open(BytesIO(response.content))
|
||||
|
||||
max_size = 1024
|
||||
if max(image.size) > max_size:
|
||||
ratio = max_size / max(image.size)
|
||||
new_size = (int(image.size[0] * ratio), int(image.size[1] * ratio))
|
||||
image = image.resize(new_size, Image.LANCZOS)
|
||||
|
||||
buffered = BytesIO()
|
||||
image_format = image.format if image.format else 'JPEG'
|
||||
image.save(buffered, format=image_format)
|
||||
img_str = base64.b64encode(buffered.getvalue()).decode()
|
||||
return img_str
|
||||
|
||||
|
||||
def encode_image_from_file(image_path):
|
||||
with open(image_path, "rb") as image_file:
|
||||
return base64.b64encode(image_file.read()).decode()
|
||||
@@ -0,0 +1,2 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
@@ -0,0 +1,12 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
from examples.common.tools.tool_action import GymAction
|
||||
from aworld.core.tool.action_factory import ActionFactory
|
||||
from aworld.core.tool.action import ExecutableAction
|
||||
|
||||
|
||||
@ActionFactory.register(name=GymAction.PLAY.value.name,
|
||||
desc=GymAction.PLAY.value.desc,
|
||||
tool_name="openai_gym")
|
||||
class Play(ExecutableAction):
|
||||
""""""
|
||||
@@ -0,0 +1,158 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
from pathlib import Path
|
||||
from typing import Dict, Any, Tuple, SupportsFloat, Union, List
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from aworld.config import ConfigDict
|
||||
from examples.common.tools.tool_action import GymAction
|
||||
from aworld.core.common import ActionModel, Observation, ActionResult
|
||||
from aworld.core.tool.base import AsyncTool, ToolFactory
|
||||
from aworld.utils.import_package import import_packages
|
||||
from aworld.tools.utils import build_observation
|
||||
|
||||
|
||||
class ActionType(object):
|
||||
DISCRETE = 'discrete'
|
||||
CONTINUOUS = 'continuous'
|
||||
|
||||
|
||||
@ToolFactory.register(name="openai_gym",
|
||||
desc="gym classic control game",
|
||||
asyn=True,
|
||||
supported_action=GymAction,
|
||||
conf_file_name=f'openai_gym_tool.yaml',
|
||||
dir=f"{Path(__file__).parent.absolute()}")
|
||||
class OpenAIGym(AsyncTool):
|
||||
def __init__(self, conf: Union[Dict[str, Any], ConfigDict, BaseModel], **kwargs) -> None:
|
||||
"""Gym environment constructor.
|
||||
|
||||
Args:
|
||||
env_id: gym environment full name
|
||||
wrappers: gym environment wrapper list
|
||||
"""
|
||||
import_packages(['pygame', 'gymnasium'])
|
||||
super(OpenAIGym, self).__init__(conf, **kwargs)
|
||||
self.env_id = self.conf.get("env_id")
|
||||
self._render = self.conf.get('render', True)
|
||||
if self._render:
|
||||
kwargs['render_mode'] = self.conf.get('render_mode', True)
|
||||
kwargs.pop('name', None)
|
||||
self.env = self._gym_env_wrappers(self.env_id, self.conf.get("wrappers", []), **kwargs)
|
||||
self.action_space = self.env.action_space
|
||||
|
||||
async def do_step(self, actions: List[ActionModel], **kwargs) -> Tuple[
|
||||
Observation, SupportsFloat, bool, bool, Dict[str, Any]]:
|
||||
if self._render:
|
||||
await self.render()
|
||||
action = actions[0].params['result']
|
||||
action = OpenAIGym.transform_action(action=action)
|
||||
state, reward, terminal, truncate, info = self.env.step(action)
|
||||
info.update(kwargs)
|
||||
self._finished = terminal
|
||||
|
||||
action_results = []
|
||||
for _ in actions:
|
||||
action_results.append(ActionResult(content=OpenAIGym.transform_state(state=state), success=True))
|
||||
return (build_observation(observer=self.name(),
|
||||
action_result=action_results,
|
||||
ability=GymAction.PLAY.value.name,
|
||||
content=OpenAIGym.transform_state(state=state),
|
||||
env_id=self.env_id,
|
||||
done=terminal,
|
||||
**kwargs),
|
||||
reward,
|
||||
terminal,
|
||||
truncate,
|
||||
info)
|
||||
|
||||
async def render(self):
|
||||
return self.env.render()
|
||||
|
||||
async def close(self):
|
||||
if self.env:
|
||||
self.env.close()
|
||||
self.env = None
|
||||
|
||||
async def reset(self, *, seed: int | None = None, options: Dict[str, str] | None = None) -> Tuple[
|
||||
Any, Dict[str, Any]]:
|
||||
state = self.env.reset()
|
||||
return build_observation(observer=self.name(),
|
||||
ability=GymAction.PLAY.value.name,
|
||||
content=OpenAIGym.transform_state(state=state),
|
||||
env_id=self.env_id,
|
||||
done=False), {}
|
||||
|
||||
def _action_dim(self):
|
||||
from gymnasium import spaces
|
||||
|
||||
if isinstance(self.env.action_space, spaces.Discrete):
|
||||
self.action_type = ActionType.DISCRETE
|
||||
return self.env.action_space.n
|
||||
elif isinstance(self.env.action_space, spaces.Box):
|
||||
self.action_type = ActionType.CONTINUOUS
|
||||
return self.env.action_space.shape[0]
|
||||
else:
|
||||
raise Exception('unsupported env.action_space: {}'.format(self.env.action_space))
|
||||
|
||||
def _state_dim(self):
|
||||
if len(self.env.observation_space.shape) == 1:
|
||||
return self.env.observation_space.shape[0]
|
||||
else:
|
||||
raise Exception('unsupported observation_space.shape: {}'.format(self.env.observation_space))
|
||||
|
||||
def _gym_env_wrappers(self, env_id, wrappers: list = [], **kwargs):
|
||||
import gymnasium
|
||||
|
||||
env = gymnasium.make(env_id, **kwargs)
|
||||
|
||||
if wrappers:
|
||||
for wrapper in wrappers:
|
||||
env = wrapper(env)
|
||||
|
||||
return env
|
||||
|
||||
@staticmethod
|
||||
def transform_state(state: Any):
|
||||
if isinstance(state, tuple):
|
||||
states = dict()
|
||||
for n, state in enumerate(state):
|
||||
state = OpenAIGym.transform_state(state=state)
|
||||
if isinstance(state, dict):
|
||||
for name, state in state.items():
|
||||
states['gym{}-{}'.format(n, name)] = state
|
||||
else:
|
||||
states['gym{}'.format(n)] = state
|
||||
return states
|
||||
elif isinstance(state, dict):
|
||||
states = dict()
|
||||
for state_name, state in state.items():
|
||||
state = OpenAIGym.transform_state(state=state)
|
||||
if isinstance(state, dict):
|
||||
for name, state in state.items():
|
||||
states['{}-{}'.format(state_name, name)] = state
|
||||
else:
|
||||
states['{}'.format(state_name)] = state
|
||||
return states
|
||||
else:
|
||||
return state
|
||||
|
||||
@staticmethod
|
||||
def transform_action(action: Any):
|
||||
if not isinstance(action, dict):
|
||||
return action
|
||||
else:
|
||||
actions = dict()
|
||||
for name, action in action.items():
|
||||
if '-' in name:
|
||||
name, inner_name = name.split('-', 1)
|
||||
if name not in actions:
|
||||
actions[name] = dict()
|
||||
actions[name][inner_name] = action
|
||||
else:
|
||||
actions[name] = action
|
||||
for name, action in actions.items():
|
||||
if isinstance(action, dict):
|
||||
actions[name] = OpenAIGym.transform_action(action=action)
|
||||
return actions
|
||||
@@ -0,0 +1,154 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
from pathlib import Path
|
||||
from typing import Dict, Any, Tuple, SupportsFloat, List, Union
|
||||
|
||||
from aworld.config import ConfigDict, ToolConfig
|
||||
from examples.common.tools.tool_action import GymAction
|
||||
from aworld.core.common import Observation, ActionModel, ActionResult
|
||||
from aworld.core.tool.base import Tool, ToolFactory
|
||||
from aworld.utils.import_package import import_packages
|
||||
from aworld.tools.utils import build_observation
|
||||
|
||||
|
||||
class ActionType(object):
|
||||
DISCRETE = 'discrete'
|
||||
CONTINUOUS = 'continuous'
|
||||
|
||||
|
||||
@ToolFactory.register(name="openai_gym",
|
||||
desc="gym classic control game",
|
||||
supported_action=GymAction,
|
||||
conf_file_name=f'openai_gym_tool.yaml',
|
||||
dir=f"{Path(__file__).parent.absolute()}")
|
||||
class OpenAIGym(Tool):
|
||||
def __init__(self, conf: Union[Dict[str, Any], ConfigDict, ToolConfig], **kwargs) -> None:
|
||||
"""Gym environment constructor.
|
||||
|
||||
Args:
|
||||
env_id: gym environment full name
|
||||
wrappers: gym environment wrapper list
|
||||
"""
|
||||
import_packages(['pygame', 'gymnasium'])
|
||||
super(OpenAIGym, self).__init__(conf, **kwargs)
|
||||
self.env_id = self.conf.get("env_id")
|
||||
self._render = self.conf.get('render', True)
|
||||
if self._render:
|
||||
kwargs['render_mode'] = self.conf.get('render_mode', 'human')
|
||||
kwargs.pop('name', None)
|
||||
self.env = self._gym_env_wrappers(self.env_id, self.conf.get("wrappers", []), **kwargs)
|
||||
self.action_space = self.env.action_space
|
||||
|
||||
def do_step(self, actions: List[ActionModel], **kwargs) -> Tuple[
|
||||
Observation, SupportsFloat, bool, bool, Dict[str, Any]]:
|
||||
if self._render:
|
||||
self.render()
|
||||
action = actions[0].params['result']
|
||||
action = OpenAIGym.transform_action(action=action)
|
||||
state, reward, terminal, truncate, info = self.env.step(action)
|
||||
info.update(kwargs)
|
||||
self._finished = terminal
|
||||
|
||||
action_results = []
|
||||
for _ in actions:
|
||||
action_results.append(ActionResult(content=OpenAIGym.transform_state(state=state), success=True))
|
||||
return (build_observation(observer=self.name(),
|
||||
action_result=action_results,
|
||||
ability=GymAction.PLAY.value.name,
|
||||
content=OpenAIGym.transform_state(state=state),
|
||||
env_id=self.env_id,
|
||||
done=terminal,
|
||||
**kwargs),
|
||||
reward,
|
||||
terminal,
|
||||
truncate,
|
||||
info)
|
||||
|
||||
def render(self):
|
||||
return self.env.render()
|
||||
|
||||
def close(self):
|
||||
if self.env:
|
||||
self.env.close()
|
||||
self.env = None
|
||||
|
||||
def reset(self, *, seed: int | None = None, options: Dict[str, str] | None = None) -> Tuple[Any, Dict[str, Any]]:
|
||||
state = self.env.reset()
|
||||
return build_observation(observer=self.name(),
|
||||
ability=GymAction.PLAY.value.name,
|
||||
content=OpenAIGym.transform_state(state=state),
|
||||
env_id=self.env_id,
|
||||
done=False), {}
|
||||
|
||||
def _action_dim(self):
|
||||
from gymnasium import spaces
|
||||
|
||||
if isinstance(self.env.action_space, spaces.Discrete):
|
||||
self.action_type = ActionType.DISCRETE
|
||||
return self.env.action_space.n
|
||||
elif isinstance(self.env.action_space, spaces.Box):
|
||||
self.action_type = ActionType.CONTINUOUS
|
||||
return self.env.action_space.shape[0]
|
||||
else:
|
||||
raise Exception('unsupported env.action_space: {}'.format(self.env.action_space))
|
||||
|
||||
def _state_dim(self):
|
||||
if len(self.env.observation_space.shape) == 1:
|
||||
return self.env.observation_space.shape[0]
|
||||
else:
|
||||
raise Exception('unsupported observation_space.shape: {}'.format(self.env.observation_space))
|
||||
|
||||
def _gym_env_wrappers(self, env_id, wrappers: list = [], **kwargs):
|
||||
import gymnasium
|
||||
|
||||
env = gymnasium.make(env_id, **kwargs)
|
||||
|
||||
if wrappers:
|
||||
for wrapper in wrappers:
|
||||
env = wrapper(env)
|
||||
|
||||
return env
|
||||
|
||||
@staticmethod
|
||||
def transform_state(state: Any):
|
||||
if isinstance(state, tuple):
|
||||
states = dict()
|
||||
for n, state in enumerate(state):
|
||||
state = OpenAIGym.transform_state(state=state)
|
||||
if isinstance(state, dict):
|
||||
for name, state in state.items():
|
||||
states['gym{}-{}'.format(n, name)] = state
|
||||
else:
|
||||
states['gym{}'.format(n)] = state
|
||||
return states
|
||||
elif isinstance(state, dict):
|
||||
states = dict()
|
||||
for state_name, state in state.items():
|
||||
state = OpenAIGym.transform_state(state=state)
|
||||
if isinstance(state, dict):
|
||||
for name, state in state.items():
|
||||
states['{}-{}'.format(state_name, name)] = state
|
||||
else:
|
||||
states['{}'.format(state_name)] = state
|
||||
return states
|
||||
else:
|
||||
return state
|
||||
|
||||
@staticmethod
|
||||
def transform_action(action: Any):
|
||||
if not isinstance(action, dict):
|
||||
return action
|
||||
else:
|
||||
actions = dict()
|
||||
for name, action in action.items():
|
||||
if '-' in name:
|
||||
name, inner_name = name.split('-', 1)
|
||||
if name not in actions:
|
||||
actions[name] = dict()
|
||||
actions[name][inner_name] = action
|
||||
else:
|
||||
actions[name] = action
|
||||
for name, action in actions.items():
|
||||
if isinstance(action, dict):
|
||||
actions[name] = OpenAIGym.transform_action(action=action)
|
||||
return actions
|
||||
@@ -0,0 +1,3 @@
|
||||
env_id: "CartPole-v1"
|
||||
render_mode: "human"
|
||||
render: True
|
||||
@@ -0,0 +1,2 @@
|
||||
gymnasium~=1.1.0
|
||||
pygame~=2.6.1
|
||||
@@ -0,0 +1,2 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
@@ -0,0 +1,68 @@
|
||||
# coding: utf-8
|
||||
|
||||
import os
|
||||
import re
|
||||
|
||||
from typing import Tuple, Any
|
||||
|
||||
from aworld.core.tool.action_factory import ActionFactory
|
||||
from aworld.core.common import ActionModel, ActionResult
|
||||
from aworld.logs.util import logger
|
||||
from aworld.core.tool.action import ExecutableAction
|
||||
from aworld.models.llm import get_llm_model, call_llm_model
|
||||
|
||||
|
||||
@ActionFactory.register(name="write_html",
|
||||
desc="a tool use for write html.",
|
||||
tool_name="html")
|
||||
class WriteHTML(ExecutableAction):
|
||||
def act(self, action: ActionModel, **kwargs) -> Tuple[ActionResult, Any]:
|
||||
logger.info("start write html!")
|
||||
goal = action.params.get("goal")
|
||||
information = action.params.get("information")
|
||||
|
||||
llm_conf = kwargs.get("llm_config")
|
||||
llm = get_llm_model(llm_conf)
|
||||
sys_prompt = "you are a helpful html writer."
|
||||
prompt = """Your task is to create a detailed and visually appealing HTML document based on the specified theme.
|
||||
The document must meet the following requirements, and you should utilize the provided reference materials to ensure accuracy and aesthetic quality.
|
||||
1) HTML Document Requirements
|
||||
Design and write the HTML document according to the following specifications:
|
||||
Theme : {goal}
|
||||
Related Info: {information}
|
||||
Structural Requirements :
|
||||
Use semantic HTML tags (e.g., <header>, <main>, <footer>, <section>) to create a clear and organized structure.
|
||||
Ensure the document includes a header, navigation bar, main content area, and footer.
|
||||
If applicable, add additional sections such as a sidebar, or call-to-action buttons.
|
||||
Styling Requirements :
|
||||
Implement a visually appealing design using CSS, including color schemes, font choices, spacing adjustments, etc.
|
||||
Ensure the page has a responsive layout that works well on different devices (use media queries or frameworks like Bootstrap).
|
||||
Add animations or interactive features (e.g., hover effects on buttons, scroll-triggered animations) to enhance user experience.
|
||||
please give me html code directly, no need other words
|
||||
"""
|
||||
|
||||
messages = [{'role': 'system', 'content': sys_prompt},
|
||||
{'role': 'user', 'content': prompt.format(goal=goal, information=information)}]
|
||||
|
||||
output = call_llm_model(llm,
|
||||
messages=messages,
|
||||
model=llm_conf.llm_model_name,
|
||||
temperature=llm_conf.llm_temperature)
|
||||
|
||||
content = output.content
|
||||
html_pattern = re.compile(r'<html.*?>.*?</html>', re.DOTALL)
|
||||
matches = html_pattern.findall(content)
|
||||
|
||||
title_pattern = re.compile(r'<title.*?>.*?</title>', re.DOTALL)
|
||||
filename = (title_pattern.findall(content)[0]
|
||||
.replace("<title>", "")
|
||||
.replace("</title>", "")
|
||||
.replace(" ", "_") + ".html")
|
||||
|
||||
with open(filename, "a", encoding='utf-8') as f:
|
||||
f.write(matches[0])
|
||||
|
||||
abs_file_path = os.path.abspath(filename)
|
||||
msg = f'Successfully wrote html to {abs_file_path}'
|
||||
|
||||
return ActionResult(content=msg, keep=True, is_done=True), None
|
||||
@@ -0,0 +1,10 @@
|
||||
# coding: utf-8
|
||||
|
||||
from aworld.tools.template_tool import TemplateTool
|
||||
from examples.common.tools.tool_action import WriteAction
|
||||
from aworld.core.tool.base import ToolFactory
|
||||
|
||||
|
||||
@ToolFactory.register(name="html", desc="html tool", supported_action=WriteAction)
|
||||
class HtmlTool(TemplateTool):
|
||||
"""Html tool"""
|
||||
@@ -0,0 +1,12 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
from examples.common.tools.tool_action import PythonToolAction
|
||||
from aworld.core.tool.action_factory import ActionFactory
|
||||
from aworld.core.tool.action import ExecutableAction
|
||||
|
||||
|
||||
@ActionFactory.register(name=PythonToolAction.EXECUTE.value.name,
|
||||
desc=PythonToolAction.EXECUTE.value.desc,
|
||||
tool_name="python_execute")
|
||||
class ExecuteAction(ExecutableAction):
|
||||
"""Only one action, define it, implemented can be omitted."""
|
||||
+5
@@ -0,0 +1,5 @@
|
||||
custom_executor: False
|
||||
enable_recording: False
|
||||
working_dir:
|
||||
max_retry: 3
|
||||
|
||||
@@ -0,0 +1,257 @@
|
||||
import sys
|
||||
import ast
|
||||
import re
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Tuple, List
|
||||
from io import StringIO
|
||||
|
||||
from aworld.logs.util import logger
|
||||
from aworld.config.conf import ToolConfig
|
||||
from examples.common.tools.tool_action import PythonToolAction
|
||||
from aworld.core.common import ActionModel, Observation, ActionResult
|
||||
from aworld.core.tool.base import Tool, AgentInput, ToolFactory
|
||||
from aworld.utils import import_package
|
||||
from aworld.tools.utils import build_observation
|
||||
|
||||
|
||||
@ToolFactory.register(name="python_execute",
|
||||
desc="python interpreter tool",
|
||||
supported_action=PythonToolAction,
|
||||
conf_file_name=f'python_execute_tool.yaml',
|
||||
dir=f"{Path(__file__).parent.absolute()}")
|
||||
class PythonTool(Tool):
|
||||
|
||||
def __init__(self,
|
||||
conf: ToolConfig,
|
||||
**kwargs) -> None:
|
||||
"""
|
||||
Initialize the PythonExecutor
|
||||
Args:
|
||||
conf: tool config
|
||||
**kwargs: -
|
||||
Return:
|
||||
None
|
||||
"""
|
||||
super(PythonTool, self).__init__(conf, **kwargs)
|
||||
self.type = "function"
|
||||
self.local_namespace = {}
|
||||
self.global_namespace = {}
|
||||
self.original_stdout = sys.stdout
|
||||
self.output_buffer = StringIO()
|
||||
self.installed_packages = set()
|
||||
import_package('langchain_experimental')
|
||||
from langchain_experimental.utilities.python import PythonREPL
|
||||
self.python_repl = PythonREPL()
|
||||
|
||||
def extract_imports(self, code: str) -> set:
|
||||
"""
|
||||
Extract import statements
|
||||
Args:
|
||||
code: python code
|
||||
Returns:
|
||||
set: import statements
|
||||
"""
|
||||
imports = set()
|
||||
|
||||
try:
|
||||
tree = ast.parse(code)
|
||||
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.Import):
|
||||
# deal import xxx or import xxx as yyy
|
||||
for name in node.names:
|
||||
package_name = name.name.split('.')[0]
|
||||
imports.add(package_name)
|
||||
|
||||
elif isinstance(node, ast.ImportFrom):
|
||||
# deal from xxx import yyy or from xxx.yyy import zzz
|
||||
if node.module:
|
||||
package_name = node.module.split('.')[0]
|
||||
imports.add(package_name)
|
||||
|
||||
except SyntaxError:
|
||||
import_pattern = r'^import\s+([\w\s,]+)|from\s+(\w+)'
|
||||
for line in code.split('\n'):
|
||||
line = line.strip()
|
||||
match = re.match(import_pattern, line)
|
||||
if match:
|
||||
if match.group(1):
|
||||
|
||||
packages = [p.strip() for p in match.group(1).split(',')]
|
||||
for package in packages:
|
||||
if package:
|
||||
package_name = package.split()[0]
|
||||
imports.add(package_name)
|
||||
elif match.group(2):
|
||||
imports.add(match.group(2))
|
||||
|
||||
return imports
|
||||
|
||||
def install_dependencies(self,
|
||||
packages: set) -> None:
|
||||
"""
|
||||
Install dependency packages
|
||||
Args:
|
||||
packages: python third packages
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
for package in packages:
|
||||
try:
|
||||
__import__(package)
|
||||
except ImportError:
|
||||
try:
|
||||
subprocess.check_call([sys.executable, "-m", "pip", "install", package])
|
||||
self.installed_packages.add(package)
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.warning(f"Failed to install {package}: {str(e)}")
|
||||
|
||||
def uninstall_dependencies(self) -> None:
|
||||
"""
|
||||
Uninstall dependency packages
|
||||
Args:
|
||||
-
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
try:
|
||||
for package in self.installed_packages:
|
||||
try:
|
||||
subprocess.check_call([sys.executable, "-m", "pip", "uninstall", "-y", package])
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.warning(f"Failed to uninstall {package}: {str(e)}")
|
||||
self.installed_packages.clear()
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to uninstall dependencies: {repr(e)}")
|
||||
|
||||
def reset(self,
|
||||
*,
|
||||
seed: int | None = None,
|
||||
options: Dict[str, str] | None = None) -> Tuple[AgentInput, dict[str, Any]]:
|
||||
"""
|
||||
Reset the executor
|
||||
Args:
|
||||
seed: -
|
||||
options: -
|
||||
Returns:
|
||||
AgentInput, dict[str, Any]: -
|
||||
"""
|
||||
self.close()
|
||||
self.local_namespace = {}
|
||||
self.global_namespace = {}
|
||||
self._finished = False
|
||||
self.installed_packages.clear()
|
||||
|
||||
return build_observation(observer=self.name(),
|
||||
ability=PythonToolAction.EXECUTE.value.name), {}
|
||||
|
||||
def close(self) -> None:
|
||||
"""
|
||||
Close the executor
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
try:
|
||||
self.uninstall_dependencies()
|
||||
sys.stdout = self.original_stdout
|
||||
self.output_buffer.close()
|
||||
self.local_namespace.clear()
|
||||
self.global_namespace.clear()
|
||||
except:
|
||||
pass
|
||||
finally:
|
||||
self._finished = True
|
||||
|
||||
def do_step(
|
||||
self,
|
||||
actions: List[ActionModel],
|
||||
**kwargs) -> Tuple[Observation, float, bool, bool, dict[str, Any]]:
|
||||
"""
|
||||
Step the executor
|
||||
Args:
|
||||
actions: actions
|
||||
**kwargs: -
|
||||
Returns:
|
||||
Observation, float, bool, bool, dict[str, Any]: -
|
||||
"""
|
||||
self.step_finished = False
|
||||
reward = 0
|
||||
fail_error = ""
|
||||
observation = build_observation(observer=self.name(),
|
||||
ability=PythonToolAction.EXECUTE.value.name)
|
||||
try:
|
||||
if not actions:
|
||||
return (observation, reward,
|
||||
kwargs.get("terminated",
|
||||
False), kwargs.get("truncated", False), {
|
||||
"exception": "actions is empty"
|
||||
})
|
||||
for action in actions:
|
||||
code = action.params.get("code", "")
|
||||
if not code:
|
||||
logger.warning(f"{action} no code to execute.")
|
||||
continue
|
||||
try:
|
||||
_, output, error = self.execute(code)
|
||||
observation.content = output
|
||||
except Exception as e:
|
||||
error = str(e)
|
||||
output = error
|
||||
|
||||
observation.action_result.append(
|
||||
ActionResult(is_done=True,
|
||||
success=False if error else True,
|
||||
content=f"{output}",
|
||||
error=f"{error}",
|
||||
keep=False))
|
||||
reward = 1
|
||||
except Exception as e:
|
||||
fail_error = str(e)
|
||||
finally:
|
||||
self._finished = True
|
||||
|
||||
info = {"exception": fail_error}
|
||||
info.update(kwargs)
|
||||
return (observation, reward, kwargs.get("terminated", False),
|
||||
kwargs.get("truncated", False), info)
|
||||
|
||||
def execute(self, code, timeout=300):
|
||||
"""
|
||||
Execute the code
|
||||
Args:
|
||||
code: python code
|
||||
timeout: timeout seconds
|
||||
Returns:
|
||||
result, output, error
|
||||
"""
|
||||
required_packages = self.extract_imports(code)
|
||||
self.install_dependencies(required_packages)
|
||||
self.python_repl.globals = self.global_namespace
|
||||
self.python_repl.locals = self.local_namespace
|
||||
error = None
|
||||
try:
|
||||
output = self.python_repl.run(code, timeout)
|
||||
except Exception as e:
|
||||
error = f'{repr(e)}'
|
||||
finally:
|
||||
self.uninstall_dependencies()
|
||||
return '', output, error
|
||||
|
||||
def get_execute_result(self):
|
||||
"""
|
||||
Get the execute result
|
||||
Returns:
|
||||
output, error
|
||||
"""
|
||||
output = None
|
||||
error = ''
|
||||
try:
|
||||
output = self.output_buffer.getvalue()
|
||||
self.output_buffer.truncate(0)
|
||||
self.output_buffer.seek(0)
|
||||
sys.stdout = self.original_stdout
|
||||
except Exception as e:
|
||||
error = f'{repr(e)}'
|
||||
logger.warning(f"Failed to get output, {repr(e)}")
|
||||
return output, error
|
||||
@@ -0,0 +1,324 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
|
||||
from aworld.core.common import ToolActionInfo, ParamInfo
|
||||
from aworld.core.tool.action import ToolAction
|
||||
|
||||
|
||||
class ChatAction(ToolAction):
|
||||
"""chat between agents """
|
||||
TASK_DONE = ToolActionInfo(name="TASK_DONE",
|
||||
desc="Complete task - with return text and if the task is finished (success=True) or not yet completly finished (success=False), because last step is reached")
|
||||
|
||||
|
||||
class SearchAction(ToolAction):
|
||||
"""Info search actions."""
|
||||
WIKI = ToolActionInfo(name="wiki",
|
||||
input_params={"query": ParamInfo(name="query",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="wiki search query input.")},
|
||||
desc="Search the entity in WikiPedia and return the summary of the required page, containing factual information about the given entity.")
|
||||
DUCK_GO = ToolActionInfo(name="duck_go",
|
||||
input_params={"query": ParamInfo(name="query",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="duckduckgo search query input"),
|
||||
"source": ParamInfo(name="source",
|
||||
type="str",
|
||||
required=False,
|
||||
desc="duckduckgo search query input.",
|
||||
default_value="text"),
|
||||
"max_results": ParamInfo(name="max_results",
|
||||
type="str",
|
||||
required=False,
|
||||
desc="duckduckgo search query input.",
|
||||
default_value=5)},
|
||||
desc="Use DuckDuckGo search engine to search information for the given query")
|
||||
GOOGLE = ToolActionInfo(name="google",
|
||||
input_params={"query": ParamInfo(name="query",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="google search query input."),
|
||||
"num_result_pages": ParamInfo(name="num_result_pages",
|
||||
type="str",
|
||||
required=False,
|
||||
desc="google search query input.",
|
||||
default_value=5)},
|
||||
desc="Use Google search engine to search information for the given query.")
|
||||
BAIDU = ToolActionInfo(name="baidu",
|
||||
input_params={"query": ParamInfo(name="query",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="baidu search query input."),
|
||||
"num_results": ParamInfo(name="num_results",
|
||||
type="str",
|
||||
required=False,
|
||||
desc="baidu search number of results.",
|
||||
default_value=5)},
|
||||
desc="Use Baidu search engine to search information for the given query.")
|
||||
|
||||
|
||||
class GymAction(ToolAction):
|
||||
PLAY = ToolActionInfo(name="play",
|
||||
input_params={"result": ParamInfo(name="result",
|
||||
type="object",
|
||||
required=True,
|
||||
desc="Agent decision result.")},
|
||||
desc="step")
|
||||
|
||||
|
||||
class BrowserAction(ToolAction):
|
||||
"""Definition of Browser tool supported action."""
|
||||
GO_TO_URL = ToolActionInfo(name="go_to_url",
|
||||
input_params={"url": ParamInfo(name="url",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="got to url in page on browser.")},
|
||||
desc="Navigate to URL in the current tab")
|
||||
INPUT_TEXT = ToolActionInfo(name="input_text",
|
||||
input_params={"text": ParamInfo(name="text",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="input text in page on browser"),
|
||||
"index": ParamInfo(name="index",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="index of click element in page on browser.")},
|
||||
desc="Input text into a input interactive element")
|
||||
SEARCH = ToolActionInfo(name="search",
|
||||
input_params={"url": ParamInfo(name="url",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="search url."),
|
||||
"query": ParamInfo(name="query",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="search query input in page on browser.")},
|
||||
desc="Search the query in search engine, Google, Baidu etc., in the current tab, the query should be a search query like humans search in search engine, concrete and not vague or super long. More the single most important items. ")
|
||||
SEARCH_GOOGLE = ToolActionInfo(name="search_google",
|
||||
input_params={"url": ParamInfo(name="url",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="search url."),
|
||||
"query": ParamInfo(name="query",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="search query input in google.")},
|
||||
desc="Search the query in Google in the current tab, the query should be a search query like humans search in Google, concrete and not vague or super long. More the single most important items. ")
|
||||
GO_BACK = ToolActionInfo(name="go_back",
|
||||
desc="Go back")
|
||||
SCROLL_DOWN = ToolActionInfo(name="scroll_down",
|
||||
input_params={"amount": ParamInfo(name="amount",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="pixel amount.")},
|
||||
desc="Scroll down the page by pixel amount - if no amount is specified, scroll down one page")
|
||||
SCROLL_UP = ToolActionInfo(name="scroll_up",
|
||||
input_params={"amount": ParamInfo(name="amount",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="Pixel amount.")},
|
||||
desc="Scroll up the page by pixel amount - if no amount is specified, scroll up one page")
|
||||
CLICK_ELEMENT = ToolActionInfo(name="click_element",
|
||||
input_params={"index": ParamInfo(name="index",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="Index of click element in page on browser.")},
|
||||
desc="Click element")
|
||||
NEW_TAB = ToolActionInfo(name="new_tab",
|
||||
input_params={"url": ParamInfo(name="url",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="Open url in new tab on browser.")},
|
||||
desc="Open url in new tab")
|
||||
SWITCH_TAB = ToolActionInfo(name="switch_tab",
|
||||
input_params={"page_id": ParamInfo(name="page_id",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="Switch tab by page id on browser.")},
|
||||
desc="Switch tab")
|
||||
WAIT = ToolActionInfo(name="wait",
|
||||
input_params={"seconds": ParamInfo(name="seconds",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="Wait some seconds.")},
|
||||
desc="Open url in new tab")
|
||||
EXTRACT_CONTENT = ToolActionInfo(name="extract_content",
|
||||
input_params={"goal": ParamInfo(name="goal",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="The goal in page content.")},
|
||||
desc="Extract page content to retrieve specific information from the page, e.g. all company names, a specifc description, all information about, links with companies in structured format or simply links")
|
||||
SEND_KEYS = ToolActionInfo(name="send_keys",
|
||||
input_params={"keys": ParamInfo(name="keys",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="Strings of special keys.")},
|
||||
desc="Send strings of special keys like Escape,Backspace, Insert, PageDown, Delete, Enter, Shortcuts such as `Control+o`, `Control+Shift+T` are supported as well. This gets used in keyboard.press. ")
|
||||
WRITE_TO_FILE = ToolActionInfo(name="write_to_file",
|
||||
input_params={
|
||||
"file_path": ParamInfo(
|
||||
name="file_path",
|
||||
type="str",
|
||||
required=False,
|
||||
default_value="tmp_result.md",
|
||||
desc="Path to the file to write to"
|
||||
),
|
||||
"content": ParamInfo(
|
||||
name="content",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="Content to write to the file"
|
||||
),
|
||||
"mode": ParamInfo(
|
||||
name="mode",
|
||||
type="str",
|
||||
required=False,
|
||||
default_value="a",
|
||||
desc="File opening mode: 'w' for write (overwrite), 'a' for append (default)"
|
||||
)
|
||||
},
|
||||
desc="Write content to a file")
|
||||
DONE = ToolActionInfo(name="done",
|
||||
desc="Complete task - with return text and if the task is finished (success=True) or not yet completly finished (success=False), because last step is reached")
|
||||
|
||||
|
||||
class AndroidAction(ToolAction):
|
||||
"""Definition of android tool supported action."""
|
||||
TAP = ToolActionInfo(name="tap",
|
||||
input_params={"tap_index": ParamInfo(name="tap_index",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="Index of tap element.")},
|
||||
desc="Tap element")
|
||||
SWIPE = ToolActionInfo(name="swipe",
|
||||
input_params={"index": ParamInfo(name="index",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="Index of swipe the screen."),
|
||||
"direction": ParamInfo(name="direction",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="Direction of swipe the screen."),
|
||||
"dist": ParamInfo(name="dist",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="Dist of swipe the screen.")},
|
||||
desc="Swipe the screen")
|
||||
LONG_PRESS = ToolActionInfo(name="long_press",
|
||||
input_params={"long_press_index": ParamInfo(name="long_press_index",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="Index of the element.")},
|
||||
desc="Long press the element")
|
||||
INPUT_TEXT = ToolActionInfo(name="input_text",
|
||||
input_params={"text": ParamInfo(name="text",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="Input text into a input interactive element.")},
|
||||
desc="Input text into a input interactive element")
|
||||
DONE = ToolActionInfo(name="done",
|
||||
input_params={"type": ParamInfo(name="type",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="Type of done."),
|
||||
"success": ParamInfo(name="success",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="Task success status.")},
|
||||
desc="task done")
|
||||
|
||||
|
||||
class FileAction(ToolAction):
|
||||
"""Definition of file supported action."""
|
||||
OPEN = ToolActionInfo(name="open",
|
||||
input_params={},
|
||||
desc="")
|
||||
|
||||
|
||||
class ImageAnalysisAction(ToolAction):
|
||||
"""Definition of image analysis supported action."""
|
||||
ANALYSIS = ToolActionInfo(name="analysis",
|
||||
input_params={},
|
||||
desc="")
|
||||
|
||||
|
||||
class CodeExecuteAction(ToolAction):
|
||||
"""Definition of code execute supported action."""
|
||||
EXECUTE_CODE = ToolActionInfo(
|
||||
name="execute_code",
|
||||
input_params={"code": ParamInfo(name="code",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="The input code to execute. Codes should be complete and runnable (like running a script), and need to explicitly use the print statement to get the output.")},
|
||||
desc="Execute the given codes. Codes should be complete and runnable (like running a script), and need to explicitly use the print statement to get the output.")
|
||||
|
||||
|
||||
class ShellAction(ToolAction):
|
||||
"""Definition of shell execute supported action."""
|
||||
EXECUTE_SCRIPT = ToolActionInfo(
|
||||
name="execute_script",
|
||||
input_params={"script": ParamInfo(name="script",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="The input script to execute. Script should be complete and runnable, and need to explicitly use the print statement to get the output.")},
|
||||
desc="Execute the given script, need to explicitly use the print statement to get the output.")
|
||||
|
||||
|
||||
class DocumentExecuteAction(ToolAction):
|
||||
"""Definition of Document execute supported action."""
|
||||
DOCUMENT_ANALYSIS = ToolActionInfo(
|
||||
name="document_analysis",
|
||||
input_params={"document_path": ParamInfo(name="document_path",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="The path of the document to be processed, either a local path or a URL. It can process image, video, audio, ppt, docx, pdf, doc, xls, xlsx and xml, etc.")},
|
||||
desc="Extract the content of a given document (or url) and return the processed text. It can process image, video, audio, ppt, docx, pdf, doc, xls, xlsx and xml, etc. It may filter out some information, resulting in inaccurate content.")
|
||||
|
||||
|
||||
class PythonToolAction(ToolAction):
|
||||
"""Definition of python code execute supported action."""
|
||||
EXECUTE = ToolActionInfo(
|
||||
name="execute",
|
||||
input_params={"code": ParamInfo(name="code",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="The input python code to execute. Python codes should be complete and runnable (like running a script), and need to explicitly use the print statement to get the output.")},
|
||||
desc="Execute the given python codes. Codes should be complete and runnable (like running a script), and need to explicitly use the print statement to get the output.")
|
||||
|
||||
|
||||
class WriteAction(ToolAction):
|
||||
"""Info Write actions."""
|
||||
WRITE_HTML = ToolActionInfo(name="write_html",
|
||||
input_params={"goal": ParamInfo(name="goal",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="the write goal, about theme, requirements for writing html file."),
|
||||
"information": ParamInfo(name="information",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="the related information for writing html file. lengths should less than 6000 words."
|
||||
)
|
||||
},
|
||||
desc="write the html file about `goal` based on `information`.")
|
||||
|
||||
class GetTraceAction(ToolAction):
|
||||
"""Definition of get trace supported action."""
|
||||
GET_TRACE = ToolActionInfo(
|
||||
name="get_trace",
|
||||
input_params={"trace_id": ParamInfo(name="trace_id",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="The trace id to get.")},
|
||||
desc="Get the trace of the current execution.")
|
||||
|
||||
class HumanExecuteAction(ToolAction):
|
||||
"""Definition of Human execute supported action."""
|
||||
HUMAN_CONFIRM = ToolActionInfo(
|
||||
name="human_confirm",
|
||||
input_params={"content": ParamInfo(name="content",
|
||||
type="str",
|
||||
required=True,
|
||||
desc="Content for user confirmation")},
|
||||
desc="The main purpose of this tool is to pass given content to the user for confirmation.")
|
||||
@@ -0,0 +1,160 @@
|
||||
import aworld.trace as trace
|
||||
import aworld.trace.instrumentation.semconv as semconv
|
||||
from aworld.trace.server import get_trace_server
|
||||
from aworld.trace.server.util import build_trace_tree
|
||||
from aworld.core.tool.base import AsyncTool, AgentInput, ToolFactory
|
||||
from examples.common.tools.tool_action import GetTraceAction
|
||||
from aworld.tools.utils import build_observation
|
||||
from aworld.config.conf import ToolConfig
|
||||
from aworld.core.common import Observation, ActionModel, ActionResult
|
||||
from typing import Tuple, Dict, Any, List
|
||||
from aworld.logs.util import logger
|
||||
|
||||
|
||||
@ToolFactory.register(name="trace",
|
||||
desc="Get the trace of the current execution.",
|
||||
supported_action=GetTraceAction,
|
||||
conf_file_name=f'trace_tool.yaml')
|
||||
class TraceTool(AsyncTool):
|
||||
def __init__(self,
|
||||
conf: ToolConfig,
|
||||
**kwargs) -> None:
|
||||
"""
|
||||
Initialize the TraceTool
|
||||
Args:
|
||||
conf: tool config
|
||||
**kwargs: -
|
||||
Return:
|
||||
None
|
||||
"""
|
||||
super(TraceTool, self).__init__(conf, **kwargs)
|
||||
self.type = "function"
|
||||
self.get_trace_url = self.conf.get('get_trace_url')
|
||||
|
||||
async def reset(self,
|
||||
*,
|
||||
seed: int | None = None,
|
||||
options: Dict[str, str] | None = None) -> Tuple[AgentInput, dict[str, Any]]:
|
||||
"""
|
||||
Reset the executor
|
||||
Args:
|
||||
seed: -
|
||||
options: -
|
||||
Returns:
|
||||
AgentInput, dict[str, Any]: -
|
||||
"""
|
||||
self._finished = False
|
||||
return build_observation(observer=self.name(),
|
||||
ability=GetTraceAction.GET_TRACE.value.name), {}
|
||||
|
||||
async def close(self) -> None:
|
||||
"""
|
||||
Close the executor
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
self._finished = True
|
||||
|
||||
async def do_step(self,
|
||||
actions: List[ActionModel],
|
||||
**kwargs) -> Tuple[Observation, float, bool, bool, dict[str, Any]]:
|
||||
reward = 0
|
||||
fail_error = ""
|
||||
observation = build_observation(observer=self.name(),
|
||||
ability=GetTraceAction.GET_TRACE.value.name)
|
||||
results = []
|
||||
try:
|
||||
if not actions:
|
||||
return (observation, reward,
|
||||
kwargs.get("terminated",
|
||||
False), kwargs.get("truncated", False), {
|
||||
"exception": "actions is empty"
|
||||
})
|
||||
for action in actions:
|
||||
trace_id = action.params.get("trace_id", "")
|
||||
if not trace_id:
|
||||
current_span = trace.get_current_span()
|
||||
if current_span:
|
||||
trace_id = current_span.get_trace_id()
|
||||
if not trace_id:
|
||||
logger.warning(f"{action} no trace_id to fetch.")
|
||||
observation.action_result.append(
|
||||
ActionResult(is_done=True,
|
||||
success=False,
|
||||
content="",
|
||||
error="no trace_id to fetch",
|
||||
keep=False))
|
||||
continue
|
||||
try:
|
||||
trace_data = self.fetch_trace_data(trace_id)
|
||||
# logger.info(f"trace_data={trace_data}")
|
||||
error = ""
|
||||
except Exception as e:
|
||||
error = str(e)
|
||||
results.append(trace_data)
|
||||
observation.action_result.append(
|
||||
ActionResult(is_done=True,
|
||||
success=False if error else True,
|
||||
content=f"{trace_data}",
|
||||
error=f"{error}",
|
||||
keep=False))
|
||||
|
||||
observation.content = f"{results}"
|
||||
reward = 1
|
||||
except Exception as e:
|
||||
fail_error = str(e)
|
||||
finally:
|
||||
self._finished = True
|
||||
|
||||
info = {"exception": fail_error}
|
||||
info.update(kwargs)
|
||||
return (observation, reward, kwargs.get("terminated", False),
|
||||
kwargs.get("truncated", False), info)
|
||||
|
||||
def fetch_trace_data(self, trace_id=None):
|
||||
'''
|
||||
fetch trace data from trace server.
|
||||
return trace data, like:
|
||||
{
|
||||
'trace_id': trace_id,
|
||||
'root_span': [],
|
||||
}
|
||||
'''
|
||||
trace_data = {"trace_id": trace_id, "root_span": []}
|
||||
try:
|
||||
if trace_id:
|
||||
trace_server = get_trace_server()
|
||||
if not trace_server:
|
||||
logger.error("No memory trace server has been set.")
|
||||
else:
|
||||
trace_storage = trace_server.get_storage()
|
||||
spans = trace_storage.get_all_spans(trace_id)
|
||||
if spans:
|
||||
trace_data["root_span"] = build_trace_tree(spans)
|
||||
return self.proccess_trace(trace_data)
|
||||
return trace_data
|
||||
except Exception as e:
|
||||
import traceback
|
||||
logger.error(
|
||||
f"Error fetching trace data traceback: {traceback.format_exc()}")
|
||||
return trace_data
|
||||
|
||||
def proccess_trace(self, trace_data):
|
||||
root_spans = trace_data.get("root_span")
|
||||
for span in root_spans:
|
||||
self.choose_attribute(span)
|
||||
return trace_data
|
||||
|
||||
def choose_attribute(self, span):
|
||||
include_attr = [semconv.GEN_AI_USAGE_INPUT_TOKENS,
|
||||
semconv.GEN_AI_USAGE_OUTPUT_TOKENS, semconv.GEN_AI_USAGE_TOTAL_TOKENS,
|
||||
semconv.GEN_AI_COMPLETION_TOOL_CALLS, "event.id"]
|
||||
result_attributes = {}
|
||||
origin_attributes = span.get("attributes") or {}
|
||||
for key, value in origin_attributes.items():
|
||||
if key in include_attr:
|
||||
result_attributes[key] = value
|
||||
span["attributes"] = result_attributes
|
||||
if span.get("children"):
|
||||
for child in span.get("children"):
|
||||
self.choose_attribute(child)
|
||||
@@ -0,0 +1,197 @@
|
||||
# GAIA Agent Setup Guide
|
||||
|
||||
## 1. Overview
|
||||
|
||||
This guide will help you set up and run the GAIA agent for the AWorld framework. GAIA is a benchmark dataset for evaluating AI agents' capabilities.
|
||||
|
||||
## 2. Prerequisites
|
||||
|
||||
### 1. System Requirements
|
||||
- **Operating System**: macOS or Linux (Windows not fully tested)
|
||||
- **Node.js**: Version 22 LTS with npm
|
||||
- **Conda**: For environment management
|
||||
|
||||
### 2. Required Software
|
||||
- `libmagic1` - File type detection
|
||||
- `libreoffice` - Document processing
|
||||
- `ffmpeg` - Media processing
|
||||
|
||||
## 3. Installation Steps
|
||||
|
||||
### 1. Clone the Repository
|
||||
|
||||
```bash
|
||||
git clone https://github.com/inclusionAI/AWorld.git
|
||||
cd AWorld
|
||||
```
|
||||
|
||||
### 2. Set Up Conda Environment
|
||||
|
||||
Create and activate a dedicated Conda environment for GAIA:
|
||||
|
||||
```bash
|
||||
conda env create -f examples/gaia/aworld-gaia.yml
|
||||
conda activate aworld-gaia
|
||||
```
|
||||
|
||||
> **Note**: If you don't have Conda installed, download Miniconda from [here](https://www.anaconda.com/docs/getting-started/miniconda/install).
|
||||
|
||||
### 3. Install AWorld Framework
|
||||
|
||||
Install the AWorld framework and build the web UI:
|
||||
|
||||
```bash
|
||||
# Install PDF processing dependencies
|
||||
pip install "marker-pdf[full]" --no-deps
|
||||
|
||||
# Build web UI
|
||||
sh -c "cd aworld/cmd/web/webui && npm install && npm run build"
|
||||
|
||||
# Install AWorld
|
||||
python setup.py install
|
||||
```
|
||||
|
||||
### 4. Install MCP Tool Dependencies
|
||||
|
||||
#### Install Playwright
|
||||
```bash
|
||||
playwright install chromium --with-deps --no-shell
|
||||
```
|
||||
|
||||
#### Install System Dependencies
|
||||
|
||||
**For macOS:**
|
||||
```bash
|
||||
brew install libmagic
|
||||
brew install ffmpeg
|
||||
brew install --cask libreoffice
|
||||
```
|
||||
|
||||
> **Note**: Install Homebrew from [brew.sh](https://brew.sh/) if not already installed.
|
||||
|
||||
**For Linux:**
|
||||
```bash
|
||||
apt-get install -y --no-install-recommends libmagic1 libreoffice ffmpeg
|
||||
```
|
||||
|
||||
### 5. Prepare GAIA Dataset
|
||||
|
||||
Download the GAIA dataset from Hugging Face:
|
||||
|
||||
```bash
|
||||
git clone git@hf.co:datasets/gaia-benchmark/GAIA examples/gaia/GAIA
|
||||
```
|
||||
|
||||
> **⚠️ Important**:
|
||||
> - You need to configure Hugging Face SSH keys to access the GAIA repository
|
||||
> - The dataset path will be used as the `GAIA_DATASET_PATH` variable in your `.env` file
|
||||
|
||||
### 6. Configure Environment Variables
|
||||
|
||||
Create the environment configuration file:
|
||||
|
||||
```bash
|
||||
cp examples/gaia/cmd/agent_deploy/gaia_agent/.env.template examples/gaia/cmd/agent_deploy/gaia_agent/.env
|
||||
```
|
||||
|
||||
Edit the `.env` file and replace all `{YOUR_CONFIG}` placeholders with your actual configuration values.
|
||||
|
||||
## 3. Running the GAIA Agent
|
||||
|
||||
### 1. Web UI Interface
|
||||
|
||||
Start the GAIA agent web interface:
|
||||
|
||||
```bash
|
||||
cd examples/gaia/cmd && aworld web
|
||||
```
|
||||
|
||||
### 2. Command Line Interface
|
||||
|
||||
Run GAIA tasks using the command line interface:
|
||||
|
||||
```bash
|
||||
python examples/gaia/run.py --split validation --q c61d22de-5f6c-4958-a7f6-5e9707bd3466
|
||||
```
|
||||
|
||||
## 4. Command Line Arguments
|
||||
|
||||
### 1. Required Arguments
|
||||
|
||||
| Argument | Type | Description | Default | Example |
|
||||
|----------|------|-------------|---------|---------|
|
||||
| `--split` | str | Dataset split to use | `validation` | `--split test` |
|
||||
| `--start` | int | Start index of the dataset | `0` | `--start 10` |
|
||||
| `--end` | int | End index of the dataset | `165` | `--end 100` |
|
||||
|
||||
### 2. Optional Arguments
|
||||
|
||||
| Argument | Type | Description | Example |
|
||||
|----------|------|-------------|---------|
|
||||
| `--q` | str | Specific question index (overrides start/end) | `--q 0-0-0-0-0` |
|
||||
| `--skip` | flag | Skip previously processed questions | `--skip` |
|
||||
| `--blacklist_file_path` | str | Path to blacklist file | `--blacklist_file_path blacklist.txt` |
|
||||
|
||||
## 5. Usage Examples
|
||||
|
||||
### 1. Basic Usage
|
||||
Process a range of questions in the validation split:
|
||||
|
||||
```bash
|
||||
python examples/gaia/run.py --start 10 --end 50 --split validation
|
||||
```
|
||||
|
||||
### 2. Process Specific Question
|
||||
Run a single question by its index:
|
||||
|
||||
```bash
|
||||
python examples/gaia/run.py --q 0-0-0-0-0
|
||||
```
|
||||
|
||||
### 3. Skip Processed Questions
|
||||
Process questions while skipping previously completed ones:
|
||||
|
||||
```bash
|
||||
python examples/gaia/run.py --start 10 --end 50 --skip
|
||||
```
|
||||
|
||||
### 4. Use Blacklist
|
||||
Skip questions listed in a blacklist file:
|
||||
|
||||
```bash
|
||||
python examples/gaia/run.py --start 10 --end 50 --blacklist_file_path blacklist.txt
|
||||
```
|
||||
|
||||
## 6. Expected Output
|
||||
|
||||
When running successfully, you should see logs similar to:
|
||||
|
||||
```
|
||||
YYYY-MM-DD HH:MM:SS - root - INFO - Agent answer: egalitarian
|
||||
YYYY-MM-DD HH:MM:SS - root - INFO - Correct answer: egalitarian
|
||||
YYYY-MM-DD HH:MM:SS - examples.gaia.utils - INFO - Evaluating egalitarian as a string.
|
||||
YYYY-MM-DD HH:MM:SS - root - INFO - Question 0 Correct!
|
||||
```
|
||||
|
||||
## 7. Troubleshooting
|
||||
|
||||
### 1. Common Issues
|
||||
|
||||
1. **Dataset Access Problems**
|
||||
- Verify your Hugging Face SSH keys are correctly configured
|
||||
- Ensure you have access to the GAIA repository
|
||||
|
||||
2. **Installation Issues**
|
||||
- Set up a pip mirror if necessary for faster downloads
|
||||
- Ensure all system dependencies are properly installed
|
||||
|
||||
3. **Environment Issues**
|
||||
- Make sure you're using the correct Conda environment (`aworld-gaia`)
|
||||
- Verify Node.js version is 22 LTS
|
||||
|
||||
### 2. Getting Help
|
||||
|
||||
If you encounter issues not covered in this guide:
|
||||
- Check the project's main documentation
|
||||
- Review the error logs for specific error messages
|
||||
- Ensure all prerequisites are properly installed
|
||||
@@ -0,0 +1,82 @@
|
||||
# GAIA Agent Guard Functionality Setup
|
||||
|
||||
## 1. Overview
|
||||
|
||||
This guide covers the additional setup required for the enhanced guard agent in the GAIA agent. For basic installation and setup, please refer to [README.md](README.md).
|
||||
|
||||
## 2. Setting Up Guard Functionality
|
||||
|
||||
### 2.1 MCP Configuration
|
||||
|
||||
To enable the guard functionality, ensure the following MCP server is registered in your `mcp.json` file:
|
||||
|
||||
```json
|
||||
{
|
||||
"mcpServers": {
|
||||
"maneuvering": {
|
||||
"command": "python",
|
||||
"args": [
|
||||
"-m",
|
||||
"examples.gaia.mcp_collections.intelligence.guard"
|
||||
],
|
||||
"env": {},
|
||||
"client_session_timeout_seconds": 9999.0
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### 2.2 Configure Environment Variables
|
||||
|
||||
Add the guard_llm API key to your `.env` file:
|
||||
|
||||
```bash
|
||||
# Add this line to examples/gaia/cmd/agent_deploy/gaia_agent/.env
|
||||
GUARD_LLM_API_KEY=your_guard_llm_api_key_here
|
||||
```
|
||||
|
||||
### 2.3 Update Prompt Configuration
|
||||
|
||||
Replace the content of `prompt.py` with the enhanced version from `prompt_w_guard.py`:
|
||||
|
||||
```bash
|
||||
cp examples/gaia/prompt_w_guard.py examples/gaia/prompt.py
|
||||
```
|
||||
|
||||
### 2.4 Configure Task Subset Processing
|
||||
|
||||
To run specific subsets of GAIA tasks, add the following code to `run.py`:
|
||||
|
||||
```python
|
||||
# load task subset from subset.txt
|
||||
subset_file_path = Path(__file__).parent / "subset.txt"
|
||||
|
||||
if subset_file_path.exists():
|
||||
with open(subset_file_path, "r", encoding="utf-8") as f:
|
||||
task_subset = set(line.strip() for line in f if line.strip())
|
||||
logging.info(f"Loaded {len(task_subset)} task IDs from subset.txt")
|
||||
else:
|
||||
task_subset = set() # Empty set if file doesn't exist
|
||||
logging.warning("subset.txt file not found, using empty task subset")
|
||||
```
|
||||
|
||||
And add the filtering logic in the main processing loop:
|
||||
|
||||
```python
|
||||
# only process tasks that are in the subset
|
||||
if dataset_i["task_id"] not in task_subset:
|
||||
continue
|
||||
```
|
||||
|
||||
|
||||
## 3. Configuration Summary
|
||||
|
||||
To enable guard functionality, ensure you have:
|
||||
|
||||
1. ✅ **MCP Configuration**: `maneuvering` server registered in `mcp.json`
|
||||
2. ✅ **Environment Variables**: `GUARD_LLM_API_KEY` added to `.env` file
|
||||
3. ✅ **Prompt Update**: `prompt_w_guard.py` content copied to `prompt.py`
|
||||
4. ✅ **Subset Processing**: Batch processing code added to `run.py`
|
||||
|
||||
|
||||
This setup provides enhanced reasoning capabilities and efficient batch processing for GAIA task evaluation.
|
||||
@@ -0,0 +1,8 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
from aworld.utils import import_packages
|
||||
|
||||
import_packages(["dotenv"])
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
@@ -0,0 +1,248 @@
|
||||
name: aworld-gaia
|
||||
channels:
|
||||
- defaults
|
||||
dependencies:
|
||||
- bzip2=1.0.8
|
||||
- ca-certificates=2025.2.25
|
||||
- expat=2.7.1
|
||||
- libcxx=17.0.6
|
||||
- libffi=3.4.4
|
||||
- ncurses=6.4
|
||||
- openssl=3.0.16
|
||||
- pip=25.1
|
||||
- python=3.12.11
|
||||
- readline=8.2
|
||||
- setuptools=78.1.1
|
||||
- sqlite=3.45.3
|
||||
- tk=8.6.14
|
||||
- wheel=0.45.1
|
||||
- xz=5.6.4
|
||||
- zlib=1.2.13
|
||||
- pip:
|
||||
- aiofiles==24.1.0
|
||||
- altair==5.5.0
|
||||
- annotated-types==0.7.0
|
||||
- anthropic
|
||||
- anyio==4.9.0
|
||||
- arxiv==2.2.0
|
||||
- asttokens==3.0.0
|
||||
- attrs==25.3.0
|
||||
- Authlib==1.6.0
|
||||
- backoff==2.2.1
|
||||
- beautifulsoup4==4.13.4
|
||||
- blinker==1.9.0
|
||||
- Brotli==1.1.0
|
||||
- browser-use==0.7.1
|
||||
- bubus==1.2.1
|
||||
- cachetools==5.5.2
|
||||
- certifi==2025.6.15
|
||||
- cffi==1.17.1
|
||||
- cfgv==3.4.0
|
||||
- chardet==5.2.0
|
||||
- charset-normalizer==3.4.2
|
||||
- chess==1.11.2
|
||||
- click==8.2.1
|
||||
- cobble==0.1.4
|
||||
- cryptography==45.0.4
|
||||
- cssselect2==0.8.0
|
||||
- Cython==3.1.2
|
||||
- decorator==5.2.1
|
||||
- defusedxml==0.7.1
|
||||
- distlib==0.3.9
|
||||
- distro==1.9.0
|
||||
- dotenv==0.9.9
|
||||
- durationpy==0.10
|
||||
- EbookLib==0.18
|
||||
- einops==0.8.1
|
||||
- et_xmlfile==2.0.0
|
||||
- executing==2.2.0
|
||||
- faiss-cpu==1.11.0
|
||||
- feedparser==6.0.11
|
||||
- filelock==3.18.0
|
||||
- filetype==1.2.0
|
||||
- Flask==3.1.1
|
||||
- fonttools==4.58.4
|
||||
- fsspec==2025.5.1
|
||||
- ftfy==6.3.1
|
||||
- gitdb==4.0.12
|
||||
- GitPython==3.1.44
|
||||
- google-ai-generativelanguage==0.6.18
|
||||
- google-api-core==2.25.1
|
||||
- google-auth==2.40.3
|
||||
- google-genai==1.21.1
|
||||
- googleapis-common-protos==1.70.0
|
||||
- greenlet==3.2.3
|
||||
- grpcio==1.73.0
|
||||
- grpcio-status==1.73.0
|
||||
- h11==0.16.0
|
||||
- h2==4.2.0
|
||||
- hf-xet==1.1.5
|
||||
- hpack==4.1.0
|
||||
- httpcore==1.0.9
|
||||
- httpx==0.28.1
|
||||
- httpx-sse==0.4.1
|
||||
- huggingface-hub==0.33.0
|
||||
- hyperframe==6.1.0
|
||||
- identify==2.6.12
|
||||
- idna==3.10
|
||||
- ipython==9.3.0
|
||||
- ipython_pygments_lexers==1.1.1
|
||||
- itsdangerous==2.2.0
|
||||
- jedi==0.19.2
|
||||
- Jinja2==3.1.6
|
||||
- jiter==0.10.0
|
||||
- joblib==1.5.1
|
||||
- jsonpatch==1.33
|
||||
- jsonpickle==4.1.1
|
||||
- jsonpointer==3.0.0
|
||||
- jsonschema==4.24.0
|
||||
- jsonschema-specifications==2025.4.1
|
||||
- kubernetes==32.0.1
|
||||
- langchain==0.3.25
|
||||
- langchain-anthropic==0.3.15
|
||||
- langchain-core==0.3.64
|
||||
- langchain-deepseek==0.1.3
|
||||
- langchain-google-genai==2.1.5
|
||||
- langchain-ollama==0.3.3
|
||||
- langchain-openai==0.3.21
|
||||
- langchain-text-splitters==0.3.8
|
||||
- langsmith==0.3.45
|
||||
- lxml==5.4.0
|
||||
- mammoth==1.9.1
|
||||
- markdown
|
||||
# - marker-pdf
|
||||
- markdown-it-py==3.0.0
|
||||
- markdown2==2.5.3
|
||||
- markdownify==1.1.0
|
||||
- MarkupSafe==3.0.2
|
||||
- matplotlib-inline==0.1.7
|
||||
- mcp==1.6.0
|
||||
- mdurl==0.1.2
|
||||
- mem0ai==0.1.111
|
||||
- mpmath==1.3.0
|
||||
- narwhals==1.44.0
|
||||
- networkx==3.5
|
||||
- nodeenv==1.9.1
|
||||
- numpy==2.3.1
|
||||
- oauthlib==3.3.1
|
||||
- ollama==0.5.1
|
||||
- openai
|
||||
- opencv-python==4.11.0.86
|
||||
- opencv-python-headless==4.11.0.86
|
||||
- openpyxl==3.1.5
|
||||
- orjson==3.10.18
|
||||
- outcome==1.3.0.post0
|
||||
- packaging==24.2
|
||||
- pandas==2.3.0
|
||||
- parso==0.8.4
|
||||
- patchright==1.52.5
|
||||
- pdftext==0.6.3
|
||||
- pexpect==4.9.0
|
||||
- pillow==10.4.0
|
||||
- platformdirs==4.3.8
|
||||
- playwright==1.52.0
|
||||
- portalocker==2.10.1
|
||||
- posthog==5.4.0
|
||||
- prompt_toolkit==3.0.51
|
||||
- proto-plus==1.26.1
|
||||
- protobuf==6.31.1
|
||||
- psutil==7.0.0
|
||||
- ptyprocess==0.7.0
|
||||
- pure_eval==0.2.3
|
||||
- pyarrow==20.0.0
|
||||
- pyasn1==0.6.1
|
||||
- pyasn1_modules==0.4.2
|
||||
- pycparser==2.22
|
||||
- pydantic==2.11.5
|
||||
- pydantic-settings
|
||||
- pydantic_core
|
||||
- pyautogui
|
||||
- rubicon-objc
|
||||
- pydeck==0.9.1
|
||||
- pydyf==0.11.0
|
||||
- pyee==13.0.0
|
||||
- Pygments==2.19.2
|
||||
- pypdfium2==4.30.0
|
||||
- pyperclip==1.9.0
|
||||
- pyphen==0.17.2
|
||||
- PySocks==1.7.1
|
||||
- pytesseract==0.3.13
|
||||
- python-dateutil==2.9.0.post0
|
||||
- python-docx==1.2.0
|
||||
- python-dotenv==1.0.1
|
||||
- python-pptx==1.0.2
|
||||
- python-magic==0.4.27
|
||||
- pytz==2025.2
|
||||
- pyvis==0.3.2
|
||||
- PyYAML==6.0.2
|
||||
- qdrant-client==1.14.3
|
||||
- RapidFuzz==3.13.0
|
||||
- referencing==0.36.2
|
||||
- regex==2024.11.6
|
||||
- requests==2.32.4
|
||||
- requests-oauthlib==2.0.0
|
||||
- requests-toolbelt==1.0.0
|
||||
- rich==14.0.0
|
||||
- rpds-py==0.25.1
|
||||
- rsa==4.9.1
|
||||
- safetensors==0.5.3
|
||||
- scikit-learn==1.7.0
|
||||
- scipy==1.16.0
|
||||
- screeninfo==0.8.1
|
||||
- selenium==4.33.0
|
||||
- sentence-transformers==4.1.0
|
||||
- setuptools==78.1.1
|
||||
- sgmllib3k==1.0.0
|
||||
- shellingham==1.5.4
|
||||
- six==1.17.0
|
||||
- smmap==5.0.2
|
||||
- sniffio==1.3.1
|
||||
- sortedcontainers==2.4.0
|
||||
- soupsieve==2.7
|
||||
- SQLAlchemy==2.0.41
|
||||
- sse-starlette==2.3.6
|
||||
- stack-data==0.6.3
|
||||
- starlette==0.47.1
|
||||
- streamlit==1.46.0
|
||||
- surya-ocr==0.14.6
|
||||
- sympy==1.14.0
|
||||
- tabulate==0.9.0
|
||||
- tenacity==8.5.0
|
||||
- threadpoolctl==3.6.0
|
||||
- tiktoken==0.9.0
|
||||
- tinycss2==1.4.0
|
||||
- tinyhtml5==2.0.0
|
||||
- tokenizers==0.21.2
|
||||
- toml==0.10.2
|
||||
- torch==2.7.1
|
||||
- tornado==6.5.1
|
||||
- tqdm==4.67.1
|
||||
- traitlets==5.14.3
|
||||
- transformers==4.52.4
|
||||
- trio==0.30.0
|
||||
- trio-websocket==0.12.2
|
||||
- typer==0.16.0
|
||||
- typing-inspection==0.4.1
|
||||
- typing_extensions==4.13.2
|
||||
- tzdata==2025.2
|
||||
- urllib3==2.4.0
|
||||
- uuid7==0.1.0
|
||||
- uvicorn==0.34.3
|
||||
- virtualenv==20.31.2
|
||||
- waybackpy==3.0.6
|
||||
- wcwidth==0.2.13
|
||||
- weasyprint==63.1
|
||||
- webencodings==0.5.1
|
||||
- websocket-client==1.8.0
|
||||
- websockets==15.0.1
|
||||
- Werkzeug==3.1.3
|
||||
- wheel==0.45.1
|
||||
- wikipedia==1.4.0
|
||||
- wsproto==1.2.0
|
||||
- xlsxwriter==3.2.5
|
||||
- youtube-transcript-api==1.1.0
|
||||
- zopfli==0.2.3.post1
|
||||
- zstandard==0.23.0
|
||||
- httpx[socks]
|
||||
- ffmpeg-python==0.2.0
|
||||
prefix: /opt/homebrew/Caskroom/miniconda/base/envs/aworld-gaia
|
||||
@@ -0,0 +1,42 @@
|
||||
import logging
|
||||
import os
|
||||
from aworld.cmd.data_model import BaseAWorldAgent, ChatCompletionRequest
|
||||
from examples.gaia.gaia_agent_runner import GaiaAgentRunner
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AWorldAgent(BaseAWorldAgent):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
os.makedirs(os.path.join(os.getcwd(), "static"), exist_ok=True)
|
||||
|
||||
async def run(self, prompt: str = None, request: ChatCompletionRequest = None):
|
||||
llm_provider = os.getenv("LLM_PROVIDER", "openai")
|
||||
llm_model_name = os.getenv("LLM_MODEL_NAME")
|
||||
llm_api_key = os.getenv("LLM_API_KEY")
|
||||
llm_base_url = os.getenv("LLM_BASE_URL")
|
||||
llm_temperature = float(os.getenv("LLM_TEMPERATURE", 0.0))
|
||||
|
||||
if not llm_model_name or not llm_api_key or not llm_base_url:
|
||||
raise ValueError(
|
||||
"LLM_MODEL_NAME, LLM_API_KEY, LLM_BASE_URL must be set in your envrionment variables"
|
||||
)
|
||||
|
||||
runner = GaiaAgentRunner(
|
||||
llm_provider=llm_provider,
|
||||
llm_model_name=llm_model_name,
|
||||
llm_base_url=llm_base_url,
|
||||
llm_api_key=llm_api_key,
|
||||
llm_temperature=llm_temperature,
|
||||
session_id=request.session_id,
|
||||
)
|
||||
|
||||
if prompt is None and request is not None:
|
||||
prompt = request.messages[-1].content
|
||||
|
||||
logger.info(f">>> Gaia Agent: prompt={prompt}, runner={runner}")
|
||||
|
||||
async for line in runner.run(prompt):
|
||||
logger.info(f">>> Gaia Agent Line: {line}")
|
||||
yield line
|
||||
@@ -0,0 +1,346 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import traceback
|
||||
from typing import AsyncGenerator
|
||||
import uuid
|
||||
from aworld.cmd.utils.agent_ui_parser import (
|
||||
AWorldWebAgentUI,
|
||||
BaseToolResultParser,
|
||||
ToolCard,
|
||||
ToolResultParserFactory,
|
||||
)
|
||||
from aworld.config.conf import AgentConfig, TaskConfig
|
||||
from aworld.agents.llm_agent import Agent
|
||||
from aworld.core.task import Task
|
||||
from aworld.output.artifact import ArtifactType
|
||||
from aworld.output.workspace import WorkSpace
|
||||
from aworld.runner import Runners
|
||||
from aworld.output.ui.base import AworldUI
|
||||
from aworld.output.base import Output, ToolResultOutput
|
||||
from .utils import (
|
||||
add_file_path,
|
||||
load_dataset_meta_dict,
|
||||
question_scorer,
|
||||
)
|
||||
from .prompt import system_prompt
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class GaiaSearchToolResultParser(BaseToolResultParser):
|
||||
async def parse(self, output: ToolResultOutput, workspace: WorkSpace):
|
||||
tool_card = ToolCard.from_tool_result(output)
|
||||
|
||||
query = ""
|
||||
try:
|
||||
args = json.loads(tool_card.arguments)
|
||||
query = args.get("query")
|
||||
# aworld search server
|
||||
if not query:
|
||||
query = args.get("query_list")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
result_items = []
|
||||
try:
|
||||
results = json.loads(tool_card.results)
|
||||
result_items = results.get("message", {}).get("results", [])
|
||||
# aworld search server return url, not link
|
||||
if result_items and isinstance(result_items, list):
|
||||
for item in result_items:
|
||||
if not item.get("link", None) and item.get("url", None):
|
||||
item["link"] = item.get("url")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if len(result_items) > 0:
|
||||
tool_card.results = ""
|
||||
|
||||
tool_card.card_type = "tool_call_card_link_list"
|
||||
tool_card.card_data = {
|
||||
"title": "🔎 Gaia Search",
|
||||
"query": query,
|
||||
"search_items": result_items,
|
||||
}
|
||||
|
||||
artifact_id = str(uuid.uuid4())
|
||||
await workspace.create_artifact(
|
||||
artifact_type=ArtifactType.WEB_PAGES,
|
||||
artifact_id=artifact_id,
|
||||
content=result_items,
|
||||
metadata={
|
||||
"query": query,
|
||||
},
|
||||
)
|
||||
tool_card.artifacts.append(
|
||||
{
|
||||
"artifact_type": ArtifactType.WEB_PAGES.value,
|
||||
"artifact_id": artifact_id,
|
||||
}
|
||||
)
|
||||
|
||||
return f"""
|
||||
\n\n**🔎 Gaia Search**\n\n
|
||||
```tool_card
|
||||
{json.dumps(tool_card.model_dump(), ensure_ascii=False, indent=2)}
|
||||
```\n
|
||||
"""
|
||||
|
||||
|
||||
class CustomToolResultParserFactory(ToolResultParserFactory):
|
||||
def get_parser(self, tool_type: str, tool_name: str):
|
||||
if tool_name in ("search_server", "search"):
|
||||
return GaiaSearchToolResultParser()
|
||||
return super().get_parser(tool_type, tool_name)
|
||||
|
||||
|
||||
# Module-level flag to ensure dependencies are installed only once per program run
|
||||
_install_dependencies_flag = False
|
||||
|
||||
|
||||
class GaiaAgentRunner:
|
||||
"""
|
||||
Gaia Agent Runner
|
||||
"""
|
||||
|
||||
def _install_dependencies(self):
|
||||
try:
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
requirements_file = os.path.join(current_dir, "requirements.txt")
|
||||
|
||||
if os.path.exists(requirements_file):
|
||||
logger.info(f"Installing dependencies from {requirements_file}")
|
||||
subprocess.check_call(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"pip",
|
||||
"install",
|
||||
"-U",
|
||||
"-r",
|
||||
requirements_file,
|
||||
]
|
||||
)
|
||||
subprocess.check_call(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"pip",
|
||||
"install",
|
||||
"--no-deps",
|
||||
"-U",
|
||||
"marker-pdf",
|
||||
"anthropic==0.46.0",
|
||||
]
|
||||
)
|
||||
logger.info("Dependencies installed successfully")
|
||||
else:
|
||||
logger.warning(f"Requirements file not found at {requirements_file}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to install dependencies: {e}")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
llm_provider: str,
|
||||
llm_model_name: str,
|
||||
llm_base_url: str,
|
||||
llm_api_key: str,
|
||||
llm_temperature: float = 0.0,
|
||||
mcp_config: dict = None,
|
||||
session_id: str = None,
|
||||
):
|
||||
global _install_dependencies_flag
|
||||
if not _install_dependencies_flag:
|
||||
self._install_dependencies()
|
||||
_install_dependencies_flag = True
|
||||
|
||||
self.session_id = session_id or str(uuid.uuid4())
|
||||
self.agent_config = AgentConfig(
|
||||
llm_provider=llm_provider,
|
||||
llm_model_name=llm_model_name,
|
||||
llm_api_key=llm_api_key,
|
||||
llm_base_url=llm_base_url,
|
||||
llm_temperature=llm_temperature,
|
||||
)
|
||||
|
||||
if mcp_config is None:
|
||||
mcp_path = os.path.join(
|
||||
os.path.dirname(os.path.abspath(__file__)), "mcp.json"
|
||||
)
|
||||
with open(mcp_path, "r") as f:
|
||||
mcp_config = json.load(f)
|
||||
logger.info(f"Gaia Agent Runner mcp_config: {mcp_config}")
|
||||
|
||||
self.super_agent = Agent(
|
||||
conf=self.agent_config,
|
||||
name="gaia_super_agent",
|
||||
system_prompt=system_prompt,
|
||||
mcp_config=mcp_config,
|
||||
mcp_servers=(
|
||||
os.getenv("GAIA_MCP_SERVERS", "").split(",") if os.getenv("GAIA_MCP_SERVERS", "") else ""
|
||||
or mcp_config.get("mcpServers", {}).keys()
|
||||
),
|
||||
)
|
||||
|
||||
self.gaia_dataset_path = os.path.abspath(
|
||||
os.getenv(
|
||||
"GAIA_DATASET_PATH",
|
||||
os.path.join(
|
||||
os.path.dirname(os.path.abspath(__file__)), "GAIA", "2023"
|
||||
),
|
||||
)
|
||||
)
|
||||
self.full_dataset = load_dataset_meta_dict(self.gaia_dataset_path)
|
||||
logger.info(
|
||||
f"Gaia Agent Runner initialized: super_agent={self.super_agent}, agent_config={self.agent_config}, gaia_dataset_path={self.gaia_dataset_path}, full_dataset={len(self.full_dataset)}"
|
||||
)
|
||||
|
||||
async def run(self, prompt: str):
|
||||
yield (f"\n### GAIA Agent Start!")
|
||||
|
||||
mcp_servers = "\n- ✅ ".join(self.super_agent.mcp_servers)
|
||||
yield (f"\n```gaia_agent_status\n- ✅ {mcp_servers}\n```\n")
|
||||
|
||||
question = None
|
||||
data_item = None
|
||||
task_id = None
|
||||
try:
|
||||
json_data = json.loads(prompt)
|
||||
task_id = json_data["task_id"]
|
||||
|
||||
data_item = self.full_dataset[task_id]
|
||||
question = add_file_path(data_item, file_path=self.gaia_dataset_path)[
|
||||
"Question"
|
||||
]
|
||||
yield (
|
||||
f"\n### Gaia Question\n```gaia_question\n{json.dumps(data_item, indent=2)}\n```\n"
|
||||
)
|
||||
except Exception as e:
|
||||
pass
|
||||
|
||||
if not question:
|
||||
logger.warning(
|
||||
"Could not find GAIA question for prompt, chat using prompt directly!"
|
||||
)
|
||||
yield (f"\n{prompt}\n")
|
||||
question = prompt
|
||||
|
||||
try:
|
||||
task = Task(
|
||||
id=task_id + "." + uuid.uuid1().hex if task_id else uuid.uuid1().hex,
|
||||
input=question,
|
||||
agent=self.super_agent,
|
||||
conf=TaskConfig(max_steps=20),
|
||||
session_id=self.session_id,
|
||||
endless_threshold=50,
|
||||
)
|
||||
|
||||
last_output: Output = None
|
||||
rich_ui = AWorldWebAgentUI(
|
||||
session_id=self.session_id,
|
||||
workspace=WorkSpace.from_local_storages(workspace_id=self.session_id),
|
||||
tool_result_parser_factory=CustomToolResultParserFactory(),
|
||||
)
|
||||
async for output in Runners.streamed_run_task(task).stream_events():
|
||||
logger.info(f"Gaia Agent Ouput: {output}")
|
||||
res = await AworldUI.parse_output(output, rich_ui)
|
||||
for item in res if isinstance(res, list) else [res]:
|
||||
if isinstance(item, AsyncGenerator):
|
||||
async for sub_item in item:
|
||||
yield sub_item
|
||||
if sub_item and str(sub_item).strip():
|
||||
last_output = sub_item
|
||||
else:
|
||||
yield item
|
||||
if item and str(item).strip():
|
||||
last_output = item
|
||||
|
||||
logger.info(f"Gaia Agent Last Output: {last_output}")
|
||||
|
||||
if data_item and last_output:
|
||||
final_response = self._judge_answer(data_item, last_output)
|
||||
yield final_response
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing {prompt}, error: {traceback.format_exc()}")
|
||||
|
||||
def _judge_answer(self, data_item: dict, result: Output):
|
||||
answer = result
|
||||
match = re.search(r"<answer>(.*?)</answer>", answer)
|
||||
if match:
|
||||
answer = match.group(1)
|
||||
logger.info(f"Agent answer: {answer}")
|
||||
logger.info(f"Correct answer: {data_item['Final answer']}")
|
||||
|
||||
if question_scorer(answer, data_item["Final answer"]):
|
||||
logger.info(f"Question {data_item['task_id']} Correct!")
|
||||
else:
|
||||
logger.info(f"Question {data_item['task_id']} Incorrect!")
|
||||
|
||||
# Create the new result record
|
||||
correct = question_scorer(answer, data_item["Final answer"])
|
||||
new_result = {
|
||||
"task_id": data_item["task_id"],
|
||||
"level": data_item["Level"],
|
||||
"question": data_item["Question"],
|
||||
"answer": data_item["Final answer"],
|
||||
"response": answer,
|
||||
"is_correct": correct,
|
||||
}
|
||||
return f"\n## Final Result: {'✅' if correct else '❌'}\n \n```gaia_result\n{json.dumps(new_result, indent=2)}\n```"
|
||||
else:
|
||||
new_result = answer
|
||||
return f"\n## Final Result:\n \n```gaia_result\n{json.dumps(new_result, indent=2)}\n```"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import asyncio
|
||||
import argparse
|
||||
from datetime import datetime
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "output")
|
||||
if not os.path.exists(output_dir):
|
||||
os.makedirs(output_dir)
|
||||
|
||||
output_file = os.path.join(
|
||||
output_dir, f"output_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md"
|
||||
)
|
||||
|
||||
async def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--prompt", type=str, default="")
|
||||
args = parser.parse_args()
|
||||
|
||||
try:
|
||||
prompt = args.prompt
|
||||
|
||||
llm_provider = os.getenv("LLM_PROVIDER", "openai")
|
||||
llm_model_name = os.getenv("LLM_MODEL_NAME")
|
||||
llm_api_key = os.getenv("LLM_API_KEY")
|
||||
llm_base_url = os.getenv("LLM_BASE_URL")
|
||||
llm_temperature = os.getenv("LLM_TEMPERATURE", 0.0)
|
||||
|
||||
def send_output(output):
|
||||
with open(output_file, "a") as f:
|
||||
f.write(f"{output}\n")
|
||||
|
||||
async for i in GaiaAgentRunner(
|
||||
llm_provider=llm_provider,
|
||||
llm_model_name=llm_model_name,
|
||||
llm_base_url=llm_base_url,
|
||||
llm_api_key=llm_api_key,
|
||||
llm_temperature=llm_temperature,
|
||||
).run(prompt):
|
||||
send_output(i)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error processing {args.prompt}, error: {traceback.format_exc()}"
|
||||
)
|
||||
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,136 @@
|
||||
from typing import AsyncGenerator
|
||||
import traceback
|
||||
import logging
|
||||
import os
|
||||
import json
|
||||
import sys
|
||||
import uuid
|
||||
import time
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class GaiaAgentServer:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def _get_model_config(self):
|
||||
try:
|
||||
llm_provider = os.getenv("LLM_PROVIDER", "openai")
|
||||
llm_model_name = os.getenv("LLM_MODEL_NAME")
|
||||
llm_api_key = os.getenv("LLM_API_KEY")
|
||||
llm_base_url = os.getenv("LLM_BASE_URL")
|
||||
llm_temperature = float(os.getenv("LLM_TEMPERATURE", 0.0))
|
||||
return {
|
||||
"provider": llm_provider,
|
||||
"model": llm_model_name,
|
||||
"api_key": llm_api_key,
|
||||
"base_url": llm_base_url,
|
||||
"temperature": llm_temperature,
|
||||
}
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f">>> Gaia Agent: GAIA_MODEL_CONFIG is not configured, using LLM"
|
||||
)
|
||||
raise e
|
||||
|
||||
def models(self):
|
||||
model = self._get_model_config()
|
||||
|
||||
return [
|
||||
{
|
||||
"id": f"{model['provider']}/{model['model']}",
|
||||
"name": f"gaia_agent@{model['provider']}/{model['model']}",
|
||||
}
|
||||
]
|
||||
|
||||
async def chat_completions(self, body: dict) -> AsyncGenerator[str, None]:
|
||||
def response_line(line: str, model: str):
|
||||
return {
|
||||
"object": "chat.completion.chunk",
|
||||
"id": str(uuid.uuid4()).replace("-", ""),
|
||||
"choices": [
|
||||
{"index": 0, "delta": {"content": line, "role": "assistant"}}
|
||||
],
|
||||
"created": int(time.time()),
|
||||
"model": model,
|
||||
}
|
||||
|
||||
try:
|
||||
logger.info(f">>> Gaia Agent: body={body}")
|
||||
|
||||
prompt = body["messages"][-1]["content"]
|
||||
model = body["model"].replace("gaia_agent.", "")
|
||||
|
||||
logger.info(f">>> Gaia Agent: prompt={prompt}, model={model}")
|
||||
|
||||
selected_model = self._get_model_config()
|
||||
|
||||
logger.info(f">>> Gaia Agent: Using model configuration: {selected_model}")
|
||||
|
||||
logger.info(f">>> Gaia Agent Python Path: sys.path={sys.path}")
|
||||
|
||||
llm_provider = selected_model.get("provider")
|
||||
llm_model_name = selected_model.get("model")
|
||||
llm_api_key = selected_model.get("api_key")
|
||||
llm_base_url = selected_model.get("base_url")
|
||||
llm_temperature = selected_model.get("temperature", 0.0)
|
||||
|
||||
from examples.gaia.gaia_agent_runner import GaiaAgentRunner
|
||||
|
||||
runner = GaiaAgentRunner(
|
||||
llm_provider=llm_provider,
|
||||
llm_model_name=llm_model_name,
|
||||
llm_base_url=llm_base_url,
|
||||
llm_api_key=llm_api_key,
|
||||
llm_temperature=llm_temperature,
|
||||
)
|
||||
|
||||
logger.info(f">>> Gaia Agent: prompt={prompt}, runner={runner}")
|
||||
|
||||
async for i in runner.run(prompt):
|
||||
line = response_line(i, model)
|
||||
logger.info(f">>> Gaia Agent Line: {line}")
|
||||
yield line
|
||||
|
||||
except Exception as e:
|
||||
emsg = traceback.format_exc()
|
||||
logger.error(f">>> Gaia Agent Error: exception {emsg}")
|
||||
yield response_line(f"Gaia Agent Error: {emsg}", model)
|
||||
|
||||
finally:
|
||||
logger.info(f">>> Gaia Agent Done")
|
||||
|
||||
|
||||
import fastapi
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
app = fastapi.FastAPI()
|
||||
|
||||
from examples.gaia.gaia_agent_server import GaiaAgentServer
|
||||
|
||||
agent_server = GaiaAgentServer()
|
||||
|
||||
|
||||
@app.get("/v1/models")
|
||||
async def models():
|
||||
return agent_server.models()
|
||||
|
||||
|
||||
@app.post("/v1/chat/completions")
|
||||
async def chat_completions(request: fastapi.Request):
|
||||
form_data = await request.json()
|
||||
logger.info(f">>> Gaia Agent Server: form_data={form_data}")
|
||||
|
||||
async def event_generator():
|
||||
async for chunk in agent_server.chat_completions(form_data):
|
||||
# Format as SSE: each line needs to start with "data: " and end with two newlines
|
||||
yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n"
|
||||
|
||||
return StreamingResponse(event_generator(), media_type="text/event-stream")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
|
||||
uvicorn.run("gaia_agent_server:app", host="0.0.0.0", port=8888)
|
||||
@@ -0,0 +1,206 @@
|
||||
{
|
||||
"mcpServers": {
|
||||
"audio": {
|
||||
"command": "python",
|
||||
"args": [
|
||||
"-m",
|
||||
"examples.gaia.mcp_collections.media.audio"
|
||||
],
|
||||
"env": {
|
||||
},
|
||||
"client_session_timeout_seconds": 9999.0
|
||||
},
|
||||
"browser": {
|
||||
"command": "python",
|
||||
"args": [
|
||||
"-m",
|
||||
"examples.gaia.mcp_collections.tools.browser"
|
||||
],
|
||||
"env": {
|
||||
"LLM_MODEL_NAME": "${LLM_MODEL_NAME}",
|
||||
"LLM_API_KEY": "${LLM_API_KEY}",
|
||||
"LLM_BASE_URL": "${LLM_BASE_URL}"
|
||||
},
|
||||
"client_session_timeout_seconds": 9999.0
|
||||
},
|
||||
"chess": {
|
||||
"command": "python",
|
||||
"args": [
|
||||
"-m",
|
||||
"examples.gaia.mcp_collections.tools.playchess"
|
||||
],
|
||||
"env": {
|
||||
},
|
||||
"client_session_timeout_seconds": 9999.0
|
||||
},
|
||||
"code": {
|
||||
"command": "npx",
|
||||
"args": [
|
||||
"-y",
|
||||
"@e2b/mcp-server"
|
||||
],
|
||||
"env": {
|
||||
"E2B_API_KEY": "${E2B_API_KEY}"
|
||||
}
|
||||
},
|
||||
"csv": {
|
||||
"command": "python",
|
||||
"args": [
|
||||
"-m",
|
||||
"examples.gaia.mcp_collections.documents.mscsv"
|
||||
],
|
||||
"env": {
|
||||
},
|
||||
"client_session_timeout_seconds": 9999.0
|
||||
},
|
||||
"docx": {
|
||||
"command": "python",
|
||||
"args": [
|
||||
"-m",
|
||||
"examples.gaia.mcp_collections.documents.msdocx"
|
||||
],
|
||||
"env": {
|
||||
},
|
||||
"client_session_timeout_seconds": 9999.0
|
||||
},
|
||||
"download": {
|
||||
"command": "python",
|
||||
"args": [
|
||||
"-m",
|
||||
"examples.gaia.mcp_collections.tools.download"
|
||||
],
|
||||
"env": {
|
||||
},
|
||||
"client_session_timeout_seconds": 9999.0
|
||||
},
|
||||
"xlsx": {
|
||||
"command": "python",
|
||||
"args": [
|
||||
"-m",
|
||||
"examples.gaia.mcp_collections.documents.msxlsx"
|
||||
],
|
||||
"env": {
|
||||
},
|
||||
"client_session_timeout_seconds": 9999.0
|
||||
},
|
||||
"image": {
|
||||
"command": "python",
|
||||
"args": [
|
||||
"-m",
|
||||
"examples.gaia.mcp_collections.media.image"
|
||||
],
|
||||
"env": {
|
||||
},
|
||||
"client_session_timeout_seconds": 9999.0
|
||||
},
|
||||
"pdf": {
|
||||
"command": "python",
|
||||
"args": [
|
||||
"-m",
|
||||
"examples.gaia.mcp_collections.documents.pdf"
|
||||
],
|
||||
"env": {
|
||||
},
|
||||
"client_session_timeout_seconds": 9999.0
|
||||
},
|
||||
"pptx": {
|
||||
"command": "python",
|
||||
"args": [
|
||||
"-m",
|
||||
"examples.gaia.mcp_collections.documents.mspptx"
|
||||
],
|
||||
"env": {
|
||||
},
|
||||
"client_session_timeout_seconds": 9999.0
|
||||
},
|
||||
"pubchem": {
|
||||
"command": "python",
|
||||
"args": [
|
||||
"-m",
|
||||
"examples.gaia.mcp_collections.tools.pubchem"
|
||||
],
|
||||
"env": {
|
||||
},
|
||||
"client_session_timeout_seconds": 9999.0
|
||||
},
|
||||
"reasoning": {
|
||||
"command": "python",
|
||||
"args": [
|
||||
"-m",
|
||||
"examples.gaia.mcp_collections.intelligence.think"
|
||||
],
|
||||
"env": {
|
||||
},
|
||||
"client_session_timeout_seconds": 9999.0
|
||||
},
|
||||
"search": {
|
||||
"command": "python",
|
||||
"args": [
|
||||
"-m",
|
||||
"examples.gaia.mcp_collections.tools.search"
|
||||
],
|
||||
"env": {
|
||||
"GOOGLE_API_KEY": "${GOOGLE_API_KEY}",
|
||||
"GOOGLE_CSE_ID": "${GOOGLE_CSE_ID}"
|
||||
},
|
||||
"client_session_timeout_seconds": 9999.0
|
||||
},
|
||||
"terminal": {
|
||||
"command": "python",
|
||||
"args": [
|
||||
"-m",
|
||||
"examples.gaia.mcp_collections.tools.terminal"
|
||||
]
|
||||
},
|
||||
"video": {
|
||||
"command": "python",
|
||||
"args": [
|
||||
"-m",
|
||||
"examples.gaia.mcp_collections.media.video"
|
||||
],
|
||||
"env": {
|
||||
},
|
||||
"client_session_timeout_seconds": 9999.0
|
||||
},
|
||||
"wayback": {
|
||||
"command": "python",
|
||||
"args": [
|
||||
"-m",
|
||||
"examples.gaia.mcp_collections.tools.wayback"
|
||||
],
|
||||
"env": {
|
||||
},
|
||||
"client_session_timeout_seconds": 9999.0
|
||||
},
|
||||
"wikipedia": {
|
||||
"command": "python",
|
||||
"args": [
|
||||
"-m",
|
||||
"examples.gaia.mcp_collections.tools.wiki"
|
||||
],
|
||||
"env": {
|
||||
},
|
||||
"client_session_timeout_seconds": 9999.0
|
||||
},
|
||||
"youtube": {
|
||||
"command": "python",
|
||||
"args": [
|
||||
"-m",
|
||||
"examples.gaia.mcp_collections.tools.youtube"
|
||||
],
|
||||
"env": {
|
||||
},
|
||||
"client_session_timeout_seconds": 9999.0
|
||||
},
|
||||
"txt": {
|
||||
"command": "python",
|
||||
"args": [
|
||||
"-m",
|
||||
"examples.gaia.mcp_collections.documents.txt"
|
||||
],
|
||||
"env": {
|
||||
},
|
||||
"client_session_timeout_seconds": 9999.0
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,102 @@
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from aworld.logs.util import Color
|
||||
from examples.gaia.utils import color_log, setup_logger
|
||||
|
||||
|
||||
class ActionArguments(BaseModel):
|
||||
r"""Protocol: MCP Action Arguments"""
|
||||
|
||||
name: str = Field(description="The name of the action")
|
||||
transport: Literal["stdio", "sse"] = Field(default="stdio", description="The transport of the action")
|
||||
workspace: str | None = Field(
|
||||
default=None,
|
||||
description="The workspace of the action."
|
||||
" If not specified or invalid, the workspace will be read from the environment variable AWORLD_WORKSPACE.",
|
||||
)
|
||||
unittest: bool = Field(default=False, description="Whether to run in unittest mode")
|
||||
|
||||
|
||||
class ActionResponse(BaseModel):
|
||||
r"""Protocol: MCP Action Response"""
|
||||
|
||||
success: bool = Field(default=False, description="Whether the action is successfully executed")
|
||||
message: Any = Field(default=None, description="The execution result of the action")
|
||||
metadata: dict[str, Any] = Field(default={}, description="The metadata of the action")
|
||||
|
||||
|
||||
class ActionCollection:
|
||||
r"""Base class for all ActionCollection."""
|
||||
|
||||
server: FastMCP
|
||||
logger: logging.Logger
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
self.unittest = arguments.unittest
|
||||
self.transport = arguments.transport
|
||||
self.supported_extensions = set()
|
||||
|
||||
self.workspace: Path = self._obtain_valid_workspace(arguments.workspace)
|
||||
|
||||
self.logger: logging.Logger = setup_logger(self.__class__.__name__, self.workspace)
|
||||
|
||||
self.server = FastMCP(arguments.name)
|
||||
for tool_name in self.__class__.__dict__:
|
||||
if tool_name.startswith("mcp_") and callable(getattr(self.__class__, tool_name)):
|
||||
tool = getattr(self, tool_name)
|
||||
self.server.add_tool(tool, description=tool.__doc__)
|
||||
|
||||
def run(self) -> None:
|
||||
if not self.unittest:
|
||||
self.server.run(transport=self.transport)
|
||||
|
||||
def _color_log(self, value: str, color: Color = None, level: str = "info"):
|
||||
return color_log(self.logger, value, color, level=level)
|
||||
|
||||
def _obtain_valid_workspace(self, workspace: str | None = None) -> Path:
|
||||
r"""
|
||||
Obtain a valid workspace path.
|
||||
Priority:
|
||||
1. user defined workspace
|
||||
2. environment variable AWORLD_WORKSPACE
|
||||
3. home directory
|
||||
"""
|
||||
path = Path(workspace) if workspace else os.getenv("AWORLD_WORKSPACE", "~")
|
||||
if path and path.expanduser().is_dir():
|
||||
return path.expanduser().resolve()
|
||||
|
||||
# self._color_log("Invalid workspace path, using home directory instead.", Color.yellow)
|
||||
return Path.home().expanduser().resolve()
|
||||
|
||||
def _validate_file_path(self, file_path: str) -> Path:
|
||||
"""Validate and resolve file path. Rely on the predefined supported_extensions class variable.
|
||||
|
||||
Args:
|
||||
file_path: Path to the document or media file
|
||||
|
||||
Returns:
|
||||
Resolved Path object
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If file doesn't exist
|
||||
ValueError: If file type is not supported
|
||||
"""
|
||||
path = Path(file_path).expanduser()
|
||||
if not path.is_absolute():
|
||||
path = self.workspace / path
|
||||
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"File not found: {path}")
|
||||
|
||||
if path.suffix.lower() not in self.supported_extensions:
|
||||
raise ValueError(
|
||||
f"Unsupported file type: {path.suffix}. Supported types: {', '.join(self.supported_extensions)}"
|
||||
)
|
||||
|
||||
return path
|
||||
@@ -0,0 +1,24 @@
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class DocumentMetadata(BaseModel):
|
||||
"""Metadata extracted from document processing."""
|
||||
|
||||
file_name: str = Field(description="Original file name")
|
||||
file_size: int = Field(description="File size in bytes")
|
||||
file_type: str = Field(description="Document file type/extension")
|
||||
absolute_path: str = Field(description="Absolute path to the document file")
|
||||
page_count: int | None = Field(default=None, description="Number of pages in document")
|
||||
processing_time: float = Field(
|
||||
description="Time taken to process the document in seconds", deprecated=True, exclude=True
|
||||
)
|
||||
extracted_images: list[str] = Field(default_factory=list, description="Paths to extracted image files")
|
||||
extracted_media: list[dict[str, str]] = Field(
|
||||
default_factory=list, description="list of extracted media files with type and path"
|
||||
)
|
||||
output_format: str = Field(description="Format of the extracted content")
|
||||
llm_enhanced: bool = Field(default=False, description="Whether LLM enhancement was used", exclude=True)
|
||||
ocr_applied: bool = Field(default=False, description="Whether OCR was applied", exclude=True)
|
||||
extracted_text_file_path: str | None = Field(
|
||||
default=None, description="Absolute path to the extracted text file (if applicable)"
|
||||
)
|
||||
@@ -0,0 +1,376 @@
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
import chardet
|
||||
import pandas as pd
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import Field
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from aworld.logs.util import Color
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
from examples.gaia.mcp_collections.documents.models import DocumentMetadata
|
||||
|
||||
|
||||
class CSVExtractionCollection(ActionCollection):
|
||||
"""MCP service for CSV document content extraction using pandas.
|
||||
|
||||
Supports extraction from CSV files with various encodings and delimiters.
|
||||
Provides LLM-friendly text output with structured metadata and data analysis.
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
self._media_output_dir = self.workspace / "extracted_media"
|
||||
self._media_output_dir.mkdir(exist_ok=True)
|
||||
|
||||
self.supported_extensions: set = {".csv", ".tsv", ".txt"}
|
||||
|
||||
self._color_log("CSV Extraction Service initialized", Color.green, "debug")
|
||||
self._color_log(f"Media output directory: {self._media_output_dir}", Color.blue, "debug")
|
||||
|
||||
def _detect_encoding(self, file_path: Path) -> str:
|
||||
"""Detect file encoding using chardet.
|
||||
|
||||
Args:
|
||||
file_path: Path to the CSV file
|
||||
|
||||
Returns:
|
||||
Detected encoding string
|
||||
"""
|
||||
try:
|
||||
with open(file_path, "rb") as f:
|
||||
raw_data = f.read(10000) # Read first 10KB for detection
|
||||
result = chardet.detect(raw_data)
|
||||
encoding = result.get("encoding", "utf-8")
|
||||
confidence = result.get("confidence", 0)
|
||||
|
||||
self._color_log(f"Detected encoding: {encoding} (confidence: {confidence:.2f})", Color.blue)
|
||||
return encoding if confidence > 0.7 else "utf-8"
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Encoding detection failed: {e}, using utf-8")
|
||||
return "utf-8"
|
||||
|
||||
def _detect_delimiter(self, file_path: Path, encoding: str) -> str:
|
||||
"""Detect CSV delimiter by analyzing the first few lines.
|
||||
|
||||
Args:
|
||||
file_path: Path to the CSV file
|
||||
encoding: File encoding
|
||||
|
||||
Returns:
|
||||
Detected delimiter character
|
||||
"""
|
||||
try:
|
||||
with open(file_path, "r", encoding=encoding) as f:
|
||||
sample = f.read(1024) # Read first 1KB
|
||||
|
||||
# Common delimiters to test
|
||||
delimiters = [",", ";", "\t", "|", ":"]
|
||||
delimiter_counts = {}
|
||||
|
||||
for delimiter in delimiters:
|
||||
count = sample.count(delimiter)
|
||||
if count > 0:
|
||||
delimiter_counts[delimiter] = count
|
||||
|
||||
if delimiter_counts:
|
||||
detected_delimiter = max(delimiter_counts, key=delimiter_counts.get)
|
||||
self._color_log(f"Detected delimiter: '{detected_delimiter}'", Color.blue)
|
||||
return detected_delimiter
|
||||
else:
|
||||
return ","
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Delimiter detection failed: {e}, using comma")
|
||||
return ","
|
||||
|
||||
def _extract_csv_content(
|
||||
self, file_path: Path, max_rows: int | None = None, encoding: str | None = None, delimiter: str | None = None
|
||||
) -> dict[str, Any]:
|
||||
"""Extract content from CSV file using pandas.
|
||||
|
||||
Args:
|
||||
file_path: Path to the CSV file
|
||||
max_rows: Maximum number of rows to read
|
||||
encoding: File encoding (auto-detected if None)
|
||||
delimiter: CSV delimiter (auto-detected if None)
|
||||
|
||||
Returns:
|
||||
Dictionary containing extracted content and metadata
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Auto-detect encoding and delimiter if not provided
|
||||
if encoding is None:
|
||||
encoding = self._detect_encoding(file_path)
|
||||
if delimiter is None:
|
||||
delimiter = self._detect_delimiter(file_path, encoding)
|
||||
|
||||
try:
|
||||
# Read CSV with pandas
|
||||
df = pd.read_csv(file_path, encoding=encoding, delimiter=delimiter, nrows=max_rows, low_memory=False)
|
||||
|
||||
# Get full file info for metadata
|
||||
full_df_info = pd.read_csv(
|
||||
file_path,
|
||||
encoding=encoding,
|
||||
delimiter=delimiter,
|
||||
nrows=0, # Just get headers and shape info
|
||||
)
|
||||
|
||||
# Count total rows efficiently
|
||||
total_rows = sum(1 for _ in open(file_path, "r", encoding=encoding)) - 1 # Subtract header
|
||||
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
return {
|
||||
"dataframe": df,
|
||||
"total_rows": total_rows,
|
||||
"total_columns": len(full_df_info.columns),
|
||||
"columns": list(df.columns),
|
||||
"encoding": encoding,
|
||||
"delimiter": delimiter,
|
||||
"processing_time": processing_time,
|
||||
"data_types": df.dtypes.to_dict(),
|
||||
"memory_usage": df.memory_usage(deep=True).sum(),
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to read CSV file: {e}")
|
||||
raise
|
||||
|
||||
def _format_content_for_llm(self, df: pd.DataFrame, output_format: str, include_stats: bool = True) -> str:
|
||||
"""Format extracted CSV content to be LLM-friendly.
|
||||
|
||||
Args:
|
||||
df: Pandas DataFrame with CSV data
|
||||
output_format: Desired output format
|
||||
include_stats: Whether to include statistical summary
|
||||
|
||||
Returns:
|
||||
Formatted content string
|
||||
"""
|
||||
if output_format.lower() == "markdown":
|
||||
# Convert to markdown table
|
||||
content = df.to_markdown(index=False, tablefmt="github")
|
||||
|
||||
if include_stats:
|
||||
# Add statistical summary
|
||||
stats_content = "\n\n## Data Summary\n\n"
|
||||
stats_content += f"- **Rows**: {len(df)}\n"
|
||||
stats_content += f"- **Columns**: {len(df.columns)}\n"
|
||||
stats_content += f"- **Column Names**: {', '.join(df.columns)}\n\n"
|
||||
|
||||
# Add data types info
|
||||
stats_content += "### Column Data Types\n\n"
|
||||
for col, dtype in df.dtypes.items():
|
||||
stats_content += f"- **{col}**: {dtype}\n"
|
||||
|
||||
# Add basic statistics for numeric columns
|
||||
numeric_cols = df.select_dtypes(include=["number"]).columns
|
||||
if len(numeric_cols) > 0:
|
||||
stats_content += "\n### Numeric Column Statistics\n\n"
|
||||
stats_df = df[numeric_cols].describe()
|
||||
stats_content += stats_df.to_markdown(tablefmt="github")
|
||||
|
||||
content += stats_content
|
||||
|
||||
elif output_format.lower() == "json":
|
||||
# Convert to JSON with metadata
|
||||
data_dict = {
|
||||
"data": df.to_dict(orient="records"),
|
||||
"metadata": {
|
||||
"rows": len(df),
|
||||
"columns": len(df.columns),
|
||||
"column_names": list(df.columns),
|
||||
"data_types": {col: str(dtype) for col, dtype in df.dtypes.items()},
|
||||
},
|
||||
}
|
||||
if include_stats:
|
||||
numeric_cols = df.select_dtypes(include=["number"]).columns
|
||||
if len(numeric_cols) > 0:
|
||||
data_dict["statistics"] = df[numeric_cols].describe().to_dict()
|
||||
|
||||
content = json.dumps(data_dict, indent=2, default=str)
|
||||
|
||||
elif output_format.lower() == "html":
|
||||
# Convert to HTML table
|
||||
content = df.to_html(index=False, classes="table table-striped")
|
||||
|
||||
else:
|
||||
# Plain text format
|
||||
content = df.to_string(index=False)
|
||||
|
||||
return content
|
||||
|
||||
def mcp_extract_csv_content(
|
||||
self,
|
||||
file_path: str = Field(description="Path to the CSV document file to extract content from"),
|
||||
output_format: Literal["markdown", "json", "html", "text"] = Field(
|
||||
default="markdown", description="Output format: 'markdown', 'json', 'html', or 'text'"
|
||||
),
|
||||
max_rows: int | None = Field(default=None, description="Maximum number of rows to read (None for all rows)"),
|
||||
include_statistics: bool = Field(default=True, description="Whether to include statistical summary in output"),
|
||||
generate_visualizations: bool = Field(
|
||||
default=False, description="Whether to generate and save data visualizations"
|
||||
),
|
||||
encoding: str | None = Field(default=None, description="File encoding (auto-detected if None)"),
|
||||
delimiter: str | None = Field(default=None, description="CSV delimiter (auto-detected if None)"),
|
||||
) -> ActionResponse:
|
||||
"""Extract content from CSV documents using pandas.
|
||||
|
||||
This tool provides comprehensive CSV document content extraction with support for:
|
||||
- CSV, TSV, and delimited text files
|
||||
- Automatic encoding and delimiter detection
|
||||
- Statistical analysis and data profiling
|
||||
- Multiple output formats (Markdown, JSON, HTML, Text)
|
||||
- Optional data visualizations
|
||||
- Memory-efficient processing for large files
|
||||
|
||||
Args:
|
||||
file_path: Path to the CSV file
|
||||
output_format: Desired output format
|
||||
max_rows: Maximum rows to process (None for all)
|
||||
include_statistics: Include statistical summary
|
||||
generate_visualizations: Generate data visualizations
|
||||
encoding: File encoding (auto-detected if None)
|
||||
delimiter: CSV delimiter (auto-detected if None)
|
||||
|
||||
Returns:
|
||||
ActionResponse with extracted content, metadata, and optional visualizations
|
||||
"""
|
||||
try:
|
||||
# Handle FieldInfo objects from pydantic
|
||||
if isinstance(file_path, FieldInfo):
|
||||
file_path = file_path.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
if isinstance(max_rows, FieldInfo):
|
||||
max_rows = max_rows.default
|
||||
if isinstance(include_statistics, FieldInfo):
|
||||
include_statistics = include_statistics.default
|
||||
if isinstance(generate_visualizations, FieldInfo):
|
||||
generate_visualizations = generate_visualizations.default
|
||||
if isinstance(encoding, FieldInfo):
|
||||
encoding = encoding.default
|
||||
if isinstance(delimiter, FieldInfo):
|
||||
delimiter = delimiter.default
|
||||
|
||||
# Validate input file
|
||||
file_path: Path = self._validate_file_path(file_path)
|
||||
self._color_log(f"Processing CSV file: {file_path.name}", Color.cyan)
|
||||
|
||||
# Extract CSV content
|
||||
extraction_result = self._extract_csv_content(
|
||||
file_path, max_rows=max_rows, encoding=encoding, delimiter=delimiter
|
||||
)
|
||||
|
||||
df: pd.DataFrame = extraction_result["dataframe"]
|
||||
|
||||
# Format content for LLM consumption
|
||||
formatted_content = self._format_content_for_llm(df, output_format, include_stats=include_statistics)
|
||||
|
||||
# Prepare metadata
|
||||
file_stats = file_path.stat()
|
||||
document_metadata = DocumentMetadata(
|
||||
file_name=file_path.name,
|
||||
file_size=file_stats.st_size,
|
||||
file_type=file_path.suffix.lower(),
|
||||
absolute_path=str(file_path.absolute()),
|
||||
page_count=None, # Not applicable for CSV
|
||||
processing_time=extraction_result["processing_time"],
|
||||
extracted_images=[], # CSV files don't contain images
|
||||
extracted_media=None,
|
||||
output_format=output_format,
|
||||
llm_enhanced=False,
|
||||
ocr_applied=False,
|
||||
)
|
||||
|
||||
# Add CSV-specific metadata
|
||||
csv_metadata = {
|
||||
"total_rows": extraction_result["total_rows"],
|
||||
"total_columns": extraction_result["total_columns"],
|
||||
"rows_processed": len(df),
|
||||
"columns_processed": len(df.columns),
|
||||
"column_names": extraction_result["columns"],
|
||||
"data_types": {k: str(v) for k, v in extraction_result["data_types"].items()},
|
||||
"encoding": extraction_result["encoding"],
|
||||
"delimiter": extraction_result["delimiter"],
|
||||
"memory_usage_bytes": int(extraction_result["memory_usage"]),
|
||||
}
|
||||
|
||||
# Merge metadata
|
||||
final_metadata = {**document_metadata.model_dump(), **csv_metadata}
|
||||
|
||||
self._color_log(
|
||||
f"Successfully extracted CSV content from {file_path.name} "
|
||||
f"({extraction_result['total_rows']} rows, {extraction_result['total_columns']} columns",
|
||||
Color.green,
|
||||
)
|
||||
|
||||
return ActionResponse(success=True, message=formatted_content, metadata=final_metadata)
|
||||
|
||||
except FileNotFoundError as e:
|
||||
self.logger.error(f"File not found: {str(e)}")
|
||||
return ActionResponse(
|
||||
success=False, message=f"File not found: {str(e)}", metadata={"error_type": "file_not_found"}
|
||||
)
|
||||
except ValueError as e:
|
||||
self.logger.error(f"Invalid input: {str(e)}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}",
|
||||
metadata={"error_type": "invalid_input"},
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"CSV extraction failed: {str(e)}: {traceback.format_exc()}")
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=f"CSV extraction failed: {str(e)}",
|
||||
metadata={"error_type": "extraction_error"},
|
||||
)
|
||||
|
||||
def mcp_list_supported_formats(self) -> ActionResponse:
|
||||
"""List all supported CSV formats for extraction.
|
||||
|
||||
Returns:
|
||||
ActionResponse with list of supported file formats and their descriptions
|
||||
"""
|
||||
supported_formats = {
|
||||
"CSV": "Comma-Separated Values files (.csv)",
|
||||
"TSV": "Tab-Separated Values files (.tsv)",
|
||||
"TXT": "Delimited text files (.txt)",
|
||||
}
|
||||
|
||||
format_list = "\n".join(
|
||||
[f"**{format_name}**: {description}" for format_name, description in supported_formats.items()]
|
||||
)
|
||||
|
||||
return ActionResponse(
|
||||
success=True,
|
||||
message=f"Supported CSV formats:\n\n{format_list}",
|
||||
metadata={"supported_formats": list(supported_formats.keys()), "total_formats": len(supported_formats)},
|
||||
)
|
||||
|
||||
|
||||
# Example usage and entry point
|
||||
if __name__ == "__main__":
|
||||
load_dotenv()
|
||||
|
||||
# Default arguments for testing
|
||||
args = ActionArguments(
|
||||
name="csv_extraction_service",
|
||||
transport="stdio",
|
||||
workspace=os.getenv("AWORLD_WORKSPACE", "~"),
|
||||
)
|
||||
|
||||
# Initialize and run the CSV extraction service
|
||||
try:
|
||||
service = CSVExtractionCollection(args)
|
||||
service.run()
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}: {traceback.format_exc()}")
|
||||
@@ -0,0 +1,617 @@
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
from docx import Document
|
||||
from docx.document import Document as DocumentType
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import Field
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from aworld.logs.util import Color
|
||||
from examples.gaia.mcp_collections.base import ActionArguments, ActionCollection, ActionResponse
|
||||
from examples.gaia.mcp_collections.documents.models import DocumentMetadata
|
||||
|
||||
|
||||
class DOCXExtractionCollection(ActionCollection):
|
||||
"""MCP service for DOCX/DOC document content extraction using python-docx.
|
||||
|
||||
Supports extraction from DOCX and DOC files with comprehensive content parsing.
|
||||
Provides LLM-friendly text output with structured metadata and media file handling.
|
||||
"""
|
||||
|
||||
def __init__(self, arguments: ActionArguments) -> None:
|
||||
super().__init__(arguments)
|
||||
self._media_output_dir = self.workspace / "extracted_media"
|
||||
self._media_output_dir.mkdir(exist_ok=True)
|
||||
|
||||
self.supported_extensions = {".docx", ".doc"}
|
||||
|
||||
self._color_log("DOCX Extraction Service initialized", Color.green, "debug")
|
||||
self._color_log(f"Media output directory: {self._media_output_dir}", Color.blue, "debug")
|
||||
|
||||
def _extract_images_from_docx(self, file_path: Path, file_stem: str) -> list[dict[str, str]]:
|
||||
"""Extract embedded images from DOCX file.
|
||||
|
||||
Args:
|
||||
file_path: Path to the DOCX file
|
||||
file_stem: Base name for saving files
|
||||
|
||||
Returns:
|
||||
List of dictionaries containing image file paths and metadata
|
||||
"""
|
||||
saved_media = []
|
||||
|
||||
try:
|
||||
# DOCX files are ZIP archives
|
||||
with zipfile.ZipFile(file_path, "r") as docx_zip:
|
||||
# Look for media files in the word/media/ directory
|
||||
media_files = [f for f in docx_zip.namelist() if f.startswith("word/media/")]
|
||||
|
||||
for idx, media_file in enumerate(media_files):
|
||||
try:
|
||||
# Extract file extension and create appropriate filename
|
||||
original_name = Path(media_file).name
|
||||
file_extension = Path(media_file).suffix
|
||||
|
||||
# Generate unique filename
|
||||
media_filename = f"{file_stem}_media_{idx}_{original_name}"
|
||||
media_path = self._media_output_dir / media_filename
|
||||
|
||||
# Extract and save the media file
|
||||
with docx_zip.open(media_file) as source:
|
||||
with open(media_path, "wb") as target:
|
||||
target.write(source.read())
|
||||
|
||||
# Determine media type based on extension
|
||||
media_type = "image"
|
||||
if file_extension.lower() in [".mp3", ".wav", ".m4a", ".ogg"]:
|
||||
media_type = "audio"
|
||||
elif file_extension.lower() in [".mp4", ".avi", ".mov", ".wmv"]:
|
||||
media_type = "video"
|
||||
elif file_extension.lower() in [".png", ".jpg", ".jpeg", ".gif", ".bmp", ".tiff"]:
|
||||
media_type = "image"
|
||||
else:
|
||||
media_type = "other"
|
||||
|
||||
saved_media.append(
|
||||
{
|
||||
"type": media_type,
|
||||
"path": str(media_path),
|
||||
"filename": media_filename,
|
||||
"original_name": original_name,
|
||||
"size_bytes": media_path.stat().st_size,
|
||||
}
|
||||
)
|
||||
|
||||
self._color_log(f"Extracted {media_type}: {media_filename}", Color.blue)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to extract media file {media_file}: {e}")
|
||||
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Could not extract media from DOCX: {e}")
|
||||
|
||||
return saved_media
|
||||
|
||||
def _extract_document_structure(self, doc: DocumentType) -> dict[str, Any]:
|
||||
"""Extract document structure and metadata.
|
||||
|
||||
Args:
|
||||
doc: python-docx Document object
|
||||
|
||||
Returns:
|
||||
Dictionary containing document structure information
|
||||
"""
|
||||
structure = {
|
||||
"paragraphs_count": len(doc.paragraphs),
|
||||
"tables_count": len(doc.tables),
|
||||
"sections_count": len(doc.sections),
|
||||
"styles_used": [],
|
||||
"has_headers_footers": False,
|
||||
"page_count": None, # Not directly available in python-docx
|
||||
}
|
||||
|
||||
# Collect unique styles used in the document
|
||||
styles_used = set()
|
||||
for paragraph in doc.paragraphs:
|
||||
if paragraph.style and paragraph.style.name:
|
||||
styles_used.add(paragraph.style.name)
|
||||
|
||||
structure["styles_used"] = list(styles_used)
|
||||
|
||||
# Check for headers and footers
|
||||
for section in doc.sections:
|
||||
if (section.header.paragraphs and any(p.text.strip() for p in section.header.paragraphs)) or (
|
||||
section.footer.paragraphs and any(p.text.strip() for p in section.footer.paragraphs)
|
||||
):
|
||||
structure["has_headers_footers"] = True
|
||||
break
|
||||
|
||||
return structure
|
||||
|
||||
def _extract_tables_content(self, doc: DocumentType) -> list[dict[str, Any]]:
|
||||
"""Extract content from all tables in the document.
|
||||
|
||||
Args:
|
||||
doc: python-docx Document object
|
||||
|
||||
Returns:
|
||||
List of dictionaries containing table data
|
||||
"""
|
||||
tables_data = []
|
||||
|
||||
for table_idx, table in enumerate(doc.tables):
|
||||
table_data = {
|
||||
"table_index": table_idx,
|
||||
"rows_count": len(table.rows),
|
||||
"columns_count": len(table.columns) if table.rows else 0,
|
||||
"data": [],
|
||||
}
|
||||
|
||||
# Extract table content
|
||||
for _, row in enumerate(table.rows):
|
||||
row_data = []
|
||||
for cell in row.cells:
|
||||
cell_text = cell.text.strip()
|
||||
row_data.append(cell_text)
|
||||
table_data["data"].append(row_data)
|
||||
|
||||
tables_data.append(table_data)
|
||||
|
||||
return tables_data
|
||||
|
||||
def _convert_doc_to_docx(self, doc_path: Path) -> Path:
|
||||
"""Convert .doc file to .docx using LibreOffice.
|
||||
|
||||
Args:
|
||||
doc_path: Path to the .doc file
|
||||
|
||||
Returns:
|
||||
Path to the converted .docx file
|
||||
"""
|
||||
import subprocess
|
||||
|
||||
output_dir = self._media_output_dir
|
||||
docx_path = output_dir / f"{doc_path.stem}.docx"
|
||||
|
||||
# Check if already converted
|
||||
if docx_path.exists():
|
||||
self._color_log(f"Using existing converted file: {docx_path.name}", Color.blue)
|
||||
return docx_path
|
||||
|
||||
self._color_log(f"Converting .doc to .docx: {doc_path.name}", Color.yellow)
|
||||
|
||||
cmd = ["libreoffice", "--headless", "--convert-to", "docx", "--outdir", str(output_dir), str(doc_path)]
|
||||
|
||||
try:
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, check=True, timeout=60) # pylint: disable=W0612
|
||||
if docx_path.exists():
|
||||
self._color_log(f"Conversion successful: {docx_path.name}", Color.green)
|
||||
return docx_path
|
||||
else:
|
||||
raise RuntimeError("Conversion completed but output file not found")
|
||||
|
||||
except subprocess.CalledProcessError as e:
|
||||
self.logger.error(f"LibreOffice conversion failed: {e.stderr}")
|
||||
raise RuntimeError(f"Failed to convert .doc to .docx: {e.stderr}") from e
|
||||
except subprocess.TimeoutExpired as e:
|
||||
self.logger.error("LibreOffice conversion timed out")
|
||||
raise RuntimeError("Conversion timed out after 60 seconds") from e
|
||||
|
||||
# Modify the _extract_content_from_docx method to handle .doc files
|
||||
def _extract_content_from_docx(
|
||||
self, file_path: Path, extract_tables: bool = True, extract_headers_footers: bool = True
|
||||
) -> dict[str, Any]:
|
||||
"""Extract content from DOCX file using python-docx.
|
||||
|
||||
Args:
|
||||
file_path: Path to the DOCX/DOC file
|
||||
extract_tables: Whether to extract table content
|
||||
extract_headers_footers: Whether to extract headers and footers
|
||||
|
||||
Returns:
|
||||
Dictionary containing extracted content and metadata
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Convert .doc to .docx if needed
|
||||
if file_path.suffix.lower() == ".doc":
|
||||
file_path = self._convert_doc_to_docx(file_path)
|
||||
|
||||
try:
|
||||
# Load the document
|
||||
doc = Document(str(file_path))
|
||||
|
||||
# Extract main text content
|
||||
paragraphs = []
|
||||
for paragraph in doc.paragraphs:
|
||||
if paragraph.text.strip(): # Skip empty paragraphs
|
||||
para_data = {"text": paragraph.text, "style": paragraph.style.name if paragraph.style else "Normal"}
|
||||
paragraphs.append(para_data)
|
||||
|
||||
# Extract document structure
|
||||
structure = self._extract_document_structure(doc)
|
||||
|
||||
# Extract tables if requested
|
||||
tables = []
|
||||
if extract_tables:
|
||||
tables = self._extract_tables_content(doc)
|
||||
|
||||
# Extract headers and footers if requested
|
||||
headers_footers = []
|
||||
if extract_headers_footers:
|
||||
for section_idx, section in enumerate(doc.sections):
|
||||
# Extract header content
|
||||
header_text = []
|
||||
for para in section.header.paragraphs:
|
||||
if para.text.strip():
|
||||
header_text.append(para.text)
|
||||
|
||||
# Extract footer content
|
||||
footer_text = []
|
||||
for para in section.footer.paragraphs:
|
||||
if para.text.strip():
|
||||
footer_text.append(para.text)
|
||||
|
||||
if header_text or footer_text:
|
||||
headers_footers.append(
|
||||
{"section_index": section_idx, "header": header_text, "footer": footer_text}
|
||||
)
|
||||
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
return {
|
||||
"paragraphs": paragraphs,
|
||||
"tables": tables,
|
||||
"headers_footers": headers_footers,
|
||||
"structure": structure,
|
||||
"processing_time": processing_time,
|
||||
"word_count": sum(len(p["text"].split()) for p in paragraphs),
|
||||
"character_count": sum(len(p["text"]) for p in paragraphs),
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to extract content from DOCX: {e}")
|
||||
raise
|
||||
|
||||
def _format_content_for_llm(
|
||||
self, extraction_result: dict[str, Any], output_format: str, include_structure: bool = True
|
||||
) -> str:
|
||||
"""Format extracted DOCX content to be LLM-friendly.
|
||||
|
||||
Args:
|
||||
extraction_result: Dictionary containing extracted content
|
||||
output_format: Desired output format
|
||||
include_structure: Whether to include document structure information
|
||||
|
||||
Returns:
|
||||
Formatted content string
|
||||
"""
|
||||
if output_format.lower() == "markdown":
|
||||
content_parts = []
|
||||
|
||||
# Add document structure info if requested
|
||||
if include_structure:
|
||||
structure = extraction_result["structure"]
|
||||
content_parts.append("# Document Structure\n")
|
||||
content_parts.append(f"- **Paragraphs**: {structure['paragraphs_count']}\n")
|
||||
content_parts.append(f"- **Tables**: {structure['tables_count']}\n")
|
||||
content_parts.append(f"- **Sections**: {structure['sections_count']}\n")
|
||||
content_parts.append(f"- **Word Count**: {extraction_result['word_count']}\n")
|
||||
content_parts.append(f"- **Character Count**: {extraction_result['character_count']}\n")
|
||||
if structure["styles_used"]:
|
||||
content_parts.append(f"- **Styles Used**: {', '.join(structure['styles_used'])}\n")
|
||||
content_parts.append("\n---\n\n")
|
||||
|
||||
# Add main content
|
||||
content_parts.append("# Document Content\n\n")
|
||||
|
||||
# Add paragraphs
|
||||
for para in extraction_result["paragraphs"]:
|
||||
# Format based on style
|
||||
text = para["text"]
|
||||
style = para["style"]
|
||||
|
||||
if "Heading" in style:
|
||||
# Convert heading styles to markdown headers
|
||||
if "Heading 1" in style:
|
||||
content_parts.append(f"# {text}\n\n")
|
||||
elif "Heading 2" in style:
|
||||
content_parts.append(f"## {text}\n\n")
|
||||
elif "Heading 3" in style:
|
||||
content_parts.append(f"### {text}\n\n")
|
||||
else:
|
||||
content_parts.append(f"#### {text}\n\n")
|
||||
else:
|
||||
content_parts.append(f"{text}\n\n")
|
||||
|
||||
# Add tables
|
||||
if extraction_result["tables"]:
|
||||
content_parts.append("\n## Tables\n\n")
|
||||
for table_idx, table in enumerate(extraction_result["tables"]):
|
||||
content_parts.append(f"### Table {table_idx + 1}\n\n")
|
||||
|
||||
if table["data"]:
|
||||
# Create markdown table
|
||||
headers = table["data"][0] if table["data"] else []
|
||||
if headers:
|
||||
content_parts.append("| " + " | ".join(headers) + " |\n")
|
||||
content_parts.append("|" + "---|" * len(headers) + "\n")
|
||||
|
||||
for row in table["data"][1:]:
|
||||
content_parts.append("| " + " | ".join(row) + " |\n")
|
||||
content_parts.append("\n")
|
||||
|
||||
# Add headers and footers
|
||||
if extraction_result["headers_footers"]:
|
||||
content_parts.append("\n## Headers and Footers\n\n")
|
||||
for hf in extraction_result["headers_footers"]:
|
||||
if hf["header"]:
|
||||
content_parts.append(f"**Header (Section {hf['section_index'] + 1}):**\n")
|
||||
for header_line in hf["header"]:
|
||||
content_parts.append(f"{header_line}\n")
|
||||
content_parts.append("\n")
|
||||
|
||||
if hf["footer"]:
|
||||
content_parts.append(f"**Footer (Section {hf['section_index'] + 1}):**\n")
|
||||
for footer_line in hf["footer"]:
|
||||
content_parts.append(f"{footer_line}\n")
|
||||
content_parts.append("\n")
|
||||
|
||||
return "".join(content_parts)
|
||||
|
||||
elif output_format.lower() == "json":
|
||||
# Return structured JSON
|
||||
return json.dumps(extraction_result, indent=2, ensure_ascii=False)
|
||||
|
||||
elif output_format.lower() == "html":
|
||||
# Convert to HTML
|
||||
html_parts = ["<html><body>"]
|
||||
|
||||
if include_structure:
|
||||
html_parts.append("<h1>Document Structure</h1>")
|
||||
structure = extraction_result["structure"]
|
||||
html_parts.append(f"<p><strong>Paragraphs:</strong> {structure['paragraphs_count']}</p>")
|
||||
html_parts.append(f"<p><strong>Tables:</strong> {structure['tables_count']}</p>")
|
||||
html_parts.append(f"<p><strong>Word Count:</strong> {extraction_result['word_count']}</p>")
|
||||
html_parts.append("<hr>")
|
||||
|
||||
html_parts.append("<h1>Document Content</h1>")
|
||||
|
||||
# Add paragraphs
|
||||
for para in extraction_result["paragraphs"]:
|
||||
text = para["text"].replace("&", "&").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()}")
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user