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:
+21
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2025 [Inclusion AI]
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
+197
@@ -0,0 +1,197 @@
|
||||
<div align="center">
|
||||
|
||||
# VirtualPC MCP Server (Incubating)
|
||||
|
||||
*A unified MCP tool runtime environment based on Debian with session-level environment isolation, environment state persistence, real-time UI visualization, distributed architecture, and extensibility*
|
||||
|
||||
[![License: MIT][license-image]][license-url]
|
||||
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
|
||||
[中文版](./README_zh.md) | [Quick Start](#quick-start) | [Development](#development) | [Contributing](#contributing)
|
||||
|
||||
</div>
|
||||
|
||||
---
|
||||
|
||||
## 1. Overview
|
||||
|
||||
VirtualPC MCP Server is a comprehensive MCP (Model Context Protocol) tool runtime environment designed to provide a unified, isolated, and scalable execution environment for AI agents. Built on Debian, it offers session-level environment isolation, persistent state management across multiple sessions, and real-time visualization capabilities.
|
||||
|
||||
### 1.1 Features
|
||||
|
||||
- **Session-Level Environment Isolation**: Each MCP session operates within its own isolated environment
|
||||
- **Multi-Session State Persistence**: Maintains environment state across multiple MCP sessions
|
||||
- **Real-Time UI Visualization**: Live monitoring and visualization of Agent MCP operations
|
||||
- **Distributed Architecture**: Supports both local Docker and Kubernetes cluster deployments
|
||||
- **Extensible Runtime**: Modular design enabling seamless integration of new MCP tool servers
|
||||
|
||||
## 2. Quick Start
|
||||
|
||||
This project supports both local Docker deployment (optimal for demos and debugging) and Kubernetes cluster deployment (recommended for production and RL training).
|
||||
|
||||
### 2.1 Local Docker Deployment
|
||||
|
||||
#### Prerequisites
|
||||
|
||||
Ensure Docker and Docker Compose are properly installed and operational:
|
||||
|
||||
```bash
|
||||
# Verify Docker installation
|
||||
docker --version
|
||||
docker compose --version
|
||||
|
||||
# Verify Docker daemon is running
|
||||
docker ps
|
||||
docker compose ps
|
||||
```
|
||||
|
||||
**Step 1: Configure Environment and Prepare Gaia Dataset**
|
||||
|
||||
1. Copy the environment template and configure your settings:
|
||||
|
||||
```bash
|
||||
cp ./gaia-mcp-server/mcp_servers/.env_template ./gaia-mcp-server/mcp_servers/.env
|
||||
```
|
||||
|
||||
Edit `./gaia-mcp-server/mcp_servers/.env` with your specific configuration values.
|
||||
|
||||
2. Download the [gaia_dataset](https://huggingface.co/datasets/gaia-benchmark/GAIA) from Hugging Face and place it in `./gaia-mcp-server/docker/gaia_dataset`
|
||||
|
||||
**Step 2: Launch VirtualPC MCP Server**
|
||||
|
||||
```bash
|
||||
sh run-docker.sh
|
||||
```
|
||||
|
||||
Monitor the terminal output for any errors during startup.
|
||||
|
||||
Generate a local bearer token and copy the printed value into `<local-debug-jwt>` below. If you changed `MCP_GATEWAY_TOKEN_SECRET` in `virtualpc-mcp/docker-compose.yaml`, export the same value before running this command.
|
||||
|
||||
```bash
|
||||
python - <<'PY'
|
||||
import base64, hashlib, hmac, json, os, time
|
||||
|
||||
def part(value):
|
||||
raw = json.dumps(value, separators=(",", ":")).encode()
|
||||
return base64.urlsafe_b64encode(raw).rstrip(b"=").decode()
|
||||
|
||||
signing_input = ".".join([
|
||||
part({"alg": "HS256", "typ": "JWT"}),
|
||||
part({"app": "local_debug", "version": 1, "time": time.time()}),
|
||||
])
|
||||
secret = os.getenv("MCP_GATEWAY_TOKEN_SECRET", "123321").encode()
|
||||
signature = hmac.new(secret, signing_input.encode(), hashlib.sha256).digest()
|
||||
print(f"{signing_input}.{base64.urlsafe_b64encode(signature).rstrip(b'=').decode()}")
|
||||
PY
|
||||
```
|
||||
|
||||
**Step 3: Connect to VirtualPC MCP Server**
|
||||
|
||||
Use the following configuration to connect to the VirtualPC MCP Server:
|
||||
|
||||
```json
|
||||
{
|
||||
"virtualpc-mcp-server": {
|
||||
"type": "streamable-http",
|
||||
"url": "http://localhost:8000/mcp",
|
||||
"headers": {
|
||||
"Authorization": "Bearer <local-debug-jwt>",
|
||||
"MCP_SERVERS": "readweb-server,browser-server"
|
||||
},
|
||||
"timeout": 6000,
|
||||
"sse_read_timeout": 6000,
|
||||
"client_session_timeout_seconds": 6000
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Note**: The Bearer token above is for local testing only. The `MCP_SERVERS` header specifies the MCP server scope for your current connection, which should be a subset of server names defined in `gaia-mcp-server/mcp_servers/mcp_config.py`.
|
||||
|
||||
### 2.2 Kubernetes Cluster Deployment
|
||||
|
||||
For production deployments and RL training scenarios, Kubernetes cluster deployment is recommended. Detailed instructions will be provided in future updates.
|
||||
|
||||
## 3. Development
|
||||
|
||||
### 3.1 Adding Custom MCP Tools to VirtualPC MCP Server
|
||||
|
||||
**Step 1: Develop MCP Tool (Optional)**
|
||||
|
||||
If you need to develop a custom MCP Tool and register it with VirtualPC MCP Server, create your MCP Tool project directory under `gaia-mcp-server/mcp_servers` and implement the MCP Tool code. Refer to the [hello_world](./gaia-mcp-server/mcp_servers/hello_world/) directory for the project structure.
|
||||
|
||||
Project specifications:
|
||||
|
||||
1. Use `pyproject.toml` to manage project dependencies for Docker image building
|
||||
|
||||
**Step 2: Register MCP Tool**
|
||||
|
||||
Register your developed MCP Tool or third-party MCP Tool with VirtualPC MCP Server.
|
||||
|
||||
Edit the [MCP Tool registration file](./gaia-mcp-server/mcp_servers/mcp_config.py):
|
||||
|
||||
```python
|
||||
"STDIO_SERVER_DEMO": {
|
||||
"type": "stdio",
|
||||
"command": "python",
|
||||
"args": ["-m", "hello_world.main"],
|
||||
"cwd": "hello_world",
|
||||
},
|
||||
"{SSE/STREAMABLE-HTTP_SERVER_NAME}": {
|
||||
"type": "sse/streamable-http",
|
||||
"url": "{URL for sse/streamable-http mcp server}",
|
||||
"headers": {
|
||||
"Authorization": f"Bearer {token}"
|
||||
}
|
||||
},
|
||||
```
|
||||
|
||||
**Step 3: Update MCP Tool Schema**
|
||||
|
||||
> **Important**: VirtualPC MCP Server utilizes pre-generated tool schema data for the `list_tools()` function, therefore you must update [mcp_tool_schema.json](./gaia-mcp-server/mcp_servers/mcp_tool_schema.json) after modifying the MCP server configuration.
|
||||
|
||||
A Python script [build_mcp_tool_schema.py](./gaia-mcp-server/mcp_servers/build_mcp_tool_schema.py) is provided to update `mcp_tool_schema.json`. Before executing this script, ensure the MCP server [.env](./gaia-mcp-server/mcp_servers/.env) file is correctly configured.
|
||||
|
||||
```bash
|
||||
cd ./gaia-mcp-server/mcp_servers/
|
||||
pip install mcp
|
||||
python build_mcp_tool_schema.py
|
||||
```
|
||||
|
||||
**Step 4: Build Docker Image and Deploy Service**
|
||||
|
||||
After completing the above steps, build the Docker image and deploy the service.
|
||||
|
||||
## 4. Contributing
|
||||
|
||||
We welcome contributions from the community! Please refer to our contributing guidelines for:
|
||||
|
||||
- Code style and standards
|
||||
- Pull request process
|
||||
- Issue reporting
|
||||
- Development setup instructions
|
||||
|
||||
## 5. References
|
||||
|
||||
### Acknowledgments
|
||||
|
||||
- **Magentic-UI Project**: We have incorporated Docker Browser source code from the [magentic-ui](https://github.com/microsoft/magentic-ui) project. Special thanks to the magentic-ui project team for their excellent work.
|
||||
|
||||
### Related Projects
|
||||
|
||||
- [Model Context Protocol (MCP)](https://modelcontextprotocol.io/)
|
||||
- [Magentic-UI](https://github.com/microsoft/magentic-ui)
|
||||
- [Debian](https://www.debian.org/)
|
||||
|
||||
---
|
||||
|
||||
<div align="center">
|
||||
|
||||
**VirtualPC MCP Server** - Empowering AI agents with robust, scalable runtime environments
|
||||
|
||||
[license-image]: https://img.shields.io/badge/License-MIT-yellow.svg
|
||||
[license-url]: https://opensource.org/licenses/MIT
|
||||
|
||||
</div>
|
||||
+197
@@ -0,0 +1,197 @@
|
||||
<div align="center">
|
||||
|
||||
# VirtualPC MCP Server (孵化中)
|
||||
|
||||
*基于 Debian 的统一 MCP 工具运行时环境，具有会话级环境隔离、环境状态持久化、实时 UI 可视化、分布式架构和可扩展性*
|
||||
|
||||
[![License: MIT][license-image]][license-url]
|
||||
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
|
||||
[English](./README.md) | [快速开始](#快速开始) | [开发](#开发) | [贡献](#贡献)
|
||||
|
||||
</div>
|
||||
|
||||
---
|
||||
|
||||
## 1. 概述
|
||||
|
||||
VirtualPC MCP Server 是一个综合性的 MCP（模型上下文协议）工具运行时环境，旨在为 AI 代理提供统一、隔离和可扩展的执行环境。基于 Debian 构建，它提供会话级环境隔离、跨多个会话的持久状态管理以及实时可视化功能。
|
||||
|
||||
### 1.1 特性
|
||||
|
||||
- **会话级环境隔离**：每个 MCP 会话在其独立的隔离环境中运行
|
||||
- **多会话状态持久化**：在多个 MCP 会话之间维护环境状态
|
||||
- **实时 UI 可视化**：Agent MCP 操作的实时监控和可视化
|
||||
- **分布式架构**：支持本地 Docker 和 Kubernetes 集群部署
|
||||
- **可扩展运行时**：模块化设计，支持无缝集成新的 MCP 工具服务器
|
||||
|
||||
## 2. 快速开始
|
||||
|
||||
本项目支持本地 Docker 部署（适用于演示和调试）和 Kubernetes 集群部署（推荐用于生产和强化学习训练）。
|
||||
|
||||
### 2.1 本地 Docker 部署
|
||||
|
||||
#### 前置要求
|
||||
|
||||
确保 Docker 和 Docker Compose 已正确安装并正常运行：
|
||||
|
||||
```bash
|
||||
# 验证 Docker 安装
|
||||
docker --version
|
||||
docker compose --version
|
||||
|
||||
# 验证 Docker 守护进程是否运行
|
||||
docker ps
|
||||
docker compose ps
|
||||
```
|
||||
|
||||
**步骤 1：配置环境并准备 Gaia 数据集**
|
||||
|
||||
1. 复制环境模板并配置您的设置：
|
||||
|
||||
```bash
|
||||
cp ./gaia-mcp-server/mcp_servers/.env_template ./gaia-mcp-server/mcp_servers/.env
|
||||
```
|
||||
|
||||
编辑 `./gaia-mcp-server/mcp_servers/.env` 文件，填入您的具体配置值。
|
||||
|
||||
2. 从 Hugging Face 下载 [gaia_dataset](https://huggingface.co/datasets/gaia-benchmark/GAIA) 并放置到 `./gaia-mcp-server/docker/gaia_dataset`
|
||||
|
||||
**步骤 2：启动 VirtualPC MCP Server**
|
||||
|
||||
```bash
|
||||
sh run-docker.sh
|
||||
```
|
||||
|
||||
监控终端输出，查看启动过程中是否有任何错误。
|
||||
|
||||
生成本地 Bearer token，并将打印出的值复制到下面的 `<local-debug-jwt>`。如果你修改了 `virtualpc-mcp/docker-compose.yaml` 中的 `MCP_GATEWAY_TOKEN_SECRET`，请在运行此命令前导出相同的值。
|
||||
|
||||
```bash
|
||||
python - <<'PY'
|
||||
import base64, hashlib, hmac, json, os, time
|
||||
|
||||
def part(value):
|
||||
raw = json.dumps(value, separators=(",", ":")).encode()
|
||||
return base64.urlsafe_b64encode(raw).rstrip(b"=").decode()
|
||||
|
||||
signing_input = ".".join([
|
||||
part({"alg": "HS256", "typ": "JWT"}),
|
||||
part({"app": "local_debug", "version": 1, "time": time.time()}),
|
||||
])
|
||||
secret = os.getenv("MCP_GATEWAY_TOKEN_SECRET", "123321").encode()
|
||||
signature = hmac.new(secret, signing_input.encode(), hashlib.sha256).digest()
|
||||
print(f"{signing_input}.{base64.urlsafe_b64encode(signature).rstrip(b'=').decode()}")
|
||||
PY
|
||||
```
|
||||
|
||||
**步骤 3：连接到 VirtualPC MCP Server**
|
||||
|
||||
使用以下配置连接到 VirtualPC MCP Server：
|
||||
|
||||
```json
|
||||
{
|
||||
"virtualpc-mcp-server": {
|
||||
"type": "streamable-http",
|
||||
"url": "http://localhost:8000/mcp",
|
||||
"headers": {
|
||||
"Authorization": "Bearer <local-debug-jwt>",
|
||||
"MCP_SERVERS": "readweb-server,browser-server"
|
||||
},
|
||||
"timeout": 6000,
|
||||
"sse_read_timeout": 6000,
|
||||
"client_session_timeout_seconds": 6000
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**注意**：上述 Bearer token 仅用于本地测试。`MCP_SERVERS` 头部指定了当前连接的 MCP 服务器范围，应该是 `gaia-mcp-server/mcp_servers/mcp_config.py` 中定义的服务器名称的子集。
|
||||
|
||||
### 2.2 Kubernetes 集群部署
|
||||
|
||||
对于生产部署和强化学习训练场景，推荐使用 Kubernetes 集群部署。详细说明将在后续更新中提供。
|
||||
|
||||
## 3. 开发
|
||||
|
||||
### 3.1 向 VirtualPC MCP Server 添加自定义 MCP 工具
|
||||
|
||||
**步骤 1：开发 MCP 工具（可选）**
|
||||
|
||||
如果您需要开发自定义 MCP 工具并将其注册到 VirtualPC MCP Server，请在 `gaia-mcp-server/mcp_servers` 下创建您的 MCP 工具项目目录并实现 MCP 工具代码。参考 [hello_world](./gaia-mcp-server/mcp_servers/hello_world/) 目录的项目结构。
|
||||
|
||||
项目规范：
|
||||
|
||||
1. 使用 `pyproject.toml` 管理项目依赖，用于 Docker 镜像构建
|
||||
|
||||
**步骤 2：注册 MCP 工具**
|
||||
|
||||
将您开发的 MCP 工具或第三方 MCP 工具注册到 VirtualPC MCP Server。
|
||||
|
||||
编辑 [MCP 工具注册文件](./gaia-mcp-server/mcp_servers/mcp_config.py)：
|
||||
|
||||
```python
|
||||
"STDIO_SERVER_DEMO": {
|
||||
"type": "stdio",
|
||||
"command": "python",
|
||||
"args": ["-m", "hello_world.main"],
|
||||
"cwd": "hello_world",
|
||||
},
|
||||
"{SSE/STREAMABLE-HTTP_SERVER_NAME}": {
|
||||
"type": "sse/streamable-http",
|
||||
"url": "{URL for sse/streamable-http mcp server}",
|
||||
"headers": {
|
||||
"Authorization": f"Bearer {token}"
|
||||
}
|
||||
},
|
||||
```
|
||||
|
||||
**步骤 3：更新 MCP 工具模式**
|
||||
|
||||
> **重要**：VirtualPC MCP Server 使用预生成的工具模式数据用于 `list_tools()` 函数，因此您必须在修改 MCP 服务器配置后更新 [mcp_tool_schema.json](./gaia-mcp-server/mcp_servers/mcp_tool_schema.json)。
|
||||
|
||||
我们提供了一个 Python 脚本 [build_mcp_tool_schema.py](./gaia-mcp-server/mcp_servers/build_mcp_tool_schema.py) 来更新 `mcp_tool_schema.json`。在执行此脚本之前，请确保 MCP 服务器 [.env](./gaia-mcp-server/mcp_servers/.env) 文件已正确配置。
|
||||
|
||||
```bash
|
||||
cd ./gaia-mcp-server/mcp_servers/
|
||||
pip install mcp
|
||||
python build_mcp_tool_schema.py
|
||||
```
|
||||
|
||||
**步骤 4：构建 Docker 镜像并部署服务**
|
||||
|
||||
完成上述步骤后，构建 Docker 镜像并部署服务。
|
||||
|
||||
## 4. 贡献
|
||||
|
||||
我们欢迎社区的贡献！请参考我们的贡献指南：
|
||||
|
||||
- 代码风格和标准
|
||||
- 拉取请求流程
|
||||
- 问题报告
|
||||
- 开发设置说明
|
||||
|
||||
## 5. 参考资料
|
||||
|
||||
### 致谢
|
||||
|
||||
- **Magentic-UI 项目**：我们已整合来自 [magentic-ui](https://github.com/microsoft/magentic-ui) 项目的 Docker Browser 源代码。特别感谢 magentic-ui 项目团队的出色工作。
|
||||
|
||||
### 相关项目
|
||||
|
||||
- [模型上下文协议 (MCP)](https://modelcontextprotocol.io/)
|
||||
- [Magentic-UI](https://github.com/microsoft/magentic-ui)
|
||||
- [Debian](https://www.debian.org/)
|
||||
|
||||
---
|
||||
|
||||
<div align="center">
|
||||
|
||||
**VirtualPC MCP Server** - 为 AI Agent 提供强大、可扩展的运行时环境
|
||||
|
||||
[license-image]: https://img.shields.io/badge/License-MIT-yellow.svg
|
||||
[license-url]: https://opensource.org/licenses/MIT
|
||||
|
||||
</div>
|
||||
@@ -0,0 +1,2 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
@@ -0,0 +1,7 @@
|
||||
FROM mcp-server-base:latest
|
||||
|
||||
COPY mcp_servers /app/mcp_servers
|
||||
|
||||
COPY docker/gaia_dataset /root/workspace/gaia_dataset
|
||||
|
||||
RUN sh /app/mcp_servers/init_env.sh
|
||||
@@ -0,0 +1,12 @@
|
||||
#!/bin/sh
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
dt=$(date +%Y%m%d%H%M%S)
|
||||
img=gaia-mcp-server
|
||||
version=$dt
|
||||
|
||||
docker build -t $img -t $img:$version -t $img:latest . && \
|
||||
|
||||
echo "✅ Build image success: $img:$version"
|
||||
|
||||
exit 0
|
||||
Vendored
+14
@@ -0,0 +1,14 @@
|
||||
[project]
|
||||
name = "audio-server"
|
||||
version = "0.1.0"
|
||||
description = "Transcribe the given audio in a list of filepaths or urls."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
aiohttp= ">=3.12.15"
|
||||
dependencies = [
|
||||
"mcp",
|
||||
"requests~=2.32.4",
|
||||
"openai~=1.93.0",
|
||||
"fastmcp~=2.11.3",
|
||||
]
|
||||
|
||||
+138
@@ -0,0 +1,138 @@
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import traceback
|
||||
from typing import List
|
||||
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
from openai import OpenAI
|
||||
from pydantic import Field
|
||||
|
||||
import logging
|
||||
|
||||
from utils import get_file_from_source
|
||||
from dotenv import load_dotenv
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Initialize MCP server
|
||||
mcp = FastMCP("audio-server")
|
||||
|
||||
|
||||
client = OpenAI(
|
||||
api_key=os.getenv("AUDIO_LLM_API_KEY"),
|
||||
base_url=os.getenv("AUDIO_LLM_BASE_URL")
|
||||
)
|
||||
|
||||
AUDIO_TRANSCRIBE = (
|
||||
"Input is a base64 encoded audio. Transcribe the audio content. "
|
||||
"Return a json string with the following format: "
|
||||
'{"audio_text": "transcribed text from audio"}'
|
||||
)
|
||||
|
||||
|
||||
def encode_audio(audio_source: str, with_header: bool = True) -> str:
|
||||
"""
|
||||
Encode audio to base64 format with robust file handling
|
||||
|
||||
Args:
|
||||
audio_source: URL or local file path of the audio
|
||||
with_header: Whether to include MIME type header
|
||||
|
||||
Returns:
|
||||
str: Base64 encoded audio string, with MIME type prefix if with_header is True
|
||||
|
||||
Raises:
|
||||
ValueError: When audio source is invalid or audio format is not supported
|
||||
IOError: When audio file cannot be read
|
||||
"""
|
||||
if not audio_source:
|
||||
raise ValueError("Audio source cannot be empty")
|
||||
|
||||
try:
|
||||
# Get file with validation (only audio files allowed)
|
||||
file_path, mime_type, content = get_file_from_source(
|
||||
audio_source,
|
||||
allowed_mime_prefixes=["audio/"],
|
||||
max_size_mb=200.0, # 200MB limit for audio files
|
||||
type="audio", # Specify type as audio to handle audio files
|
||||
)
|
||||
|
||||
# Encode to base64
|
||||
audio_base64 = base64.b64encode(content).decode()
|
||||
|
||||
# Format with header if requested
|
||||
final_audio = (
|
||||
f"data:{mime_type};base64,{audio_base64}" if with_header else audio_base64
|
||||
)
|
||||
|
||||
# Clean up temporary file if it was created for a URL
|
||||
if file_path != os.path.abspath(audio_source) and os.path.exists(file_path):
|
||||
os.unlink(file_path)
|
||||
|
||||
return final_audio
|
||||
|
||||
except Exception:
|
||||
logger.error(
|
||||
f"Error encoding audio from {audio_source}: {traceback.format_exc()}"
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
@mcp.tool(description="Transcribe the given audio in a list of filepaths or urls.")
|
||||
async def mcp_transcribe_audio(
|
||||
audio_urls: List[str] = Field(
|
||||
description="The input audio in given a list of filepaths or urls."
|
||||
),
|
||||
) -> str:
|
||||
"""
|
||||
Transcribe the given audio in a list of filepaths or urls.
|
||||
|
||||
Args:
|
||||
audio_urls: List of audio file paths or URLs
|
||||
|
||||
Returns:
|
||||
str: JSON string containing transcriptions
|
||||
"""
|
||||
transcriptions = []
|
||||
for audio_url in audio_urls:
|
||||
try:
|
||||
# Get file with validation (only audio files allowed)
|
||||
file_path, _, _ = get_file_from_source(
|
||||
audio_url,
|
||||
allowed_mime_prefixes=["audio/"],
|
||||
max_size_mb=200.0, # 200MB limit for audio files
|
||||
type="audio", # Specify type as audio to handle audio files
|
||||
)
|
||||
|
||||
# Use the file for transcription
|
||||
with open(file_path, "rb") as audio_file:
|
||||
transcription = client.audio.transcriptions.create(
|
||||
file=audio_file,
|
||||
model=os.getenv("AUDIO_LLM_MODEL_NAME"),
|
||||
response_format="text",
|
||||
)
|
||||
transcriptions.append(transcription)
|
||||
|
||||
# Clean up temporary file if it was created for a URL
|
||||
if file_path != os.path.abspath(audio_url) and os.path.exists(file_path):
|
||||
os.unlink(file_path)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error transcribing {audio_url}: {traceback.format_exc()}")
|
||||
transcriptions.append(f"Error: {str(e)}")
|
||||
|
||||
logger.info(f"---get_text_by_transcribe-transcription:{transcriptions}")
|
||||
return json.dumps(transcriptions, ensure_ascii=False)
|
||||
|
||||
|
||||
# Run the server when the script is executed directly
|
||||
if __name__ == "__main__":
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
|
||||
datefmt='%Y-%m-%d %H:%M:%S'
|
||||
)
|
||||
load_dotenv()
|
||||
logger.info("Starting Audio MCP Server...")
|
||||
mcp.run(transport="stdio")
|
||||
+141
@@ -0,0 +1,141 @@
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from typing import List, Optional, Tuple
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import requests
|
||||
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def is_url(path_or_url: str) -> bool:
|
||||
"""
|
||||
Check if the given string is a URL.
|
||||
|
||||
Args:
|
||||
path_or_url: String to check
|
||||
|
||||
Returns:
|
||||
bool: True if the string is a URL, False otherwise
|
||||
"""
|
||||
parsed = urlparse(path_or_url)
|
||||
return bool(parsed.scheme and parsed.netloc)
|
||||
|
||||
|
||||
def get_file_from_source(
|
||||
source: str,
|
||||
allowed_mime_prefixes: List[str] = None,
|
||||
max_size_mb: float = 100.0,
|
||||
timeout: int = 60,
|
||||
type: str = "image",
|
||||
) -> Tuple[str, str, bytes]:
|
||||
"""
|
||||
Unified function to get file content from a URL or local path with validation.
|
||||
|
||||
Args:
|
||||
source: URL or local file path
|
||||
allowed_mime_prefixes: List of allowed MIME type prefixes (e.g., ['audio/', 'video/'])
|
||||
max_size_mb: Maximum allowed file size in MB
|
||||
timeout: Timeout for URL requests in seconds
|
||||
|
||||
Returns:
|
||||
Tuple[str, str, bytes]: (file_path, mime_type, file_content)
|
||||
- For URLs, file_path will be a temporary file path
|
||||
- For local files, file_path will be the original path
|
||||
|
||||
Raises:
|
||||
ValueError: When file doesn't exist, exceeds size limit, or has invalid MIME type
|
||||
IOError: When file cannot be read
|
||||
requests.RequestException: When URL request fails
|
||||
"""
|
||||
max_size_bytes = max_size_mb * 1024 * 1024
|
||||
temp_file = None
|
||||
|
||||
try:
|
||||
if is_url(source):
|
||||
# Handle URL
|
||||
logger.info(f"Downloading file from URL: {source}")
|
||||
response = requests.get(source, stream=True, timeout=timeout)
|
||||
response.raise_for_status()
|
||||
|
||||
# Check Content-Length if available
|
||||
content_length = response.headers.get("Content-Length")
|
||||
if content_length and int(content_length) > max_size_bytes:
|
||||
raise ValueError(f"File size exceeds limit of {max_size_mb}MB")
|
||||
|
||||
# Create a temporary file
|
||||
temp_file = tempfile.NamedTemporaryFile(delete=False)
|
||||
file_path = temp_file.name
|
||||
|
||||
# Download content in chunks to avoid memory issues
|
||||
content = bytearray()
|
||||
downloaded_size = 0
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
downloaded_size += len(chunk)
|
||||
if downloaded_size > max_size_bytes:
|
||||
raise ValueError(f"File size exceeds limit of {max_size_mb}MB")
|
||||
temp_file.write(chunk)
|
||||
content.extend(chunk)
|
||||
|
||||
temp_file.close()
|
||||
|
||||
# Get MIME type
|
||||
if type == "audio":
|
||||
mime_type = "audio/mpeg"
|
||||
elif type == "image":
|
||||
mime_type = "image/jpeg"
|
||||
elif type == "video":
|
||||
mime_type = "video/mp4"
|
||||
|
||||
|
||||
# For URLs where magic fails, try to use Content-Type header
|
||||
if mime_type == "application/octet-stream":
|
||||
content_type = response.headers.get("Content-Type", "").split(";")[0]
|
||||
if content_type:
|
||||
mime_type = content_type
|
||||
else:
|
||||
# Handle local file
|
||||
file_path = os.path.abspath(source)
|
||||
|
||||
# Check if file exists
|
||||
if not os.path.exists(file_path):
|
||||
raise ValueError(f"File not found: {file_path}")
|
||||
|
||||
# Check file size
|
||||
file_size = os.path.getsize(file_path)
|
||||
if file_size > max_size_bytes:
|
||||
raise ValueError(f"File size exceeds limit of {max_size_mb}MB")
|
||||
|
||||
# Get MIME type
|
||||
if type == "audio":
|
||||
mime_type = "audio/mpeg"
|
||||
elif type == "image":
|
||||
mime_type = "image/jpeg"
|
||||
elif type == "video":
|
||||
mime_type = "video/mp4"
|
||||
|
||||
# Read file content
|
||||
with open(file_path, "rb") as f:
|
||||
content = f.read()
|
||||
|
||||
# Validate MIME type if allowed_mime_prefixes is provided
|
||||
if allowed_mime_prefixes:
|
||||
if not any(
|
||||
mime_type.startswith(prefix) for prefix in allowed_mime_prefixes
|
||||
):
|
||||
allowed_types = ", ".join(allowed_mime_prefixes)
|
||||
raise ValueError(
|
||||
f"Invalid file type: {mime_type}. Allowed types: {allowed_types}"
|
||||
)
|
||||
|
||||
return file_path, mime_type, content
|
||||
|
||||
except Exception as e:
|
||||
# Clean up temporary file if an error occurs
|
||||
if temp_file and os.path.exists(temp_file.name):
|
||||
os.unlink(temp_file.name)
|
||||
raise e
|
||||
Vendored
+16
@@ -0,0 +1,16 @@
|
||||
[project]
|
||||
name = "browser-server"
|
||||
version = "0.1.0"
|
||||
description = "Add your description here"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
aiohttp= ">=3.12.15"
|
||||
dependencies = [
|
||||
"mcp",
|
||||
"aiohttp~=3.12.15",
|
||||
"requests~=2.32.4",
|
||||
"fastmcp~=2.11.3",
|
||||
"tavily-python~=0.7.10",
|
||||
"browser-use~=0.5.5"
|
||||
|
||||
]
|
||||
+224
@@ -0,0 +1,224 @@
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import requests
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import Any, Literal
|
||||
|
||||
|
||||
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)"
|
||||
)
|
||||
|
||||
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")
|
||||
|
||||
|
||||
def _validate_file_path(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.exists():
|
||||
raise FileNotFoundError(f"File not found: {path}")
|
||||
|
||||
return path
|
||||
|
||||
def is_url(path_or_url: str) -> bool:
|
||||
"""
|
||||
Check if the given string is a URL.
|
||||
|
||||
Args:
|
||||
path_or_url: String to check
|
||||
|
||||
Returns:
|
||||
bool: True if the string is a URL, False otherwise
|
||||
"""
|
||||
parsed = urlparse(path_or_url)
|
||||
return bool(parsed.scheme and parsed.netloc)
|
||||
|
||||
|
||||
def get_mime_type(file_path: str, default_mime: str | None = None) -> str:
|
||||
"""
|
||||
Detect MIME type of a file using python-magic if available,
|
||||
otherwise fallback to extension-based detection.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
default_mime: Default MIME type to return if detection fails
|
||||
|
||||
Returns:
|
||||
str: Detected MIME type
|
||||
"""
|
||||
# Try using python-magic for accurate MIME type detection
|
||||
try:
|
||||
import magic
|
||||
mime = magic.Magic(mime=True)
|
||||
return mime.from_file(file_path)
|
||||
except (AttributeError, IOError):
|
||||
# Fallback to extension-based detection
|
||||
extension_mime_map = {
|
||||
# Audio formats
|
||||
".mp3": "audio/mpeg",
|
||||
".wav": "audio/wav",
|
||||
".ogg": "audio/ogg",
|
||||
".m4a": "audio/mp4",
|
||||
".flac": "audio/flac",
|
||||
# Image formats
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".png": "image/png",
|
||||
".gif": "image/gif",
|
||||
".webp": "image/webp",
|
||||
".bmp": "image/bmp",
|
||||
".tiff": "image/tiff",
|
||||
# Video formats
|
||||
".mp4": "video/mp4",
|
||||
".avi": "video/x-msvideo",
|
||||
".mov": "video/quicktime",
|
||||
".mkv": "video/x-matroska",
|
||||
".webm": "video/webm",
|
||||
}
|
||||
|
||||
ext = Path(file_path).suffix.lower()
|
||||
return extension_mime_map.get(ext, default_mime or "application/octet-stream")
|
||||
|
||||
|
||||
def get_file_from_source(
|
||||
source: str,
|
||||
max_size_mb: float = 100.0,
|
||||
timeout: int = 60,
|
||||
) -> tuple[str, str, bytes]:
|
||||
"""
|
||||
Unified function to get file content from a URL or local path with validation.
|
||||
|
||||
Args:
|
||||
source: URL or local file path
|
||||
max_size_mb: Maximum allowed file size in MB
|
||||
timeout: Timeout for URL requests in seconds
|
||||
|
||||
Returns:
|
||||
Tuple[str, str, bytes]: (file_path, mime_type, file_content)
|
||||
- For URLs, file_path will be a temporary file path
|
||||
- For local files, file_path will be the original path
|
||||
|
||||
Raises:
|
||||
ValueError: When file doesn't exist, exceeds size limit, or has invalid MIME type
|
||||
IOError: When file cannot be read
|
||||
requests.RequestException: When URL request fails
|
||||
"""
|
||||
max_size_bytes = max_size_mb * 1024 * 1024
|
||||
|
||||
if is_url(source):
|
||||
# Handle URL source
|
||||
try:
|
||||
# Make a HEAD request first to check content length
|
||||
head_response = requests.head(source, timeout=timeout, allow_redirects=True)
|
||||
head_response.raise_for_status()
|
||||
|
||||
# Check content length if available
|
||||
content_length = head_response.headers.get("content-length")
|
||||
if content_length and int(content_length) > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({int(content_length) / (1024 * 1024):.2f} MB) "
|
||||
f"exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Download the file
|
||||
response = requests.get(source, timeout=timeout, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
# Read content with size checking
|
||||
content = b""
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
if len(content) + len(chunk) > max_size_bytes:
|
||||
raise ValueError(f"File size exceeds maximum allowed size ({max_size_mb} MB)")
|
||||
content += chunk
|
||||
|
||||
# Create temporary file
|
||||
parsed_url = urlparse(source)
|
||||
filename = os.path.basename(parsed_url.path) or "downloaded_file"
|
||||
|
||||
# Create temporary file with proper extension
|
||||
suffix = Path(filename).suffix or ".tmp"
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
|
||||
temp_file.write(content)
|
||||
temp_path = temp_file.name
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(temp_path)
|
||||
|
||||
return temp_path, mime_type, content
|
||||
|
||||
except requests.RequestException as e:
|
||||
raise requests.RequestException(f"Failed to download file from URL: {e}: {traceback.format_exc()}")
|
||||
except Exception as e:
|
||||
raise IOError(f"Error processing URL: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
else:
|
||||
# Handle local file path
|
||||
file_path = Path(source)
|
||||
|
||||
# Check if file exists
|
||||
if not file_path.exists():
|
||||
raise ValueError(f"File does not exist: {source}")
|
||||
|
||||
if not file_path.is_file():
|
||||
raise ValueError(f"Path is not a file: {source}")
|
||||
|
||||
# Check file size
|
||||
file_size = file_path.stat().st_size
|
||||
if file_size > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({file_size / (1024 * 1024):.2f} MB) exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Read file content
|
||||
try:
|
||||
with open(file_path, "rb") as f:
|
||||
content = f.read()
|
||||
except Exception as e:
|
||||
raise IOError(f"Cannot read file {source}: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(str(file_path))
|
||||
|
||||
return str(file_path), mime_type, content
|
||||
+427
@@ -0,0 +1,427 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from typing import Union
|
||||
|
||||
from playwright.async_api import async_playwright
|
||||
from pydantic import BaseModel, Field
|
||||
from browser_use import Agent, AgentHistoryList, BrowserProfile, BrowserSession
|
||||
from browser_use.llm import ChatOpenAI
|
||||
from dotenv import load_dotenv
|
||||
from mcp.server import FastMCP
|
||||
from mcp.types import TextContent
|
||||
from mcp.server.fastmcp import Context
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
|
||||
)
|
||||
|
||||
workspace = Path.home()
|
||||
logs_path = workspace / "logs"
|
||||
logs_path.mkdir(parents=True, exist_ok=True)
|
||||
trace_log_dir = str(logs_path)
|
||||
|
||||
extended_browser_system_prompt = """
|
||||
# 效率指南
|
||||
0. 如果有下载选项，尽可能**下载**！下载根目录必须在在~/wokspace目录下面，可以根据场景适当的创建文件夹来存放下载的文件，同时，在结果中报告要包含存放文件的完整路径。
|
||||
1. 使用包含任务关键词的特定搜索查询
|
||||
2. 避免被无关信息分散注意力
|
||||
3. 如果被付费墙阻挡，尝试使用archive.org或类似替代方案
|
||||
4. 清晰简洁地记录每个重要发现
|
||||
5. 以最少的浏览步骤精确提取必要信息。
|
||||
6. ***重要****如果操作浏览器过程中出现需要人工干预的过程，比如：登录、验证码输入、输入密码、支付等操作，就不能继续往下操作，需要返回习惯人工干预的提示信息从而等待人工干预之后继续操作(但需要保持当前操作浏览器窗口)
|
||||
样例：
|
||||
1、遇到登录页面
|
||||
返回：当前操作需要用户进行登录，请你在页面上进行相关登录操作，之后继续执行
|
||||
2、遇到输入验证码页面
|
||||
返回：当前操作需要用户进行输入验证码，请你在页面上进行输入验证码，之后继续执行
|
||||
3、遇到输入密码页面
|
||||
返回：当前操作需要用户进行输入密码，请你在页面上进行输入密码操作，之后继续执行
|
||||
4、遇到支付页面
|
||||
返回：当前操作需要用户进行支付，请你在页面上进行相关支付操作，之后继续执行
|
||||
"""
|
||||
|
||||
# Initialize LLM configuration
|
||||
llm_config = ChatOpenAI(
|
||||
model=os.getenv("LLM_MODEL_NAME"),
|
||||
api_key=os.getenv("LLM_API_KEY"),
|
||||
base_url=os.getenv("LLM_BASE_URL"),
|
||||
temperature=1.0,
|
||||
)
|
||||
|
||||
browser_profile = BrowserProfile(
|
||||
cookies_file=os.getenv("COOKIES_FILE_PATH"),
|
||||
downloads_dir=str(workspace),
|
||||
downloads_path=str(workspace),
|
||||
save_recording_path=str(workspace),
|
||||
save_downloads_path=str(workspace),
|
||||
chromium_sandbox=False,
|
||||
headless=False,
|
||||
keep_alive=True,
|
||||
)
|
||||
|
||||
from base import ActionResponse
|
||||
|
||||
|
||||
class BrowserMetadata(BaseModel):
|
||||
"""Metadata for browser automation results."""
|
||||
|
||||
task: str
|
||||
execution_successful: bool
|
||||
steps_taken: int | None = None
|
||||
downloaded_files: list[str] = Field(default_factory=list)
|
||||
visited_urls: list[str] = Field(default_factory=list)
|
||||
execution_time: float | None = None
|
||||
error_type: str | None = None
|
||||
trace_log_path: str | None = None
|
||||
|
||||
|
||||
browser_session = BrowserSession(
|
||||
browser_profile=browser_profile
|
||||
)
|
||||
|
||||
mcp = FastMCP("browser-server")
|
||||
|
||||
|
||||
async def show_vnc_window(
|
||||
ctx: Context,
|
||||
):
|
||||
"""Show the VNC window"""
|
||||
await ctx.report_progress(
|
||||
progress=0.0,
|
||||
total=1.0,
|
||||
message="tool_call_card_novnc_window",
|
||||
)
|
||||
|
||||
|
||||
def _extract_visited_urls(extracted_content: list[str]) -> list[str]:
|
||||
"""Inner method to extract URLs from content using regex.
|
||||
|
||||
Args:
|
||||
content_list: List of content strings to search for URLs
|
||||
|
||||
Returns:
|
||||
List of unique URLs found in the content
|
||||
"""
|
||||
url_pattern = r'https?://[^\s<>"\[\]{}|\\^`]+'
|
||||
visited_urls = set()
|
||||
|
||||
for content in extracted_content:
|
||||
if content and isinstance(content, str):
|
||||
urls = re.findall(url_pattern, content)
|
||||
visited_urls.update(urls)
|
||||
|
||||
return list(visited_urls)
|
||||
|
||||
|
||||
def _format_extracted_content(extracted_content: list[str]) -> str:
|
||||
"""Format extracted content to be LLM-friendly.
|
||||
|
||||
Args:
|
||||
extracted_content: List of extracted content strings from browser execution
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if not extracted_content:
|
||||
return "No content extracted from browser execution."
|
||||
|
||||
# Handle list of strings
|
||||
if len(extracted_content) == 1:
|
||||
# Single item - return it directly with formatting
|
||||
return f"**Extracted Content:**\n{extracted_content[0]}"
|
||||
else:
|
||||
# Multiple items - format as numbered list
|
||||
formatted_parts = ["**Extracted Content:**"]
|
||||
for i, content in enumerate(extracted_content, 1):
|
||||
if content.strip(): # Only include non-empty content
|
||||
formatted_parts.append(f"{i}. {content}")
|
||||
|
||||
return (
|
||||
"\n".join(formatted_parts)
|
||||
if len(formatted_parts) > 1
|
||||
else "No meaningful content extracted from browser execution."
|
||||
)
|
||||
|
||||
|
||||
async def _create_browser_agent(task: str) -> Agent:
|
||||
"""Create a browser agent instance with configured settings.
|
||||
|
||||
Args:
|
||||
task: The task description for the browser agent
|
||||
|
||||
Returns:
|
||||
Configured Agent instance
|
||||
"""
|
||||
playwright = await async_playwright().start()
|
||||
ws_remote_url: str = f"ws://localhost:37367/default"
|
||||
browser = await playwright.chromium.connect(ws_endpoint=ws_remote_url)
|
||||
return Agent(
|
||||
task=task,
|
||||
llm=llm_config,
|
||||
extend_system_message=extended_browser_system_prompt,
|
||||
use_vision=True,
|
||||
enable_memory=False,
|
||||
browser=browser,
|
||||
#browser_profile=browser_profile,
|
||||
browser_session=browser_session,
|
||||
save_conversation_path=trace_log_dir + "/trace.log",
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Get information about browser automation capabilities and configuration.
|
||||
"""
|
||||
)
|
||||
async def get_browser_capabilities() -> Union[str, TextContent]:
|
||||
"""Get information about browser automation capabilities and configuration.
|
||||
|
||||
Returns:
|
||||
ActionResponse with browser service capabilities and current configuration
|
||||
"""
|
||||
capabilities = {
|
||||
"automation_features": [
|
||||
"Web scraping and content extraction",
|
||||
"Form submission and interaction",
|
||||
"File downloads and media handling",
|
||||
"LLM-enhanced browsing with vision",
|
||||
"Memory-enabled browsing sessions",
|
||||
"Robot detection and paywall handling",
|
||||
],
|
||||
"supported_formats": ["markdown", "json", "text"],
|
||||
"configuration": {
|
||||
"llm_model": os.getenv("LLM_MODEL_NAME", "Not configured"),
|
||||
"downloads_directory": browser_profile.downloads_path,
|
||||
"cookies_enabled": bool(os.getenv("COOKIES_FILE_PATH")),
|
||||
"trace_logging": True,
|
||||
"vision_enabled": True,
|
||||
"headless": False,
|
||||
},
|
||||
}
|
||||
|
||||
formatted_info = f"""# Browser Automation Service Capabilities
|
||||
|
||||
## Features
|
||||
{chr(10).join(f"- {feature}" for feature in capabilities["automation_features"])}
|
||||
|
||||
## Supported Output Formats
|
||||
{chr(10).join(f"- {fmt}" for fmt in capabilities["supported_formats"])}
|
||||
|
||||
## Current Configuration
|
||||
- **LLM Model:** {capabilities["configuration"]["llm_model"]}
|
||||
- **Downloads Directory:** {capabilities["configuration"]["downloads_directory"]}
|
||||
- **Cookies Enabled:** {capabilities["configuration"]["cookies_enabled"]}
|
||||
- **Vision Enabled:** {capabilities["configuration"]["vision_enabled"]}
|
||||
- **Memory Enabled:** {capabilities["configuration"]["memory_enabled"]}
|
||||
- **Trace Logging:** {capabilities["configuration"]["trace_logging"]}
|
||||
"""
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=formatted_info, metadata=capabilities
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Perform browser automation tasks using the browser-use package.
|
||||
|
||||
This tool provides comprehensive browser automation capabilities including:
|
||||
- Web scraping and content extraction
|
||||
- Form submission and automated interactions
|
||||
- File downloads and media handling
|
||||
- LLM-enhanced browsing with memory and vision
|
||||
- Automatic handling of robot detection and paywalls
|
||||
"""
|
||||
)
|
||||
async def browser_use(
|
||||
context: Context,
|
||||
task: str = Field(
|
||||
description="The task to perform using the browser automation agent"
|
||||
),
|
||||
max_steps: int = Field(
|
||||
default=50, description="Maximum number of steps for browser execution"
|
||||
),
|
||||
extract_format: str = Field(
|
||||
default="markdown",
|
||||
description="Format for extracted content: 'markdown', 'json', or 'text'",
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
try:
|
||||
logging.info(f"🎯 Starting browser task: {task}")
|
||||
|
||||
# Create browser agent
|
||||
agent = await _create_browser_agent(task)
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
await show_vnc_window(context)
|
||||
|
||||
browser_execution: AgentHistoryList = await agent.run(max_steps=max_steps)
|
||||
|
||||
execution_time = time.time() - start_time
|
||||
result_content = ""
|
||||
|
||||
if browser_execution is not None and hasattr(browser_execution, 'history') and browser_execution.history is not None:
|
||||
latest_history = browser_execution.history[-1]
|
||||
if latest_history is not None and hasattr(latest_history, 'result') and latest_history.result is not None:
|
||||
first_result = latest_history.result[0]
|
||||
if first_result is not None and hasattr(first_result, 'extracted_content'):
|
||||
result_content = first_result.extracted_content
|
||||
|
||||
if result_content:
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=result_content, # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
else:
|
||||
# Handle execution failure
|
||||
error_msg = "Browser execution failed or was not completed successfully"
|
||||
|
||||
metadata = BrowserMetadata(
|
||||
task=task,
|
||||
execution_successful=False,
|
||||
execution_time=execution_time,
|
||||
error_type="execution_failure",
|
||||
trace_log_path=trace_log_dir + "/browser_log/trace.log",
|
||||
)
|
||||
|
||||
logging.info(f"❌ {error_msg}")
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=False, message=error_msg, metadata=metadata.model_dump()
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
# if (
|
||||
# browser_execution is not None
|
||||
# and browser_execution.is_done()
|
||||
# and browser_execution.is_successful()
|
||||
# ):
|
||||
# # Extract and format content
|
||||
# extracted_content = browser_execution.extracted_content()
|
||||
# final_result = browser_execution.final_result()
|
||||
#
|
||||
# # Format content based on requested format
|
||||
# if extract_format.lower() == "json":
|
||||
# formatted_content = json.dumps(
|
||||
# {"summary": final_result, "extracted_data": extracted_content},
|
||||
# indent=2,
|
||||
# )
|
||||
# elif extract_format.lower() == "text":
|
||||
# formatted_content = (
|
||||
# f"{final_result}\n\n{_format_extracted_content(extracted_content)}"
|
||||
# )
|
||||
# else: # markdown (default)
|
||||
# formatted_content = (
|
||||
# f"## Browser Automation Result\n\n**Summary:** {final_result}\n\n"
|
||||
# f"{_format_extracted_content(extracted_content)}"
|
||||
# )
|
||||
#
|
||||
# # Prepare metadata
|
||||
# metadata = BrowserMetadata(
|
||||
# task=task,
|
||||
# execution_successful=True,
|
||||
# steps_taken=(
|
||||
# len(browser_execution.history)
|
||||
# if hasattr(browser_execution, "history")
|
||||
# else None
|
||||
# ),
|
||||
# downloaded_files=[],
|
||||
# visited_urls=_extract_visited_urls(extracted_content),
|
||||
# execution_time=execution_time,
|
||||
# trace_log_path=trace_log_dir + "/browser_log/trace.log",
|
||||
# )
|
||||
#
|
||||
# logging.info(f"🗒️ Detail: {extracted_content}")
|
||||
# logging.info(f"🌏 Result: {final_result}")
|
||||
#
|
||||
# action_response = ActionResponse(
|
||||
# success=True,
|
||||
# message=formatted_content,
|
||||
# metadata=metadata.model_dump(),
|
||||
# )
|
||||
# return TextContent(
|
||||
# type="text",
|
||||
# text=json.dumps(
|
||||
# action_response.model_dump()
|
||||
# ), # Empty string instead of None
|
||||
# **{"metadata": {}}, # Pass as additional fields
|
||||
# )
|
||||
#
|
||||
# else:
|
||||
# # Handle execution failure
|
||||
# error_msg = "Browser execution failed or was not completed successfully"
|
||||
#
|
||||
# metadata = BrowserMetadata(
|
||||
# task=task,
|
||||
# execution_successful=False,
|
||||
# execution_time=execution_time,
|
||||
# error_type="execution_failure",
|
||||
# trace_log_path=trace_log_dir + "/browser_log/trace.log",
|
||||
# )
|
||||
#
|
||||
# logging.info(f"❌ {error_msg}")
|
||||
#
|
||||
# action_response = ActionResponse(
|
||||
# success=False, message=error_msg, metadata=metadata.model_dump()
|
||||
# )
|
||||
# return TextContent(
|
||||
# type="text",
|
||||
# text=json.dumps(
|
||||
# action_response.model_dump()
|
||||
# ), # Empty string instead of None
|
||||
# **{"metadata": {}}, # Pass as additional fields
|
||||
# )
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Browser automation failed: {str(e)}"
|
||||
error_trace = traceback.format_exc()
|
||||
|
||||
logging.info(f"Browser execution error: {error_trace}")
|
||||
|
||||
metadata = BrowserMetadata(
|
||||
task=task,
|
||||
execution_successful=False,
|
||||
error_type="exception",
|
||||
trace_log_path=trace_log_dir + "/browser_log/trace.log",
|
||||
)
|
||||
|
||||
logging.info(f"❌ {error_msg}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"{error_msg}\n\nError details: {error_trace}",
|
||||
metadata=metadata.model_dump(),
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
load_dotenv(override=True)
|
||||
logging.info("Starting browser-server MCP server!")
|
||||
mcp.run(transport="stdio")
|
||||
Vendored
Vendored
+16
@@ -0,0 +1,16 @@
|
||||
[project]
|
||||
name = "browseruse-server"
|
||||
version = "0.1.0"
|
||||
description = "Add your description here"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
aiohttp= ">=3.12.15"
|
||||
dependencies = [
|
||||
"mcp",
|
||||
"aiohttp~=3.12.15",
|
||||
"requests~=2.32.4",
|
||||
"fastmcp~=2.11.3",
|
||||
"tavily-python~=0.7.10",
|
||||
"browser-use~=0.5.5"
|
||||
|
||||
]
|
||||
Vendored
+99
@@ -0,0 +1,99 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
from typing import Union
|
||||
|
||||
from browser_use import Agent, BrowserSession
|
||||
from browser_use.llm import ChatOpenAI
|
||||
from dotenv import load_dotenv
|
||||
from fastmcp.server.server import FastMCP
|
||||
from mcp.types import TextContent
|
||||
from pydantic import Field
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
|
||||
|
||||
mcp = FastMCP("browseruse-server")
|
||||
|
||||
extended_browser_system_prompt = """
|
||||
|
||||
# 效率指南
|
||||
0. 如果用户问句里面有明确的URL地址，可以直接访问
|
||||
1. 使用包含任务关键术语的特定搜索查询
|
||||
2. 避免被无关信息分散注意力
|
||||
3. 如果被付费墙阻挡，尝试使用 archive.org 或类似替代方案
|
||||
4. 清晰简洁地记录每个重要发现
|
||||
5. 用最少的浏览步骤精确提取必要信息。
|
||||
|
||||
## 输出规则
|
||||
1、如果任务要求查找相关资料内容，可以返回相关资料内容的总结
|
||||
2、如果任务要求查询相关下载，那尽量的找出可以下载的链接，返回可以下载的链接(比如github上可以下载的链接一般是：https://raw.githubusercontent.com/，如果是huggingface相关的，要找到对应文件页面里面raw标签里面的地址)，下载的链接需要根据任务选择最匹配的地址，示例：
|
||||
Example 1:
|
||||
```json
|
||||
{
|
||||
"url": "https://"
|
||||
}
|
||||
```
|
||||
"""
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""Use browser to visit a web page, extract content,
|
||||
and optionally download files/images, ...
|
||||
|
||||
Returns a dict with execution trace, answer (extracted content),
|
||||
and downloaded file/image paths."""
|
||||
)
|
||||
async def complete_browser_task(
|
||||
task: str = Field(
|
||||
...,
|
||||
description=(
|
||||
"任务相关描述"
|
||||
),
|
||||
)
|
||||
)-> Union[str, TextContent]:
|
||||
browser_session = BrowserSession(
|
||||
# headless=True, # 关键参数：设置为 True 启用无头模式
|
||||
headless=False, # 关键参数：设置为 True 启用无头模式
|
||||
)
|
||||
try:
|
||||
load_dotenv()
|
||||
model = os.environ['LLM_MODEL_NAME']
|
||||
base_url = os.environ['LLM_BASE_URL']
|
||||
api_key = os.environ['LLM_API_KEY']
|
||||
if not model or not base_url or not api_key:
|
||||
logging.warning(f"Query failed: LLM_MODEL_NAME, LLM_BASE_URL, LLM_API_KEY parameters incomplete")
|
||||
return None
|
||||
llm = ChatOpenAI(
|
||||
model=model,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
temperature=float(0.1),
|
||||
)
|
||||
agent = Agent(
|
||||
task=task,
|
||||
llm=llm,
|
||||
extend_system_message=extended_browser_system_prompt,
|
||||
browser_session=browser_session
|
||||
)
|
||||
final_result = ""
|
||||
result = await agent.run()
|
||||
logging.info(f"complete_browser_task result: {result}")
|
||||
if result and result.history[-1] and result.history[-1].result and result.history[-1].result[0]:
|
||||
final_result = result.history[-1].result[0].extracted_content
|
||||
|
||||
return final_result
|
||||
except BaseException as e:
|
||||
logging.warning(f"complete_browser_task error: {e}")
|
||||
return None
|
||||
except Exception:
|
||||
logging.warning(f"complete_browser_task error: {traceback.format_exc()}")
|
||||
return None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
load_dotenv(override=True)
|
||||
logging.info("Starting browseruse-server MCP server!")
|
||||
mcp.run(transport="stdio")
|
||||
|
||||
+118
@@ -0,0 +1,118 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Script to generate mcp_tool_schema.json"""
|
||||
|
||||
import asyncio
|
||||
from contextlib import AsyncExitStack
|
||||
from datetime import timedelta
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from mcp.client.sse import sse_client
|
||||
from mcp.client.stdio import stdio_client, StdioServerParameters
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.client.streamable_http import streamablehttp_client
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def _list_tools(name: str, config: dict):
|
||||
logger.info(f"Starting tool server {name} with config {config}")
|
||||
async with AsyncExitStack() as exit_stack:
|
||||
try:
|
||||
if config.get("type") == "sse":
|
||||
read_stream, write_stream = await exit_stack.enter_async_context(
|
||||
sse_client(
|
||||
url=config.get("url", ""),
|
||||
headers=config.get("headers", {}),
|
||||
timeout=config.get("timeout", 5),
|
||||
sse_read_timeout=config.get("sse_read_timeout", 60 * 5),
|
||||
auth=config.get("auth", None),
|
||||
)
|
||||
)
|
||||
|
||||
elif config.get("type") == "streamable_http":
|
||||
read_stream, write_stream, _ = await exit_stack.enter_async_context(
|
||||
streamablehttp_client(
|
||||
url=config.get("url", ""),
|
||||
headers=config.get("headers", {}),
|
||||
timeout=config.get("timeout", 60),
|
||||
sse_read_timeout=config.get("sse_read_timeout", 60 * 5),
|
||||
auth=config.get("auth", None),
|
||||
)
|
||||
)
|
||||
|
||||
else: # stdio
|
||||
base_folder = (
|
||||
Path(__file__).resolve().parent
|
||||
)
|
||||
server_params = StdioServerParameters(
|
||||
command=config.get("command", ""),
|
||||
args=config.get("args", []),
|
||||
env=config.get("env", {}),
|
||||
cwd=str(base_folder / config.get("cwd", "")),
|
||||
)
|
||||
read_stream, write_stream = await exit_stack.enter_async_context(
|
||||
stdio_client(server=server_params)
|
||||
)
|
||||
|
||||
# Create session and tool manager
|
||||
session = await exit_stack.enter_async_context(
|
||||
ClientSession(
|
||||
read_stream,
|
||||
write_stream,
|
||||
read_timeout_seconds=timedelta(
|
||||
seconds=config.get("read_timeout", 60)
|
||||
),
|
||||
)
|
||||
)
|
||||
await session.initialize()
|
||||
result = await session.list_tools()
|
||||
return name, result.tools if result else []
|
||||
except Exception as e:
|
||||
logger.error(f"Error starting tool server {name}: {config}, {e}")
|
||||
return name, []
|
||||
|
||||
|
||||
async def list_mcp_server_tools():
|
||||
from mcp_config import mcp_config as config
|
||||
|
||||
json_tools = {}
|
||||
for server_name, server_config in config.get("mcpServers", {}).items():
|
||||
result = await _list_tools(server_name, server_config)
|
||||
name, tools = result
|
||||
logger.info(f"Result for {name}: {tools}")
|
||||
|
||||
tools_dict = [
|
||||
{
|
||||
"name": tool.name,
|
||||
"title": tool.title,
|
||||
"description": tool.description,
|
||||
"inputSchema": tool.inputSchema,
|
||||
"outputSchema": tool.outputSchema,
|
||||
"annotations": (
|
||||
{
|
||||
"title": tool.annotations.title,
|
||||
"readOnlyHint": tool.annotations.readOnlyHint,
|
||||
"destructiveHint": tool.annotations.destructiveHint,
|
||||
"idempotentHint": tool.annotations.idempotentHint,
|
||||
"openWorldHint": tool.annotations.openWorldHint,
|
||||
}
|
||||
if tool.annotations
|
||||
else None
|
||||
),
|
||||
"meta": tool.meta,
|
||||
}
|
||||
for tool in tools
|
||||
]
|
||||
json_tools[name] = tools_dict
|
||||
|
||||
if json_tools:
|
||||
with open(Path(__file__).resolve().parent / "mcp_tool_schema.json", "w") as f:
|
||||
f.write(json.dumps(json_tools, indent=4, ensure_ascii=False))
|
||||
else:
|
||||
logger.error("No results for list_results!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(list_mcp_server_tools())
|
||||
Vendored
+22
@@ -0,0 +1,22 @@
|
||||
[project]
|
||||
name = "documents-server"
|
||||
version = "0.1.0"
|
||||
description = "Add your description here"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
aiohttp= ">=3.12.15"
|
||||
dependencies = [
|
||||
"mcp",
|
||||
"aiohttp~=3.12.15",
|
||||
"requests~=2.32.4",
|
||||
"fastmcp~=2.11.3",
|
||||
"chardet~=3.0.4",
|
||||
"pandas~=2.3.0",
|
||||
"python-magic~=0.4.27",
|
||||
"python-docx~=1.2.0",
|
||||
"filetype~=1.2.0",
|
||||
"datalab-python-sdk~=0.1.4",
|
||||
"python-pptx~=1.0.2",
|
||||
|
||||
|
||||
]
|
||||
Vendored
+226
@@ -0,0 +1,226 @@
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import requests
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import Any, Literal
|
||||
|
||||
|
||||
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)"
|
||||
)
|
||||
|
||||
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")
|
||||
|
||||
|
||||
def _validate_file_path(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)
|
||||
if not path.is_absolute():
|
||||
path = path.expanduser().resolve()
|
||||
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"File not found: {path}")
|
||||
|
||||
return path
|
||||
|
||||
def is_url(path_or_url: str) -> bool:
|
||||
"""
|
||||
Check if the given string is a URL.
|
||||
|
||||
Args:
|
||||
path_or_url: String to check
|
||||
|
||||
Returns:
|
||||
bool: True if the string is a URL, False otherwise
|
||||
"""
|
||||
parsed = urlparse(path_or_url)
|
||||
return bool(parsed.scheme and parsed.netloc)
|
||||
|
||||
|
||||
def get_mime_type(file_path: str, default_mime: str | None = None) -> str:
|
||||
"""
|
||||
Detect MIME type of a file using python-magic if available,
|
||||
otherwise fallback to extension-based detection.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
default_mime: Default MIME type to return if detection fails
|
||||
|
||||
Returns:
|
||||
str: Detected MIME type
|
||||
"""
|
||||
# Try using python-magic for accurate MIME type detection
|
||||
try:
|
||||
import magic
|
||||
mime = magic.Magic(mime=True)
|
||||
return mime.from_file(file_path)
|
||||
except (AttributeError, IOError):
|
||||
# Fallback to extension-based detection
|
||||
extension_mime_map = {
|
||||
# Audio formats
|
||||
".mp3": "audio/mpeg",
|
||||
".wav": "audio/wav",
|
||||
".ogg": "audio/ogg",
|
||||
".m4a": "audio/mp4",
|
||||
".flac": "audio/flac",
|
||||
# Image formats
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".png": "image/png",
|
||||
".gif": "image/gif",
|
||||
".webp": "image/webp",
|
||||
".bmp": "image/bmp",
|
||||
".tiff": "image/tiff",
|
||||
# Video formats
|
||||
".mp4": "video/mp4",
|
||||
".avi": "video/x-msvideo",
|
||||
".mov": "video/quicktime",
|
||||
".mkv": "video/x-matroska",
|
||||
".webm": "video/webm",
|
||||
}
|
||||
|
||||
ext = Path(file_path).suffix.lower()
|
||||
return extension_mime_map.get(ext, default_mime or "application/octet-stream")
|
||||
|
||||
|
||||
def get_file_from_source(
|
||||
source: str,
|
||||
max_size_mb: float = 100.0,
|
||||
timeout: int = 60,
|
||||
) -> tuple[str, str, bytes]:
|
||||
"""
|
||||
Unified function to get file content from a URL or local path with validation.
|
||||
|
||||
Args:
|
||||
source: URL or local file path
|
||||
max_size_mb: Maximum allowed file size in MB
|
||||
timeout: Timeout for URL requests in seconds
|
||||
|
||||
Returns:
|
||||
Tuple[str, str, bytes]: (file_path, mime_type, file_content)
|
||||
- For URLs, file_path will be a temporary file path
|
||||
- For local files, file_path will be the original path
|
||||
|
||||
Raises:
|
||||
ValueError: When file doesn't exist, exceeds size limit, or has invalid MIME type
|
||||
IOError: When file cannot be read
|
||||
requests.RequestException: When URL request fails
|
||||
"""
|
||||
max_size_bytes = max_size_mb * 1024 * 1024
|
||||
|
||||
if is_url(source):
|
||||
# Handle URL source
|
||||
try:
|
||||
# Make a HEAD request first to check content length
|
||||
head_response = requests.head(source, timeout=timeout, allow_redirects=True)
|
||||
head_response.raise_for_status()
|
||||
|
||||
# Check content length if available
|
||||
content_length = head_response.headers.get("content-length")
|
||||
if content_length and int(content_length) > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({int(content_length) / (1024 * 1024):.2f} MB) "
|
||||
f"exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Download the file
|
||||
response = requests.get(source, timeout=timeout, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
# Read content with size checking
|
||||
content = b""
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
if len(content) + len(chunk) > max_size_bytes:
|
||||
raise ValueError(f"File size exceeds maximum allowed size ({max_size_mb} MB)")
|
||||
content += chunk
|
||||
|
||||
# Create temporary file
|
||||
parsed_url = urlparse(source)
|
||||
filename = os.path.basename(parsed_url.path) or "downloaded_file"
|
||||
|
||||
# Create temporary file with proper extension
|
||||
suffix = Path(filename).suffix or ".tmp"
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
|
||||
temp_file.write(content)
|
||||
temp_path = temp_file.name
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(temp_path)
|
||||
|
||||
return temp_path, mime_type, content
|
||||
|
||||
except requests.RequestException as e:
|
||||
raise requests.RequestException(f"Failed to download file from URL: {e}: {traceback.format_exc()}")
|
||||
except Exception as e:
|
||||
raise IOError(f"Error processing URL: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
else:
|
||||
# Handle local file path
|
||||
file_path = Path(source)
|
||||
|
||||
# Check if file exists
|
||||
if not file_path.exists():
|
||||
raise ValueError(f"File does not exist: {source}")
|
||||
|
||||
if not file_path.is_file():
|
||||
raise ValueError(f"Path is not a file: {source}")
|
||||
|
||||
# Check file size
|
||||
file_size = file_path.stat().st_size
|
||||
if file_size > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({file_size / (1024 * 1024):.2f} MB) exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Read file content
|
||||
try:
|
||||
with open(file_path, "rb") as f:
|
||||
content = f.read()
|
||||
except Exception as e:
|
||||
raise IOError(f"Cannot read file {source}: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(str(file_path))
|
||||
|
||||
return str(file_path), mime_type, content
|
||||
Vendored
+446
@@ -0,0 +1,446 @@
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from typing import Any, Union, Literal
|
||||
|
||||
import chardet
|
||||
import pandas as pd
|
||||
from dotenv import load_dotenv
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.types import TextContent
|
||||
from pydantic import Field
|
||||
|
||||
from base import (
|
||||
ActionResponse,
|
||||
_validate_file_path,
|
||||
get_file_from_source,
|
||||
DocumentMetadata,
|
||||
)
|
||||
|
||||
load_dotenv()
|
||||
workspace = Path.home()
|
||||
_media_output_dir = workspace / "extracted_media"
|
||||
_media_output_dir.mkdir(exist_ok=True, parents=True)
|
||||
|
||||
supported_extensions: set = {".csv", ".tsv", ".txt"}
|
||||
|
||||
|
||||
mcp = FastMCP(
|
||||
"documents-csv-server",
|
||||
instructions="""
|
||||
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.
|
||||
""",
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
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)
|
||||
- Optional data visualizations
|
||||
- Memory-efficient processing for large files
|
||||
"""
|
||||
)
|
||||
async def extract_csv_content(
|
||||
file_path: str = Field(
|
||||
description="Path to the CSV document file to extract content from"
|
||||
),
|
||||
output_format: Literal["markdown"] = Field(
|
||||
default="markdown", description="Output format: 'markdown'"
|
||||
),
|
||||
# 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)"
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
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 = _validate_file_path(file_path)
|
||||
logging.info(f"Processing CSV file: {file_path.name}")
|
||||
|
||||
# Extract CSV content
|
||||
extraction_result = _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 = _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=[], # CSV files don't contain media files
|
||||
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}
|
||||
|
||||
logging.info(
|
||||
f"Successfully extracted CSV content from {file_path.name}-({extraction_result['total_rows']} rows, {extraction_result['total_columns']} columns "
|
||||
)
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=formatted_content, metadata=final_metadata
|
||||
)
|
||||
output_dict = {"artifact_type": "MARKDOWN", "artifact_data": formatted_content}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except FileNotFoundError as e:
|
||||
logging.error(f"File not found: {str(e)}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"File not found: {str(e)}",
|
||||
metadata={"error_type": "file_not_found"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
except ValueError as e:
|
||||
logging.error(f"Invalid input: {str(e)}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}",
|
||||
metadata={"error_type": "invalid_input"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
except Exception as e:
|
||||
logging.error(f"CSV extraction failed: {str(e)}: {traceback.format_exc()}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"CSV extraction failed: {str(e)}",
|
||||
metadata={"error_type": "extraction_error"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
List all supported CSV formats for extraction.
|
||||
"""
|
||||
)
|
||||
async def list_supported_formats() -> Union[str, TextContent]:
|
||||
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()
|
||||
]
|
||||
)
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True,
|
||||
message=f"Supported CSV formats:\n\n{format_list}",
|
||||
metadata={
|
||||
"supported_formats": list(supported_formats.keys()),
|
||||
"total_formats": len(supported_formats),
|
||||
},
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": {
|
||||
"supported_formats": list(supported_formats.keys()),
|
||||
"total_formats": len(supported_formats),
|
||||
}
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
def _detect_encoding(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)
|
||||
|
||||
logging.info(
|
||||
f"Detected encoding: {encoding} (confidence: {confidence:.2f})"
|
||||
)
|
||||
return encoding if confidence > 0.7 else "utf-8"
|
||||
except Exception as e:
|
||||
logging.warning(f"Encoding detection failed: {e}, using utf-8")
|
||||
return "utf-8"
|
||||
|
||||
|
||||
def _detect_delimiter(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)
|
||||
logging.info(f"Detected delimiter: '{detected_delimiter}'")
|
||||
return detected_delimiter
|
||||
else:
|
||||
return ","
|
||||
except Exception as e:
|
||||
logging.warning(f"Delimiter detection failed: {e}, using comma")
|
||||
return ","
|
||||
|
||||
|
||||
def _extract_csv_content(
|
||||
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 = _detect_encoding(file_path)
|
||||
if delimiter is None:
|
||||
delimiter = _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:
|
||||
logging.error(f"Failed to read CSV file: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def _format_content_for_llm(
|
||||
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
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
load_dotenv(override=True)
|
||||
logging.info("Starting documents-csv-server MCP server!")
|
||||
mcp.run(transport="stdio")
|
||||
Vendored
+783
@@ -0,0 +1,783 @@
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import traceback
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from typing import Any, Union, Literal
|
||||
|
||||
from docx import Document
|
||||
from docx.document import Document as DocumentType
|
||||
from dotenv import load_dotenv
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.types import TextContent
|
||||
from pydantic import Field
|
||||
|
||||
from base import (
|
||||
ActionResponse,
|
||||
_validate_file_path,
|
||||
get_file_from_source,
|
||||
DocumentMetadata,
|
||||
)
|
||||
|
||||
load_dotenv()
|
||||
workspace = Path.home()
|
||||
_media_output_dir = workspace / "extracted_media"
|
||||
_media_output_dir.mkdir(exist_ok=True, parents=True)
|
||||
|
||||
supported_extensions = {".docx", ".doc"}
|
||||
|
||||
|
||||
mcp = FastMCP(
|
||||
"documents-docx-server",
|
||||
instructions="""
|
||||
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.
|
||||
""",
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
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)
|
||||
- Document structure analysis
|
||||
"""
|
||||
)
|
||||
async def extract_docx_content(
|
||||
file_path: str = Field(
|
||||
description="Path to the DOCX/DOC document file to extract content from"
|
||||
),
|
||||
output_format: Literal["markdown"] = Field(
|
||||
default="markdown", description="Output format: 'markdown'"
|
||||
),
|
||||
# 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",
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
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 = _validate_file_path(file_path)
|
||||
logging.info(f"Processing DOCX document: {file_path.name}")
|
||||
|
||||
# Extract embedded media if requested
|
||||
saved_media = []
|
||||
if extract_images and file_path.suffix.lower() == ".docx":
|
||||
saved_media = _extract_images_from_docx(file_path, file_path.stem)
|
||||
|
||||
# Extract document content
|
||||
extraction_result = _extract_content_from_docx(
|
||||
file_path,
|
||||
extract_tables=extract_tables,
|
||||
extract_headers_footers=extract_headers_footers,
|
||||
)
|
||||
|
||||
# Format content for LLM consumption
|
||||
formatted_content = _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}
|
||||
|
||||
logging.info(
|
||||
f"Successfully extracted DOCX content from {file_path.name} -({extraction_result['word_count']} words, {extraction_result['structure']['tables_count']} tables, {len(saved_media)} media files)"
|
||||
)
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=formatted_content, metadata=final_metadata
|
||||
)
|
||||
output_dict = {"artifact_type": "MARKDOWN", "artifact_data": formatted_content}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except FileNotFoundError as e:
|
||||
logging.error(f"File not found: {str(e)}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"File not found: {str(e)}",
|
||||
metadata={"error_type": "file_not_found"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
except ValueError as e:
|
||||
logging.error(f"Invalid input: {str(e)}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}",
|
||||
metadata={"error_type": "invalid_input"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
except ImportError as e:
|
||||
logging.error(f"Missing dependency: {str(e)}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Missing dependency: {str(e)}. Please install python-docx: pip install python-docx",
|
||||
metadata={"error_type": "missing_dependency"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
except Exception as e:
|
||||
logging.error(f"DOCX extraction failed: {str(e)}: {traceback.format_exc()}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"DOCX extraction failed: {str(e)}",
|
||||
metadata={"error_type": "extraction_error"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except FileNotFoundError as e:
|
||||
logging.error(f"File not found: {str(e)}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"File not found: {str(e)}",
|
||||
metadata={"error_type": "file_not_found"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
except ValueError as e:
|
||||
logging.error(f"Invalid input: {str(e)}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}",
|
||||
metadata={"error_type": "invalid_input"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
except Exception as e:
|
||||
logging.error(f"CSV extraction failed: {str(e)}: {traceback.format_exc()}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"CSV extraction failed: {str(e)}",
|
||||
metadata={"error_type": "extraction_error"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
List all supported document formats for extraction.
|
||||
"""
|
||||
)
|
||||
async def list_supported_formats() -> Union[str, TextContent]:
|
||||
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()
|
||||
]
|
||||
)
|
||||
|
||||
action_response = 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),
|
||||
},
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": {
|
||||
"supported_formats": list(supported_formats.keys()),
|
||||
"total_formats": len(supported_formats),
|
||||
}
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
def _extract_images_from_docx(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 = _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,
|
||||
}
|
||||
)
|
||||
|
||||
logging.info(f"Extracted {media_type}: {media_filename}")
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to extract media file {media_file}: {e}")
|
||||
|
||||
except Exception as e:
|
||||
logging.warning(f"Could not extract media from DOCX: {e}")
|
||||
|
||||
return saved_media
|
||||
|
||||
|
||||
def _extract_document_structure(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(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(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 = _media_output_dir
|
||||
docx_path = output_dir / f"{doc_path.stem}.docx"
|
||||
|
||||
# Check if already converted
|
||||
if docx_path.exists():
|
||||
logging.info(f"Using existing converted file: {docx_path.name}")
|
||||
return docx_path
|
||||
|
||||
logging.info(f"Converting .doc to .docx: {doc_path.name}")
|
||||
|
||||
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():
|
||||
logging.info(f"Conversion successful: {docx_path.name}")
|
||||
return docx_path
|
||||
else:
|
||||
raise RuntimeError("Conversion completed but output file not found")
|
||||
|
||||
except subprocess.CalledProcessError as e:
|
||||
logging.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:
|
||||
logging.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(
|
||||
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 = _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 = _extract_document_structure(doc)
|
||||
|
||||
# Extract tables if requested
|
||||
tables = []
|
||||
if extract_tables:
|
||||
tables = _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:
|
||||
logging.error(f"Failed to extract content from DOCX: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def _format_content_for_llm(
|
||||
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)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
load_dotenv(override=True)
|
||||
logging.info("Starting documents-docx-server MCP server!")
|
||||
mcp.run(transport="stdio")
|
||||
Vendored
+258
@@ -0,0 +1,258 @@
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
|
||||
import filetype
|
||||
from pathlib import Path
|
||||
from typing import Literal, Union, Any
|
||||
|
||||
import requests
|
||||
from datalab_sdk.models import ConversionResult, ConvertOptions, OCROptions
|
||||
from dotenv import load_dotenv
|
||||
from mcp.server import FastMCP
|
||||
from mcp.types import TextContent
|
||||
from pydantic import Field, BaseModel
|
||||
from pydantic.fields import FieldInfo
|
||||
from requests import Response
|
||||
|
||||
|
||||
|
||||
|
||||
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")
|
||||
|
||||
# 设置Python路径，确保子进程能找到core模块
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
project_root = os.path.abspath(os.path.join(current_dir, '..', '..'))
|
||||
if project_root not in sys.path:
|
||||
sys.path.insert(0, project_root)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
load_dotenv()
|
||||
_models_loaded = False
|
||||
_marker_models = None
|
||||
workspace = Path.home()
|
||||
_extracted_texts_dir = workspace / "processed_documents"
|
||||
_extracted_texts_dir.mkdir(exist_ok=True, parents=True)
|
||||
DATALAB_URL = "https://www.datalab.to/api/v1"
|
||||
|
||||
supported_extensions = {".pdf"}
|
||||
|
||||
class DocumentResult(BaseModel):
|
||||
file_path: str = Field(..., description="Path to the processed document")
|
||||
conversion_result: ConversionResult | None = (
|
||||
Field(None, description="Conversion result"),
|
||||
)
|
||||
errors: list[str] | None = Field(
|
||||
None, description="Error messages if processing failed"
|
||||
)
|
||||
|
||||
|
||||
class DocumentEntity(BaseModel):
|
||||
"""Represents a document entity with its metadata and content."""
|
||||
|
||||
file_name: str = Field(..., description="Name of the document file")
|
||||
file_path: str | Path = Field(..., description="Absolute path to the document file")
|
||||
file_types: Literal[
|
||||
# pdf
|
||||
"application/pdf",
|
||||
# spreadsheet
|
||||
"application/vnd.ms-excel", # xls
|
||||
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", # xlsx
|
||||
"application/vnd.oasis.opendocument.spreadsheet", # ods
|
||||
# word
|
||||
"application/msword", # doc
|
||||
"application/vnd.openxmlformats-officedocument.wordprocessingml.document", # docx
|
||||
"application/vnd.oasis.opendocument.text", # odt
|
||||
# presentation
|
||||
"application/vnd.ms-powerpoint", # ppt
|
||||
"application/vnd.openxmlformats-officedocument.presentationml.presentation", # pptx
|
||||
"application/vnd.oasis.opendocument.presentation", # odp
|
||||
# html
|
||||
"text/html",
|
||||
# image
|
||||
"image/jpeg",
|
||||
"image/png",
|
||||
"image/gif",
|
||||
"image/tiff",
|
||||
"image/webp",
|
||||
# epub
|
||||
"application/epub+zip",
|
||||
] = Field(..., description="MIME type of the document file")
|
||||
|
||||
|
||||
mcp = FastMCP("documents-pdf-server", instructions="""
|
||||
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.
|
||||
""")
|
||||
|
||||
@mcp.tool(
|
||||
description=(
|
||||
"Convert PDF document to markdown foramt. "
|
||||
)
|
||||
)
|
||||
async def convert_document_to_markdown(
|
||||
file_path: str = Field(..., description="Path to the document file"),
|
||||
paginate: bool = Field(False, description="Add page delimiters to the output"),
|
||||
) -> Union[str, TextContent]:
|
||||
#-> DocumentResult:
|
||||
"""
|
||||
Process document using Datalab SDK with advanced OCR capabilities.
|
||||
|
||||
Convert document to markdown foramt. "
|
||||
Support PDFs, DOCX, XLSX, PPTX, HTML, and images.
|
||||
|
||||
Returns DocumentResult with conversion results.
|
||||
"""
|
||||
if isinstance(file_path, FieldInfo):
|
||||
file_path: str = file_path.default
|
||||
if isinstance(paginate, FieldInfo):
|
||||
paginate: bool = paginate.default
|
||||
|
||||
result: DocumentResult = DocumentResult(file_path=file_path)
|
||||
try:
|
||||
file_entity: DocumentEntity = _prepare_file_entity(file_path)
|
||||
session: dict = await _establish_document_session(
|
||||
file_entity,
|
||||
options=ConvertOptions(
|
||||
output_format="markdown",
|
||||
paginate=paginate,
|
||||
use_llm=True,
|
||||
max_pages=None,
|
||||
),
|
||||
)
|
||||
conversion_result: ConversionResult = await _poll_result(
|
||||
session["request_check_url"]
|
||||
)
|
||||
result.conversion_result = conversion_result
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(result.model_dump()),
|
||||
}
|
||||
action_response = ActionResponse(
|
||||
success=True,
|
||||
message=json.dumps(result.model_dump()),
|
||||
metadata=output_dict,
|
||||
)
|
||||
except Exception as e:
|
||||
result.errors = [traceback.format_exc(), str(e)]
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=json.dumps(result.model_dump()),
|
||||
metadata={},
|
||||
)
|
||||
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": {}} # Pass as additional fields
|
||||
)
|
||||
|
||||
async def _poll_result(
|
||||
check_url: str,
|
||||
max_polls: int = 300,
|
||||
poll_interval: int = 2,
|
||||
) -> ConversionResult:
|
||||
"""
|
||||
Poll the Datalab API for the result of the document processing.
|
||||
"""
|
||||
try:
|
||||
for _ in range(max_polls):
|
||||
await asyncio.sleep(poll_interval)
|
||||
response: Response = requests.get(
|
||||
check_url,
|
||||
headers={"X-Api-Key": os.getenv("DATALAB_API_KEY")},
|
||||
timeout=5,
|
||||
)
|
||||
result_data: dict = response.json()
|
||||
if result_data["status"] == "complete":
|
||||
return ConversionResult(
|
||||
success=result_data.get("success", False),
|
||||
output_format="markdown",
|
||||
markdown=result_data.get("markdown"),
|
||||
html=result_data.get("html"),
|
||||
json=result_data.get("json"),
|
||||
#images=result_data.get("images"),
|
||||
metadata=result_data.get("metadata"),
|
||||
error=result_data.get("error"),
|
||||
page_count=result_data.get("page_count"),
|
||||
status=result_data.get("status", "complete"),
|
||||
)
|
||||
raise TimeoutError("Document processing timed out.")
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
f"Failed to poll document result: {traceback.format_exc()}"
|
||||
) from e
|
||||
|
||||
def _prepare_file_entity(file_path: [str | Path]) -> DocumentEntity:
|
||||
"""
|
||||
Prepare a DocumentEntity from the given file path.
|
||||
"""
|
||||
file_path: Path = Path(file_path)
|
||||
if not file_path.exists():
|
||||
raise FileNotFoundError(f"File not found: {file_path}")
|
||||
|
||||
return DocumentEntity(
|
||||
file_name=file_path.resolve().name,
|
||||
file_path=file_path.resolve(),
|
||||
file_types=filetype.guess(file_path).mime or "application/pdf",
|
||||
)
|
||||
|
||||
async def _establish_document_session(
|
||||
file_entity: DocumentEntity,
|
||||
options: [ConvertOptions | OCROptions],
|
||||
endpoint: str = "/marker",
|
||||
) -> dict:
|
||||
"""
|
||||
Establish a session for document processing with Datalab API.
|
||||
"""
|
||||
if not os.getenv("DATALAB_API_KEY"):
|
||||
raise ValueError("DATALAB_API_KEY environment variable is not set.")
|
||||
|
||||
try:
|
||||
response: Response = requests.post(
|
||||
url=DATALAB_URL + endpoint,
|
||||
files={
|
||||
"file": (
|
||||
file_entity.file_name,
|
||||
open(file_entity.file_path, "rb"),
|
||||
file_entity.file_types,
|
||||
),
|
||||
"force_ocr": (None, False),
|
||||
"paginate": (
|
||||
None,
|
||||
options.paginate if isinstance(options, ConvertOptions) else False,
|
||||
),
|
||||
"output_format": (None, "markdown"),
|
||||
"use_llm": (
|
||||
None,
|
||||
options.use_llm if isinstance(options, ConvertOptions) else False,
|
||||
),
|
||||
"strip_existing_ocr": (None, False),
|
||||
"disable_image_extraction": (None, False),
|
||||
},
|
||||
headers={"X-Api-Key": os.getenv("DATALAB_API_KEY")},
|
||||
timeout=120,
|
||||
)
|
||||
return response.json()
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
f"Failed to establish document session: {traceback.format_exc()}"
|
||||
) from e
|
||||
|
||||
if __name__ == "__main__":
|
||||
load_dotenv(override=True)
|
||||
logging.info("Starting documents-pdf-server MCP server!")
|
||||
mcp.run(transport="stdio")
|
||||
Vendored
+630
@@ -0,0 +1,630 @@
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import traceback
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from typing import Any, Union, Literal
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from pptx import Presentation
|
||||
from pptx.presentation import Presentation as PresentationType
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.types import TextContent
|
||||
from pydantic import Field
|
||||
|
||||
from base import (
|
||||
ActionResponse,
|
||||
_validate_file_path,
|
||||
get_file_from_source,
|
||||
DocumentMetadata,
|
||||
)
|
||||
|
||||
|
||||
load_dotenv()
|
||||
workspace = Path.home()
|
||||
_media_output_dir = workspace / "extracted_media"
|
||||
_media_output_dir.mkdir(exist_ok=True, parents=True)
|
||||
|
||||
supported_extensions = {".pptx", ".ppt"}
|
||||
|
||||
|
||||
mcp = FastMCP(
|
||||
"documents-pptx-server",
|
||||
instructions="""
|
||||
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.
|
||||
""",
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
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)
|
||||
- LLM-optimized formatting
|
||||
"""
|
||||
)
|
||||
async def extract_pptx_content(
|
||||
file_path: str = Field(
|
||||
description="Path to the PPTX/PPT presentation file to extract content from"
|
||||
),
|
||||
output_format: Literal["markdown"] = Field(
|
||||
default="markdown", description="Output format: 'markdown'"
|
||||
),
|
||||
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",
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
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 = _validate_file_path(file_path)
|
||||
logging.info(f"Processing PPTX presentation: {file_path.name}")
|
||||
|
||||
# Extract embedded media if requested
|
||||
saved_media = []
|
||||
if extract_images and file_path.suffix.lower() == ".pptx":
|
||||
saved_media = _extract_images_from_pptx(file_path, file_path.stem)
|
||||
|
||||
# Extract presentation content
|
||||
extraction_result = _extract_content_from_pptx(
|
||||
file_path, extract_notes=extract_notes
|
||||
)
|
||||
|
||||
# Format content for LLM consumption
|
||||
formatted_content = _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}
|
||||
|
||||
logging.info(
|
||||
f"Successfully extracted content from {file_path.name} -({extraction_result['slide_count']} slides, {len(formatted_content)} characters, {len(saved_media)} media files)"
|
||||
)
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=formatted_content, metadata=final_metadata
|
||||
)
|
||||
output_dict = {"artifact_type": "MARKDOWN", "artifact_data": formatted_content}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except FileNotFoundError as e:
|
||||
logging.error(f"File not found: {str(e)}: {traceback.format_exc()}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"File not found: {str(e)}",
|
||||
metadata={"error_type": "file_not_found"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
except ValueError as e:
|
||||
logging.error(f"Invalid input: {str(e)}: {traceback.format_exc()}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}",
|
||||
metadata={"error_type": "invalid_input"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
except Exception as e:
|
||||
logging.error(f"PPTX extraction failed: {str(e)}: {traceback.format_exc()}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"PPTX extraction failed: {str(e)}",
|
||||
metadata={"error_type": "extraction_error"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
List all supported presentation formats for extraction.
|
||||
"""
|
||||
)
|
||||
async def list_supported_formats() -> Union[str, TextContent]:
|
||||
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()
|
||||
]
|
||||
)
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True,
|
||||
message=f"Supported presentation formats:\n\n{format_list}",
|
||||
metadata={
|
||||
"supported_formats": list(supported_formats.keys()),
|
||||
"total_formats": len(supported_formats),
|
||||
},
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": {
|
||||
"supported_formats": list(supported_formats.keys()),
|
||||
"total_formats": len(supported_formats),
|
||||
},
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
def _extract_images_from_pptx(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 = _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,
|
||||
}
|
||||
)
|
||||
|
||||
logging.info(f"Saved media: {media_filename}")
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to extract media {media_file}: {e}")
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to extract media from PPTX: {e}")
|
||||
|
||||
return saved_media
|
||||
|
||||
|
||||
def _extract_slide_structure(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(
|
||||
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:
|
||||
logging.warning(
|
||||
f"Failed to extract notes from slide {slide_idx + 1}: {e}"
|
||||
)
|
||||
|
||||
slides_content.append(slide_data)
|
||||
|
||||
# Extract presentation structure
|
||||
structure = _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:
|
||||
logging.error(f"Failed to extract content from PPTX: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def _format_content_for_llm(
|
||||
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)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
load_dotenv(override=True)
|
||||
logging.info("Starting documents-pptx-server MCP server!")
|
||||
mcp.run(transport="stdio")
|
||||
Vendored
+721
@@ -0,0 +1,721 @@
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from typing import Any, Union, Literal
|
||||
|
||||
import chardet
|
||||
from dotenv import load_dotenv
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.types import TextContent
|
||||
from pydantic import Field
|
||||
|
||||
from base import (
|
||||
ActionResponse,
|
||||
_validate_file_path,
|
||||
get_file_from_source,
|
||||
DocumentMetadata,
|
||||
get_mime_type,
|
||||
)
|
||||
|
||||
load_dotenv()
|
||||
|
||||
workspace = Path.home()
|
||||
_media_output_dir = workspace / "extracted_media"
|
||||
_media_output_dir.mkdir(exist_ok=True, parents=True)
|
||||
|
||||
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",
|
||||
}
|
||||
|
||||
|
||||
mcp = FastMCP(
|
||||
"documents-txt-server",
|
||||
instructions="""
|
||||
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.
|
||||
""",
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
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
|
||||
"""
|
||||
)
|
||||
async def extract_text_content(
|
||||
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'"
|
||||
),
|
||||
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)",
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
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 = _validate_file_path(file_path)
|
||||
logging.info(f"Processing text document: {file_path.name}")
|
||||
|
||||
# Extract content from text file
|
||||
extraction_result = _extract_text_content(file_path, encoding)
|
||||
|
||||
# Check if file appears to be binary
|
||||
if extraction_result["encoding_info"]["is_binary"]:
|
||||
logging.info("Warning: File appears to contain binary data")
|
||||
|
||||
# Format content for LLM consumption
|
||||
formatted_content = _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)
|
||||
|
||||
logging.info(
|
||||
f"Successfully extracted content from {file_path.name} -({extraction_result['statistics']['character_count']:,} characters,{extraction_result['statistics']['line_count']:,} lines,encoding: {extraction_result['used_encoding']})"
|
||||
)
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=formatted_content, metadata=combined_metadata
|
||||
)
|
||||
output_dict = {"artifact_type": "MARKDOWN", "artifact_data": formatted_content}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except FileNotFoundError as e:
|
||||
logging.error(f"File not found: {str(e)}: {traceback.format_exc()}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"File not found: {str(e)}",
|
||||
metadata={"error_type": "file_not_found"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
except ValueError as e:
|
||||
logging.error(f"Invalid input: {str(e)}: {traceback.format_exc()}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}",
|
||||
metadata={"error_type": "invalid_input"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
except Exception as e:
|
||||
logging.error(f"Text extraction failed: {str(e)}: {traceback.format_exc()}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Text extraction failed: {str(e)}",
|
||||
metadata={"error_type": "extraction_error"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
List all supported text formats for extraction.
|
||||
"""
|
||||
)
|
||||
async def list_supported_formats() -> Union[str, TextContent]:
|
||||
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()
|
||||
]
|
||||
)
|
||||
|
||||
action_response = 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,
|
||||
},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
def _validate_file_path(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 = path.expanduser().resolve()
|
||||
#path = 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 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 _is_likely_text(sample):
|
||||
logging.info(
|
||||
f"Detected text file without standard extension: {path.suffix}"
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported file type: {path.suffix}. "
|
||||
f"Supported types: {', '.join(sorted(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(supported_extensions))}"
|
||||
) from e
|
||||
|
||||
return path
|
||||
|
||||
|
||||
def _is_likely_text(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(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 _is_likely_text(raw_data[:1024])
|
||||
|
||||
except Exception as e:
|
||||
logging.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(
|
||||
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 = _detect_encoding(file_path)
|
||||
|
||||
if encoding:
|
||||
# Use specified encoding
|
||||
target_encoding = encoding
|
||||
logging.info(f"Using specified encoding: {encoding}")
|
||||
else:
|
||||
# Use detected encoding
|
||||
target_encoding = encoding_info["detected_encoding"]
|
||||
logging.info(
|
||||
f"Detected encoding: {target_encoding} (confidence: {encoding_info['confidence']:.2f})"
|
||||
)
|
||||
|
||||
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 = _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:
|
||||
logging.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()
|
||||
|
||||
logging.info(
|
||||
f"Successfully read with fallback encoding: {fallback_encoding}"
|
||||
)
|
||||
|
||||
# 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": _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(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(
|
||||
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)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
load_dotenv(override=True)
|
||||
logging.info("Starting documents-txt-server MCP server!")
|
||||
mcp.run(transport="stdio")
|
||||
Vendored
+15
@@ -0,0 +1,15 @@
|
||||
[project]
|
||||
name = "download-server"
|
||||
version = "0.1.0"
|
||||
description = "Add your description here"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
aiohttp= ">=3.12.15"
|
||||
dependencies = [
|
||||
"mcp",
|
||||
"aiohttp~=3.12.15",
|
||||
"requests~=2.32.4",
|
||||
"fastmcp~=2.11.3",
|
||||
"python-magic~=0.4.27",
|
||||
|
||||
]
|
||||
Vendored
+225
@@ -0,0 +1,225 @@
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import magic
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import requests
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import Any, Literal
|
||||
|
||||
|
||||
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)"
|
||||
)
|
||||
|
||||
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")
|
||||
|
||||
|
||||
def _validate_file_path(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)
|
||||
if not path.is_absolute():
|
||||
path = path.expanduser().resolve()
|
||||
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"File not found: {path}")
|
||||
|
||||
return path
|
||||
|
||||
def is_url(path_or_url: str) -> bool:
|
||||
"""
|
||||
Check if the given string is a URL.
|
||||
|
||||
Args:
|
||||
path_or_url: String to check
|
||||
|
||||
Returns:
|
||||
bool: True if the string is a URL, False otherwise
|
||||
"""
|
||||
parsed = urlparse(path_or_url)
|
||||
return bool(parsed.scheme and parsed.netloc)
|
||||
|
||||
|
||||
def get_mime_type(file_path: str, default_mime: str | None = None) -> str:
|
||||
"""
|
||||
Detect MIME type of a file using python-magic if available,
|
||||
otherwise fallback to extension-based detection.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
default_mime: Default MIME type to return if detection fails
|
||||
|
||||
Returns:
|
||||
str: Detected MIME type
|
||||
"""
|
||||
# Try using python-magic for accurate MIME type detection
|
||||
try:
|
||||
mime = magic.Magic(mime=True)
|
||||
return mime.from_file(file_path)
|
||||
except (AttributeError, IOError):
|
||||
# Fallback to extension-based detection
|
||||
extension_mime_map = {
|
||||
# Audio formats
|
||||
".mp3": "audio/mpeg",
|
||||
".wav": "audio/wav",
|
||||
".ogg": "audio/ogg",
|
||||
".m4a": "audio/mp4",
|
||||
".flac": "audio/flac",
|
||||
# Image formats
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".png": "image/png",
|
||||
".gif": "image/gif",
|
||||
".webp": "image/webp",
|
||||
".bmp": "image/bmp",
|
||||
".tiff": "image/tiff",
|
||||
# Video formats
|
||||
".mp4": "video/mp4",
|
||||
".avi": "video/x-msvideo",
|
||||
".mov": "video/quicktime",
|
||||
".mkv": "video/x-matroska",
|
||||
".webm": "video/webm",
|
||||
}
|
||||
|
||||
ext = Path(file_path).suffix.lower()
|
||||
return extension_mime_map.get(ext, default_mime or "application/octet-stream")
|
||||
|
||||
|
||||
def get_file_from_source(
|
||||
source: str,
|
||||
max_size_mb: float = 100.0,
|
||||
timeout: int = 60,
|
||||
) -> tuple[str, str, bytes]:
|
||||
"""
|
||||
Unified function to get file content from a URL or local path with validation.
|
||||
|
||||
Args:
|
||||
source: URL or local file path
|
||||
max_size_mb: Maximum allowed file size in MB
|
||||
timeout: Timeout for URL requests in seconds
|
||||
|
||||
Returns:
|
||||
Tuple[str, str, bytes]: (file_path, mime_type, file_content)
|
||||
- For URLs, file_path will be a temporary file path
|
||||
- For local files, file_path will be the original path
|
||||
|
||||
Raises:
|
||||
ValueError: When file doesn't exist, exceeds size limit, or has invalid MIME type
|
||||
IOError: When file cannot be read
|
||||
requests.RequestException: When URL request fails
|
||||
"""
|
||||
max_size_bytes = max_size_mb * 1024 * 1024
|
||||
|
||||
if is_url(source):
|
||||
# Handle URL source
|
||||
try:
|
||||
# Make a HEAD request first to check content length
|
||||
head_response = requests.head(source, timeout=timeout, allow_redirects=True)
|
||||
head_response.raise_for_status()
|
||||
|
||||
# Check content length if available
|
||||
content_length = head_response.headers.get("content-length")
|
||||
if content_length and int(content_length) > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({int(content_length) / (1024 * 1024):.2f} MB) "
|
||||
f"exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Download the file
|
||||
response = requests.get(source, timeout=timeout, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
# Read content with size checking
|
||||
content = b""
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
if len(content) + len(chunk) > max_size_bytes:
|
||||
raise ValueError(f"File size exceeds maximum allowed size ({max_size_mb} MB)")
|
||||
content += chunk
|
||||
|
||||
# Create temporary file
|
||||
parsed_url = urlparse(source)
|
||||
filename = os.path.basename(parsed_url.path) or "downloaded_file"
|
||||
|
||||
# Create temporary file with proper extension
|
||||
suffix = Path(filename).suffix or ".tmp"
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
|
||||
temp_file.write(content)
|
||||
temp_path = temp_file.name
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(temp_path)
|
||||
|
||||
return temp_path, mime_type, content
|
||||
|
||||
except requests.RequestException as e:
|
||||
raise requests.RequestException(f"Failed to download file from URL: {e}: {traceback.format_exc()}")
|
||||
except Exception as e:
|
||||
raise IOError(f"Error processing URL: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
else:
|
||||
# Handle local file path
|
||||
file_path = Path(source)
|
||||
|
||||
# Check if file exists
|
||||
if not file_path.exists():
|
||||
raise ValueError(f"File does not exist: {source}")
|
||||
|
||||
if not file_path.is_file():
|
||||
raise ValueError(f"Path is not a file: {source}")
|
||||
|
||||
# Check file size
|
||||
file_size = file_path.stat().st_size
|
||||
if file_size > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({file_size / (1024 * 1024):.2f} MB) exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Read file content
|
||||
try:
|
||||
with open(file_path, "rb") as f:
|
||||
content = f.read()
|
||||
except Exception as e:
|
||||
raise IOError(f"Cannot read file {source}: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(str(file_path))
|
||||
|
||||
return str(file_path), mime_type, content
|
||||
Vendored
+510
@@ -0,0 +1,510 @@
|
||||
import json
|
||||
import logging
|
||||
import shutil
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Union
|
||||
from urllib.parse import urlparse
|
||||
import requests
|
||||
from dotenv import load_dotenv
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.types import TextContent
|
||||
from pydantic import Field, BaseModel
|
||||
|
||||
from base import ActionResponse
|
||||
|
||||
load_dotenv()
|
||||
workspace = Path.home()
|
||||
|
||||
default_timeout = 60 * 3 # 3 minutes timeout
|
||||
max_file_size = 1024 * 1024 * 1024 # 1GB limit
|
||||
supported_schemes = {"http", "https"}
|
||||
headers = {
|
||||
"User-Agent": (
|
||||
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
|
||||
"AppleWebKit/537.36 (KHTML, like Gecko) "
|
||||
"Chrome/91.0.4472.124 Safari/537.36"
|
||||
),
|
||||
"Accept-Language": "en-US,en;q=0.9",
|
||||
}
|
||||
|
||||
|
||||
class DownloadResult(BaseModel):
|
||||
"""Individual download operation result with structured data."""
|
||||
|
||||
url: str
|
||||
file_path: str
|
||||
success: bool
|
||||
file_size: int | None = None
|
||||
duration: str
|
||||
timestamp: str
|
||||
error_message: str | None = None
|
||||
|
||||
|
||||
class DownloadMetadata(BaseModel):
|
||||
"""Metadata for download operation results."""
|
||||
|
||||
url: str
|
||||
output_path: str
|
||||
timeout_seconds: int
|
||||
overwrite_enabled: bool
|
||||
execution_time: float | None = None
|
||||
file_size_bytes: int | None = None
|
||||
content_type: str | None = None
|
||||
status_code: int | None = None
|
||||
error_type: str | None = None
|
||||
headers_used: bool = False
|
||||
|
||||
|
||||
mcp = FastMCP(
|
||||
"download-server",
|
||||
instructions="""
|
||||
MCP service for file download operations with comprehensive controls.
|
||||
|
||||
Provides secure file download capabilities including:
|
||||
- HTTP/HTTPS URL support
|
||||
- Configurable timeout controls
|
||||
- Custom headers for authentication
|
||||
- Path validation and directory creation
|
||||
- LLM-friendly result formatting
|
||||
- Error handling and logging
|
||||
""",
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Download a file from a URL with comprehensive options and controls.
|
||||
|
||||
This tool provides secure file download capabilities with:
|
||||
- HTTP/HTTPS URL support
|
||||
- Configurable timeout controls
|
||||
- Path validation and directory creation
|
||||
- File size limits and safety checks
|
||||
- LLM-optimized result formatting
|
||||
"""
|
||||
)
|
||||
async def download_file(
|
||||
url: str = Field(description="HTTP/HTTPS URL of the file to download"),
|
||||
output_file_path: str = Field(
|
||||
description="Local path where the file should be saved (absolute or relative to workspace)"
|
||||
),
|
||||
overwrite: bool = Field(
|
||||
default=False,
|
||||
description="Whether to overwrite existing files (default: False)",
|
||||
),
|
||||
timeout: int = Field(
|
||||
default=60, description="Download timeout in seconds (default: 60)"
|
||||
),
|
||||
output_format: str = Field(
|
||||
default="markdown", description="Output format: 'markdown', 'json', or 'text'"
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
if isinstance(url, FieldInfo):
|
||||
url = url.default
|
||||
if isinstance(output_file_path, FieldInfo):
|
||||
output_file_path = output_file_path.default
|
||||
if isinstance(overwrite, FieldInfo):
|
||||
overwrite = overwrite.default
|
||||
if isinstance(timeout, FieldInfo):
|
||||
timeout = timeout.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
try:
|
||||
# Validate URL
|
||||
url_valid, url_error = _validate_url(url)
|
||||
if not url_valid:
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid URL: {url_error}",
|
||||
metadata=DownloadMetadata(
|
||||
url=url,
|
||||
output_path=output_file_path,
|
||||
timeout_seconds=timeout,
|
||||
overwrite_enabled=overwrite,
|
||||
error_type="invalid_url",
|
||||
).model_dump(),
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
# Resolve output path
|
||||
output_path = _resolve_output_path(output_file_path)
|
||||
|
||||
# Check if file exists and overwrite setting
|
||||
if output_path.exists() and not overwrite:
|
||||
existing_size = output_path.stat().st_size
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"File already exists at {output_path} ({existing_size:,} bytes) and overwrite is disabled",
|
||||
metadata=DownloadMetadata(
|
||||
url=url,
|
||||
output_path=str(output_path),
|
||||
timeout_seconds=timeout,
|
||||
overwrite_enabled=overwrite,
|
||||
file_size_bytes=existing_size,
|
||||
error_type="file_exists",
|
||||
).model_dump(),
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
# Perform download
|
||||
start_time = time.time()
|
||||
result = await _download_file_async(url, output_path, timeout, headers)
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
# Format output
|
||||
formatted_output = _format_download_output(result, output_format)
|
||||
|
||||
# Create metadata
|
||||
metadata = DownloadMetadata(
|
||||
url=url,
|
||||
output_path=str(output_path),
|
||||
timeout_seconds=timeout,
|
||||
overwrite_enabled=overwrite,
|
||||
execution_time=execution_time,
|
||||
file_size_bytes=result.file_size,
|
||||
headers_used=headers is not None,
|
||||
)
|
||||
|
||||
if not result.success:
|
||||
metadata.error_type = "download_failure"
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=result.success,
|
||||
message=formatted_output,
|
||||
metadata=metadata.model_dump(),
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
),
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to download file: {str(e)}"
|
||||
logging.error(f"Download error: {traceback.format_exc()}")
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=error_msg,
|
||||
metadata=DownloadMetadata(
|
||||
url=url,
|
||||
output_path=output_file_path,
|
||||
timeout_seconds=timeout,
|
||||
overwrite_enabled=overwrite,
|
||||
error_type="internal_error",
|
||||
).model_dump(),
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Get information about download service capabilities and configuration."""
|
||||
)
|
||||
async def get_download_capabilities() -> Union[str, TextContent]:
|
||||
capabilities = {
|
||||
"requests_available": requests is not None,
|
||||
"supported_schemes": list(supported_schemes),
|
||||
"supported_features": [
|
||||
"HTTP/HTTPS URL downloads",
|
||||
"Configurable timeout controls",
|
||||
"Custom headers support",
|
||||
"Path validation and directory creation",
|
||||
"File size limits and safety checks",
|
||||
"Multiple output formats (markdown, json, text)",
|
||||
"LLM-optimized result formatting",
|
||||
"Comprehensive error handling",
|
||||
],
|
||||
"supported_formats": ["markdown", "json", "text"],
|
||||
"configuration": {
|
||||
"default_timeout": default_timeout,
|
||||
"max_file_size_bytes": max_file_size,
|
||||
"workspace": str(workspace),
|
||||
},
|
||||
"safety_features": [
|
||||
"URL validation",
|
||||
"File size limits",
|
||||
"Timeout controls",
|
||||
"Path validation",
|
||||
"Overwrite protection",
|
||||
"Error handling and logging",
|
||||
],
|
||||
}
|
||||
|
||||
max_size_mb = max_file_size / (1024 * 1024)
|
||||
formatted_info = f"""# Download Service Capabilities
|
||||
|
||||
## Status
|
||||
- **Workspace:** `{workspace}`
|
||||
|
||||
## Supported Features
|
||||
{chr(10).join(f"- {feature}" for feature in capabilities["supported_features"])}
|
||||
|
||||
## Supported URL Schemes
|
||||
{chr(10).join(f"- {scheme}://" for scheme in capabilities["supported_schemes"])}
|
||||
|
||||
## Supported Output Formats
|
||||
{chr(10).join(f"- {fmt}" for fmt in capabilities["supported_formats"])}
|
||||
|
||||
## Configuration
|
||||
- **Default Timeout:** {capabilities["configuration"]["default_timeout"]} seconds
|
||||
- **Max File Size:** {max_file_size:,} bytes ({max_size_mb:.1f} MB)
|
||||
|
||||
## Safety Features
|
||||
{chr(10).join(f"- {feature}" for feature in capabilities["safety_features"])}
|
||||
"""
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True,
|
||||
message=formatted_info,
|
||||
metadata=capabilities,
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
),
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
def _validate_url(url: str) -> tuple[bool, str | None]:
|
||||
"""Validate URL format and scheme.
|
||||
|
||||
Args:
|
||||
url: URL to validate
|
||||
|
||||
Returns:
|
||||
Tuple of (is_valid, error_message)
|
||||
"""
|
||||
try:
|
||||
parsed = urlparse(url)
|
||||
|
||||
if not parsed.scheme:
|
||||
return False, "URL must include a scheme (http:// or https://)"
|
||||
|
||||
if parsed.scheme.lower() not in supported_schemes:
|
||||
return (
|
||||
False,
|
||||
f"Unsupported URL scheme: {parsed.scheme}. Supported: {', '.join(supported_schemes)}",
|
||||
)
|
||||
|
||||
if not parsed.netloc:
|
||||
return False, "URL must include a valid domain"
|
||||
|
||||
return True, None
|
||||
|
||||
except Exception as e:
|
||||
return False, f"Invalid URL format: {str(e)}"
|
||||
|
||||
|
||||
def _resolve_output_path(output_path: str) -> Path:
|
||||
"""Resolve and validate output file path.
|
||||
|
||||
Args:
|
||||
output_path: Output file path (absolute or relative)
|
||||
|
||||
Returns:
|
||||
Resolved Path object
|
||||
"""
|
||||
path = Path(output_path).expanduser()
|
||||
|
||||
if not path.is_absolute():
|
||||
path = workspace / path
|
||||
|
||||
# Ensure parent directory exists
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
return path.resolve()
|
||||
|
||||
|
||||
def _format_download_output(
|
||||
result: DownloadResult, output_format: str = "markdown"
|
||||
) -> str:
|
||||
"""Format download results for LLM consumption.
|
||||
|
||||
Args:
|
||||
result: Download execution result
|
||||
output_format: Format type ('markdown', 'json', 'text')
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if output_format == "json":
|
||||
return json.dumps(result.model_dump(), indent=2)
|
||||
|
||||
elif output_format == "text":
|
||||
output_parts = [
|
||||
f"URL: {result.url}",
|
||||
f"File Path: {result.file_path}",
|
||||
f"Status: {'SUCCESS' if result.success else 'FAILED'}",
|
||||
f"Duration: {result.duration}",
|
||||
f"Timestamp: {result.timestamp}",
|
||||
]
|
||||
|
||||
if result.file_size is not None:
|
||||
output_parts.append(f"File Size: {result.file_size:,} bytes")
|
||||
|
||||
if result.error_message:
|
||||
output_parts.append(f"Error: {result.error_message}")
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
else: # markdown (default)
|
||||
status_emoji = "✅" if result.success else "❌"
|
||||
|
||||
output_parts = [
|
||||
f"# File Download {status_emoji}",
|
||||
f"**URL:** `{result.url}`",
|
||||
f"**File Path:** `{result.file_path}`",
|
||||
f"**Status:** {'SUCCESS' if result.success else 'FAILED'}",
|
||||
f"**Duration:** {result.duration}",
|
||||
f"**Timestamp:** {result.timestamp}",
|
||||
]
|
||||
|
||||
if result.file_size is not None:
|
||||
size_mb = result.file_size / (1024 * 1024)
|
||||
output_parts.append(
|
||||
f"**File Size:** {result.file_size:,} bytes ({size_mb:.2f} MB)"
|
||||
)
|
||||
|
||||
if result.error_message:
|
||||
output_parts.extend(
|
||||
["\n## Error Details", f"```\n{result.error_message}\n```"]
|
||||
)
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
|
||||
async def _download_file_async(
|
||||
url: str, output_path: Path, timeout: int, headers: dict[str, str] | None
|
||||
) -> DownloadResult:
|
||||
"""Download file asynchronously with comprehensive error handling.
|
||||
|
||||
Args:
|
||||
url: URL to download from
|
||||
output_path: Local path to save file
|
||||
timeout: Request timeout in seconds
|
||||
headers: Optional custom headers
|
||||
|
||||
Returns:
|
||||
DownloadResult with execution details
|
||||
"""
|
||||
start_time = datetime.now()
|
||||
|
||||
try:
|
||||
logging.info(f"📥 Starting download: {url}")
|
||||
|
||||
with requests.get(
|
||||
url, stream=True, timeout=timeout, headers=headers
|
||||
) as response:
|
||||
response.raise_for_status()
|
||||
|
||||
# Check content length if available
|
||||
content_length = response.headers.get("content-length")
|
||||
if content_length and int(content_length) > max_file_size:
|
||||
raise ValueError(
|
||||
f"File too large: {content_length} bytes (max: {max_file_size})"
|
||||
)
|
||||
|
||||
# Download file
|
||||
with open(output_path, "wb") as f:
|
||||
shutil.copyfileobj(response.raw, f)
|
||||
|
||||
file_size = output_path.stat().st_size
|
||||
duration = str(datetime.now() - start_time)
|
||||
|
||||
logging.info(f"✅ Download completed: {file_size:,} bytes")
|
||||
|
||||
return DownloadResult(
|
||||
url=url,
|
||||
file_path=str(output_path),
|
||||
success=True,
|
||||
file_size=file_size,
|
||||
duration=duration,
|
||||
timestamp=start_time.isoformat(),
|
||||
)
|
||||
|
||||
except requests.exceptions.Timeout:
|
||||
duration = str(datetime.now() - start_time)
|
||||
error_msg = f"Download timed out after {timeout} seconds"
|
||||
logging.info(f"⏰ {error_msg}")
|
||||
|
||||
return DownloadResult(
|
||||
url=url,
|
||||
file_path=str(output_path),
|
||||
success=False,
|
||||
duration=duration,
|
||||
timestamp=start_time.isoformat(),
|
||||
error_message=error_msg,
|
||||
)
|
||||
|
||||
except requests.exceptions.RequestException as e:
|
||||
duration = str(datetime.now() - start_time)
|
||||
error_msg = f"Request failed: {str(e)}"
|
||||
logging.info(f"❌ {error_msg}")
|
||||
|
||||
return DownloadResult(
|
||||
url=url,
|
||||
file_path=str(output_path),
|
||||
success=False,
|
||||
duration=duration,
|
||||
timestamp=start_time.isoformat(),
|
||||
error_message=error_msg,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
duration = str(datetime.now() - start_time)
|
||||
error_msg = f"Unexpected error: {str(e)}"
|
||||
logging.info(f"💥 {error_msg}")
|
||||
|
||||
return DownloadResult(
|
||||
url=url,
|
||||
file_path=str(output_path),
|
||||
success=False,
|
||||
duration=duration,
|
||||
timestamp=start_time.isoformat(),
|
||||
error_message=error_msg,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
load_dotenv(override=True)
|
||||
logging.info("Starting download-server MCP server!")
|
||||
mcp.run(transport="stdio")
|
||||
Vendored
+16
@@ -0,0 +1,16 @@
|
||||
[project]
|
||||
name = "e2b-code-server"
|
||||
version = "0.1.0"
|
||||
description = "Run code in a specified e2b sandbox."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
aiohttp= ">=3.12.15"
|
||||
dependencies = [
|
||||
"mcp",
|
||||
"Pillow~=10.4.0",
|
||||
"requests~=2.32.4",
|
||||
"openai~=1.93.0",
|
||||
"fastmcp~=2.11.3",
|
||||
"e2b-code-interpreter~=1.2.0",
|
||||
]
|
||||
|
||||
Vendored
+86
@@ -0,0 +1,86 @@
|
||||
from e2b_code_interpreter import Sandbox
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
from dotenv import load_dotenv
|
||||
import os
|
||||
import sys
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Initialize MCP server
|
||||
mcp = FastMCP("e2b-code-server")
|
||||
|
||||
|
||||
@mcp.tool(description="Upload local file to e2b sandbox.")
|
||||
async def e2b_upload_file(
|
||||
path: str = Field(
|
||||
description="The local file path to upload."
|
||||
)
|
||||
) -> str:
|
||||
"""
|
||||
Upload local file to e2b sandbox.
|
||||
|
||||
Args:
|
||||
path (str): The local file path to upload.
|
||||
|
||||
Returns:
|
||||
str: E2b file path and sandbox_id.
|
||||
|
||||
"""
|
||||
try:
|
||||
os.environ["E2B_API_KEY"] = os.getenv("E2B_API_KEY")
|
||||
sbx = Sandbox()
|
||||
local_file_name = os.path.basename(path)
|
||||
e2b_file_path = f"/home/user/{local_file_name}"
|
||||
# Read local file relative to the current working directory
|
||||
with open(path, "rb") as file:
|
||||
# Upload file to the sandbox to absolute path
|
||||
sbx.files.write(e2b_file_path, file)
|
||||
return f"{e2b_file_path}, {sbx.sandbox_id}"
|
||||
except Exception as e:
|
||||
return f"Upload failed. Error: {str(e)}"
|
||||
|
||||
|
||||
@mcp.tool(description="Run code in a specified e2b sandbox.")
|
||||
async def e2b_run_code(
|
||||
sandbox_id: str = Field(
|
||||
default=None,
|
||||
description="The sandbox id to run code in, if you have uploaded a file, you should use the sandbox_id returned by the e2b_upload_file function."
|
||||
),
|
||||
code_block: str = Field(
|
||||
default=None,
|
||||
description="The code block to run in e2b sandbox."
|
||||
),
|
||||
) -> str:
|
||||
"""
|
||||
Run code in a specified e2b sandbox.
|
||||
|
||||
Args:
|
||||
sandbox_id (str): The sandbox id to run code in.
|
||||
code_block (str): The code block to run in e2b sandbox.
|
||||
|
||||
Returns:
|
||||
str: The result of running the code block.
|
||||
"""
|
||||
try:
|
||||
os.environ["E2B_API_KEY"] = os.getenv("E2B_API_KEY")
|
||||
sbx = Sandbox(
|
||||
sandbox_id=sandbox_id,
|
||||
)
|
||||
execution = sbx.run_code(code_block)
|
||||
return execution.logs
|
||||
except Exception as e:
|
||||
return f"Run code failed. Error: {str(e)}"
|
||||
|
||||
|
||||
# Run the server when the script is executed directly
|
||||
if __name__ == "__main__":
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
|
||||
datefmt='%Y-%m-%d %H:%M:%S'
|
||||
)
|
||||
load_dotenv()
|
||||
logger.info("Starting E2b Code MCP Server...")
|
||||
mcp.run(transport='stdio')
|
||||
Vendored
Vendored
+16
@@ -0,0 +1,16 @@
|
||||
[project]
|
||||
name = "googlesearch-server"
|
||||
version = "0.1.0"
|
||||
description = "Add your description here"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
aiohttp= ">=3.12.15"
|
||||
dependencies = [
|
||||
"mcp",
|
||||
"aiohttp~=3.12.15",
|
||||
"requests~=2.32.4",
|
||||
"fastmcp~=2.11.3",
|
||||
"arxiv~=2.2.0",
|
||||
"python-magic~=0.4.27",
|
||||
|
||||
]
|
||||
Vendored
+224
@@ -0,0 +1,224 @@
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import magic
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import requests
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import Any, Literal
|
||||
|
||||
|
||||
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)"
|
||||
)
|
||||
|
||||
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")
|
||||
|
||||
|
||||
def _validate_file_path(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.exists():
|
||||
raise FileNotFoundError(f"File not found: {path}")
|
||||
|
||||
return path
|
||||
|
||||
def is_url(path_or_url: str) -> bool:
|
||||
"""
|
||||
Check if the given string is a URL.
|
||||
|
||||
Args:
|
||||
path_or_url: String to check
|
||||
|
||||
Returns:
|
||||
bool: True if the string is a URL, False otherwise
|
||||
"""
|
||||
parsed = urlparse(path_or_url)
|
||||
return bool(parsed.scheme and parsed.netloc)
|
||||
|
||||
|
||||
def get_mime_type(file_path: str, default_mime: str | None = None) -> str:
|
||||
"""
|
||||
Detect MIME type of a file using python-magic if available,
|
||||
otherwise fallback to extension-based detection.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
default_mime: Default MIME type to return if detection fails
|
||||
|
||||
Returns:
|
||||
str: Detected MIME type
|
||||
"""
|
||||
# Try using python-magic for accurate MIME type detection
|
||||
try:
|
||||
mime = magic.Magic(mime=True)
|
||||
return mime.from_file(file_path)
|
||||
except (AttributeError, IOError):
|
||||
# Fallback to extension-based detection
|
||||
extension_mime_map = {
|
||||
# Audio formats
|
||||
".mp3": "audio/mpeg",
|
||||
".wav": "audio/wav",
|
||||
".ogg": "audio/ogg",
|
||||
".m4a": "audio/mp4",
|
||||
".flac": "audio/flac",
|
||||
# Image formats
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".png": "image/png",
|
||||
".gif": "image/gif",
|
||||
".webp": "image/webp",
|
||||
".bmp": "image/bmp",
|
||||
".tiff": "image/tiff",
|
||||
# Video formats
|
||||
".mp4": "video/mp4",
|
||||
".avi": "video/x-msvideo",
|
||||
".mov": "video/quicktime",
|
||||
".mkv": "video/x-matroska",
|
||||
".webm": "video/webm",
|
||||
}
|
||||
|
||||
ext = Path(file_path).suffix.lower()
|
||||
return extension_mime_map.get(ext, default_mime or "application/octet-stream")
|
||||
|
||||
|
||||
def get_file_from_source(
|
||||
source: str,
|
||||
max_size_mb: float = 100.0,
|
||||
timeout: int = 60,
|
||||
) -> tuple[str, str, bytes]:
|
||||
"""
|
||||
Unified function to get file content from a URL or local path with validation.
|
||||
|
||||
Args:
|
||||
source: URL or local file path
|
||||
max_size_mb: Maximum allowed file size in MB
|
||||
timeout: Timeout for URL requests in seconds
|
||||
|
||||
Returns:
|
||||
Tuple[str, str, bytes]: (file_path, mime_type, file_content)
|
||||
- For URLs, file_path will be a temporary file path
|
||||
- For local files, file_path will be the original path
|
||||
|
||||
Raises:
|
||||
ValueError: When file doesn't exist, exceeds size limit, or has invalid MIME type
|
||||
IOError: When file cannot be read
|
||||
requests.RequestException: When URL request fails
|
||||
"""
|
||||
max_size_bytes = max_size_mb * 1024 * 1024
|
||||
|
||||
if is_url(source):
|
||||
# Handle URL source
|
||||
try:
|
||||
# Make a HEAD request first to check content length
|
||||
head_response = requests.head(source, timeout=timeout, allow_redirects=True)
|
||||
head_response.raise_for_status()
|
||||
|
||||
# Check content length if available
|
||||
content_length = head_response.headers.get("content-length")
|
||||
if content_length and int(content_length) > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({int(content_length) / (1024 * 1024):.2f} MB) "
|
||||
f"exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Download the file
|
||||
response = requests.get(source, timeout=timeout, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
# Read content with size checking
|
||||
content = b""
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
if len(content) + len(chunk) > max_size_bytes:
|
||||
raise ValueError(f"File size exceeds maximum allowed size ({max_size_mb} MB)")
|
||||
content += chunk
|
||||
|
||||
# Create temporary file
|
||||
parsed_url = urlparse(source)
|
||||
filename = os.path.basename(parsed_url.path) or "downloaded_file"
|
||||
|
||||
# Create temporary file with proper extension
|
||||
suffix = Path(filename).suffix or ".tmp"
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
|
||||
temp_file.write(content)
|
||||
temp_path = temp_file.name
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(temp_path)
|
||||
|
||||
return temp_path, mime_type, content
|
||||
|
||||
except requests.RequestException as e:
|
||||
raise requests.RequestException(f"Failed to download file from URL: {e}: {traceback.format_exc()}")
|
||||
except Exception as e:
|
||||
raise IOError(f"Error processing URL: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
else:
|
||||
# Handle local file path
|
||||
file_path = Path(source)
|
||||
|
||||
# Check if file exists
|
||||
if not file_path.exists():
|
||||
raise ValueError(f"File does not exist: {source}")
|
||||
|
||||
if not file_path.is_file():
|
||||
raise ValueError(f"Path is not a file: {source}")
|
||||
|
||||
# Check file size
|
||||
file_size = file_path.stat().st_size
|
||||
if file_size > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({file_size / (1024 * 1024):.2f} MB) exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Read file content
|
||||
try:
|
||||
with open(file_path, "rb") as f:
|
||||
content = f.read()
|
||||
except Exception as e:
|
||||
raise IOError(f"Cannot read file {source}: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(str(file_path))
|
||||
|
||||
return str(file_path), mime_type, content
|
||||
+443
@@ -0,0 +1,443 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from typing import Union
|
||||
|
||||
import requests
|
||||
from dotenv import load_dotenv
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.types import TextContent
|
||||
from pydantic import Field, BaseModel
|
||||
|
||||
from base import ActionResponse
|
||||
|
||||
load_dotenv()
|
||||
workspace = Path.home()
|
||||
google_api_key = os.getenv("GOOGLE_API_KEY")
|
||||
google_cse_id = os.getenv("GOOGLE_CSE_ID")
|
||||
|
||||
supported_extensions: set = {".csv", ".tsv", ".txt"}
|
||||
|
||||
|
||||
class SearchResult(BaseModel):
|
||||
"""Individual search result with structured data."""
|
||||
|
||||
id: str
|
||||
title: str
|
||||
url: str
|
||||
snippet: str
|
||||
source: str
|
||||
display_link: str | None = None
|
||||
formatted_url: str | None = None
|
||||
|
||||
|
||||
class SearchMetadata(BaseModel):
|
||||
"""Metadata for search operation results."""
|
||||
|
||||
query: str
|
||||
search_engine: str
|
||||
total_results: int
|
||||
search_time: float | None = None
|
||||
language: str = "en"
|
||||
country: str = "us"
|
||||
safe_search: bool = True
|
||||
error_type: str | None = None
|
||||
api_quota_used: bool = False
|
||||
|
||||
|
||||
mcp = FastMCP("google-search-server", instructions="""
|
||||
Search MCP Server
|
||||
|
||||
This module provides MCP server functionality for performing web searches using various search engines.
|
||||
It supports structured queries and returns LLM-friendly formatted search results.
|
||||
|
||||
Key features:
|
||||
- Perform web searches using Google Custom Search API
|
||||
- Filter and format search results for LLM consumption
|
||||
- Validate and process search queries with metadata tracking
|
||||
|
||||
Main functions:
|
||||
- mcp_search_google: Searches the web using Google Custom Search API
|
||||
- mcp_get_search_capabilities: Returns information about search service capabilities
|
||||
""")
|
||||
|
||||
|
||||
@mcp.tool(description="""
|
||||
Search the web using Google Custom Search API.
|
||||
|
||||
This tool provides comprehensive web search capabilities with:
|
||||
- Google Custom Search API integration
|
||||
- Configurable result count and filtering
|
||||
- Safe search and localization options
|
||||
- LLM-optimized result formatting
|
||||
- Detailed metadata tracking
|
||||
""")
|
||||
async def search_google(
|
||||
query: str = Field(description="The search query string to search for"),
|
||||
num_results: int = Field(default=5, description="Number of search results to return (1-10, default: 5)"),
|
||||
safe_search: bool = Field(default=True, description="Whether to enable safe search filtering"),
|
||||
language: str = Field(default="en", description="Language code for search results (e.g., 'en', 'es', 'fr')"),
|
||||
country: str = Field(default="us", description="Country code for search results (e.g., 'us', 'uk', 'ca')"),
|
||||
output_format: str = Field(default="json", description="Output format: 'markdown', 'json', or 'text'"),
|
||||
) -> Union[str, TextContent]:
|
||||
if isinstance(query, FieldInfo):
|
||||
query = query.default
|
||||
if isinstance(num_results, FieldInfo):
|
||||
num_results = num_results.default
|
||||
if isinstance(safe_search, FieldInfo):
|
||||
safe_search = safe_search.default
|
||||
if isinstance(language, FieldInfo):
|
||||
language = language.default
|
||||
if isinstance(country, FieldInfo):
|
||||
country = country.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
try:
|
||||
# Validate API credentials
|
||||
if not google_api_key or not google_cse_id:
|
||||
return ActionResponse(
|
||||
success=False,
|
||||
message=(
|
||||
"Google Search API credentials not configured. "
|
||||
"Please set GOOGLE_API_KEY and GOOGLE_CSE_ID environment variables."
|
||||
),
|
||||
metadata={"error_type": "missing_credentials"},
|
||||
)
|
||||
|
||||
# Validate parameters
|
||||
validated_query, validated_num_results = _validate_search_parameters(query, num_results)
|
||||
|
||||
logging.info(f"🔍 Searching Google for: '{validated_query}'")
|
||||
|
||||
# Prepare API request
|
||||
start_time = time.time()
|
||||
|
||||
url = "https://www.googleapis.com/customsearch/v1"
|
||||
params = {
|
||||
"key": google_api_key,
|
||||
"cx": google_cse_id,
|
||||
"q": validated_query,
|
||||
"num": validated_num_results,
|
||||
"safe": "active" if safe_search else "off",
|
||||
"hl": language,
|
||||
"gl": country,
|
||||
}
|
||||
|
||||
# Make API request
|
||||
response = requests.get(url, params=params, timeout=10)
|
||||
response.raise_for_status()
|
||||
|
||||
search_time = time.time() - start_time
|
||||
data = response.json()
|
||||
|
||||
# Parse search results
|
||||
search_results = []
|
||||
search_items = [] # 新增：用于存储加工后的数据
|
||||
|
||||
if "items" in data:
|
||||
# 使用字典去重，避免重复URL
|
||||
url_dict = {}
|
||||
|
||||
for i, item in enumerate(data["items"]):
|
||||
# 原有的SearchResult处理逻辑保持不变
|
||||
result = SearchResult(
|
||||
id=f"google-{i}",
|
||||
title=item.get("title", ""),
|
||||
url=item.get("link", ""),
|
||||
snippet=item.get("snippet", ""),
|
||||
source="google",
|
||||
display_link=item.get("displayLink", ""),
|
||||
formatted_url=item.get("formattedUrl", ""),
|
||||
)
|
||||
search_results.append(result)
|
||||
|
||||
# 新增：同时加工search_items数据
|
||||
url = item.get("link", "")
|
||||
if url not in url_dict:
|
||||
url_dict[url] = {
|
||||
"title": item.get("title", ""),
|
||||
"url": url,
|
||||
"snippet": item.get("snippet", "")[:100] + "..." if len(
|
||||
item.get("snippet", "")) > 100 else item.get("snippet", ""),
|
||||
"content": item.get("snippet", "")[:1000] + "..." if len(
|
||||
item.get("snippet", "")) > 1000 else item.get("snippet", ""),
|
||||
}
|
||||
|
||||
# 将字典值转换为列表
|
||||
search_items = list(url_dict.values())
|
||||
|
||||
# Format results based on requested format
|
||||
if "json" == "json":
|
||||
formatted_content = {
|
||||
"query": validated_query,
|
||||
"results": [result.model_dump() for result in search_results],
|
||||
"count": len(search_results),
|
||||
}
|
||||
|
||||
message_content = formatted_content
|
||||
elif output_format.lower() == "text":
|
||||
if search_results:
|
||||
result_lines = []
|
||||
for i, result in enumerate(search_results, 1):
|
||||
result_lines.append(f"{i}. {result.title}")
|
||||
result_lines.append(f" URL: {result.url}")
|
||||
result_lines.append(f" Summary: {result.snippet}")
|
||||
result_lines.append("") # Empty line
|
||||
message_content = "\n".join(result_lines)
|
||||
else:
|
||||
message_content = f"No results found for: {validated_query}"
|
||||
else: # markdown (default)
|
||||
message_content = _format_search_results_for_llm(search_results, validated_query)
|
||||
|
||||
# Prepare metadata
|
||||
metadata = SearchMetadata(
|
||||
query=validated_query,
|
||||
search_engine="google",
|
||||
total_results=len(search_results),
|
||||
search_time=search_time,
|
||||
language=language,
|
||||
country=country,
|
||||
safe_search=safe_search,
|
||||
api_quota_used=True,
|
||||
)
|
||||
|
||||
logging.info(f"✅ Found {len(search_results)} results in {search_time:.2f}s")
|
||||
|
||||
action_response = ActionResponse(success=True, message=message_content, metadata=metadata.model_dump())
|
||||
search_output_dict = {
|
||||
"artifact_type": "WEB_PAGES",
|
||||
"artifact_data": {
|
||||
"query": query,
|
||||
"results": search_items
|
||||
}
|
||||
}
|
||||
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": search_output_dict} # Pass as additional fields
|
||||
)
|
||||
|
||||
except requests.exceptions.RequestException as e:
|
||||
error_msg = f"Google Search API request failed: {str(e)}"
|
||||
logging.error(f"Search API error: {traceback.format_exc()}")
|
||||
|
||||
metadata = SearchMetadata(
|
||||
query=query, search_engine="google", total_results=0, error_type="api_request_failed"
|
||||
)
|
||||
|
||||
logging.info(f"❌ {error_msg}")
|
||||
|
||||
action_response = ActionResponse(success=False, message=error_msg, metadata=metadata.model_dump())
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": {}} # Pass as additional fields
|
||||
)
|
||||
|
||||
except ValueError as e:
|
||||
error_msg = f"Invalid search parameters: {str(e)}"
|
||||
|
||||
metadata = SearchMetadata(
|
||||
query=query, search_engine="google", total_results=0, error_type="invalid_parameters"
|
||||
)
|
||||
|
||||
logging.info(f"❌ {error_msg}")
|
||||
|
||||
action_response = ActionResponse(success=False, message=error_msg, metadata=metadata.model_dump())
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": {}} # Pass as additional fields
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Search operation failed: {str(e)}"
|
||||
error_trace = traceback.format_exc()
|
||||
|
||||
logging.error(f"Unexpected search error: {error_trace}")
|
||||
|
||||
metadata = SearchMetadata(
|
||||
query=query, search_engine="google", total_results=0, error_type="unexpected_error"
|
||||
)
|
||||
|
||||
logging.info(f"❌ {error_msg}")
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=False, message=f"{error_msg}\n\nError details: {error_trace}", metadata=metadata.model_dump()
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": {}} # Pass as additional fields
|
||||
)
|
||||
|
||||
except FileNotFoundError as e:
|
||||
logging.error(f"File not found: {str(e)}")
|
||||
action_response = ActionResponse(
|
||||
success=False, message=f"File not found: {str(e)}", metadata={"error_type": "file_not_found"}
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": {}} # Pass as additional fields
|
||||
)
|
||||
except ValueError as e:
|
||||
logging.error(f"Invalid input: {str(e)}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}",
|
||||
metadata={"error_type": "invalid_input"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": {}} # Pass as additional fields
|
||||
)
|
||||
except Exception as e:
|
||||
logging.error(f"CSV extraction failed: {str(e)}: {traceback.format_exc()}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"CSV extraction failed: {str(e)}",
|
||||
metadata={"error_type": "extraction_error"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": {}} # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(description="""
|
||||
Get information about search service capabilities and configuration.
|
||||
""")
|
||||
async def get_search_capabilities(
|
||||
) -> Union[str, TextContent]:
|
||||
capabilities = {
|
||||
"search_engines": ["Google Custom Search API"],
|
||||
"supported_features": [
|
||||
"Web search with customizable result count",
|
||||
"Safe search filtering",
|
||||
"Language and country localization",
|
||||
"Multiple output formats (markdown, json, text)",
|
||||
"LLM-optimized result formatting",
|
||||
"Detailed metadata tracking",
|
||||
],
|
||||
"supported_formats": ["markdown", "json", "text"],
|
||||
"configuration": {
|
||||
"google_api_configured": bool(google_api_key and google_cse_id),
|
||||
"max_results_per_query": 10,
|
||||
"default_language": "en",
|
||||
"default_country": "us",
|
||||
"safe_search_default": True,
|
||||
},
|
||||
"limitations": [
|
||||
"Google CSE has daily quota limits",
|
||||
"Maximum 10 results per query",
|
||||
"Requires valid API credentials",
|
||||
],
|
||||
}
|
||||
|
||||
formatted_info = f"""# Search Service Capabilities
|
||||
|
||||
## Available Search Engines
|
||||
{chr(10).join(f"- {engine}" for engine in capabilities["search_engines"])}
|
||||
|
||||
## Features
|
||||
{chr(10).join(f"- {feature}" for feature in capabilities["supported_features"])}
|
||||
|
||||
## Supported Output Formats
|
||||
{chr(10).join(f"- {fmt}" for fmt in capabilities["supported_formats"])}
|
||||
|
||||
## Current Configuration
|
||||
- **Google API Configured:** {capabilities["configuration"]["google_api_configured"]}
|
||||
- **Max Results Per Query:** {capabilities["configuration"]["max_results_per_query"]}
|
||||
- **Default Language:** {capabilities["configuration"]["default_language"]}
|
||||
- **Default Country:** {capabilities["configuration"]["default_country"]}
|
||||
- **Safe Search Default:** {capabilities["configuration"]["safe_search_default"]}
|
||||
|
||||
## Limitations
|
||||
{chr(10).join(f"- {limitation}" for limitation in capabilities["limitations"])}
|
||||
"""
|
||||
|
||||
action_response = ActionResponse(success=True, message=formatted_info, metadata=capabilities)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
),
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": output_dict} # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
def _format_search_results_for_llm(results: list[SearchResult], query: str) -> str:
|
||||
"""Format search results to be LLM-friendly.
|
||||
|
||||
Args:
|
||||
results: List of search results
|
||||
query: Original search query
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if not results:
|
||||
return f"No search results found for query: '{query}'"
|
||||
|
||||
formatted_parts = [f"# Search Results for: '{query}'", f"Found {len(results)} results:\n"]
|
||||
|
||||
for i, result in enumerate(results, 1):
|
||||
result_section = [
|
||||
f"## Result {i}: {result.title}",
|
||||
f"**URL:** {result.url}",
|
||||
f"**Source:** {result.source}",
|
||||
]
|
||||
|
||||
if result.display_link:
|
||||
result_section.append(f"**Domain:** {result.display_link}")
|
||||
|
||||
result_section.append(f"**Summary:** {result.snippet}")
|
||||
result_section.append("") # Empty line for spacing
|
||||
|
||||
formatted_parts.append("\n".join(result_section))
|
||||
|
||||
return "\n".join(formatted_parts)
|
||||
|
||||
|
||||
def _validate_search_parameters(query: str, num_results: int) -> tuple[str, int]:
|
||||
"""Validate and normalize search parameters.
|
||||
|
||||
Args:
|
||||
query: Search query string
|
||||
num_results: Number of results requested
|
||||
|
||||
Returns:
|
||||
Tuple of (validated_query, validated_num_results)
|
||||
|
||||
Raises:
|
||||
ValueError: If parameters are invalid
|
||||
"""
|
||||
if not query or not query.strip():
|
||||
raise ValueError("Search query cannot be empty")
|
||||
|
||||
# Normalize query
|
||||
validated_query = query.strip()
|
||||
|
||||
# Validate and clamp num_results
|
||||
validated_num_results = max(1, min(num_results, 10)) # Google CSE limit is 10
|
||||
|
||||
return validated_query, validated_num_results
|
||||
|
||||
if __name__ == "__main__":
|
||||
load_dotenv(override=True)
|
||||
logging.info("Starting googlesearch-server MCP server!")
|
||||
mcp.run(transport="stdio")
|
||||
+11
@@ -0,0 +1,11 @@
|
||||
[project]
|
||||
name = "hello-world"
|
||||
version = "0.1.0"
|
||||
description = "Add your description here"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
dependencies = ["mcp"]
|
||||
|
||||
[build-system]
|
||||
requires = ["setuptools"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
+23
@@ -0,0 +1,23 @@
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
import logging
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.DEBUG, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
mcp = FastMCP(name="hello-world", log_level="DEBUG")
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def hello_world(name: str) -> str:
|
||||
"""
|
||||
Say hello to the world.
|
||||
"""
|
||||
return f"Hello, {name}!"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
logger.info("Starting hello-world MCP server!")
|
||||
mcp.run(transport="stdio")
|
||||
Vendored
+17
@@ -0,0 +1,17 @@
|
||||
[project]
|
||||
name = "image-server"
|
||||
version = "0.1.0"
|
||||
description = "Solve the question by careful reasoning given the image(s) in given local filepath or url, including reasoning, ocr, etc."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
aiohttp= ">=3.12.15"
|
||||
dependencies = [
|
||||
"mcp",
|
||||
"Pillow~=10.4.0",
|
||||
"requests~=2.32.4",
|
||||
"openai~=1.93.0",
|
||||
"fastmcp~=2.11.3",
|
||||
]
|
||||
|
||||
|
||||
|
||||
+230
@@ -0,0 +1,230 @@
|
||||
"""
|
||||
Image MCP Server
|
||||
|
||||
This module provides MCP server functionality for image processing and analysis.
|
||||
It handles image encoding, optimization, and various image analysis tasks such as
|
||||
OCR (Optical Character Recognition) and visual reasoning.
|
||||
|
||||
The server supports both local image files and remote image URLs with proper validation
|
||||
and handles various image formats including JPEG, PNG, GIF, and others.
|
||||
|
||||
Main functions:
|
||||
- encode_images: Encodes images to base64 format with optimization
|
||||
- optimize_image: Resizes and optimizes images for better performance
|
||||
- Various MCP tools for image analysis and processing
|
||||
"""
|
||||
|
||||
|
||||
import base64
|
||||
import os
|
||||
from io import BytesIO
|
||||
from typing import Any, Dict, List
|
||||
|
||||
import logging
|
||||
|
||||
from PIL import Image
|
||||
from pydantic import Field
|
||||
|
||||
from utils import get_file_from_source
|
||||
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
from openai import OpenAI
|
||||
from dotenv import load_dotenv
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Initialize MCP server
|
||||
mcp = FastMCP("image-server")
|
||||
|
||||
|
||||
IMAGE_OCR = (
|
||||
"Input is a base64 encoded image. Read text from image if present. "
|
||||
"Return a json string with the following format: "
|
||||
'{"image_text": "text from image"}'
|
||||
)
|
||||
|
||||
IMAGE_REASONING = (
|
||||
"Input is a base64 encoded image. Given user's task: {task}, "
|
||||
"solve it following the guide line:\n"
|
||||
"1. Careful visual inspection\n"
|
||||
"2. Contextual reasoning\n"
|
||||
"3. Text transcription where relevant\n"
|
||||
"4. Logical deduction from visual evidence\n"
|
||||
"Return a json string with the following format: "
|
||||
'{"image_reasoning_result": "reasoning result given task and image"}'
|
||||
)
|
||||
|
||||
|
||||
def optimize_image(image_data: bytes, max_size: int = 1024) -> bytes:
|
||||
"""
|
||||
Optimize image by resizing if needed
|
||||
|
||||
Args:
|
||||
image_data: Raw image data
|
||||
max_size: Maximum dimension size in pixels
|
||||
|
||||
Returns:
|
||||
bytes: Optimized image data
|
||||
|
||||
Raises:
|
||||
ValueError: When image cannot be processed
|
||||
"""
|
||||
try:
|
||||
image = Image.open(BytesIO(image_data))
|
||||
|
||||
# Resize if image is too large
|
||||
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.Resampling.LANCZOS)
|
||||
|
||||
# Save to buffer
|
||||
buffered = BytesIO()
|
||||
image_format = image.format if image.format else "JPEG"
|
||||
image.save(buffered, format=image_format)
|
||||
return buffered.getvalue()
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to optimize image: {str(e)}")
|
||||
return image_data # Return original data if optimization fails
|
||||
|
||||
|
||||
def encode_images(image_sources: List[str], with_header: bool = True) -> List[str]:
|
||||
"""
|
||||
Encode images to base64 format with robust file handling
|
||||
|
||||
Args:
|
||||
image_sources: List of URLs or local file paths of images
|
||||
with_header: Whether to include MIME type header
|
||||
|
||||
Returns:
|
||||
List[str]: Base64 encoded image strings, with MIME type prefix if with_header is True
|
||||
|
||||
Raises:
|
||||
ValueError: When image source is invalid or image format is not supported
|
||||
"""
|
||||
if not image_sources:
|
||||
raise ValueError("Image sources cannot be empty")
|
||||
|
||||
images = []
|
||||
for image_source in image_sources:
|
||||
try:
|
||||
# Get file with validation (only image files allowed)
|
||||
file_path, mime_type, content = get_file_from_source(
|
||||
image_source,
|
||||
allowed_mime_prefixes=["image/"],
|
||||
max_size_mb=30.0, # 10MB limit for images
|
||||
type="image",
|
||||
)
|
||||
|
||||
# Optimize image
|
||||
optimized_content = optimize_image(content)
|
||||
|
||||
# Encode to base64
|
||||
image_base64 = base64.b64encode(optimized_content).decode()
|
||||
|
||||
# Format with header if requested
|
||||
final_image = (
|
||||
f"data:{mime_type};base64,{image_base64}"
|
||||
if with_header
|
||||
else image_base64
|
||||
)
|
||||
|
||||
images.append(final_image)
|
||||
|
||||
# Clean up temporary file if it was created for a URL
|
||||
if file_path != os.path.abspath(image_source) and os.path.exists(file_path):
|
||||
os.unlink(file_path)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error encoding image from {image_source}: {str(e)}")
|
||||
raise
|
||||
|
||||
return images
|
||||
|
||||
def image_to_base64(image_path):
|
||||
try:
|
||||
with Image.open(image_path) as image:
|
||||
buffered = BytesIO()
|
||||
image_format = image.format if image.format else "JPEG"
|
||||
image.save(buffered, format=image_format)
|
||||
image_bytes = buffered.getvalue()
|
||||
base64_encoded = base64.b64encode(image_bytes).decode('utf-8')
|
||||
return base64_encoded
|
||||
except Exception as e:
|
||||
logger.error(f"Base64 error: {e}", exc_info=True)
|
||||
return None
|
||||
|
||||
|
||||
def create_image_contents(prompt: str, image_base64: List[str]) -> List[Dict[str, Any]]:
|
||||
"""Create uniform image format for querying llm."""
|
||||
content = [
|
||||
{"type": "text", "text": prompt},
|
||||
]
|
||||
content.extend(
|
||||
[{"type": "image_url", "image_url": {"url": url}} for url in image_base64]
|
||||
)
|
||||
return content
|
||||
|
||||
|
||||
@mcp.tool(description="Solve the question by careful reasoning given the image(s) in given local filepath or url, including reasoning, ocr, etc.")
|
||||
def mcp_image_recognition(
|
||||
image_urls: List[str] = Field(
|
||||
description="The input image(s) in given a list of local filepaths or urls."
|
||||
),
|
||||
question: str = Field(description="The question to ask."),
|
||||
) -> str:
|
||||
"""solve the question by careful reasoning given the image(s) in given filepath or url."""
|
||||
|
||||
try:
|
||||
image_base64 = image_to_base64(image_urls[0])
|
||||
logger.info(f"image_url: {image_urls[0]}")
|
||||
reasoning_prompt = question
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content":
|
||||
[
|
||||
{"type": "text", "text": reasoning_prompt},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/jpeg;base64,{image_base64}"
|
||||
}
|
||||
},
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
client = OpenAI(
|
||||
api_key=os.getenv("IMAGE_LLM_API_KEY"),
|
||||
base_url=os.getenv("IMAGE_LLM_BASE_URL")
|
||||
)
|
||||
response = client.chat.completions.create(
|
||||
model=os.getenv("IMAGE_LLM_MODEL_NAME"),
|
||||
messages=messages,
|
||||
)
|
||||
|
||||
logger.info(f"response: {response}")
|
||||
image_reasoning_result = response.choices[0].message.content
|
||||
|
||||
except Exception as e:
|
||||
image_reasoning_result = ""
|
||||
logger.error(f"image_reasoning_result-Execute error: {e}", exc_info=True)
|
||||
|
||||
logger.info(
|
||||
f"---get_reasoning_by_image-image_reasoning_result:{image_reasoning_result}"
|
||||
)
|
||||
|
||||
return image_reasoning_result
|
||||
|
||||
|
||||
# Run the server when the script is executed directly
|
||||
if __name__ == "__main__":
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
|
||||
datefmt='%Y-%m-%d %H:%M:%S'
|
||||
)
|
||||
load_dotenv()
|
||||
logger.info("Starting Image MCP Server...")
|
||||
mcp.run(transport='stdio')
|
||||
+142
@@ -0,0 +1,142 @@
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from typing import List, Optional, Tuple
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import requests
|
||||
from mcp.server import FastMCP
|
||||
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def is_url(path_or_url: str) -> bool:
|
||||
"""
|
||||
Check if the given string is a URL.
|
||||
|
||||
Args:
|
||||
path_or_url: String to check
|
||||
|
||||
Returns:
|
||||
bool: True if the string is a URL, False otherwise
|
||||
"""
|
||||
parsed = urlparse(path_or_url)
|
||||
return bool(parsed.scheme and parsed.netloc)
|
||||
|
||||
|
||||
def get_file_from_source(
|
||||
source: str,
|
||||
allowed_mime_prefixes: List[str] = None,
|
||||
max_size_mb: float = 100.0,
|
||||
timeout: int = 60,
|
||||
type: str = "image",
|
||||
) -> Tuple[str, str, bytes]:
|
||||
"""
|
||||
Unified function to get file content from a URL or local path with validation.
|
||||
|
||||
Args:
|
||||
source: URL or local file path
|
||||
allowed_mime_prefixes: List of allowed MIME type prefixes (e.g., ['audio/', 'video/'])
|
||||
max_size_mb: Maximum allowed file size in MB
|
||||
timeout: Timeout for URL requests in seconds
|
||||
|
||||
Returns:
|
||||
Tuple[str, str, bytes]: (file_path, mime_type, file_content)
|
||||
- For URLs, file_path will be a temporary file path
|
||||
- For local files, file_path will be the original path
|
||||
|
||||
Raises:
|
||||
ValueError: When file doesn't exist, exceeds size limit, or has invalid MIME type
|
||||
IOError: When file cannot be read
|
||||
requests.RequestException: When URL request fails
|
||||
"""
|
||||
max_size_bytes = max_size_mb * 1024 * 1024
|
||||
temp_file = None
|
||||
|
||||
try:
|
||||
if is_url(source):
|
||||
# Handle URL
|
||||
logger.info(f"Downloading file from URL: {source}")
|
||||
response = requests.get(source, stream=True, timeout=timeout)
|
||||
response.raise_for_status()
|
||||
|
||||
# Check Content-Length if available
|
||||
content_length = response.headers.get("Content-Length")
|
||||
if content_length and int(content_length) > max_size_bytes:
|
||||
raise ValueError(f"File size exceeds limit of {max_size_mb}MB")
|
||||
|
||||
# Create a temporary file
|
||||
temp_file = tempfile.NamedTemporaryFile(delete=False)
|
||||
file_path = temp_file.name
|
||||
|
||||
# Download content in chunks to avoid memory issues
|
||||
content = bytearray()
|
||||
downloaded_size = 0
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
downloaded_size += len(chunk)
|
||||
if downloaded_size > max_size_bytes:
|
||||
raise ValueError(f"File size exceeds limit of {max_size_mb}MB")
|
||||
temp_file.write(chunk)
|
||||
content.extend(chunk)
|
||||
|
||||
temp_file.close()
|
||||
|
||||
# Get MIME type
|
||||
if type == "audio":
|
||||
mime_type = "audio/mpeg"
|
||||
elif type == "image":
|
||||
mime_type = "image/jpeg"
|
||||
elif type == "video":
|
||||
mime_type = "video/mp4"
|
||||
|
||||
|
||||
# For URLs where magic fails, try to use Content-Type header
|
||||
if mime_type == "application/octet-stream":
|
||||
content_type = response.headers.get("Content-Type", "").split(";")[0]
|
||||
if content_type:
|
||||
mime_type = content_type
|
||||
else:
|
||||
# Handle local file
|
||||
file_path = os.path.abspath(source)
|
||||
|
||||
# Check if file exists
|
||||
if not os.path.exists(file_path):
|
||||
raise ValueError(f"File not found: {file_path}")
|
||||
|
||||
# Check file size
|
||||
file_size = os.path.getsize(file_path)
|
||||
if file_size > max_size_bytes:
|
||||
raise ValueError(f"File size exceeds limit of {max_size_mb}MB")
|
||||
|
||||
# Get MIME type
|
||||
if type == "audio":
|
||||
mime_type = "audio/mpeg"
|
||||
elif type == "image":
|
||||
mime_type = "image/jpeg"
|
||||
elif type == "video":
|
||||
mime_type = "video/mp4"
|
||||
|
||||
# Read file content
|
||||
with open(file_path, "rb") as f:
|
||||
content = f.read()
|
||||
|
||||
# Validate MIME type if allowed_mime_prefixes is provided
|
||||
if allowed_mime_prefixes:
|
||||
if not any(
|
||||
mime_type.startswith(prefix) for prefix in allowed_mime_prefixes
|
||||
):
|
||||
allowed_types = ", ".join(allowed_mime_prefixes)
|
||||
raise ValueError(
|
||||
f"Invalid file type: {mime_type}. Allowed types: {allowed_types}"
|
||||
)
|
||||
|
||||
return file_path, mime_type, content
|
||||
|
||||
except Exception as e:
|
||||
# Clean up temporary file if an error occurs
|
||||
if temp_file and os.path.exists(temp_file.name):
|
||||
os.unlink(temp_file.name)
|
||||
raise e
|
||||
@@ -0,0 +1,12 @@
|
||||
#!/bin/bash
|
||||
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
for i in */; do
|
||||
if [ -d "$i" ] && [ -f "$i/pyproject.toml" ]; then
|
||||
echo "Installing dependencies for $i"
|
||||
(cd "$i" && uv sync) &
|
||||
fi
|
||||
done
|
||||
|
||||
wait
|
||||
Vendored
Vendored
+22
@@ -0,0 +1,22 @@
|
||||
[project]
|
||||
name = "documents-server"
|
||||
version = "0.1.0"
|
||||
description = "Add your description here"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
aiohttp= ">=3.12.15"
|
||||
dependencies = [
|
||||
"mcp",
|
||||
"aiohttp~=3.12.15",
|
||||
"requests~=2.32.4",
|
||||
"fastmcp~=2.11.3",
|
||||
"python-magic~=0.4.27",
|
||||
"chardet~=3.0.4",
|
||||
"pandas~=2.3.0",
|
||||
#"aworld~=0.2.5",
|
||||
"opencv-python~=4.12.0.88",
|
||||
"opencv-python-headless~=4.12.0.88",
|
||||
"numpy~=2.2.3",
|
||||
"openai~=1.93.0",
|
||||
|
||||
]
|
||||
Vendored
+226
@@ -0,0 +1,226 @@
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import magic
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import requests
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import Any, Literal
|
||||
|
||||
|
||||
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)"
|
||||
)
|
||||
|
||||
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")
|
||||
|
||||
|
||||
def _validate_file_path(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)
|
||||
if not path.is_absolute():
|
||||
path = path.expanduser().resolve()
|
||||
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"File not found: {path}")
|
||||
|
||||
return path
|
||||
|
||||
def is_url(path_or_url: str) -> bool:
|
||||
"""
|
||||
Check if the given string is a URL.
|
||||
|
||||
Args:
|
||||
path_or_url: String to check
|
||||
|
||||
Returns:
|
||||
bool: True if the string is a URL, False otherwise
|
||||
"""
|
||||
parsed = urlparse(path_or_url)
|
||||
return bool(parsed.scheme and parsed.netloc)
|
||||
|
||||
|
||||
def get_mime_type(file_path: str, default_mime: str | None = None) -> str:
|
||||
"""
|
||||
Detect MIME type of a file using python-magic if available,
|
||||
otherwise fallback to extension-based detection.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
default_mime: Default MIME type to return if detection fails
|
||||
|
||||
Returns:
|
||||
str: Detected MIME type
|
||||
"""
|
||||
# Try using python-magic for accurate MIME type detection
|
||||
try:
|
||||
mime = magic.Magic(mime=True)
|
||||
return mime.from_file(file_path)
|
||||
except (AttributeError, IOError):
|
||||
# Fallback to extension-based detection
|
||||
extension_mime_map = {
|
||||
# Audio formats
|
||||
".mp3": "audio/mpeg",
|
||||
".wav": "audio/wav",
|
||||
".ogg": "audio/ogg",
|
||||
".m4a": "audio/mp4",
|
||||
".flac": "audio/flac",
|
||||
# Image formats
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".png": "image/png",
|
||||
".gif": "image/gif",
|
||||
".webp": "image/webp",
|
||||
".bmp": "image/bmp",
|
||||
".tiff": "image/tiff",
|
||||
# Video formats
|
||||
".mp4": "video/mp4",
|
||||
".avi": "video/x-msvideo",
|
||||
".mov": "video/quicktime",
|
||||
".mkv": "video/x-matroska",
|
||||
".webm": "video/webm",
|
||||
}
|
||||
|
||||
ext = Path(file_path).suffix.lower()
|
||||
return extension_mime_map.get(ext, default_mime or "application/octet-stream")
|
||||
|
||||
|
||||
def get_file_from_source(
|
||||
source: str,
|
||||
max_size_mb: float = 100.0,
|
||||
timeout: int = 60,
|
||||
) -> tuple[str, str, bytes]:
|
||||
"""
|
||||
Unified function to get file content from a URL or local path with validation.
|
||||
|
||||
Args:
|
||||
source: URL or local file path
|
||||
max_size_mb: Maximum allowed file size in MB
|
||||
timeout: Timeout for URL requests in seconds
|
||||
|
||||
Returns:
|
||||
Tuple[str, str, bytes]: (file_path, mime_type, file_content)
|
||||
- For URLs, file_path will be a temporary file path
|
||||
- For local files, file_path will be the original path
|
||||
|
||||
Raises:
|
||||
ValueError: When file doesn't exist, exceeds size limit, or has invalid MIME type
|
||||
IOError: When file cannot be read
|
||||
requests.RequestException: When URL request fails
|
||||
"""
|
||||
max_size_bytes = max_size_mb * 1024 * 1024
|
||||
|
||||
if is_url(source):
|
||||
# Handle URL source
|
||||
try:
|
||||
# Make a HEAD request first to check content length
|
||||
head_response = requests.head(source, timeout=timeout, allow_redirects=True)
|
||||
head_response.raise_for_status()
|
||||
|
||||
# Check content length if available
|
||||
content_length = head_response.headers.get("content-length")
|
||||
if content_length and int(content_length) > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({int(content_length) / (1024 * 1024):.2f} MB) "
|
||||
f"exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Download the file
|
||||
response = requests.get(source, timeout=timeout, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
# Read content with size checking
|
||||
content = b""
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
if len(content) + len(chunk) > max_size_bytes:
|
||||
raise ValueError(f"File size exceeds maximum allowed size ({max_size_mb} MB)")
|
||||
content += chunk
|
||||
|
||||
# Create temporary file
|
||||
parsed_url = urlparse(source)
|
||||
filename = os.path.basename(parsed_url.path) or "downloaded_file"
|
||||
|
||||
# Create temporary file with proper extension
|
||||
suffix = Path(filename).suffix or ".tmp"
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
|
||||
temp_file.write(content)
|
||||
temp_path = temp_file.name
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(temp_path)
|
||||
|
||||
return temp_path, mime_type, content
|
||||
|
||||
except requests.RequestException as e:
|
||||
raise requests.RequestException(f"Failed to download file from URL: {e}: {traceback.format_exc()}")
|
||||
except Exception as e:
|
||||
raise IOError(f"Error processing URL: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
else:
|
||||
# Handle local file path
|
||||
file_path = Path(source)
|
||||
|
||||
# Check if file exists
|
||||
if not file_path.exists():
|
||||
raise ValueError(f"File does not exist: {source}")
|
||||
|
||||
if not file_path.is_file():
|
||||
raise ValueError(f"Path is not a file: {source}")
|
||||
|
||||
# Check file size
|
||||
file_size = file_path.stat().st_size
|
||||
if file_size > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({file_size / (1024 * 1024):.2f} MB) exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Read file content
|
||||
try:
|
||||
with open(file_path, "rb") as f:
|
||||
content = f.read()
|
||||
except Exception as e:
|
||||
raise IOError(f"Cannot read file {source}: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(str(file_path))
|
||||
|
||||
return str(file_path), mime_type, content
|
||||
Vendored
+367
@@ -0,0 +1,367 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from typing import Union, Literal
|
||||
|
||||
from pydantic import BaseModel
|
||||
from dotenv import load_dotenv
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.types import TextContent
|
||||
from pydantic import Field
|
||||
from openai import OpenAI
|
||||
|
||||
from base import (
|
||||
ActionResponse,
|
||||
_validate_file_path
|
||||
)
|
||||
|
||||
load_dotenv()
|
||||
|
||||
mcp = FastMCP(
|
||||
"intelligence-code-server",
|
||||
instructions="""
|
||||
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
|
||||
""",
|
||||
)
|
||||
|
||||
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
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
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
|
||||
"""
|
||||
)
|
||||
async def generate_python_code(
|
||||
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'",
|
||||
)
|
||||
) -> Union[str, TextContent]:
|
||||
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")
|
||||
|
||||
logging.info(f"Generating code for: {task_description[:100]}...")
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Prepare the code generation prompt
|
||||
prompt = _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 = _call_code_model(prompt, temperature)
|
||||
|
||||
# Extract clean Python code
|
||||
generated_code = _extract_python_code(raw_response)
|
||||
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
# Populate metadata fields
|
||||
metadata = CodeGenerationMetadata(
|
||||
model_name=os.getenv("CODE_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 = _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)
|
||||
logging.info(f"Generated code also saved to: {output_file_path_obj}")
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to save code to file '{save_to_file_path}': {str(e)}")
|
||||
metadata.file_save_error = str(e)
|
||||
|
||||
logging.info(
|
||||
f"Successfully generated code ({metadata.code_length} characters, "
|
||||
f"{metadata.processing_time_seconds:.2f}s)"
|
||||
)
|
||||
|
||||
action_response = ActionResponse(success=True, message=generated_code,
|
||||
metadata=metadata.model_dump(exclude_none=True))
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(action_response.model_dump()),
|
||||
}
|
||||
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": output_dict} # Pass as additional fields
|
||||
)
|
||||
|
||||
except ValueError as e:
|
||||
logging.error(f"Invalid input: {str(e)}")
|
||||
metadata.error_type = "invalid_input"
|
||||
metadata.error_message = str(e)
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}",
|
||||
metadata=metadata.model_dump(exclude_none=True),
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": {}} # Pass as additional fields
|
||||
)
|
||||
except Exception as e:
|
||||
logging.error(f"Code generation failed: {str(e)}: {traceback.format_exc()}")
|
||||
metadata.error_type = "generation_error"
|
||||
metadata.error_message = str(e)
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Code generation failed: {str(e)}",
|
||||
metadata=metadata.model_dump(exclude_none=True),
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": {}} # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Get information about the reasoning service capabilities.
|
||||
"""
|
||||
)
|
||||
async def get_reasoning_capabilities() -> Union[str, TextContent]:
|
||||
capabilities = {
|
||||
"Mathematical Problems": "Advanced mathematical reasoning, proofs, and calculations",
|
||||
"Code Contests": "Programming challenges, algorithm design, and optimization",
|
||||
"Logic Puzzles": "Brain teasers, riddles, and logical reasoning problems",
|
||||
"STEM Problems": "Competition-level science, technology, engineering, and math",
|
||||
"Multi-step Analysis": "Complex analytical reasoning with multiple interconnected steps",
|
||||
}
|
||||
|
||||
capability_list = "\n".join(
|
||||
[
|
||||
f"**{capability}**: {description}"
|
||||
for capability, description in capabilities.items()
|
||||
]
|
||||
)
|
||||
|
||||
metadata = {
|
||||
"model_name": os.getenv("THINK_LLM_MODEL_NAME", ""),
|
||||
"provider": "openai",
|
||||
"supported_capabilities": list(capabilities.keys()),
|
||||
"total_capabilities": len(capabilities),
|
||||
"reasoning_styles": ["detailed", "concise", "step-by-step"],
|
||||
}
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True,
|
||||
message=f"Intelligence Reasoning Service Capabilities:\n\n{capability_list}",
|
||||
metadata=metadata,
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": f"Intelligence Reasoning Service Capabilities:\n\n{capability_list}"
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
def _prepare_code_prompt(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(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
|
||||
"""
|
||||
openai_params = {
|
||||
"model": os.getenv("TCODE_LLM_MODEL_NAME", ""),
|
||||
"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,
|
||||
}
|
||||
try:
|
||||
client: OpenAI = OpenAI(
|
||||
api_key=os.getenv("CODE_LLM_API_KEY"),
|
||||
base_url=os.getenv("TCODE_LLM_BASE_URL"),
|
||||
)
|
||||
response = client.chat.completions.create(**openai_params)
|
||||
content = ""
|
||||
if response and hasattr(response, 'choices') and response.choices:
|
||||
content = response.choices[0].message.content
|
||||
return content
|
||||
except BaseException as e:
|
||||
logging.warn(f"coding failed: {str(e)}: {traceback.format_exc()}")
|
||||
return f"coding failed: {str(e)}"
|
||||
|
||||
return content
|
||||
|
||||
def _extract_python_code(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()
|
||||
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
load_dotenv(override=True)
|
||||
logging.info("Starting intelligence-think-server MCP server!")
|
||||
mcp.run(transport="stdio")
|
||||
Vendored
+294
@@ -0,0 +1,294 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from typing import Any, Union, Literal
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.types import TextContent
|
||||
from pydantic import Field
|
||||
|
||||
from base import ActionResponse
|
||||
from openai import OpenAI
|
||||
|
||||
load_dotenv()
|
||||
|
||||
mcp = FastMCP(
|
||||
"intelligence-guard-server",
|
||||
instructions="""
|
||||
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.
|
||||
""",
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
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.
|
||||
|
||||
"""
|
||||
)
|
||||
async def guarding_reasoning_process(
|
||||
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)",
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
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"
|
||||
)
|
||||
|
||||
logging.info(f"Processing guarding request: {question[:100]}...")
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Prepare the guarding prompt
|
||||
prompt = _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 = _call_guarding_model(prompt, temperature)
|
||||
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
# Prepare metadata
|
||||
metadata = {
|
||||
"model_name": os.getenv("GUARD_LLM_MODEL_NAME", ""),
|
||||
"guarding_style": guarding_style,
|
||||
"response_length": len(guarding_result),
|
||||
}
|
||||
|
||||
logging.info(
|
||||
f"Successfully completed guarding ({len(guarding_result)} characters, {processing_time:.2f}s)"
|
||||
)
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=guarding_result, metadata=metadata
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": guarding_result
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except ValueError as e:
|
||||
logging.error(f"Invalid input: {str(e)}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}",
|
||||
metadata={"error_type": "invalid_input", "error_message": str(e)},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
except Exception as e:
|
||||
logging.error(f"Guarding failed: {str(e)}: {traceback.format_exc()}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Guarding failed: {str(e)}",
|
||||
metadata={"error_type": "guarding_error", "error_message": str(e)},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Get information about the guarding reasoning process service capabilities.
|
||||
"""
|
||||
)
|
||||
async def get_guarding_capabilities() -> Union[str, TextContent]:
|
||||
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": os.getenv("GUARD_LLM_MODEL_NAME", ""),
|
||||
"provider": "openai",
|
||||
"supported_capabilities": list(capabilities.keys()),
|
||||
"total_capabilities": len(capabilities),
|
||||
"guarding_styles": ["detailed", "concise", "step-by-step"],
|
||||
}
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True,
|
||||
message=f"Intelligence guarding Service Capabilities:\n\n{capability_list}",
|
||||
metadata=metadata,
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": f"Intelligence guarding Service Capabilities:\n\n{capability_list}"
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
def _prepare_guarding_prompt(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(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
|
||||
"""
|
||||
client: OpenAI = OpenAI(
|
||||
api_key=os.getenv("GUARD_LLM_API_KEY"),
|
||||
base_url=os.getenv("GUARD_LLM_BASE_URL"),
|
||||
)
|
||||
openai_params = {
|
||||
"model": os.getenv("GUARD_LLM_MODEL_NAME", ""),
|
||||
"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,
|
||||
}
|
||||
try:
|
||||
response = client.chat.completions.create(**openai_params)
|
||||
content = ""
|
||||
if response and hasattr(response, 'choices') and response.choices:
|
||||
content = response.choices[0].message.content
|
||||
return content
|
||||
except BaseException as e:
|
||||
logging.warn(f"Reasoning failed: {str(e)}: {traceback.format_exc()}")
|
||||
return f"Reasoning failed: {str(e)}"
|
||||
|
||||
return content
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
load_dotenv(override=True)
|
||||
logging.info("Starting intelligence-guard-server MCP server!")
|
||||
mcp.run(transport="stdio")
|
||||
Vendored
+271
@@ -0,0 +1,271 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from typing import Union, Literal
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from pydantic.fields import FieldInfo
|
||||
from openai import OpenAI
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.types import TextContent
|
||||
from pydantic import Field
|
||||
|
||||
from base import ActionResponse
|
||||
|
||||
load_dotenv()
|
||||
|
||||
mcp = FastMCP(
|
||||
"intelligence-think-server",
|
||||
instructions="""
|
||||
MCP service for complex problem reasoning using powerful reasoning models.
|
||||
|
||||
Supports advanced reasoning for:
|
||||
- Mathematical problems and proofs
|
||||
- Code contest and programming challenges
|
||||
- Logic puzzles and riddles
|
||||
- Competition-level STEM problems
|
||||
- Multi-step analytical reasoning
|
||||
""",
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
his tool provides comprehensive reasoning capabilities for:
|
||||
- Mathematical problems and proofs
|
||||
- Programming and algorithm challenges
|
||||
- Logic puzzles, brain teasers, and fun riddles
|
||||
- Competition-level STEM problems
|
||||
- Multi-step analytical reasoning tasks
|
||||
|
||||
Weakness:
|
||||
- Inability to process media types: image, audio, or video
|
||||
- Require precise description of problem context and settings
|
||||
"""
|
||||
)
|
||||
async def complex_problem_reasoning(
|
||||
question: str = Field(
|
||||
description="The input question for complex problem reasoning, such as math and code contest problems"
|
||||
),
|
||||
original_task: str = Field(
|
||||
default="", description="The original task description."
|
||||
),
|
||||
temperature: float = Field(
|
||||
default=0.3,
|
||||
description="Model temperature for response variability (0.0-1.0)",
|
||||
ge=0.0,
|
||||
le=1.0,
|
||||
),
|
||||
reasoning_style: Literal["detailed", "concise", "step-by-step"] = Field(
|
||||
default="detailed",
|
||||
description="Style of reasoning output: detailed(analysis), concise(summary), or step-by-step(breakdown)",
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
try:
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(question, FieldInfo):
|
||||
question = question.default
|
||||
if isinstance(original_task, FieldInfo):
|
||||
original_task = original_task.default
|
||||
if isinstance(temperature, FieldInfo):
|
||||
temperature = temperature.default
|
||||
if isinstance(reasoning_style, FieldInfo):
|
||||
reasoning_style = reasoning_style.default
|
||||
|
||||
# Validate input
|
||||
if not question or not question.strip():
|
||||
raise ValueError("Question is required for complex problem reasoning")
|
||||
|
||||
logging.info(f"Processing reasoning request: {question[:100]}...")
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Prepare the reasoning prompt
|
||||
prompt = _prepare_reasoning_prompt(question, original_task)
|
||||
|
||||
# Enhance prompt based on reasoning style
|
||||
if reasoning_style == "step-by-step":
|
||||
prompt += "\n\nPlease provide a clear step-by-step breakdown of your reasoning process."
|
||||
elif reasoning_style == "concise":
|
||||
prompt += (
|
||||
"\n\nPlease provide a concise but complete reasoning and final answer."
|
||||
)
|
||||
elif reasoning_style == "detailed":
|
||||
prompt += (
|
||||
"\n\nPlease provide detailed analysis with comprehensive reasoning."
|
||||
)
|
||||
|
||||
# Call the reasoning model
|
||||
reasoning_result = _call_reasoning_model(prompt, temperature)
|
||||
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
# Prepare metadata
|
||||
metadata = {
|
||||
"model_name": os.getenv("THINK_LLM_MODEL_NAME", ""),
|
||||
"reasoning_style": reasoning_style,
|
||||
"response_length": len(reasoning_result),
|
||||
}
|
||||
|
||||
logging.info(
|
||||
f"Successfully completed reasoning ({len(reasoning_result)} characters, {processing_time:.2f}s)"
|
||||
)
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=reasoning_result, metadata=metadata
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": reasoning_result
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except ValueError as e:
|
||||
logging.error(f"Invalid input: {str(e)}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Invalid input: {str(e)}",
|
||||
metadata={"error_type": "invalid_input", "error_message": str(e)},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
except Exception as e:
|
||||
logging.error(f"Reasoning failed: {str(e)}: {traceback.format_exc()}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Reasoning failed: {str(e)}",
|
||||
metadata={"error_type": "reasoning_error", "error_message": str(e)},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Get information about the reasoning service capabilities.
|
||||
"""
|
||||
)
|
||||
async def get_reasoning_capabilities() -> Union[str, TextContent]:
|
||||
capabilities = {
|
||||
"Mathematical Problems": "Advanced mathematical reasoning, proofs, and calculations",
|
||||
"Code Contests": "Programming challenges, algorithm design, and optimization",
|
||||
"Logic Puzzles": "Brain teasers, riddles, and logical reasoning problems",
|
||||
"STEM Problems": "Competition-level science, technology, engineering, and math",
|
||||
"Multi-step Analysis": "Complex analytical reasoning with multiple interconnected steps",
|
||||
}
|
||||
|
||||
capability_list = "\n".join(
|
||||
[
|
||||
f"**{capability}**: {description}"
|
||||
for capability, description in capabilities.items()
|
||||
]
|
||||
)
|
||||
|
||||
metadata = {
|
||||
"model_name": os.getenv("THINK_LLM_MODEL_NAME", ""),
|
||||
"provider": "openai",
|
||||
"supported_capabilities": list(capabilities.keys()),
|
||||
"total_capabilities": len(capabilities),
|
||||
"reasoning_styles": ["detailed", "concise", "step-by-step"],
|
||||
}
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True,
|
||||
message=f"Intelligence Reasoning Service Capabilities:\n\n{capability_list}",
|
||||
metadata=metadata,
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": f"Intelligence Reasoning Service Capabilities:\n\n{capability_list}"
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
def _prepare_reasoning_prompt(question: str, original_task: str = "") -> str:
|
||||
"""Prepare the reasoning prompt with question and optional context.
|
||||
|
||||
Args:
|
||||
question: The main question for reasoning
|
||||
original_task: Optional original task description for context
|
||||
|
||||
Returns:
|
||||
Formatted prompt string
|
||||
"""
|
||||
if original_task:
|
||||
return f"Original Task: {original_task}\n\nQuestion: {question}"
|
||||
return f"Question: {question}"
|
||||
|
||||
|
||||
def _call_reasoning_model(prompt: str, temperature: float = 0.3) -> str:
|
||||
"""Call the reasoning model with the prepared prompt.
|
||||
|
||||
Args:
|
||||
prompt: The formatted prompt for reasoning
|
||||
temperature: Model temperature for response variability
|
||||
|
||||
Returns:
|
||||
Reasoning result from the model
|
||||
|
||||
Raises:
|
||||
Exception: If model call fails
|
||||
"""
|
||||
openai_params = {
|
||||
"model": os.getenv("THINK_LLM_MODEL_NAME", ""),
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"You are an expert at solving complex problems including math, "
|
||||
"code contests, riddles, and puzzles. "
|
||||
"Provide detailed step-by-step reasoning and a clear final answer."
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
"temperature": temperature,
|
||||
}
|
||||
try:
|
||||
client: OpenAI = OpenAI(
|
||||
api_key=os.getenv("THINK_LLM_API_KEY"),
|
||||
base_url=os.getenv("THINK_LLM_BASE_URL"),
|
||||
)
|
||||
response = client.chat.completions.create(**openai_params)
|
||||
content = ""
|
||||
if response and hasattr(response, 'choices') and response.choices:
|
||||
content = response.choices[0].message.content
|
||||
return content
|
||||
except BaseException as e:
|
||||
logging.warn(f"Reasoning failed: {str(e)}: {traceback.format_exc()}")
|
||||
return f"Reasoning failed: {str(e)}"
|
||||
|
||||
return content
|
||||
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
load_dotenv(override=True)
|
||||
logging.info("Starting intelligence-think-server MCP server!")
|
||||
mcp.run(transport="stdio")
|
||||
+262
@@ -0,0 +1,262 @@
|
||||
import os
|
||||
|
||||
import dotenv
|
||||
|
||||
dotenv.load_dotenv(verbose=True, override=True, interpolate=True)
|
||||
|
||||
mcp_config = {
|
||||
"mcpServers": {
|
||||
"readweb-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/main.py"],
|
||||
"cwd": "readweb_server",
|
||||
"env": {
|
||||
"PIC_SEARCH_URL": os.getenv("PIC_SEARCH_URL"),
|
||||
"PIC_SEARCH_TOTAL_NUM": os.getenv("PIC_SEARCH_TOTAL_NUM"),
|
||||
"PIC_SEARCH_SLICE_NUM": os.getenv("PIC_SEARCH_SLICE_NUM"),
|
||||
"PIC_SEARCH_DOMAIN": os.getenv("PIC_SEARCH_DOMAIN"),
|
||||
"PIC_SEARCH_SEARCHMODE": os.getenv("PIC_SEARCH_SEARCHMODE"),
|
||||
"PIC_SEARCH_SOURCE": os.getenv("PIC_SEARCH_SOURCE"),
|
||||
"PIC_SEARCH_UID": os.getenv("PIC_SEARCH_UID"),
|
||||
"JINA_API_KEY": os.getenv("JINA_API_KEY"),
|
||||
"TAVILY_API_KEY": os.getenv("TAVILY_API_KEY"),
|
||||
"GOOGLE_API_KEY": os.getenv("GOOGLE_API_KEY"),
|
||||
"GOOGLE_CSE_ID": os.getenv("GOOGLE_CSE_ID"),
|
||||
},
|
||||
},
|
||||
"browser-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/main.py"],
|
||||
"cwd": "browser_server",
|
||||
"env": {
|
||||
"LLM_BASE_URL": os.getenv("BROWSERUSE_LLM_BASE_URL"),
|
||||
"LLM_MODEL_NAME": os.getenv("BROWSERUSE_LLM_MODEL_NAME"),
|
||||
"LLM_API_KEY": os.getenv("BROWSERUSE_LLM_API_KEY"),
|
||||
},
|
||||
},
|
||||
"documents-csv-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/csv_server.py"],
|
||||
"cwd": "documents_server",
|
||||
},
|
||||
"documents-docx-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/docx_server.py"],
|
||||
"cwd": "documents_server",
|
||||
},
|
||||
"documents-pptx-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/pptx_server.py"],
|
||||
"cwd": "documents_server",
|
||||
},
|
||||
"documents-pdf-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/pdf_server.py"],
|
||||
"cwd": "documents_server",
|
||||
"env": {
|
||||
"DATALAB_API_KEY": os.getenv("DATALAB_API_KEY"),
|
||||
},
|
||||
},
|
||||
"documents-txt-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/txt_server.py"],
|
||||
"cwd": "documents_server",
|
||||
},
|
||||
"download-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/download.py"],
|
||||
"cwd": "download_server",
|
||||
},
|
||||
"intelligence-code-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/code.py"],
|
||||
"cwd": "intelligence_server",
|
||||
"env": {
|
||||
"CODE_LLM_BASE_URL": os.getenv("CODE_LLM_BASE_URL"),
|
||||
"CODE_LLM_MODEL_NAME": os.getenv("CODE_LLM_MODEL_NAME"),
|
||||
"CODE_LLM_API_KEY": os.getenv("CODE_LLM_API_KEY"),
|
||||
},
|
||||
},
|
||||
"intelligence-think-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/think.py"],
|
||||
"cwd": "intelligence_server",
|
||||
"env": {
|
||||
"THINK_LLM_BASE_URL": os.getenv("THINK_LLM_BASE_URL"),
|
||||
"THINK_LLM_MODEL_NAME": os.getenv("THINK_LLM_MODEL_NAME"),
|
||||
"THINK_LLM_API_KEY": os.getenv("THINK_LLM_API_KEY"),
|
||||
},
|
||||
},
|
||||
"intelligence-guard-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/guard.py"],
|
||||
"cwd": "intelligence_server",
|
||||
"env": {
|
||||
"GUARD_LLM_BASE_URL": os.getenv("GUARD_LLM_BASE_URL"),
|
||||
"GUARD_LLM_MODEL_NAME": os.getenv("GUARD_LLM_MODEL_NAME"),
|
||||
"GUARD_LLM_API_KEY": os.getenv("GUARD_LLM_API_KEY"),
|
||||
},
|
||||
},
|
||||
"media-audio-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/audio.py"],
|
||||
"cwd": "media_server",
|
||||
"env": {
|
||||
"AUDIO_LLM_BASE_URL": os.getenv("AUDIO_LLM_BASE_URL"),
|
||||
"AUDIO_LLM_MODEL_NAME": os.getenv("AUDIO_LLM_MODEL_NAME"),
|
||||
"AUDIO_LLM_API_KEY": os.getenv("AUDIO_LLM_API_KEY"),
|
||||
},
|
||||
},
|
||||
"media-image-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/image.py"],
|
||||
"cwd": "media_server",
|
||||
"env": {
|
||||
"IMAGE_LLM_BASE_URL": os.getenv("IMAGE_LLM_BASE_URL"),
|
||||
"IMAGE_LLM_MODEL_NAME": os.getenv("IMAGE_LLM_MODEL_NAME"),
|
||||
"IMAGE_LLM_API_KEY": os.getenv("IMAGE_LLM_API_KEY"),
|
||||
},
|
||||
},
|
||||
"media-video-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/video.py"],
|
||||
"cwd": "media_server",
|
||||
"env": {
|
||||
"VIDEO_LLM_BASE_URL": os.getenv("VIDEO_LLM_BASE_URL"),
|
||||
"VIDEO_LLM_MODEL_NAME": os.getenv("VIDEO_LLM_MODEL_NAME"),
|
||||
"VIDEO_LLM_API_KEY": os.getenv("VIDEO_LLM_API_KEY"),
|
||||
},
|
||||
},
|
||||
"parxiv-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/parxiv.py"],
|
||||
"cwd": "parxiv_server",
|
||||
},
|
||||
"terminal-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/terminal.py"],
|
||||
"cwd": "terminal_server",
|
||||
},
|
||||
"wayback-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/wayback.py"],
|
||||
"cwd": "wayback_server",
|
||||
},
|
||||
"wiki-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/wiki.py"],
|
||||
"cwd": "wiki_server",
|
||||
},
|
||||
"googlesearch-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/googlesearch.py"],
|
||||
"cwd": "googlesearch_server",
|
||||
"env": {
|
||||
"GOOGLE_API_KEY": os.getenv("GOOGLE_API_KEY"),
|
||||
"GOOGLE_CSE_ID": os.getenv("GOOGLE_CSE_ID"),
|
||||
},
|
||||
},
|
||||
"filesystem-server": {
|
||||
"type": "stdio",
|
||||
"command": "npx",
|
||||
"args": ["-y", "@modelcontextprotocol/server-filesystem", "~/workspace"],
|
||||
},
|
||||
"terminal-controller": {
|
||||
"command": "uvx",
|
||||
"args": ["terminal_controller"],
|
||||
"env": {"SESSION_REQUEST_CONNECT_TIMEOUT": "300"},
|
||||
},
|
||||
"excel": {
|
||||
"command": "uvx",
|
||||
"args": ["excel-mcp-server", "stdio"],
|
||||
"env": {
|
||||
"EXCEL_MCP_PAGING_CELLS_LIMIT": "4000",
|
||||
"SESSION_REQUEST_CONNECT_TIMEOUT": "120",
|
||||
},
|
||||
},
|
||||
"google-search": {
|
||||
"command": "npx",
|
||||
"args": ["-y", "@adenot/mcp-google-search"],
|
||||
"env": {
|
||||
"GOOGLE_API_KEY": os.environ["GOOGLE_API_KEY"],
|
||||
"GOOGLE_SEARCH_ENGINE_ID": os.environ["GOOGLE_CSE_ID"],
|
||||
"SESSION_REQUEST_CONNECT_TIMEOUT": "60",
|
||||
},
|
||||
},
|
||||
"audio-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/main.py"],
|
||||
"cwd": "audio_server",
|
||||
"env": {
|
||||
"AUDIO_LLM_API_KEY": os.environ["AUDIO_LLM_API_KEY"],
|
||||
"AUDIO_LLM_BASE_URL": os.environ["AUDIO_LLM_BASE_URL"],
|
||||
"AUDIO_LLM_MODEL_NAME": os.environ["AUDIO_LLM_MODEL_NAME"],
|
||||
"SESSION_REQUEST_CONNECT_TIMEOUT": "60",
|
||||
},
|
||||
},
|
||||
"image-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/main.py"],
|
||||
"cwd": "image_server",
|
||||
"env": {
|
||||
"IMAGE_LLM_API_KEY": os.environ.get("IMAGE_LLM_API_KEY"),
|
||||
"IMAGE_LLM_MODEL_NAME": os.environ.get("IMAGE_LLM_MODEL_NAME"),
|
||||
"IMAGE_LLM_BASE_URL": os.environ.get("IMAGE_LLM_BASE_URL"),
|
||||
"SESSION_REQUEST_CONNECT_TIMEOUT": "60",
|
||||
},
|
||||
},
|
||||
"e2b-code-server": {
|
||||
"type": "stdio",
|
||||
"command": "uv",
|
||||
"args": ["run", "src/main.py"],
|
||||
"cwd": "e2b_code_server",
|
||||
"env": {
|
||||
"E2B_API_KEY": os.environ["E2B_API_KEY"],
|
||||
"SESSION_REQUEST_CONNECT_TIMEOUT": "120",
|
||||
},
|
||||
},
|
||||
"ms-playwright": {
|
||||
"command": "npx",
|
||||
"args": [
|
||||
"@playwright/mcp@latest",
|
||||
"--no-sandbox",
|
||||
"--headless",
|
||||
"--isolated",
|
||||
],
|
||||
"env": {
|
||||
"PLAYWRIGHT_TIMEOUT": "120000",
|
||||
"SESSION_REQUEST_CONNECT_TIMEOUT": "120",
|
||||
},
|
||||
},
|
||||
# "calculator": {
|
||||
# "command": "uvx",
|
||||
# "args": [
|
||||
# "mcp_server_calculator"
|
||||
# ],
|
||||
# "env": {
|
||||
# "SESSION_REQUEST_CONNECT_TIMEOUT": "20"
|
||||
# }
|
||||
# },
|
||||
}
|
||||
}
|
||||
+10046
File diff suppressed because it is too large
Load Diff
Vendored
+23
@@ -0,0 +1,23 @@
|
||||
[project]
|
||||
name = "media-server"
|
||||
version = "0.1.0"
|
||||
description = "Add your description here"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
aiohttp= ">=3.12.15"
|
||||
dependencies = [
|
||||
"mcp",
|
||||
"aiohttp~=3.12.15",
|
||||
"requests~=2.32.4",
|
||||
"fastmcp~=2.11.3",
|
||||
"python-magic~=0.4.27",
|
||||
"chardet~=3.0.4",
|
||||
"pandas~=2.3.0",
|
||||
"numpy~=2.2.3",
|
||||
"openai~=1.93.0",
|
||||
"pytesseract~=0.3.10",
|
||||
"pillow~=10.4.0",
|
||||
"opencv-python~=4.12.0.88",
|
||||
"opencv-python-headless~=4.12.0.88",
|
||||
|
||||
]
|
||||
+667
@@ -0,0 +1,667 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from typing import Any, Union, Literal
|
||||
from openai import OpenAI
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.types import TextContent
|
||||
from pydantic import Field, BaseModel
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from base import ActionResponse, _validate_file_path
|
||||
|
||||
workspace = Path.home()
|
||||
_audio_output_dir = workspace / "processed_audio"
|
||||
_audio_output_dir.mkdir(exist_ok=True, parents=True)
|
||||
|
||||
mcp = FastMCP(
|
||||
"media-audio-server",
|
||||
instructions="""
|
||||
MCP service for comprehensive audio processing using ffmpeg.
|
||||
|
||||
Supports various audio operations including:
|
||||
- Audio format conversion
|
||||
- Audio transcription (speech-to-text)
|
||||
- Audio metadata extraction
|
||||
- Audio quality enhancement
|
||||
- Audio trimming and editing
|
||||
- Audio analysis and feature extraction
|
||||
""",
|
||||
)
|
||||
|
||||
|
||||
class AudioMetadata(BaseModel):
|
||||
"""Metadata extracted from audio processing."""
|
||||
|
||||
file_name: str = Field(description="Original audio file name")
|
||||
file_size: int = Field(description="File size in bytes")
|
||||
file_type: str = Field(description="Audio file type/extension")
|
||||
absolute_path: str = Field(description="Absolute path to the audio file")
|
||||
duration: float | None = Field(
|
||||
default=None, description="Duration of audio in seconds"
|
||||
)
|
||||
sample_rate: int | None = Field(default=None, description="Audio sample rate in Hz")
|
||||
channels: int | None = Field(default=None, description="Number of audio channels")
|
||||
bitrate: int | None = Field(default=None, description="Audio bitrate in kbps")
|
||||
codec: str | None = Field(default=None, description="Audio codec used")
|
||||
processing_time: float = Field(
|
||||
description="Time taken to process the audio in seconds"
|
||||
)
|
||||
output_files: list[str] = Field(
|
||||
default_factory=list, description="Paths to generated output files"
|
||||
)
|
||||
transcription: str | None = Field(
|
||||
default=None, description="Transcribed text from audio"
|
||||
)
|
||||
word_count: int | None = Field(
|
||||
default=None, description="Number of words in transcription"
|
||||
)
|
||||
output_format: str = Field(description="Format of the processed output")
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Transcribe audio file to text using OpenAI Whisper.
|
||||
|
||||
This tool converts speech in audio files to text with high accuracy.
|
||||
Supports multiple languages and provides various output formats including
|
||||
timestamped segments for detailed analysis.
|
||||
"""
|
||||
)
|
||||
async def transcribe_audio(
|
||||
file_path: str = Field(description="Path to the audio file to transcribe"),
|
||||
model_size: Literal["tiny", "base", "small", "medium", "large"] = Field(
|
||||
default="base",
|
||||
description="Whisper model size: tiny (fastest), base (balanced), small, medium, large (most accurate)",
|
||||
),
|
||||
output_format: Literal["text", "detailed", "segments"] = Field(
|
||||
default="text",
|
||||
description="Output format: 'text' (plain text), 'detailed' (with metadata), 'segments' (timestamped)",
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
try:
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(file_path, FieldInfo):
|
||||
file_path = file_path.default
|
||||
if isinstance(model_size, FieldInfo):
|
||||
model_size = model_size.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Validate input file
|
||||
file_path: Path = _validate_file_path(file_path)
|
||||
logging.info(f"Transcribing audio: {file_path.name}")
|
||||
|
||||
# Get original metadata
|
||||
original_metadata = _get_audio_metadata(file_path)
|
||||
|
||||
# Prepare audio for transcription
|
||||
prepared_audio = _prepare_audio_for_transcription(file_path)
|
||||
|
||||
# Perform transcription
|
||||
transcription_result = _transcribe_with_whisper(prepared_audio)
|
||||
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
# Prepare file statistics
|
||||
file_stats = file_path.stat()
|
||||
|
||||
# Count words in transcription
|
||||
word_count = (
|
||||
len(transcription_result["text"].split())
|
||||
if transcription_result["text"]
|
||||
else 0
|
||||
)
|
||||
|
||||
# Create metadata object
|
||||
audio_metadata = AudioMetadata(
|
||||
file_name=file_path.name,
|
||||
file_size=file_stats.st_size,
|
||||
file_type=file_path.suffix.lower(),
|
||||
absolute_path=str(file_path.absolute()),
|
||||
duration=original_metadata.get("duration"),
|
||||
sample_rate=original_metadata.get("sample_rate"),
|
||||
channels=original_metadata.get("channels"),
|
||||
bitrate=original_metadata.get("bitrate"),
|
||||
codec=original_metadata.get("codec"),
|
||||
processing_time=processing_time,
|
||||
output_files=[str(prepared_audio)],
|
||||
transcription=transcription_result["text"],
|
||||
word_count=word_count,
|
||||
output_format=f"transcription_{output_format}",
|
||||
)
|
||||
|
||||
# Format output based on requested format
|
||||
if output_format == "text":
|
||||
result_message = transcription_result["text"]
|
||||
elif output_format == "detailed":
|
||||
confidence_str = (
|
||||
f"{transcription_result['confidence']:.2f}"
|
||||
if transcription_result["confidence"]
|
||||
else "N/A"
|
||||
)
|
||||
result_message = (
|
||||
f"Transcription Results for {file_path.name}:\n\n"
|
||||
f"**Text:** {transcription_result['text']}\n\n"
|
||||
f"**Confidence:** {confidence_str}\n"
|
||||
f"**Word Count:** {word_count}\n"
|
||||
f"**Duration:** {original_metadata.get('duration', 0):.2f} seconds\n"
|
||||
f"**Model:** {model_size}\n"
|
||||
f"**Processing Time:** {processing_time:.2f} seconds"
|
||||
)
|
||||
elif output_format == "segments":
|
||||
segments_text = "\n".join(
|
||||
[
|
||||
f"[{seg.get('start', 0):.2f}s - {seg.get('end', 0):.2f}s]: {seg.get('text', '').strip()}"
|
||||
for seg in transcription_result.get("segments", [])
|
||||
]
|
||||
)
|
||||
result_message = (
|
||||
f"Timestamped Transcription for {file_path.name}:\n\n"
|
||||
f"{segments_text}\n\n"
|
||||
f"**Full Text:** {transcription_result['text']}"
|
||||
)
|
||||
else:
|
||||
result_message = transcription_result["text"]
|
||||
|
||||
# Clean up temporary file
|
||||
try:
|
||||
prepared_audio.unlink()
|
||||
except Exception:
|
||||
pass # Ignore cleanup errors
|
||||
|
||||
logging.info(
|
||||
f"Transcription completed: {word_count} words, {processing_time:.2f}s"
|
||||
)
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=result_message, metadata=audio_metadata.model_dump()
|
||||
)
|
||||
# output_dict = {
|
||||
# "artifact_type": "MARKDOWN",
|
||||
# "artifact_data": formatted_content
|
||||
# }
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
)
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Audio transcription failed: {str(e)}: {traceback.format_exc()}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Audio transcription failed: {str(e)}",
|
||||
metadata={"error_type": "transcription_error"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Extract comprehensive metadata from audio files.
|
||||
|
||||
This tool analyzes audio files and extracts detailed metadata including
|
||||
duration, sample rate, channels, bitrate, codec, and other technical information.
|
||||
"""
|
||||
)
|
||||
async def extract_audio_metadata(
|
||||
file_path: str = Field(description="Path to the audio file to analyze"),
|
||||
) -> Union[str, TextContent]:
|
||||
try:
|
||||
if isinstance(file_path, FieldInfo):
|
||||
file_path = file_path.default
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Validate input file
|
||||
file_path: Path = _validate_file_path(file_path)
|
||||
logging.info(f"Extracting metadata from: {file_path.name}")
|
||||
|
||||
# Extract metadata
|
||||
metadata = _get_audio_metadata(file_path)
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
# Prepare file statistics
|
||||
file_stats = file_path.stat()
|
||||
|
||||
# Create metadata object
|
||||
audio_metadata = AudioMetadata(
|
||||
file_name=file_path.name,
|
||||
file_size=file_stats.st_size,
|
||||
file_type=file_path.suffix.lower(),
|
||||
absolute_path=str(file_path.absolute()),
|
||||
duration=metadata.get("duration"),
|
||||
sample_rate=metadata.get("sample_rate"),
|
||||
channels=metadata.get("channels"),
|
||||
bitrate=metadata.get("bitrate"),
|
||||
codec=metadata.get("codec"),
|
||||
processing_time=processing_time,
|
||||
output_files=[],
|
||||
output_format="metadata",
|
||||
)
|
||||
|
||||
# Format metadata for LLM consumption
|
||||
result_message = (
|
||||
f"Audio Metadata for {file_path.name}:\n"
|
||||
f"Duration: {metadata.get('duration', 'Unknown'):.2f} seconds\n"
|
||||
f"Sample Rate: {metadata.get('sample_rate', 'Unknown')} Hz\n"
|
||||
f"Channels: {metadata.get('channels', 'Unknown')}\n"
|
||||
f"Bitrate: {metadata.get('bitrate', 'Unknown')} kbps\n"
|
||||
f"Codec: {metadata.get('codec', 'Unknown')}\n"
|
||||
f"File Size: {file_stats.st_size / 1024 / 1024:.2f} MB\n"
|
||||
f"Format: {file_path.suffix.upper()}"
|
||||
)
|
||||
|
||||
logging.info(f"Metadata extraction completed in {processing_time:.2f}s")
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=result_message, metadata=audio_metadata.model_dump()
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
)
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Metadata extraction failed: {str(e)}: {traceback.format_exc()}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Metadata extraction failed: {str(e)}",
|
||||
metadata={"error_type": "metadata_error"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Trim audio file to specified time range.
|
||||
|
||||
This tool cuts audio files to extract specific segments based on start time
|
||||
and duration. Useful for creating clips or removing unwanted sections.
|
||||
"""
|
||||
)
|
||||
async def trim_audio(
|
||||
file_path: str = Field(description="Path to the audio file to trim"),
|
||||
start_time: float = Field(description="Start time in seconds"),
|
||||
duration: float | None = Field(
|
||||
default=None, description="Duration in seconds (if None, trim to end)"
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
try:
|
||||
if isinstance(file_path, FieldInfo):
|
||||
file_path = file_path.default
|
||||
if isinstance(start_time, FieldInfo):
|
||||
start_time = start_time.default
|
||||
if isinstance(duration, FieldInfo):
|
||||
duration = duration.default
|
||||
|
||||
process_start = time.time()
|
||||
|
||||
# Validate input file
|
||||
file_path: Path = _validate_file_path(file_path)
|
||||
logging.info(f"Trimming audio: {file_path.name}")
|
||||
|
||||
# Get original metadata
|
||||
original_metadata = _get_audio_metadata(file_path)
|
||||
|
||||
# Validate time parameters
|
||||
if start_time < 0:
|
||||
raise ValueError("Start time cannot be negative")
|
||||
if duration is not None and duration <= 0:
|
||||
raise ValueError("Duration must be positive")
|
||||
if (
|
||||
original_metadata.get("duration")
|
||||
and start_time >= original_metadata["duration"]
|
||||
):
|
||||
raise ValueError("Start time exceeds audio duration")
|
||||
|
||||
# Trim audio
|
||||
output_path = _trim_audio(file_path, start_time, duration)
|
||||
|
||||
# Get trimmed file metadata
|
||||
trimmed_metadata = _get_audio_metadata(output_path)
|
||||
processing_time = time.time() - process_start
|
||||
|
||||
# Prepare file statistics
|
||||
file_stats = file_path.stat()
|
||||
|
||||
# Create metadata object
|
||||
audio_metadata = AudioMetadata(
|
||||
file_name=file_path.name,
|
||||
file_size=file_stats.st_size,
|
||||
file_type=file_path.suffix.lower(),
|
||||
absolute_path=str(file_path.absolute()),
|
||||
duration=trimmed_metadata.get("duration"),
|
||||
sample_rate=trimmed_metadata.get("sample_rate"),
|
||||
channels=trimmed_metadata.get("channels"),
|
||||
bitrate=trimmed_metadata.get("bitrate"),
|
||||
codec=trimmed_metadata.get("codec"),
|
||||
processing_time=processing_time,
|
||||
output_files=[str(output_path)],
|
||||
output_format="trimmed_audio",
|
||||
)
|
||||
|
||||
end_time = start_time + (
|
||||
duration or (original_metadata.get("duration", 0) - start_time)
|
||||
)
|
||||
result_message = (
|
||||
f"Successfully trimmed {file_path.name}\n"
|
||||
f"Original duration: {original_metadata.get('duration', 0):.2f} seconds\n"
|
||||
f"Trimmed segment: {start_time:.2f}s - {end_time:.2f}s\n"
|
||||
f"New duration: {trimmed_metadata.get('duration', 0):.2f} seconds\n"
|
||||
f"Output file: {output_path.name}"
|
||||
)
|
||||
|
||||
logging.info(f"Audio trimming completed in {processing_time:.2f}s")
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=result_message, metadata=audio_metadata.model_dump()
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
)
|
||||
}
|
||||
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Audio trimming failed: {str(e)}: {traceback.format_exc()}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Audio trimming failed: {str(e)}",
|
||||
metadata={"error_type": "trimming_error"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
List all supported audio formats for processing.
|
||||
"""
|
||||
)
|
||||
async def list_supported_formats() -> Union[str, TextContent]:
|
||||
supported_formats = {
|
||||
"MP3": "MPEG Audio Layer III (.mp3) - Most common compressed format",
|
||||
"WAV": "Waveform Audio File Format (.wav) - Uncompressed, high quality",
|
||||
"FLAC": "Free Lossless Audio Codec (.flac) - Lossless compression",
|
||||
"AAC": "Advanced Audio Coding (.aac) - Efficient compression",
|
||||
"OGG": "Ogg Vorbis (.ogg) - Open source compressed format",
|
||||
"M4A": "MPEG-4 Audio (.m4a) - Apple's preferred format",
|
||||
"WMA": "Windows Media Audio (.wma) - Microsoft format",
|
||||
"OPUS": "Opus Audio (.opus) - Modern, efficient codec",
|
||||
"AIFF": "Audio Interchange File Format (.aiff) - Apple's uncompressed format",
|
||||
"AU": "Sun Audio (.au) - Unix audio format",
|
||||
"RA": "RealAudio (.ra) - Streaming audio format",
|
||||
"AMR": "Adaptive Multi-Rate (.amr) - Mobile audio format",
|
||||
}
|
||||
format_list = "\n".join(
|
||||
[
|
||||
f"**{format_name}**: {description}"
|
||||
for format_name, description in supported_formats.items()
|
||||
]
|
||||
)
|
||||
action_response = ActionResponse(
|
||||
success=True,
|
||||
message=f"Supported audio formats:\n\n{format_list}",
|
||||
metadata={
|
||||
"supported_formats": list(supported_formats.keys()),
|
||||
"total_formats": len(supported_formats),
|
||||
"ffmpeg_available": _check_ffmpeg_availability(),
|
||||
},
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(action_response.model_dump())
|
||||
}
|
||||
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
def _check_ffmpeg_availability() -> bool:
|
||||
"""Check if ffmpeg is available in the system.
|
||||
|
||||
Returns:
|
||||
bool: True if ffmpeg is available, False otherwise
|
||||
"""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["ffmpeg", "-version"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=10,
|
||||
check=False,
|
||||
)
|
||||
if result.returncode == 0:
|
||||
logging.info("FFmpeg is available")
|
||||
else:
|
||||
logging.info("FFmpeg not found in system PATH")
|
||||
return result.returncode == 0
|
||||
except (subprocess.TimeoutExpired, FileNotFoundError) as e:
|
||||
logging.info(
|
||||
"FFmpeg not available or timeout, Please refer to https://github.com/inclusionAI/AWorld/blob/main/examples/gaia/README.md#system-tools-setup for more details"
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def _get_audio_metadata(file_path: Path) -> dict[str, Any]:
|
||||
"""Extract audio metadata using ffprobe.
|
||||
|
||||
Args:
|
||||
file_path: Path to the audio file
|
||||
|
||||
Returns:
|
||||
Dictionary containing audio metadata
|
||||
"""
|
||||
try:
|
||||
cmd = [
|
||||
"ffprobe",
|
||||
"-v",
|
||||
"quiet",
|
||||
"-print_format",
|
||||
"json",
|
||||
"-show_format",
|
||||
"-show_streams",
|
||||
str(file_path),
|
||||
]
|
||||
result = subprocess.run(
|
||||
cmd, capture_output=True, text=True, timeout=30, check=False
|
||||
)
|
||||
|
||||
if result.returncode == 0:
|
||||
metadata = json.loads(result.stdout)
|
||||
|
||||
# Extract relevant audio information
|
||||
format_info = metadata.get("format", {})
|
||||
streams = metadata.get("streams", [])
|
||||
audio_stream = next(
|
||||
(s for s in streams if s.get("codec_type") == "audio"), {}
|
||||
)
|
||||
|
||||
return {
|
||||
"duration": float(format_info.get("duration", 0)),
|
||||
"sample_rate": (
|
||||
int(audio_stream.get("sample_rate", 0))
|
||||
if audio_stream.get("sample_rate")
|
||||
else None
|
||||
),
|
||||
"channels": (
|
||||
int(audio_stream.get("channels", 0))
|
||||
if audio_stream.get("channels")
|
||||
else None
|
||||
),
|
||||
"bitrate": (
|
||||
int(format_info.get("bit_rate", 0)) // 1000
|
||||
if format_info.get("bit_rate")
|
||||
else None
|
||||
),
|
||||
"codec": audio_stream.get("codec_name"),
|
||||
}
|
||||
else:
|
||||
logging.warning(f"Failed to extract metadata: {result.stderr}")
|
||||
return {}
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Error extracting audio metadata: {str(e)}")
|
||||
return {}
|
||||
|
||||
|
||||
def _trim_audio(
|
||||
input_path: Path, start_time: float, duration: float | None = None
|
||||
) -> Path:
|
||||
"""Trim audio file to specified time range.
|
||||
|
||||
Args:
|
||||
input_path: Path to input audio file
|
||||
start_time: Start time in seconds
|
||||
duration: Duration in seconds (if None, trim to end)
|
||||
|
||||
Returns:
|
||||
Path to trimmed audio file
|
||||
"""
|
||||
output_path = _audio_output_dir / f"{input_path.stem}_trimmed{input_path.suffix}"
|
||||
|
||||
cmd = ["ffmpeg", "-i", str(input_path), "-ss", str(start_time), "-y"]
|
||||
|
||||
if duration is not None:
|
||||
cmd.extend(["-t", str(duration)])
|
||||
|
||||
cmd.extend(["-c", "copy", str(output_path)])
|
||||
|
||||
result = subprocess.run(
|
||||
cmd, capture_output=True, text=True, timeout=300, check=False
|
||||
)
|
||||
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(f"Audio trimming failed: {result.stderr}")
|
||||
|
||||
return output_path
|
||||
|
||||
|
||||
def _prepare_audio_for_transcription(file_path: Path) -> Path:
|
||||
"""Prepare audio file for transcription by converting to optimal format.
|
||||
|
||||
Args:
|
||||
file_path: Path to the original audio file
|
||||
|
||||
Returns:
|
||||
Path to the prepared audio file (WAV format, 16kHz)
|
||||
"""
|
||||
output_path = _audio_output_dir / f"{file_path.stem}_for_transcription.wav"
|
||||
|
||||
# Convert to WAV format with 16kHz sample rate for optimal transcription
|
||||
cmd = [
|
||||
"ffmpeg",
|
||||
"-i",
|
||||
str(file_path),
|
||||
"-ar",
|
||||
"16000", # 16kHz sample rate
|
||||
"-ac",
|
||||
"1", # Mono channel
|
||||
"-c:a",
|
||||
"pcm_s16le", # 16-bit PCM
|
||||
"-y",
|
||||
str(output_path),
|
||||
]
|
||||
|
||||
result = subprocess.run(
|
||||
cmd, capture_output=True, text=True, timeout=300, check=False
|
||||
)
|
||||
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(
|
||||
f"Audio preparation for transcription failed: {result.stderr}"
|
||||
)
|
||||
|
||||
return output_path
|
||||
|
||||
|
||||
def _transcribe_with_whisper(audio_path: Path) -> dict[str, Any]:
|
||||
"""Transcribe audio using OpenAI Whisper.
|
||||
|
||||
Args:
|
||||
audio_path: Path to the audio file
|
||||
|
||||
Returns:
|
||||
Dictionary containing transcription results
|
||||
"""
|
||||
try:
|
||||
client: OpenAI = OpenAI(
|
||||
api_key=os.getenv("AUDIO_LLM_API_KEY"),
|
||||
base_url=os.getenv("AUDIO_LLM_BASE_URL"),
|
||||
)
|
||||
|
||||
# Use the file for transcription
|
||||
with open(audio_path, "rb") as audio_file:
|
||||
transcription: str = client.audio.transcriptions.create(
|
||||
file=audio_file,
|
||||
model=os.getenv("AUDIO_LLM_MODEL_NAME"),
|
||||
response_format="text",
|
||||
)
|
||||
|
||||
return {"text": transcription.strip() if transcription else ""}
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
f"Audio transcription failed: {e}: {traceback.format_exc()}"
|
||||
) from e
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
load_dotenv(override=True)
|
||||
logging.info("Starting media-audio-server MCP server!")
|
||||
mcp.run(transport="stdio")
|
||||
+226
@@ -0,0 +1,226 @@
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import magic
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import requests
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import Any, Literal
|
||||
|
||||
|
||||
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)"
|
||||
)
|
||||
|
||||
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")
|
||||
|
||||
|
||||
def _validate_file_path(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)
|
||||
if not path.is_absolute():
|
||||
path = path.expanduser().resolve()
|
||||
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"File not found: {path}")
|
||||
|
||||
return path
|
||||
|
||||
def is_url(path_or_url: str) -> bool:
|
||||
"""
|
||||
Check if the given string is a URL.
|
||||
|
||||
Args:
|
||||
path_or_url: String to check
|
||||
|
||||
Returns:
|
||||
bool: True if the string is a URL, False otherwise
|
||||
"""
|
||||
parsed = urlparse(path_or_url)
|
||||
return bool(parsed.scheme and parsed.netloc)
|
||||
|
||||
|
||||
def get_mime_type(file_path: str, default_mime: str | None = None) -> str:
|
||||
"""
|
||||
Detect MIME type of a file using python-magic if available,
|
||||
otherwise fallback to extension-based detection.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
default_mime: Default MIME type to return if detection fails
|
||||
|
||||
Returns:
|
||||
str: Detected MIME type
|
||||
"""
|
||||
# Try using python-magic for accurate MIME type detection
|
||||
try:
|
||||
mime = magic.Magic(mime=True)
|
||||
return mime.from_file(file_path)
|
||||
except (AttributeError, IOError):
|
||||
# Fallback to extension-based detection
|
||||
extension_mime_map = {
|
||||
# Audio formats
|
||||
".mp3": "audio/mpeg",
|
||||
".wav": "audio/wav",
|
||||
".ogg": "audio/ogg",
|
||||
".m4a": "audio/mp4",
|
||||
".flac": "audio/flac",
|
||||
# Image formats
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".png": "image/png",
|
||||
".gif": "image/gif",
|
||||
".webp": "image/webp",
|
||||
".bmp": "image/bmp",
|
||||
".tiff": "image/tiff",
|
||||
# Video formats
|
||||
".mp4": "video/mp4",
|
||||
".avi": "video/x-msvideo",
|
||||
".mov": "video/quicktime",
|
||||
".mkv": "video/x-matroska",
|
||||
".webm": "video/webm",
|
||||
}
|
||||
|
||||
ext = Path(file_path).suffix.lower()
|
||||
return extension_mime_map.get(ext, default_mime or "application/octet-stream")
|
||||
|
||||
|
||||
def get_file_from_source(
|
||||
source: str,
|
||||
max_size_mb: float = 100.0,
|
||||
timeout: int = 60,
|
||||
) -> tuple[str, str, bytes]:
|
||||
"""
|
||||
Unified function to get file content from a URL or local path with validation.
|
||||
|
||||
Args:
|
||||
source: URL or local file path
|
||||
max_size_mb: Maximum allowed file size in MB
|
||||
timeout: Timeout for URL requests in seconds
|
||||
|
||||
Returns:
|
||||
Tuple[str, str, bytes]: (file_path, mime_type, file_content)
|
||||
- For URLs, file_path will be a temporary file path
|
||||
- For local files, file_path will be the original path
|
||||
|
||||
Raises:
|
||||
ValueError: When file doesn't exist, exceeds size limit, or has invalid MIME type
|
||||
IOError: When file cannot be read
|
||||
requests.RequestException: When URL request fails
|
||||
"""
|
||||
max_size_bytes = max_size_mb * 1024 * 1024
|
||||
|
||||
if is_url(source):
|
||||
# Handle URL source
|
||||
try:
|
||||
# Make a HEAD request first to check content length
|
||||
head_response = requests.head(source, timeout=timeout, allow_redirects=True)
|
||||
head_response.raise_for_status()
|
||||
|
||||
# Check content length if available
|
||||
content_length = head_response.headers.get("content-length")
|
||||
if content_length and int(content_length) > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({int(content_length) / (1024 * 1024):.2f} MB) "
|
||||
f"exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Download the file
|
||||
response = requests.get(source, timeout=timeout, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
# Read content with size checking
|
||||
content = b""
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
if len(content) + len(chunk) > max_size_bytes:
|
||||
raise ValueError(f"File size exceeds maximum allowed size ({max_size_mb} MB)")
|
||||
content += chunk
|
||||
|
||||
# Create temporary file
|
||||
parsed_url = urlparse(source)
|
||||
filename = os.path.basename(parsed_url.path) or "downloaded_file"
|
||||
|
||||
# Create temporary file with proper extension
|
||||
suffix = Path(filename).suffix or ".tmp"
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
|
||||
temp_file.write(content)
|
||||
temp_path = temp_file.name
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(temp_path)
|
||||
|
||||
return temp_path, mime_type, content
|
||||
|
||||
except requests.RequestException as e:
|
||||
raise requests.RequestException(f"Failed to download file from URL: {e}: {traceback.format_exc()}")
|
||||
except Exception as e:
|
||||
raise IOError(f"Error processing URL: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
else:
|
||||
# Handle local file path
|
||||
file_path = Path(source)
|
||||
|
||||
# Check if file exists
|
||||
if not file_path.exists():
|
||||
raise ValueError(f"File does not exist: {source}")
|
||||
|
||||
if not file_path.is_file():
|
||||
raise ValueError(f"Path is not a file: {source}")
|
||||
|
||||
# Check file size
|
||||
file_size = file_path.stat().st_size
|
||||
if file_size > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({file_size / (1024 * 1024):.2f} MB) exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Read file content
|
||||
try:
|
||||
with open(file_path, "rb") as f:
|
||||
content = f.read()
|
||||
except Exception as e:
|
||||
raise IOError(f"Cannot read file {source}: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(str(file_path))
|
||||
|
||||
return str(file_path), mime_type, content
|
||||
+560
@@ -0,0 +1,560 @@
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import Any, Union
|
||||
import pytesseract
|
||||
from PIL import Image, ImageEnhance, ImageFilter
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.types import TextContent
|
||||
from pydantic import Field, BaseModel
|
||||
from openai import OpenAI
|
||||
|
||||
from base import ActionResponse, _validate_file_path
|
||||
|
||||
|
||||
load_dotenv()
|
||||
workspace = Path.home()
|
||||
_image_output_dir = workspace / "processed_images"
|
||||
_image_output_dir.mkdir(exist_ok=True, parents=True)
|
||||
|
||||
|
||||
supported_extensions = {
|
||||
".jpg",
|
||||
".jpeg",
|
||||
".png",
|
||||
".gif",
|
||||
".webp",
|
||||
".bmp",
|
||||
".tiff",
|
||||
".tif",
|
||||
".ico",
|
||||
".svg",
|
||||
}
|
||||
|
||||
mcp = FastMCP(
|
||||
"media-image-server",
|
||||
instructions="""
|
||||
MCP service for comprehensive image processing and analysis.
|
||||
|
||||
Supports various image operations including:
|
||||
- Metadata extraction
|
||||
- OCR (Optical Character Recognition)
|
||||
- AI-powered image analysis and reasoning
|
||||
""",
|
||||
)
|
||||
|
||||
|
||||
class ImageMetadata(BaseModel):
|
||||
"""Metadata extracted from image processing."""
|
||||
|
||||
file_name: str = Field(description="Original image file name")
|
||||
file_size: int = Field(description="File size in bytes")
|
||||
file_type: str = Field(description="Image file type/extension")
|
||||
absolute_path: str = Field(description="Absolute path to the image file")
|
||||
width: int | None = Field(default=None, description="Image width in pixels")
|
||||
height: int | None = Field(default=None, description="Image height in pixels")
|
||||
mode: str | None = Field(
|
||||
default=None, description="Image color mode (RGB, RGBA, L, etc.)"
|
||||
)
|
||||
format: str | None = Field(
|
||||
default=None, description="Image format (JPEG, PNG, etc.)"
|
||||
)
|
||||
has_transparency: bool = Field(
|
||||
default=False, description="Whether image has transparency"
|
||||
)
|
||||
processing_time: float = Field(
|
||||
description="Time taken to process the image in seconds", exclude=True
|
||||
)
|
||||
output_files: list[str] = Field(
|
||||
default_factory=list, description="Paths to generated output files"
|
||||
)
|
||||
extracted_text: str | None = Field(
|
||||
default=None, description="Text extracted via OCR"
|
||||
)
|
||||
analysis_result: str | None = Field(default=None, description="AI analysis result")
|
||||
compression_ratio: float | None = Field(
|
||||
default=None, description="Compression ratio if optimized"
|
||||
)
|
||||
output_format: str = Field(description="Format of the processed output")
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Extract text from images using Optical Character Recognition (OCR).
|
||||
|
||||
This tool uses Tesseract OCR to extract text content from images,
|
||||
with optional preprocessing to improve recognition accuracy.
|
||||
"""
|
||||
)
|
||||
async def extract_text_ocr(
|
||||
file_path: str = Field(description="Path to the image file for OCR"),
|
||||
language: str = Field(
|
||||
default="eng", description="OCR language code (e.g., 'eng', 'spa', 'fra')"
|
||||
),
|
||||
preprocess: bool = Field(
|
||||
default=True, description="Whether to preprocess image for better OCR"
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
try:
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(file_path, FieldInfo):
|
||||
file_path = file_path.default
|
||||
if isinstance(language, FieldInfo):
|
||||
language = language.default
|
||||
if isinstance(preprocess, FieldInfo):
|
||||
preprocess = preprocess.default
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Validate input file
|
||||
file_path: Path = _validate_file_path(file_path)
|
||||
logging.info(f"Performing OCR on: {file_path.name}")
|
||||
|
||||
# Load image
|
||||
image = _load_image(file_path)
|
||||
original_metadata = _get_image_metadata(image, file_path)
|
||||
|
||||
# Preprocess image for better OCR if requested
|
||||
if preprocess:
|
||||
# Convert to grayscale
|
||||
if image.mode != "L":
|
||||
image = image.convert("L")
|
||||
|
||||
# Enhance contrast
|
||||
enhancer = ImageEnhance.Contrast(image)
|
||||
image = enhancer.enhance(2.0)
|
||||
|
||||
# Apply slight sharpening
|
||||
image = image.filter(ImageFilter.SHARPEN)
|
||||
|
||||
# Perform OCR
|
||||
extracted_text = _perform_ocr(image)
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
# Count words and characters
|
||||
word_count = len(extracted_text.split()) if extracted_text else 0
|
||||
char_count = len(extracted_text) if extracted_text else 0
|
||||
|
||||
# Create metadata object
|
||||
image_metadata = ImageMetadata(
|
||||
file_name=file_path.name,
|
||||
file_size=file_path.stat().st_size,
|
||||
file_type=file_path.suffix.lower(),
|
||||
absolute_path=str(file_path.absolute()),
|
||||
width=original_metadata["width"],
|
||||
height=original_metadata["height"],
|
||||
mode=original_metadata["mode"],
|
||||
format=original_metadata["format"],
|
||||
has_transparency=original_metadata["has_transparency"],
|
||||
processing_time=processing_time,
|
||||
output_files=[],
|
||||
extracted_text=extracted_text,
|
||||
output_format="ocr_text",
|
||||
)
|
||||
|
||||
if extracted_text:
|
||||
result_message = (
|
||||
f"OCR Results for {file_path.name}:\n\n"
|
||||
f"**Extracted Text:**\n{extracted_text}\n\n"
|
||||
f"**Statistics:**\n"
|
||||
f"- Words: {word_count}\n"
|
||||
f"- Characters: {char_count}\n"
|
||||
f"- Language: {language}\n"
|
||||
f"- Processing time: {processing_time:.2f}s"
|
||||
)
|
||||
else:
|
||||
result_message = (
|
||||
f"No text detected in {file_path.name}."
|
||||
" The image may not contain readable text or OCR preprocessing may be needed."
|
||||
)
|
||||
|
||||
logging.info(
|
||||
f"OCR completed: {word_count} words extracted in {processing_time:.2f}s"
|
||||
)
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=result_message, metadata=image_metadata.model_dump()
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
)
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"OCR failed: {str(e)}: {traceback.format_exc()}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"OCR failed: {str(e)}",
|
||||
metadata={"error_type": "ocr_error"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Analyze image content using AI vision models.
|
||||
|
||||
This tool uses advanced AI models to analyze and describe image content,
|
||||
answer questions about images, or perform specific visual reasoning tasks.
|
||||
"""
|
||||
)
|
||||
async def analyze_image_ai(
|
||||
file_path: str = Field(description="Path to the image file for AI analysis"),
|
||||
task: str = Field(
|
||||
default="Describe what you see in this image",
|
||||
description="Specific analysis task or question about the image",
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
try:
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(file_path, FieldInfo):
|
||||
file_path = file_path.default
|
||||
if isinstance(task, FieldInfo):
|
||||
task = task.default
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Validate input file
|
||||
file_path: Path = _validate_file_path(file_path)
|
||||
logging.info(f"Analyzing image with AI: {file_path.name}")
|
||||
|
||||
# Load image
|
||||
image = _load_image(file_path)
|
||||
original_metadata = _get_image_metadata(image, file_path)
|
||||
|
||||
# Convert to base64 for AI analysis
|
||||
image_base64 = _image_to_base64(image, "JPEG")
|
||||
|
||||
# Perform AI analysis
|
||||
analysis_result = _analyze_with_ai(image_base64, task)
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
# Create metadata object
|
||||
metadata_dict = {
|
||||
"file_name": file_path.name,
|
||||
"file_size": file_path.stat().st_size,
|
||||
"file_type": file_path.suffix.lower(),
|
||||
"absolute_path": str(file_path.absolute()),
|
||||
"width": original_metadata["width"],
|
||||
"height": original_metadata["height"],
|
||||
"mode": original_metadata["mode"],
|
||||
"format": original_metadata["format"],
|
||||
"has_transparency": original_metadata["has_transparency"],
|
||||
"processing_time": processing_time,
|
||||
"output_files": [],
|
||||
"analysis_result": analysis_result,
|
||||
"output_format": "ai_analysis",
|
||||
}
|
||||
|
||||
image_metadata = ImageMetadata(**metadata_dict)
|
||||
|
||||
result_message = (
|
||||
f"AI Analysis Results for {file_path.name}:\n\n"
|
||||
f"**Task:** {task}\n\n"
|
||||
f"**Analysis:**\n{analysis_result}\n\n"
|
||||
f"**Image Info:**\n"
|
||||
f"- Dimensions: {original_metadata['width']}x{original_metadata['height']}\n"
|
||||
f"- Format: {original_metadata['format']}\n"
|
||||
f"- Processing time: {processing_time:.2f}s"
|
||||
)
|
||||
|
||||
logging.infor_log(f"AI analysis completed in {processing_time:.2f}s")
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=result_message, metadata=image_metadata.model_dump()
|
||||
)
|
||||
output_dict = {"artifact_type": "MARKDOWN", "artifact_data": analysis_result}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"AI image analysis failed: {str(e)}: {traceback.format_exc()}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"AI image analysis failed: {str(e)}",
|
||||
metadata={"error_type": "ai_analysis_error"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Extract comprehensive metadata from image files.
|
||||
|
||||
This tool analyzes image files and extracts detailed metadata including
|
||||
dimensions, format, color mode, file size, and other technical information.
|
||||
"""
|
||||
)
|
||||
async def get_image_metadata(
|
||||
file_path: str = Field(description="Path to the image file to analyze"),
|
||||
) -> Union[str, TextContent]:
|
||||
try:
|
||||
if isinstance(file_path, FieldInfo):
|
||||
file_path = file_path.default
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Validate input file
|
||||
file_path: Path = _validate_file_path(file_path)
|
||||
logging.info(f"Extracting metadata from: {file_path.name}")
|
||||
|
||||
# Load image
|
||||
image = _load_image(file_path)
|
||||
metadata = _get_image_metadata(image, file_path)
|
||||
processing_time = time.time() - start_time
|
||||
|
||||
# Get file statistics
|
||||
file_stats = file_path.stat()
|
||||
|
||||
# Create metadata object
|
||||
image_metadata = ImageMetadata(
|
||||
file_name=file_path.name,
|
||||
file_size=file_stats.st_size,
|
||||
file_type=file_path.suffix.lower(),
|
||||
absolute_path=str(file_path.absolute()),
|
||||
width=metadata["width"],
|
||||
height=metadata["height"],
|
||||
mode=metadata["mode"],
|
||||
format=metadata["format"],
|
||||
has_transparency=metadata["has_transparency"],
|
||||
processing_time=processing_time,
|
||||
output_files=[],
|
||||
output_format="metadata",
|
||||
)
|
||||
|
||||
# Calculate additional info
|
||||
aspect_ratio = (
|
||||
metadata["width"] / metadata["height"] if metadata["height"] > 0 else 0
|
||||
)
|
||||
megapixels = (metadata["width"] * metadata["height"]) / 1_000_000
|
||||
|
||||
result_message = (
|
||||
f"Image Metadata for {file_path.name}:\n"
|
||||
f"Dimensions: {metadata['width']}x{metadata['height']} pixels\n"
|
||||
f"Aspect Ratio: {aspect_ratio:.2f}:1\n"
|
||||
f"Megapixels: {megapixels:.2f} MP\n"
|
||||
f"Color Mode: {metadata['mode']}\n"
|
||||
f"Format: {metadata['format']}\n"
|
||||
f"Has Transparency: {metadata['has_transparency']}\n"
|
||||
f"File Size: {file_stats.st_size / 1024 / 1024:.2f} MB\n"
|
||||
f"File Type: {file_path.suffix.upper()}"
|
||||
)
|
||||
|
||||
logging.info(f"Metadata extraction completed in {processing_time:.2f}s")
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=result_message, metadata=image_metadata.model_dump()
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
)
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Metadata extraction failed: {str(e)}: {traceback.format_exc()}")
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Metadata extraction failed: {str(e)}",
|
||||
metadata={"error_type": "metadata_error"},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
def _load_image(file_path: Path) -> Image.Image:
|
||||
"""Load image from file path.
|
||||
|
||||
Args:
|
||||
file_path: Path to the image file
|
||||
|
||||
Returns:
|
||||
PIL Image object
|
||||
"""
|
||||
try:
|
||||
return Image.open(file_path)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to load image {file_path}: {str(e)}") from e
|
||||
|
||||
|
||||
def _get_image_metadata(image: Image.Image, file_path: Path) -> dict[str, Any]:
|
||||
"""Extract metadata from PIL Image object.
|
||||
|
||||
Args:
|
||||
image: PIL Image object
|
||||
file_path: Path to the original file
|
||||
|
||||
Returns:
|
||||
Dictionary containing image metadata
|
||||
"""
|
||||
return {
|
||||
"width": image.width,
|
||||
"height": image.height,
|
||||
"mode": image.mode,
|
||||
"format": image.format or file_path.suffix.upper().lstrip("."),
|
||||
"has_transparency": image.mode in ("RGBA", "LA")
|
||||
or "transparency" in image.info,
|
||||
}
|
||||
|
||||
|
||||
def _optimize_image(
|
||||
image: Image.Image, max_size: tuple[int, int] | None = None
|
||||
) -> Image.Image:
|
||||
"""Optimize image for size and quality.
|
||||
|
||||
Args:
|
||||
image: PIL Image object
|
||||
max_size: Maximum dimensions (width, height)
|
||||
|
||||
Returns:
|
||||
Optimized PIL Image object
|
||||
"""
|
||||
optimized = image.copy()
|
||||
|
||||
# Resize if max_size specified
|
||||
if max_size:
|
||||
optimized.thumbnail(max_size, Image.Resampling.LANCZOS)
|
||||
|
||||
return optimized
|
||||
|
||||
|
||||
def _perform_ocr(image: Image.Image) -> str:
|
||||
"""Perform OCR on image to extract text.
|
||||
|
||||
Args:
|
||||
image: PIL Image object
|
||||
|
||||
Returns:
|
||||
Extracted text string
|
||||
"""
|
||||
try:
|
||||
return pytesseract.image_to_string(image).strip()
|
||||
except ImportError:
|
||||
return "OCR not available - pytesseract not installed"
|
||||
except Exception as e:
|
||||
return f"OCR failed: {str(e)}"
|
||||
|
||||
|
||||
def _analyze_with_ai(image_base64: str, task: str) -> str:
|
||||
"""Analyze image using AI model.
|
||||
|
||||
Args:
|
||||
image_base64: Base64 encoded image
|
||||
task: Analysis task description
|
||||
|
||||
Returns:
|
||||
AI analysis result
|
||||
"""
|
||||
try:
|
||||
openai_params = {
|
||||
"model": os.getenv("IMAGE_LLM_MODEL_NAME", ""),
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "content": task},
|
||||
{"type": "image_url", "image_url": {"url": image_base64}},
|
||||
],
|
||||
},
|
||||
]
|
||||
}
|
||||
client: OpenAI = OpenAI(
|
||||
api_key=os.getenv("IMAGE_LLM_API_KEY"),
|
||||
base_url=os.getenv("IMAGE_LLM_BASE_URL"),
|
||||
)
|
||||
response = client.chat.completions.create(**openai_params)
|
||||
content = ""
|
||||
if response and hasattr(response, 'choices') and response.choices:
|
||||
content = response.choices[0].message.content
|
||||
return content
|
||||
except Exception as e:
|
||||
return f"AI analysis failed: {str(e)}"
|
||||
|
||||
|
||||
def _image_to_base64(image: Image.Image, output_format: str = "JPEG") -> str:
|
||||
"""Convert PIL Image to base64 string.
|
||||
|
||||
Args:
|
||||
image: PIL Image object
|
||||
output_format: Output format (JPEG, PNG, etc.)
|
||||
|
||||
Returns:
|
||||
Base64 encoded image string
|
||||
"""
|
||||
buffer = BytesIO()
|
||||
|
||||
# Convert RGBA or P to RGB for JPEG
|
||||
if output_format.upper() == "JPEG":
|
||||
if image.mode in ("RGBA", "LA"):
|
||||
background = Image.new("RGB", image.size, (255, 255, 255))
|
||||
background.paste(
|
||||
image, mask=image.split()[-1] if image.mode == "RGBA" else None
|
||||
)
|
||||
image = background
|
||||
elif image.mode == "P":
|
||||
image = image.convert("RGB")
|
||||
|
||||
image.save(
|
||||
buffer,
|
||||
format=output_format,
|
||||
quality=85 if output_format.upper() == "JPEG" else None,
|
||||
)
|
||||
|
||||
mime_type = f"image/{output_format.lower()}"
|
||||
img_base64 = base64.b64encode(buffer.getvalue()).decode()
|
||||
|
||||
return f"data:{mime_type};base64,{img_base64}"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
load_dotenv(override=True)
|
||||
logging.info("Starting media-image-server MCP server!")
|
||||
mcp.run(transport="stdio")
|
||||
+986
@@ -0,0 +1,986 @@
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from pathlib import Path
|
||||
from typing import Union
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.types import TextContent
|
||||
from pydantic import Field, BaseModel
|
||||
from openai import OpenAI
|
||||
|
||||
from base import ActionResponse, _validate_file_path, get_file_from_source
|
||||
|
||||
load_dotenv()
|
||||
workspace = Path.home()
|
||||
|
||||
supported_extensions = {".mp4", ".avi", ".mov", ".mkv", ".webm", ".flv"}
|
||||
video_analyze_prompt = (
|
||||
"Input is a sequence of video frames. Given user's task: {task}. "
|
||||
"analyze the video content following these steps:\n"
|
||||
"1. Temporal sequence understanding\n"
|
||||
"2. Motion and action analysis\n"
|
||||
"3. Scene context interpretation\n"
|
||||
"4. Object and person tracking\n"
|
||||
)
|
||||
|
||||
video_summarize_prompt = (
|
||||
"Input is a sequence of video frames. "
|
||||
"Summarize the main content of the video. "
|
||||
"Include key points, main topics, and important visual elements. "
|
||||
)
|
||||
|
||||
|
||||
mcp = FastMCP(
|
||||
"media-video-server",
|
||||
instructions="""
|
||||
MCP service for video operations with AI-powered analysis.
|
||||
|
||||
Provides video processing capabilities including:
|
||||
- AI-powered video content analysis
|
||||
- Video summarization with key insights
|
||||
- Keyframe extraction with scene detection
|
||||
- Subtitle extraction from video content
|
||||
- LLM-friendly result formatting
|
||||
- Error handling and logging
|
||||
""",
|
||||
)
|
||||
|
||||
|
||||
class VideoAnalysisResult(BaseModel):
|
||||
"""Video analysis result model with structured data"""
|
||||
|
||||
video_source: str
|
||||
analysis_result: str
|
||||
frame_count: int
|
||||
duration_analyzed: float
|
||||
success: bool
|
||||
error: str | None = None
|
||||
|
||||
|
||||
class VideoSummaryResult(BaseModel):
|
||||
"""Video summary result model with structured data"""
|
||||
|
||||
video_source: str
|
||||
summary: str
|
||||
frame_count: int
|
||||
duration_analyzed: float
|
||||
success: bool
|
||||
error: str | None = None
|
||||
|
||||
|
||||
class KeyframeResult(BaseModel):
|
||||
"""Keyframe extraction result model with file information"""
|
||||
|
||||
frame_paths: list[str]
|
||||
frame_timestamps: list[float]
|
||||
output_directory: str
|
||||
frame_count: int
|
||||
target_time: float
|
||||
window_size: float
|
||||
success: bool
|
||||
error: str | None = None
|
||||
|
||||
|
||||
class VideoMetadata(BaseModel):
|
||||
"""Metadata for video operation results"""
|
||||
|
||||
operation: str
|
||||
video_source: str | None = None
|
||||
sample_rate: int | None = None
|
||||
start_time: float | None = None
|
||||
end_time: float | None = None
|
||||
target_time: float | None = None
|
||||
window_size: float | None = None
|
||||
output_directory: str | None = None
|
||||
frame_count: int | None = None
|
||||
execution_time: float | None = None
|
||||
error_type: str | None = None
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Analyze video content using AI with parallel processing.
|
||||
|
||||
This tool provides comprehensive video analysis capabilities including:
|
||||
- Content understanding and description
|
||||
- Object and scene detection
|
||||
- Action and movement analysis
|
||||
- Temporal event tracking
|
||||
- Question-answering about video content
|
||||
- Parallel processing for faster analysis
|
||||
"""
|
||||
)
|
||||
async def analyze_video(
|
||||
video_url: str = Field(description="Path or URL to the video file to analyze"),
|
||||
question: str = Field(description="Question or task for video analysis"),
|
||||
sample_rate: float = Field(
|
||||
default=1.0, description="Frame sampling rate (frames per second)"
|
||||
),
|
||||
start_time: float = Field(default=0.0, description="Start time in seconds"),
|
||||
end_time: float | None = Field(
|
||||
default=None, description="End time in seconds (None for full video)"
|
||||
),
|
||||
output_format: str = Field(
|
||||
default="markdown",
|
||||
description="Output format: 'markdown', 'json', or 'text'",
|
||||
),
|
||||
max_workers: int = Field(
|
||||
default=4, description="Maximum number of parallel workers for analysis"
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
start_exec_time = time.time()
|
||||
|
||||
try:
|
||||
# Validate video file
|
||||
video_path = _validate_file_path(video_url)
|
||||
|
||||
logging.info(f"🎬 Analyzing video: {video_url}")
|
||||
logging.info(
|
||||
f"📋 Question: {question}",
|
||||
)
|
||||
|
||||
# Extract video frames
|
||||
video_frames = _get_video_frames(
|
||||
str(video_path), sample_rate, start_time, end_time
|
||||
)
|
||||
logging.info(
|
||||
f"📸 Extracted {len(video_frames)} frames",
|
||||
)
|
||||
|
||||
# Process frames in chunks of 64 frames for parallel analysis
|
||||
chunk_size = 128
|
||||
chunks = []
|
||||
|
||||
# Create chunks of `chunk_size` continuous frames
|
||||
for i in range(0, len(video_frames), chunk_size):
|
||||
chunk_frames = video_frames[i : i + chunk_size]
|
||||
chunks.append((i // chunk_size, chunk_frames, question))
|
||||
|
||||
logging.info(
|
||||
f"🔄 Processing {len(chunks)} chunks with {max_workers} parallel workers"
|
||||
)
|
||||
|
||||
# Process chunks in parallel
|
||||
all_results = [None] * len(chunks) # Pre-allocate to maintain order
|
||||
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
# Submit all chunk analysis tasks
|
||||
future_to_chunk = {
|
||||
executor.submit(_analyze_frame_chunk, chunk_data): chunk_data[0]
|
||||
for chunk_data in chunks
|
||||
}
|
||||
|
||||
# Collect results as they complete
|
||||
for future in as_completed(future_to_chunk):
|
||||
try:
|
||||
chunk_index, result = future.result()
|
||||
all_results[chunk_index] = (
|
||||
f"Result of video part {chunk_index + 1}: {result}"
|
||||
)
|
||||
except Exception as e:
|
||||
chunk_index = future_to_chunk[future]
|
||||
logging.info(
|
||||
f"❌ Error processing chunk {chunk_index + 1}: {str(e)}"
|
||||
)
|
||||
all_results[chunk_index] = (
|
||||
f"Result of video part {chunk_index + 1}: Analysis failed - {str(e)}"
|
||||
)
|
||||
|
||||
# Filter out None results and join
|
||||
analysis_result = "\n".join(
|
||||
[result for result in all_results if result is not None]
|
||||
)
|
||||
duration_analyzed = (
|
||||
end_time - start_time if end_time else len(video_frames) / sample_rate
|
||||
)
|
||||
|
||||
# Create result
|
||||
result = VideoAnalysisResult(
|
||||
video_source=video_url,
|
||||
analysis_result=analysis_result,
|
||||
frame_count=len(video_frames),
|
||||
duration_analyzed=duration_analyzed,
|
||||
success=True,
|
||||
error=None,
|
||||
)
|
||||
|
||||
# Format output for LLM
|
||||
message = _format_analysis_output(result, output_format)
|
||||
execution_time = time.time() - start_exec_time
|
||||
|
||||
# Create metadata
|
||||
metadata = {
|
||||
"video_source": video_url,
|
||||
"frame_count": len(video_frames),
|
||||
"chunks_processed": len(chunks),
|
||||
"chunk_size": chunk_size,
|
||||
"parallel_workers": max_workers,
|
||||
"duration_analyzed": duration_analyzed,
|
||||
"sample_rate": sample_rate,
|
||||
"start_time": start_time,
|
||||
"end_time": end_time,
|
||||
"execution_time": execution_time,
|
||||
"output_format": output_format,
|
||||
"success": True,
|
||||
}
|
||||
|
||||
logging.info(
|
||||
f"✅ Video analysis completed in {execution_time:.2f}s chunks, {len(video_frames)} frames)"
|
||||
)
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=message, metadata=metadata
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
)
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
execution_time = time.time() - start_exec_time
|
||||
error_msg = f"Video analysis failed: {str(e)}"
|
||||
logging.info(f"❌ {error_msg}")
|
||||
logging.error(f"{error_msg}: {traceback.format_exc()}")
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=error_msg,
|
||||
metadata={
|
||||
"video_source": video_url,
|
||||
"execution_time": execution_time,
|
||||
"error": str(e),
|
||||
"success": False,
|
||||
},
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Summarize the main content of a video using AI analysis.
|
||||
|
||||
This tool provides AI-powered video summarization with:
|
||||
- Key point extraction
|
||||
- Main topic identification
|
||||
- Important visual element recognition
|
||||
- LLM-optimized result formatting
|
||||
"""
|
||||
)
|
||||
async def summarize_video(
|
||||
video_url: str = Field(description="The input video filepath or URL to summarize."),
|
||||
sample_rate: int = Field(
|
||||
default=1, description="Sample n frames per second (default: 1)."
|
||||
),
|
||||
start_time: float = Field(
|
||||
default=0,
|
||||
description="Start time of the video segment in seconds (default: 0).",
|
||||
),
|
||||
end_time: float | None = Field(
|
||||
default=None,
|
||||
description="End time of the video segment in seconds (default: None).",
|
||||
),
|
||||
output_format: str = Field(
|
||||
default="markdown",
|
||||
description="Output format: 'markdown', 'json', or 'text' (default: markdown).",
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
start_exec_time = time.time()
|
||||
|
||||
try:
|
||||
# Validate video file
|
||||
video_path = _validate_file_path(video_url)
|
||||
|
||||
logging.info(
|
||||
f"🎬 Summarizing video: {video_url}",
|
||||
)
|
||||
|
||||
# Extract video frames
|
||||
video_frames = _get_video_frames(
|
||||
str(video_path), sample_rate, start_time, end_time
|
||||
)
|
||||
logging.info(f"📸 Extracted {len(video_frames)} frames")
|
||||
|
||||
# Process frames in larger chunks for summarization
|
||||
interval = 490
|
||||
frame_nums = 500
|
||||
all_results = []
|
||||
|
||||
for i in range(0, len(video_frames), interval):
|
||||
cur_frames = video_frames[i : i + frame_nums]
|
||||
content = _create_video_content(video_summarize_prompt, cur_frames)
|
||||
inputs = [{"role": "user", "content": content}]
|
||||
|
||||
try:
|
||||
openai_params = {
|
||||
"model": os.getenv("VIDEO_LLM_MODEL_NAME", ""),
|
||||
"messages": inputs,
|
||||
"temperature" :float(os.getenv("VIDEO_LLM_TEMPERATURE", "1.0")),
|
||||
}
|
||||
client: OpenAI = OpenAI(
|
||||
api_key=os.getenv("VIDEO_LLM_API_KEY"),
|
||||
base_url=os.getenv("VIDEO_LLM_BASE_URL"),
|
||||
)
|
||||
response = client.chat.completions.create(**openai_params)
|
||||
content = ""
|
||||
if response and hasattr(response, 'choices') and response.choices:
|
||||
content = response.choices[0].message.content
|
||||
|
||||
cur_summary = content
|
||||
except Exception as e:
|
||||
logging.info(
|
||||
f"LLM summary error for chunk {i // interval + 1}: {str(e)}"
|
||||
)
|
||||
cur_summary = f"Summary failed for video segment {i // interval + 1}"
|
||||
|
||||
all_results.append(
|
||||
f"Summary of video part {i // interval + 1}: {cur_summary}"
|
||||
)
|
||||
|
||||
if i + frame_nums >= len(video_frames):
|
||||
break
|
||||
|
||||
summary_result = "\n".join(all_results)
|
||||
duration_analyzed = (
|
||||
end_time - start_time if end_time else len(video_frames) / sample_rate
|
||||
)
|
||||
|
||||
# Create result
|
||||
result = VideoSummaryResult(
|
||||
video_source=video_url,
|
||||
summary=summary_result,
|
||||
frame_count=len(video_frames),
|
||||
duration_analyzed=duration_analyzed,
|
||||
success=True,
|
||||
error=None,
|
||||
)
|
||||
|
||||
# Format output for LLM
|
||||
message = _format_summary_output(result, output_format)
|
||||
execution_time = time.time() - start_exec_time
|
||||
|
||||
# Create metadata
|
||||
metadata = VideoMetadata(
|
||||
operation="summarize",
|
||||
video_source=video_url,
|
||||
sample_rate=sample_rate,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
frame_count=len(video_frames),
|
||||
execution_time=execution_time,
|
||||
).model_dump()
|
||||
|
||||
logging.info("✅ Video summarization completed successfully")
|
||||
action_response = ActionResponse(
|
||||
success=True, message=message, metadata=metadata
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
)
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
logging.info(f"❌ Video summarization error: {traceback.format_exc()}")
|
||||
|
||||
# Format error for LLM
|
||||
message = f"Failed to summarize video: {error_msg}"
|
||||
execution_time = time.time() - start_exec_time
|
||||
|
||||
# Create metadata
|
||||
metadata = VideoMetadata(
|
||||
operation="summarize",
|
||||
video_source=video_url,
|
||||
error_type="summarization_failure",
|
||||
execution_time=execution_time,
|
||||
).model_dump()
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=False, message=message, metadata=metadata
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Extract key frames around a target time with scene detection.
|
||||
|
||||
This tool provides keyframe extraction with:
|
||||
- Scene detection for significant frame changes
|
||||
- Configurable time windows
|
||||
- Automatic output directory management
|
||||
- LLM-optimized result formatting
|
||||
"""
|
||||
)
|
||||
async def extract_keyframes(
|
||||
video_path: str = Field(description="The input video filepath or URL."),
|
||||
target_time: int = Field(
|
||||
description="The specific time point for extraction (in seconds), centered within the window_size."
|
||||
),
|
||||
window_size: int = Field(
|
||||
default=5,
|
||||
description="The window size for extraction (in seconds, default: 5).",
|
||||
),
|
||||
output_dir: str = Field(
|
||||
default=None,
|
||||
description="Directory where extracted frames will be saved (default: workspace/keyframes).",
|
||||
),
|
||||
output_format: str = Field(
|
||||
default="markdown",
|
||||
description="Output format: 'markdown', 'json', or 'text' (default: markdown).",
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
start_exec_time = time.time()
|
||||
|
||||
try:
|
||||
# Validate video file
|
||||
validated_path = _validate_file_path(video_path)
|
||||
|
||||
# Set default output directory
|
||||
output_dir = str(workspace / "keyframes") if output_dir is None else output_dir
|
||||
|
||||
output_path = Path(output_dir)
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
logging.info(f"🎬 Extracting keyframes from: {video_path}")
|
||||
logging.info(f"🎯 Target time: {target_time}s, Window: {window_size}s")
|
||||
|
||||
# Extract keyframes with scene detection
|
||||
frames, frame_times = _extract_keyframes_with_scene_detection(
|
||||
str(validated_path), target_time, window_size
|
||||
)
|
||||
|
||||
# Save frames to disk
|
||||
frame_paths, frame_timestamps = _save_keyframes(
|
||||
frames, frame_times, str(output_path)
|
||||
)
|
||||
|
||||
# Cleanup if requested
|
||||
# if cleanup and validated_path.exists():
|
||||
# validated_path.unlink()
|
||||
# self._color_log("🗑️ Cleaned up original video file", Color.yellow)
|
||||
|
||||
# Create result
|
||||
result = KeyframeResult(
|
||||
frame_paths=frame_paths,
|
||||
frame_timestamps=frame_timestamps,
|
||||
output_directory=str(output_path),
|
||||
frame_count=len(frame_paths),
|
||||
target_time=float(target_time),
|
||||
window_size=float(window_size),
|
||||
success=True,
|
||||
error=None,
|
||||
)
|
||||
|
||||
# Format output for LLM
|
||||
message = _format_keyframe_output(result, output_format)
|
||||
execution_time = time.time() - start_exec_time
|
||||
|
||||
# Create metadata
|
||||
metadata = VideoMetadata(
|
||||
operation="extract_keyframes",
|
||||
video_source=video_path,
|
||||
target_time=float(target_time),
|
||||
window_size=float(window_size),
|
||||
output_directory=str(output_path),
|
||||
frame_count=len(frame_paths),
|
||||
execution_time=execution_time,
|
||||
).model_dump()
|
||||
|
||||
logging.info(f"✅ Extracted {len(frame_paths)} keyframes successfully")
|
||||
action_response = ActionResponse(
|
||||
success=True, message=message, metadata=metadata
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
)
|
||||
}
|
||||
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
logging.info(f"❌ Keyframe extraction error: {traceback.format_exc()}")
|
||||
|
||||
# Format error for LLM
|
||||
message = f"Failed to extract keyframes: {error_msg}"
|
||||
execution_time = time.time() - start_exec_time
|
||||
|
||||
# Create metadata
|
||||
metadata = VideoMetadata(
|
||||
operation="extract_keyframes",
|
||||
video_source=video_path,
|
||||
target_time=float(target_time) if target_time else None,
|
||||
window_size=float(window_size) if window_size else None,
|
||||
error_type="keyframe_extraction_failure",
|
||||
execution_time=execution_time,
|
||||
).model_dump()
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=False, message=message, metadata=metadata
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
def _get_video_frames(
|
||||
video_source: str,
|
||||
sample_rate: int = 2,
|
||||
start_time: float = 0,
|
||||
end_time: float | None = None,
|
||||
) -> list[dict[str, any]]:
|
||||
"""Extract frames from video with given sample rate.
|
||||
|
||||
Args:
|
||||
video_source: Path or URL to the video file
|
||||
sample_rate: Number of frames to sample per second
|
||||
start_time: Start time of the video segment in seconds
|
||||
end_time: End time of the video segment in seconds
|
||||
|
||||
Returns:
|
||||
List of dictionaries containing frame data and timestamp
|
||||
|
||||
Raises:
|
||||
ValueError: When video file cannot be opened or is not valid
|
||||
"""
|
||||
try:
|
||||
# Get file with validation (only video files allowed)
|
||||
file_path, _, _ = get_file_from_source(
|
||||
video_source,
|
||||
max_size_mb=2500.0, # 2500MB limit for videos
|
||||
)
|
||||
|
||||
# Open video file
|
||||
video = cv2.VideoCapture(file_path) # pylint: disable=E1101
|
||||
if not video.isOpened():
|
||||
raise ValueError(f"Could not open video file: {file_path}")
|
||||
|
||||
fps = video.get(cv2.CAP_PROP_FPS) # pylint: disable=E1101
|
||||
frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) # pylint: disable=E1101
|
||||
video_duration = frame_count / fps
|
||||
|
||||
if end_time is None:
|
||||
end_time = video_duration
|
||||
|
||||
if start_time > end_time:
|
||||
raise ValueError("Start time cannot be greater than end time.")
|
||||
|
||||
if start_time < 0:
|
||||
start_time = 0
|
||||
|
||||
if end_time > video_duration:
|
||||
end_time = video_duration
|
||||
|
||||
start_frame = int(start_time * fps)
|
||||
end_frame = int(end_time * fps)
|
||||
|
||||
all_frames = []
|
||||
frames = []
|
||||
|
||||
# Calculate frame interval based on sample rate
|
||||
frame_interval = max(1, int(fps / sample_rate))
|
||||
|
||||
# Set the video capture to the start frame
|
||||
video.set(cv2.CAP_PROP_POS_FRAMES, start_frame) # pylint: disable=E1101
|
||||
|
||||
for i in range(start_frame, end_frame):
|
||||
ret, frame = video.read()
|
||||
if not ret:
|
||||
break
|
||||
|
||||
# Convert frame to JPEG format
|
||||
_, buffer = cv2.imencode(".jpg", frame) # pylint: disable=E1101
|
||||
frame_data = base64.b64encode(buffer).decode("utf-8")
|
||||
|
||||
# Add data URL prefix for JPEG image
|
||||
frame_data = f"data:image/jpeg;base64,{frame_data}"
|
||||
|
||||
all_frames.append({"data": frame_data, "time": i / fps})
|
||||
|
||||
for i in range(0, len(all_frames), frame_interval):
|
||||
frames.append(all_frames[i])
|
||||
|
||||
video.release()
|
||||
|
||||
# Clean up temporary file if it was created for a URL
|
||||
if file_path != str(Path(video_source).resolve()) and Path(file_path).exists():
|
||||
Path(file_path).unlink()
|
||||
|
||||
if not frames:
|
||||
raise ValueError(f"Could not extract any frames from video: {video_source}")
|
||||
|
||||
return frames
|
||||
|
||||
except Exception as e:
|
||||
logging.info(f"Error extracting frames from {video_source}: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
def _create_video_content(
|
||||
prompt: str, video_frames: list[dict[str, any]]
|
||||
) -> list[dict[str, any]]:
|
||||
"""Create uniform video format for querying LLM."""
|
||||
content = [{"type": "text", "text": prompt}]
|
||||
content.extend(
|
||||
[
|
||||
{"type": "image_url", "image_url": {"url": frame["data"]}}
|
||||
for frame in video_frames
|
||||
]
|
||||
)
|
||||
return content
|
||||
|
||||
|
||||
def _format_analysis_output(
|
||||
result: VideoAnalysisResult, format_type: str = "markdown"
|
||||
) -> str:
|
||||
"""Format video analysis results for LLM consumption.
|
||||
|
||||
Args:
|
||||
result: Video analysis result
|
||||
format_type: Output format ('markdown', 'json', 'text')
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if not result.success:
|
||||
return f"Failed to analyze video: {result.error}"
|
||||
|
||||
if format_type == "json":
|
||||
return result.model_dump_json(indent=2)
|
||||
|
||||
elif format_type == "text":
|
||||
output_parts = [
|
||||
"Video Analysis Results",
|
||||
f"Source: {result.video_source}",
|
||||
f"Frames Analyzed: {result.frame_count}",
|
||||
f"Duration: {result.duration_analyzed:.2f} seconds",
|
||||
"",
|
||||
"Analysis:",
|
||||
result.analysis_result,
|
||||
]
|
||||
return "\n".join(output_parts)
|
||||
|
||||
else: # markdown (default)
|
||||
output_parts = [
|
||||
"# Video Analysis Results ✅",
|
||||
"",
|
||||
"## Video Information",
|
||||
f"**Source:** `{result.video_source}`",
|
||||
f"**Frames Analyzed:** {result.frame_count}",
|
||||
f"**Duration:** {result.duration_analyzed:.2f} seconds",
|
||||
"",
|
||||
"## Analysis Results",
|
||||
result.analysis_result,
|
||||
]
|
||||
return "\n".join(output_parts)
|
||||
|
||||
|
||||
def _format_summary_output(
|
||||
result: VideoSummaryResult, format_type: str = "markdown"
|
||||
) -> str:
|
||||
"""Format video summary results for LLM consumption.
|
||||
|
||||
Args:
|
||||
result: Video summary result
|
||||
format_type: Output format ('markdown', 'json', 'text')
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if not result.success:
|
||||
return f"Failed to summarize video: {result.error}"
|
||||
|
||||
if format_type == "json":
|
||||
return result.model_dump_json(indent=2)
|
||||
|
||||
elif format_type == "text":
|
||||
output_parts = [
|
||||
"Video Summary",
|
||||
f"Source: {result.video_source}",
|
||||
f"Frames Analyzed: {result.frame_count}",
|
||||
f"Duration: {result.duration_analyzed:.2f} seconds",
|
||||
"",
|
||||
"Summary:",
|
||||
result.summary,
|
||||
]
|
||||
return "\n".join(output_parts)
|
||||
|
||||
else: # markdown (default)
|
||||
output_parts = [
|
||||
"# Video Summary ✅",
|
||||
"",
|
||||
"## Video Information",
|
||||
f"**Source:** `{result.video_source}`",
|
||||
f"**Frames Analyzed:** {result.frame_count}",
|
||||
f"**Duration:** {result.duration_analyzed:.2f} seconds",
|
||||
"",
|
||||
"## Summary",
|
||||
result.summary,
|
||||
]
|
||||
return "\n".join(output_parts)
|
||||
|
||||
|
||||
def _format_keyframe_output(
|
||||
result: KeyframeResult, format_type: str = "markdown"
|
||||
) -> str:
|
||||
"""Format keyframe extraction results for LLM consumption.
|
||||
|
||||
Args:
|
||||
result: Keyframe extraction result
|
||||
format_type: Output format ('markdown', 'json', 'text')
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if not result.success:
|
||||
return f"Failed to extract keyframes: {result.error}"
|
||||
|
||||
if format_type == "json":
|
||||
return result.model_dump_json(indent=2)
|
||||
|
||||
elif format_type == "text":
|
||||
output_parts = [
|
||||
"Keyframe Extraction Results",
|
||||
f"Target Time: {result.target_time}s",
|
||||
f"Window Size: {result.window_size}s",
|
||||
f"Frames Extracted: {result.frame_count}",
|
||||
f"Output Directory: {result.output_directory}",
|
||||
"",
|
||||
"Frame Files:",
|
||||
]
|
||||
for i, (path, timestamp) in enumerate(
|
||||
zip(result.frame_paths, result.frame_timestamps), 1
|
||||
):
|
||||
output_parts.append(f"{i}. {path} (at {timestamp:.2f}s)")
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
else: # markdown (default)
|
||||
output_parts = [
|
||||
"# Keyframe Extraction Results ✅",
|
||||
"",
|
||||
"## Extraction Parameters",
|
||||
f"**Target Time:** {result.target_time}s",
|
||||
f"**Window Size:** {result.window_size}s",
|
||||
f"**Frames Extracted:** {result.frame_count}",
|
||||
f"**Output Directory:** `{result.output_directory}`",
|
||||
"",
|
||||
"## Extracted Frames",
|
||||
]
|
||||
|
||||
for i, (path, timestamp) in enumerate(
|
||||
zip(result.frame_paths, result.frame_timestamps), 1
|
||||
):
|
||||
output_parts.append(f"{i}. `{path}` (at {timestamp:.2f}s)")
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
|
||||
def _analyze_frame_chunk(chunk_data: tuple[int, list, str]) -> tuple[int, str]:
|
||||
"""Analyze a chunk of video frames using LLM.
|
||||
|
||||
Args:
|
||||
chunk_data: Tuple containing (chunk_index, frames, question)
|
||||
|
||||
Returns:
|
||||
Tuple of (chunk_index, analysis_result)
|
||||
"""
|
||||
chunk_index, frames, question = chunk_data
|
||||
|
||||
try:
|
||||
content = _create_video_content(
|
||||
video_analyze_prompt.format(task=question), frames
|
||||
)
|
||||
inputs = [{"role": "user", "content": content}]
|
||||
openai_params = {
|
||||
"model": os.getenv("VIDEO_LLM_MODEL_NAME", ""),
|
||||
"messages": inputs,
|
||||
"temperature" :float(os.getenv("VIDEO_LLM_TEMPERATURE", "1.0")),
|
||||
|
||||
}
|
||||
client: OpenAI = OpenAI(
|
||||
api_key=os.getenv("IMAGE_LLM_API_KEY"),
|
||||
base_url=os.getenv("IMAGE_LLM_BASE_URL"),
|
||||
)
|
||||
response = client.chat.completions.create(**openai_params)
|
||||
content = ""
|
||||
if response and hasattr(response, 'choices') and response.choices:
|
||||
content = response.choices[0].message.content
|
||||
|
||||
analysis_result = content
|
||||
logging.info(f"✅ Completed analysis for chunk {chunk_index + 1}")
|
||||
|
||||
except Exception as e:
|
||||
logging.info(f"❌ LLM analysis error for chunk {chunk_index + 1}: {str(e)}")
|
||||
analysis_result = (
|
||||
f"Analysis failed for video segment {chunk_index + 1}: {str(e)}"
|
||||
)
|
||||
|
||||
return chunk_index, analysis_result
|
||||
|
||||
|
||||
def _extract_keyframes_with_scene_detection(
|
||||
video_path: str, target_time: int, window_size: int
|
||||
) -> tuple[list[any], list[float]]:
|
||||
"""Extract key frames around the target time with scene detection.
|
||||
|
||||
Args:
|
||||
video_path: Path to the video file
|
||||
target_time: Target time in seconds
|
||||
window_size: Window size in seconds
|
||||
|
||||
Returns:
|
||||
Tuple of (frames, frame_times)
|
||||
"""
|
||||
cap = cv2.VideoCapture(video_path) # pylint: disable=E1101
|
||||
fps = cap.get(cv2.CAP_PROP_FPS) # pylint: disable=E1101
|
||||
|
||||
# Calculate frame numbers for the time window
|
||||
start_frame = int((target_time - window_size / 2) * fps)
|
||||
end_frame = int((target_time + window_size / 2) * fps)
|
||||
total_frames_in_window = end_frame - start_frame
|
||||
|
||||
max_frames = 384 # Maximum allowed frames to prevent memory issues
|
||||
|
||||
# Calculate sampling interval for even distribution
|
||||
if total_frames_in_window <= max_frames:
|
||||
# If total frames is within limit, use scene detection normally
|
||||
frame_interval = 1
|
||||
use_scene_detection = True
|
||||
else:
|
||||
# If exceeds limit, sample evenly across the window
|
||||
frame_interval = total_frames_in_window // max_frames
|
||||
use_scene_detection = False # Skip scene detection for even sampling
|
||||
|
||||
frames = []
|
||||
frame_times = []
|
||||
|
||||
# Set video position to start_frame
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, max(0, start_frame)) # pylint: disable=E1101
|
||||
|
||||
prev_frame = None
|
||||
frame_count = 0
|
||||
|
||||
while cap.isOpened() and len(frames) < max_frames:
|
||||
frame_pos = cap.get(cv2.CAP_PROP_POS_FRAMES) # pylint: disable=E1101
|
||||
if frame_pos >= end_frame:
|
||||
break
|
||||
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
break
|
||||
|
||||
if use_scene_detection:
|
||||
# Use original scene detection logic
|
||||
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) # pylint: disable=E1101
|
||||
|
||||
# If this is the first frame, save it
|
||||
if prev_frame is None:
|
||||
frames.append(frame)
|
||||
frame_times.append(frame_pos / fps)
|
||||
else:
|
||||
# Calculate difference between current and previous frame
|
||||
diff = cv2.absdiff(gray, prev_frame) # pylint: disable=E1101
|
||||
mean_diff = np.mean(diff)
|
||||
|
||||
# If significant change detected, save frame
|
||||
if mean_diff > 20: # Threshold for scene change
|
||||
frames.append(frame)
|
||||
frame_times.append(frame_pos / fps)
|
||||
|
||||
prev_frame = gray
|
||||
else:
|
||||
# Use even sampling for large windows
|
||||
if frame_count % frame_interval == 0:
|
||||
frames.append(frame)
|
||||
frame_times.append(frame_pos / fps)
|
||||
|
||||
frame_count += 1
|
||||
|
||||
cap.release()
|
||||
return frames, frame_times
|
||||
|
||||
|
||||
def _save_keyframes(
|
||||
frames: list[any], frame_times: list[float], output_dir: str
|
||||
) -> tuple[list[str], list[float]]:
|
||||
"""Save extracted frames to disk.
|
||||
|
||||
Args:
|
||||
frames: List of frame objects
|
||||
frame_times: List of frame timestamps
|
||||
output_dir: Output directory path
|
||||
|
||||
Returns:
|
||||
Tuple of (saved_paths, saved_timestamps)
|
||||
"""
|
||||
output_path = Path(output_dir)
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
saved_paths = []
|
||||
saved_timestamps = []
|
||||
|
||||
for _, (frame, timestamp) in enumerate(zip(frames, frame_times)):
|
||||
filename = output_path / f"frame_{timestamp:.2f}s.jpg"
|
||||
cv2.imwrite(str(filename), frame) # pylint: disable=E1101
|
||||
saved_paths.append(str(filename))
|
||||
saved_timestamps.append(timestamp)
|
||||
|
||||
return saved_paths, saved_timestamps
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
load_dotenv(override=True)
|
||||
logging.info("Starting media-video-server MCP server!")
|
||||
mcp.run(transport="stdio")
|
||||
Vendored
+16
@@ -0,0 +1,16 @@
|
||||
[project]
|
||||
name = "parxiv-server"
|
||||
version = "0.1.0"
|
||||
description = "Add your description here"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
aiohttp= ">=3.12.15"
|
||||
dependencies = [
|
||||
"mcp",
|
||||
"aiohttp~=3.12.15",
|
||||
"requests~=2.32.4",
|
||||
"fastmcp~=2.11.3",
|
||||
"arxiv~=2.2.0",
|
||||
"python-magic~=0.4.27",
|
||||
|
||||
]
|
||||
+226
@@ -0,0 +1,226 @@
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import magic
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import requests
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import Any, Literal
|
||||
|
||||
|
||||
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)"
|
||||
)
|
||||
|
||||
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")
|
||||
|
||||
|
||||
def _validate_file_path(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)
|
||||
if not path.is_absolute():
|
||||
path = path.expanduser().resolve()
|
||||
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"File not found: {path}")
|
||||
|
||||
return path
|
||||
|
||||
def is_url(path_or_url: str) -> bool:
|
||||
"""
|
||||
Check if the given string is a URL.
|
||||
|
||||
Args:
|
||||
path_or_url: String to check
|
||||
|
||||
Returns:
|
||||
bool: True if the string is a URL, False otherwise
|
||||
"""
|
||||
parsed = urlparse(path_or_url)
|
||||
return bool(parsed.scheme and parsed.netloc)
|
||||
|
||||
|
||||
def get_mime_type(file_path: str, default_mime: str | None = None) -> str:
|
||||
"""
|
||||
Detect MIME type of a file using python-magic if available,
|
||||
otherwise fallback to extension-based detection.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
default_mime: Default MIME type to return if detection fails
|
||||
|
||||
Returns:
|
||||
str: Detected MIME type
|
||||
"""
|
||||
# Try using python-magic for accurate MIME type detection
|
||||
try:
|
||||
mime = magic.Magic(mime=True)
|
||||
return mime.from_file(file_path)
|
||||
except (AttributeError, IOError):
|
||||
# Fallback to extension-based detection
|
||||
extension_mime_map = {
|
||||
# Audio formats
|
||||
".mp3": "audio/mpeg",
|
||||
".wav": "audio/wav",
|
||||
".ogg": "audio/ogg",
|
||||
".m4a": "audio/mp4",
|
||||
".flac": "audio/flac",
|
||||
# Image formats
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".png": "image/png",
|
||||
".gif": "image/gif",
|
||||
".webp": "image/webp",
|
||||
".bmp": "image/bmp",
|
||||
".tiff": "image/tiff",
|
||||
# Video formats
|
||||
".mp4": "video/mp4",
|
||||
".avi": "video/x-msvideo",
|
||||
".mov": "video/quicktime",
|
||||
".mkv": "video/x-matroska",
|
||||
".webm": "video/webm",
|
||||
}
|
||||
|
||||
ext = Path(file_path).suffix.lower()
|
||||
return extension_mime_map.get(ext, default_mime or "application/octet-stream")
|
||||
|
||||
|
||||
def get_file_from_source(
|
||||
source: str,
|
||||
max_size_mb: float = 100.0,
|
||||
timeout: int = 60,
|
||||
) -> tuple[str, str, bytes]:
|
||||
"""
|
||||
Unified function to get file content from a URL or local path with validation.
|
||||
|
||||
Args:
|
||||
source: URL or local file path
|
||||
max_size_mb: Maximum allowed file size in MB
|
||||
timeout: Timeout for URL requests in seconds
|
||||
|
||||
Returns:
|
||||
Tuple[str, str, bytes]: (file_path, mime_type, file_content)
|
||||
- For URLs, file_path will be a temporary file path
|
||||
- For local files, file_path will be the original path
|
||||
|
||||
Raises:
|
||||
ValueError: When file doesn't exist, exceeds size limit, or has invalid MIME type
|
||||
IOError: When file cannot be read
|
||||
requests.RequestException: When URL request fails
|
||||
"""
|
||||
max_size_bytes = max_size_mb * 1024 * 1024
|
||||
|
||||
if is_url(source):
|
||||
# Handle URL source
|
||||
try:
|
||||
# Make a HEAD request first to check content length
|
||||
head_response = requests.head(source, timeout=timeout, allow_redirects=True)
|
||||
head_response.raise_for_status()
|
||||
|
||||
# Check content length if available
|
||||
content_length = head_response.headers.get("content-length")
|
||||
if content_length and int(content_length) > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({int(content_length) / (1024 * 1024):.2f} MB) "
|
||||
f"exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Download the file
|
||||
response = requests.get(source, timeout=timeout, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
# Read content with size checking
|
||||
content = b""
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
if len(content) + len(chunk) > max_size_bytes:
|
||||
raise ValueError(f"File size exceeds maximum allowed size ({max_size_mb} MB)")
|
||||
content += chunk
|
||||
|
||||
# Create temporary file
|
||||
parsed_url = urlparse(source)
|
||||
filename = os.path.basename(parsed_url.path) or "downloaded_file"
|
||||
|
||||
# Create temporary file with proper extension
|
||||
suffix = Path(filename).suffix or ".tmp"
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
|
||||
temp_file.write(content)
|
||||
temp_path = temp_file.name
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(temp_path)
|
||||
|
||||
return temp_path, mime_type, content
|
||||
|
||||
except requests.RequestException as e:
|
||||
raise requests.RequestException(f"Failed to download file from URL: {e}: {traceback.format_exc()}")
|
||||
except Exception as e:
|
||||
raise IOError(f"Error processing URL: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
else:
|
||||
# Handle local file path
|
||||
file_path = Path(source)
|
||||
|
||||
# Check if file exists
|
||||
if not file_path.exists():
|
||||
raise ValueError(f"File does not exist: {source}")
|
||||
|
||||
if not file_path.is_file():
|
||||
raise ValueError(f"Path is not a file: {source}")
|
||||
|
||||
# Check file size
|
||||
file_size = file_path.stat().st_size
|
||||
if file_size > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({file_size / (1024 * 1024):.2f} MB) exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Read file content
|
||||
try:
|
||||
with open(file_path, "rb") as f:
|
||||
content = f.read()
|
||||
except Exception as e:
|
||||
raise IOError(f"Cannot read file {source}: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(str(file_path))
|
||||
|
||||
return str(file_path), mime_type, content
|
||||
Vendored
+851
@@ -0,0 +1,851 @@
|
||||
import json
|
||||
import logging
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Union
|
||||
|
||||
import arxiv
|
||||
from dotenv import load_dotenv
|
||||
from pydantic.fields import FieldInfo
|
||||
from mcp.server import FastMCP
|
||||
from mcp.types import TextContent
|
||||
from pydantic import Field, BaseModel
|
||||
|
||||
from base import ActionResponse
|
||||
|
||||
|
||||
load_dotenv()
|
||||
workspace = Path.home()
|
||||
supported_extensions = {".pdf"}
|
||||
|
||||
# ArXiv client configuration
|
||||
client = arxiv.Client(
|
||||
page_size=100,
|
||||
delay_seconds=3.0, # Be respectful to ArXiv servers
|
||||
num_retries=3,
|
||||
)
|
||||
|
||||
# Create downloads directory
|
||||
_downloads_dir = workspace / "arxiv_downloads"
|
||||
_downloads_dir.mkdir(exist_ok=True)
|
||||
|
||||
# ArXiv subject categories mapping
|
||||
subject_categories = {
|
||||
"cs": "Computer Science",
|
||||
"math": "Mathematics",
|
||||
"physics": "Physics",
|
||||
"astro-ph": "Astrophysics",
|
||||
"cond-mat": "Condensed Matter",
|
||||
"gr-qc": "General Relativity and Quantum Cosmology",
|
||||
"hep-ex": "High Energy Physics - Experiment",
|
||||
"hep-lat": "High Energy Physics - Lattice",
|
||||
"hep-ph": "High Energy Physics - Phenomenology",
|
||||
"hep-th": "High Energy Physics - Theory",
|
||||
"math-ph": "Mathematical Physics",
|
||||
"nlin": "Nonlinear Sciences",
|
||||
"nucl-ex": "Nuclear Experiment",
|
||||
"nucl-th": "Nuclear Theory",
|
||||
"quant-ph": "Quantum Physics",
|
||||
"q-bio": "Quantitative Biology",
|
||||
"q-fin": "Quantitative Finance",
|
||||
"stat": "Statistics",
|
||||
"econ": "Economics",
|
||||
"eess": "Electrical Engineering and Systems Science",
|
||||
}
|
||||
|
||||
|
||||
class PaperResult(BaseModel):
|
||||
"""Individual paper search result with structured data."""
|
||||
|
||||
entry_id: str
|
||||
title: str
|
||||
authors: list[str]
|
||||
summary: str
|
||||
published: str
|
||||
updated: str | None = None
|
||||
categories: list[str]
|
||||
primary_category: str
|
||||
pdf_url: str | None = None
|
||||
doi: str | None = None
|
||||
journal_ref: str | None = None
|
||||
comment: str | None = None
|
||||
|
||||
|
||||
class ArxivMetadata(BaseModel):
|
||||
"""Metadata for ArXiv operation results."""
|
||||
|
||||
operation: str
|
||||
query: str | None = None
|
||||
max_results: int | None = None
|
||||
sort_by: str | None = None
|
||||
sort_order: str | None = None
|
||||
total_results: int | None = None
|
||||
execution_time: float | None = None
|
||||
error_type: str | None = None
|
||||
paper_id: str | None = None
|
||||
download_path: str | None = None
|
||||
file_size: int | None = None
|
||||
|
||||
|
||||
mcp = FastMCP(
|
||||
"parxiv-server",
|
||||
instructions="""
|
||||
ArXiv MCP Server
|
||||
|
||||
This module provides MCP server functionality for ArXiv academic paper operations.
|
||||
It supports paper search, metadata extraction, and content retrieval with LLM-friendly formatting.
|
||||
|
||||
Key features:
|
||||
- Search ArXiv papers by query, author, category, or ID
|
||||
- Extract paper metadata (title, authors, abstract, etc.)
|
||||
- Download and process paper PDFs
|
||||
- LLM-optimized content formatting
|
||||
- Comprehensive error handling and logging
|
||||
|
||||
Main functions:
|
||||
- mcp_search_papers: Search ArXiv papers with flexible criteria
|
||||
- mcp_get_paper_details: Get detailed information about specific papers
|
||||
- mcp_download_paper: Download paper PDF and extract text content
|
||||
- mcp_get_categories: Get available ArXiv subject categories
|
||||
- mcp_get_arxiv_capabilities: Get service capabilities and configuration
|
||||
""",
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Search ArXiv papers with flexible criteria.
|
||||
|
||||
This tool provides comprehensive ArXiv paper search with:
|
||||
- Keyword, title, and author search capabilities
|
||||
- Category filtering for specific subject areas
|
||||
- Flexible sorting options (relevance, date)
|
||||
- Configurable result limits
|
||||
- LLM-optimized result formatting
|
||||
"""
|
||||
)
|
||||
async def search_papers(
|
||||
query: str = Field(description="Search query (keywords, title, author, etc.)"),
|
||||
sort_by: str = Field(
|
||||
default="relevance",
|
||||
description="Sort by: 'relevance', 'lastUpdatedDate', 'submittedDate'",
|
||||
),
|
||||
sort_order: str = Field(
|
||||
default="descending", description="Sort order: 'ascending' or 'descending'"
|
||||
),
|
||||
category: str | None = Field(
|
||||
default=None, description="Filter by ArXiv category (e.g., 'cs.AI', 'math.CO')"
|
||||
),
|
||||
output_format: str = Field(
|
||||
default="markdown", description="Output format: 'markdown', 'json', or 'text'"
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
if isinstance(query, FieldInfo):
|
||||
query = query.default
|
||||
if isinstance(sort_by, FieldInfo):
|
||||
sort_by = sort_by.default
|
||||
if isinstance(sort_order, FieldInfo):
|
||||
sort_order = sort_order.default
|
||||
if isinstance(category, FieldInfo):
|
||||
category = category.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
try:
|
||||
logging.info(f"🔍 Searching ArXiv for: {query}")
|
||||
|
||||
start_time = datetime.now()
|
||||
|
||||
# Build search query
|
||||
search_query = query
|
||||
if category:
|
||||
search_query = f"cat:{category} AND ({query})"
|
||||
|
||||
# Configure sort criteria
|
||||
sort_criterion = arxiv.SortCriterion.Relevance
|
||||
if sort_by == "lastUpdatedDate":
|
||||
sort_criterion = arxiv.SortCriterion.LastUpdatedDate
|
||||
elif sort_by == "submittedDate":
|
||||
sort_criterion = arxiv.SortCriterion.SubmittedDate
|
||||
|
||||
sort_order_enum = arxiv.SortOrder.Descending
|
||||
if sort_order == "ascending":
|
||||
sort_order_enum = arxiv.SortOrder.Ascending
|
||||
|
||||
# Perform search
|
||||
search = arxiv.Search(
|
||||
query=search_query,
|
||||
max_results=300000,
|
||||
sort_by=sort_criterion,
|
||||
sort_order=sort_order_enum,
|
||||
)
|
||||
|
||||
# Execute search and collect results
|
||||
results = []
|
||||
for paper in client.results(search):
|
||||
results.append(_format_paper_result(paper))
|
||||
|
||||
execution_time = (datetime.now() - start_time).total_seconds()
|
||||
|
||||
# Format output
|
||||
formatted_output = _format_search_results(results, output_format)
|
||||
|
||||
# Create metadata
|
||||
metadata = ArxivMetadata(
|
||||
operation="search_papers",
|
||||
query=query,
|
||||
sort_by=sort_by,
|
||||
sort_order=sort_order,
|
||||
total_results=len(results),
|
||||
execution_time=execution_time,
|
||||
)
|
||||
|
||||
logging.info(f"✅ Found {len(results)} papers in {execution_time:.2f}s")
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=formatted_output, metadata=metadata.model_dump()
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
)
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to search ArXiv papers: {str(e)}"
|
||||
logging.error(f"ArXiv search error: {traceback.format_exc()}")
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=error_msg,
|
||||
metadata=ArxivMetadata(
|
||||
operation="search_papers", query=query, error_type="search_error"
|
||||
).model_dump(),
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Get detailed information about a specific ArXiv paper.
|
||||
"""
|
||||
)
|
||||
async def get_paper_details(
|
||||
paper_id: str = Field(
|
||||
description="ArXiv paper ID (e.g., '2301.07041' or 'arxiv:2301.07041')"
|
||||
),
|
||||
output_format: str = Field(
|
||||
default="markdown", description="Output format: 'markdown', 'json', or 'text'"
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
if isinstance(paper_id, FieldInfo):
|
||||
paper_id = paper_id.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
try:
|
||||
# Clean paper ID
|
||||
clean_id = paper_id.replace("arxiv:", "").strip()
|
||||
|
||||
logging.info(f"📄 Getting details for paper: {clean_id}")
|
||||
|
||||
start_time = datetime.now()
|
||||
|
||||
# Search for the specific paper
|
||||
search = arxiv.Search(id_list=[clean_id])
|
||||
|
||||
paper = next(client.results(search), None)
|
||||
if not paper:
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Paper not found: {clean_id}",
|
||||
metadata=ArxivMetadata(
|
||||
operation="get_paper_details",
|
||||
paper_id=clean_id,
|
||||
error_type="paper_not_found",
|
||||
).model_dump(),
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
execution_time = (datetime.now() - start_time).total_seconds()
|
||||
|
||||
# Format paper details
|
||||
paper_result = _format_paper_result(paper)
|
||||
formatted_output = _format_paper_details(paper_result, output_format)
|
||||
|
||||
# Create metadata
|
||||
metadata = ArxivMetadata(
|
||||
operation="get_paper_details",
|
||||
paper_id=clean_id,
|
||||
execution_time=execution_time,
|
||||
)
|
||||
|
||||
logging.info(f"✅ Retrieved paper details in {execution_time:.2f}s")
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=formatted_output, metadata=metadata.model_dump()
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
)
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to get paper details: {str(e)}"
|
||||
logging.error(f"Paper details error: {traceback.format_exc()}")
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=error_msg,
|
||||
metadata=ArxivMetadata(
|
||||
operation="get_paper_details",
|
||||
paper_id=paper_id,
|
||||
error_type="retrieval_error",
|
||||
).model_dump(),
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Download ArXiv paper PDF and optionally extract text content.
|
||||
"""
|
||||
)
|
||||
async def download_paper(
|
||||
paper_id: str = Field(
|
||||
description="ArXiv paper ID (e.g., '2301.07041' or 'arxiv:2301.07041')"
|
||||
),
|
||||
extract_text: bool = Field(
|
||||
default=True, description="Whether to extract text content from PDF"
|
||||
),
|
||||
output_format: str = Field(
|
||||
default="markdown", description="Output format: 'markdown', 'json', or 'text'"
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
if isinstance(paper_id, FieldInfo):
|
||||
paper_id = paper_id.default
|
||||
if isinstance(extract_text, FieldInfo):
|
||||
extract_text = extract_text.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
try:
|
||||
# Clean paper ID
|
||||
clean_id = paper_id.replace("arxiv:", "").strip()
|
||||
|
||||
logging.info(f"📥 Downloading paper: {clean_id}")
|
||||
|
||||
start_time = datetime.now()
|
||||
|
||||
# Search for the paper
|
||||
search = arxiv.Search(id_list=[clean_id])
|
||||
paper = next(client.results(search), None)
|
||||
|
||||
if not paper:
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Paper not found: {clean_id}",
|
||||
metadata=ArxivMetadata(
|
||||
operation="download_paper",
|
||||
paper_id=clean_id,
|
||||
error_type="paper_not_found",
|
||||
).model_dump(),
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
# Download PDF
|
||||
filename = f"{clean_id.replace('/', '_')}.pdf"
|
||||
download_path = _downloads_dir / filename
|
||||
|
||||
paper.download_pdf(dirpath=str(_downloads_dir), filename=filename)
|
||||
|
||||
execution_time = (datetime.now() - start_time).total_seconds()
|
||||
file_size = download_path.stat().st_size if download_path.exists() else 0
|
||||
|
||||
# Prepare response message
|
||||
if output_format == "json":
|
||||
response_data = {
|
||||
"paper_id": clean_id,
|
||||
"title": paper.title,
|
||||
"download_path": str(download_path),
|
||||
"file_size": file_size,
|
||||
"download_time": execution_time,
|
||||
}
|
||||
|
||||
if extract_text:
|
||||
try:
|
||||
# Basic text extraction (would need additional libraries like PyPDF2 or pdfplumber)
|
||||
response_data["text_extraction"] = (
|
||||
"Text extraction requires additional PDF processing libraries"
|
||||
)
|
||||
except Exception:
|
||||
response_data["text_extraction"] = "Text extraction failed"
|
||||
|
||||
formatted_output = json.dumps(response_data, indent=2)
|
||||
|
||||
elif output_format == "text":
|
||||
output_parts = [
|
||||
"Paper Downloaded Successfully",
|
||||
f"Paper ID: {clean_id}",
|
||||
f"Title: {paper.title}",
|
||||
f"Download Path: {download_path}",
|
||||
f"File Size: {file_size:,} bytes",
|
||||
f"Download Time: {execution_time:.2f} seconds",
|
||||
]
|
||||
|
||||
if extract_text:
|
||||
output_parts.append(
|
||||
"\nNote: Text extraction requires additional PDF processing libraries"
|
||||
)
|
||||
|
||||
formatted_output = "\n".join(output_parts)
|
||||
|
||||
else: # markdown (default)
|
||||
output_parts = [
|
||||
"# 📥 Paper Download Complete",
|
||||
"",
|
||||
f"**Paper ID:** `{clean_id}`",
|
||||
f"**Title:** {paper.title}",
|
||||
f"**Download Path:** `{download_path}`",
|
||||
f"**File Size:** {file_size:,} bytes",
|
||||
f"**Download Time:** {execution_time:.2f} seconds",
|
||||
]
|
||||
|
||||
if extract_text:
|
||||
output_parts.extend(
|
||||
[
|
||||
"",
|
||||
"## 📄 Text Extraction",
|
||||
(
|
||||
"*Note: Text extraction requires additional "
|
||||
"PDF processing libraries like PyPDF2 or pdfplumber*"
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
formatted_output = "\n".join(output_parts)
|
||||
|
||||
# Create metadata
|
||||
metadata = ArxivMetadata(
|
||||
operation="download_paper",
|
||||
paper_id=clean_id,
|
||||
download_path=str(download_path),
|
||||
file_size=file_size,
|
||||
execution_time=execution_time,
|
||||
)
|
||||
|
||||
logging.info(
|
||||
f"✅ Downloaded paper in {execution_time:.2f}s ({file_size:,} bytes)"
|
||||
)
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=formatted_output, metadata=metadata.model_dump()
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
)
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to download paper: {str(e)}"
|
||||
logging.error(f"Paper download error: {traceback.format_exc()}")
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=error_msg,
|
||||
metadata=ArxivMetadata(
|
||||
operation="download_paper",
|
||||
paper_id=paper_id,
|
||||
error_type="download_error",
|
||||
).model_dump(),
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Get information about ArXiv service capabilities and configuration.
|
||||
"""
|
||||
)
|
||||
async def get_arxiv_capabilities() -> Union[str, TextContent]:
|
||||
capabilities = {
|
||||
"supported_operations": [
|
||||
"Paper search with flexible criteria",
|
||||
"Detailed paper metadata retrieval",
|
||||
"PDF download and storage",
|
||||
"Subject category filtering",
|
||||
"Multiple output formats (markdown, json, text)",
|
||||
"LLM-optimized result formatting",
|
||||
],
|
||||
"search_capabilities": [
|
||||
"Keyword and phrase search",
|
||||
"Author name search",
|
||||
"Title search",
|
||||
"Category filtering",
|
||||
"Date-based sorting",
|
||||
"Relevance-based sorting",
|
||||
],
|
||||
"supported_formats": ["markdown", "json", "text"],
|
||||
"configuration": {
|
||||
"downloads_directory": str(_downloads_dir),
|
||||
"client_page_size": 100,
|
||||
"client_delay_seconds": 3.0,
|
||||
"client_num_retries": 3,
|
||||
"supported_categories_count": len(subject_categories),
|
||||
},
|
||||
"rate_limiting": {
|
||||
"delay_between_requests": "3.0 seconds",
|
||||
"retry_attempts": 3,
|
||||
"respectful_usage": "Configured for ArXiv server guidelines",
|
||||
},
|
||||
}
|
||||
|
||||
formatted_info = f"""# ArXiv Service Capabilities
|
||||
|
||||
## Supported Operations
|
||||
{chr(10).join(f"- {op}" for op in capabilities["supported_operations"])}
|
||||
|
||||
## Search Capabilities
|
||||
{chr(10).join(f"- {cap}" for cap in capabilities["search_capabilities"])}
|
||||
|
||||
## Supported Output Formats
|
||||
{chr(10).join(f"- {fmt}" for fmt in capabilities["supported_formats"])}
|
||||
|
||||
## Current Configuration
|
||||
- **Downloads Directory:** {capabilities["configuration"]["downloads_directory"]}
|
||||
- **Client Page Size:** {capabilities["configuration"]["client_page_size"]}
|
||||
- **Request Delay:** {capabilities["configuration"]["client_delay_seconds"]} seconds
|
||||
- **Retry Attempts:** {capabilities["configuration"]["client_num_retries"]}
|
||||
- **Available Categories:** {capabilities["configuration"]["supported_categories_count"]}
|
||||
|
||||
## Rate Limiting & Ethics
|
||||
- **Delay Between Requests:** {capabilities["rate_limiting"]["delay_between_requests"]}
|
||||
- **Retry Policy:** {capabilities["rate_limiting"]["retry_attempts"]} attempts
|
||||
- **Server Respect:** {capabilities["rate_limiting"]["respectful_usage"]}
|
||||
"""
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=formatted_info, metadata=capabilities
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
)
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Get available ArXiv subject categories.
|
||||
"""
|
||||
)
|
||||
async def get_categories(
|
||||
output_format: str = Field(
|
||||
default="markdown", description="Output format: 'markdown', 'json', or 'text'"
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
try:
|
||||
if output_format == "json":
|
||||
formatted_output = json.dumps(subject_categories, indent=2)
|
||||
|
||||
elif output_format == "text":
|
||||
output_parts = ["ArXiv Subject Categories:\n"]
|
||||
for code, name in subject_categories.items():
|
||||
output_parts.append(f"{code}: {name}")
|
||||
formatted_output = "\n".join(output_parts)
|
||||
|
||||
else: # markdown (default)
|
||||
output_parts = [
|
||||
"# ArXiv Subject Categories",
|
||||
"",
|
||||
"Available categories for filtering search results:",
|
||||
"",
|
||||
]
|
||||
|
||||
for code, name in subject_categories.items():
|
||||
output_parts.append(f"- **`{code}`**: {name}")
|
||||
|
||||
output_parts.extend(
|
||||
[
|
||||
"",
|
||||
"## Usage Examples",
|
||||
"- `cs.AI` - Artificial Intelligence",
|
||||
"- `cs.LG` - Machine Learning",
|
||||
"- `math.CO` - Combinatorics",
|
||||
"- `physics.gen-ph` - General Physics",
|
||||
]
|
||||
)
|
||||
|
||||
formatted_output = "\n".join(output_parts)
|
||||
|
||||
metadata = ArxivMetadata(
|
||||
operation="get_categories", total_results=len(subject_categories)
|
||||
)
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=formatted_output, metadata=metadata.model_dump()
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to get categories: {str(e)}"
|
||||
logging.error(f"Categories error: {traceback.format_exc()}")
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=error_msg,
|
||||
metadata=ArxivMetadata(
|
||||
operation="get_categories", error_type="internal_error"
|
||||
).model_dump(),
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
)
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
def _format_paper_result(paper: arxiv.Result) -> PaperResult:
|
||||
"""Convert arxiv.Result to structured PaperResult.
|
||||
|
||||
Args:
|
||||
paper: ArXiv paper result object
|
||||
|
||||
Returns:
|
||||
Structured PaperResult object
|
||||
"""
|
||||
return PaperResult(
|
||||
entry_id=paper.entry_id,
|
||||
title=paper.title.strip(),
|
||||
authors=[author.name for author in paper.authors],
|
||||
summary=paper.summary.strip(),
|
||||
published=paper.published.isoformat(),
|
||||
updated=paper.updated.isoformat() if paper.updated else None,
|
||||
categories=paper.categories,
|
||||
primary_category=paper.primary_category,
|
||||
pdf_url=paper.pdf_url,
|
||||
doi=paper.doi,
|
||||
journal_ref=paper.journal_ref,
|
||||
comment=paper.comment,
|
||||
)
|
||||
|
||||
|
||||
def _format_search_results(
|
||||
results: list[PaperResult], output_format: str = "markdown"
|
||||
) -> str:
|
||||
"""Format paper search results for LLM consumption.
|
||||
|
||||
Args:
|
||||
results: List of paper results
|
||||
output_format: Format type ('markdown', 'json', 'text')
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if not results:
|
||||
return "No papers found matching the search criteria."
|
||||
|
||||
if output_format == "json":
|
||||
return json.dumps([result.model_dump() for result in results], indent=2)
|
||||
|
||||
elif output_format == "text":
|
||||
output_parts = [f"Found {len(results)} papers:\n"]
|
||||
|
||||
for i, paper in enumerate(results, 1):
|
||||
authors_str = ", ".join(paper.authors[:3])
|
||||
if len(paper.authors) > 3:
|
||||
authors_str += f" et al. ({len(paper.authors)} total)"
|
||||
|
||||
output_parts.extend(
|
||||
[
|
||||
f"{i}. {paper.title}",
|
||||
f" Authors: {authors_str}",
|
||||
f" Published: {paper.published[:10]}",
|
||||
f" Categories: {', '.join(paper.categories)}",
|
||||
f" ArXiv ID: {paper.entry_id.split('/')[-1]}",
|
||||
f" Abstract: {paper.summary[:200]}...",
|
||||
"",
|
||||
]
|
||||
)
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
else: # markdown (default)
|
||||
output_parts = [f"# ArXiv Search Results\n\nFound **{len(results)}** papers:\n"]
|
||||
|
||||
for i, paper in enumerate(results, 1):
|
||||
authors_str = ", ".join(paper.authors[:3])
|
||||
if len(paper.authors) > 3:
|
||||
authors_str += f" *et al.* ({len(paper.authors)} total)"
|
||||
|
||||
arxiv_id = paper.entry_id.split("/")[-1]
|
||||
|
||||
output_parts.extend(
|
||||
[
|
||||
f"## {i}. {paper.title}",
|
||||
f"**Authors:** {authors_str}",
|
||||
f"**Published:** {paper.published[:10]}",
|
||||
f"**Categories:** {', '.join(paper.categories)}",
|
||||
f"**ArXiv ID:** `{arxiv_id}`",
|
||||
f"**PDF:** [Download]({paper.pdf_url})" if paper.pdf_url else "",
|
||||
"",
|
||||
f"**Abstract:** {paper.summary}",
|
||||
"",
|
||||
"---",
|
||||
"",
|
||||
]
|
||||
)
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
|
||||
def _format_paper_details(paper: PaperResult, output_format: str = "markdown") -> str:
|
||||
"""Format detailed paper information for LLM consumption.
|
||||
|
||||
Args:
|
||||
paper: Paper result object
|
||||
output_format: Format type ('markdown', 'json', 'text')
|
||||
|
||||
Returns:
|
||||
Formatted string with detailed paper information
|
||||
"""
|
||||
if output_format == "json":
|
||||
return json.dumps(paper.model_dump(), indent=2)
|
||||
|
||||
elif output_format == "text":
|
||||
output_parts = [
|
||||
f"Title: {paper.title}",
|
||||
f"Authors: {', '.join(paper.authors)}",
|
||||
f"Published: {paper.published}",
|
||||
f"Updated: {paper.updated or 'N/A'}",
|
||||
f"Primary Category: {paper.primary_category}",
|
||||
f"All Categories: {', '.join(paper.categories)}",
|
||||
f"ArXiv ID: {paper.entry_id.split('/')[-1]}",
|
||||
f"PDF URL: {paper.pdf_url or 'N/A'}",
|
||||
f"DOI: {paper.doi or 'N/A'}",
|
||||
f"Journal Reference: {paper.journal_ref or 'N/A'}",
|
||||
f"Comment: {paper.comment or 'N/A'}",
|
||||
"",
|
||||
"Abstract:",
|
||||
paper.summary,
|
||||
]
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
else: # markdown (default)
|
||||
arxiv_id = paper.entry_id.split("/")[-1]
|
||||
|
||||
output_parts = [
|
||||
f"# {paper.title}",
|
||||
"",
|
||||
f"**Authors:** {', '.join(paper.authors)}",
|
||||
f"**Published:** {paper.published[:10]}",
|
||||
f"**Updated:** {paper.updated[:10] if paper.updated else 'N/A'}",
|
||||
f"**Primary Category:** {paper.primary_category}",
|
||||
f"**All Categories:** {', '.join(paper.categories)}",
|
||||
f"**ArXiv ID:** `{arxiv_id}`",
|
||||
(
|
||||
f"**PDF:** [Download]({paper.pdf_url})"
|
||||
if paper.pdf_url
|
||||
else "**PDF:** N/A"
|
||||
),
|
||||
f"**DOI:** {paper.doi}" if paper.doi else "**DOI:** N/A",
|
||||
(
|
||||
f"**Journal Reference:** {paper.journal_ref}"
|
||||
if paper.journal_ref
|
||||
else "**Journal Reference:** N/A"
|
||||
),
|
||||
f"**Comment:** {paper.comment}" if paper.comment else "**Comment:** N/A",
|
||||
"",
|
||||
"## Abstract",
|
||||
"",
|
||||
paper.summary,
|
||||
]
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
load_dotenv(override=True)
|
||||
logging.info("Starting parxiv-server MCP server!")
|
||||
mcp.run(transport="stdio")
|
||||
Vendored
+15
@@ -0,0 +1,15 @@
|
||||
[project]
|
||||
name = "aworldsearch-server"
|
||||
version = "0.1.0"
|
||||
description = "Add your description here"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
aiohttp= ">=3.12.15"
|
||||
dependencies = [
|
||||
"mcp",
|
||||
"aiohttp~=3.12.15",
|
||||
"requests~=2.32.4",
|
||||
"fastmcp~=2.11.3",
|
||||
"tavily-python~=0.7.10"
|
||||
|
||||
]
|
||||
+217
@@ -0,0 +1,217 @@
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import traceback
|
||||
from typing import List, Dict, Any, Optional, Union
|
||||
|
||||
import aiohttp
|
||||
import requests
|
||||
from mcp.server import FastMCP
|
||||
from mcp.types import TextContent
|
||||
from pydantic import Field
|
||||
from tavily import TavilyClient
|
||||
from dotenv import load_dotenv
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
mcp = FastMCP("readweb-server")
|
||||
|
||||
|
||||
def filter_valid_images(result: Optional[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""Filter valid document results, returns empty list if input is None"""
|
||||
if result is None:
|
||||
return []
|
||||
|
||||
try:
|
||||
valid_docs = []
|
||||
|
||||
# Check success field
|
||||
if not result.get("success"):
|
||||
return valid_docs
|
||||
|
||||
# Check searchDocs field
|
||||
search_docs = result.get("searchImages", [])
|
||||
if not search_docs:
|
||||
return valid_docs
|
||||
|
||||
# Extract required fields
|
||||
required_fields = ["title", "picUrl"]
|
||||
|
||||
for doc in search_docs:
|
||||
# Check if all required fields exist and are not empty
|
||||
is_valid = True
|
||||
for field in required_fields:
|
||||
if field not in doc or not doc[field]:
|
||||
is_valid = False
|
||||
break
|
||||
|
||||
if is_valid:
|
||||
# Keep only required fields
|
||||
filtered_doc = {field: doc[field] for field in required_fields}
|
||||
valid_docs.append({
|
||||
"type": "IMAGE",
|
||||
"title": filtered_doc.get("title", ""),
|
||||
"url": filtered_doc.get("picUrl", "")
|
||||
})
|
||||
|
||||
return valid_docs
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
def filter_valid_docs(result: Optional[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""Filter valid document results, returns empty list if input is None"""
|
||||
if result is None:
|
||||
return []
|
||||
|
||||
try:
|
||||
valid_docs = []
|
||||
|
||||
# Check success field
|
||||
if not result.get("success"):
|
||||
return valid_docs
|
||||
|
||||
# Check searchDocs field
|
||||
search_docs = result.get("searchDocs", [])
|
||||
if not search_docs:
|
||||
return valid_docs
|
||||
|
||||
# Extract required fields
|
||||
required_fields = ["title", "docAbstract", "url", "doc"]
|
||||
|
||||
for doc in search_docs:
|
||||
# Check if all required fields exist and are not empty
|
||||
is_valid = True
|
||||
for field in required_fields:
|
||||
if field not in doc or not doc[field]:
|
||||
is_valid = False
|
||||
break
|
||||
|
||||
if is_valid:
|
||||
# Keep only required fields
|
||||
filtered_doc = {field: doc[field] for field in required_fields}
|
||||
valid_docs.append(filtered_doc)
|
||||
|
||||
return valid_docs
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
@mcp.tool(
|
||||
description="提取并将网页内容转换为清晰、可读的markdown格式。非常适合阅读文章、文档、博客文章或任何网页内容。当您需要分析网站的文本内容、绕过付费墙或获取结构化数据时，请使用此工具。")
|
||||
async def read_url(
|
||||
url: str = Field(
|
||||
description="一个强大的网页内容提取工具，可以从指定URL检索和处理原始内容，非常适合数据收集、内容分析和研究任务。"
|
||||
),
|
||||
include_images: bool = Field(
|
||||
False,
|
||||
description="在响应中包含从URL提取的图片列表"
|
||||
)
|
||||
) -> Union[str, TextContent]:
|
||||
try:
|
||||
if not url:
|
||||
return TextContent(
|
||||
type="text",
|
||||
text="", # Empty string instead of None
|
||||
**{"metadata": {}} # Pass as additional fields
|
||||
)
|
||||
urls = [url]
|
||||
TAVILY_API_KEY = os.getenv('TAVILY_API_KEY')
|
||||
if not TAVILY_API_KEY:
|
||||
return TextContent(
|
||||
type="text",
|
||||
text="", # Empty string instead of None
|
||||
**{"metadata": {}} # Pass as additional fields
|
||||
)
|
||||
tavily_client = TavilyClient(api_key=TAVILY_API_KEY)
|
||||
response = tavily_client.extract(urls=urls, include_images=include_images,format="text")
|
||||
text_extracted_content = ""
|
||||
|
||||
if response and isinstance(response, dict):
|
||||
results = response.get("results", [])
|
||||
if results and isinstance(results, list) and len(results) > 0:
|
||||
first_result = results[0]
|
||||
if isinstance(first_result, dict):
|
||||
raw_content = first_result.get("raw_content")
|
||||
if raw_content and isinstance(raw_content, str):
|
||||
text_extracted_content = raw_content
|
||||
markdown_response = tavily_client.extract(urls=urls, include_images=include_images, include_favicon=False,
|
||||
format="markdown")
|
||||
markdown_extracted_content = ""
|
||||
if markdown_response and isinstance(markdown_response, dict):
|
||||
results = markdown_response.get("results", [])
|
||||
if results and isinstance(results, list) and len(results) > 0:
|
||||
first_result = results[0]
|
||||
if isinstance(first_result, dict):
|
||||
raw_content = first_result.get("raw_content")
|
||||
if raw_content and isinstance(raw_content, str):
|
||||
markdown_extracted_content = raw_content
|
||||
|
||||
search_output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": markdown_extracted_content
|
||||
}
|
||||
# Initialize TextContent with additional parameters
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=text_extracted_content,
|
||||
**{"metadata": search_output_dict} # Pass processed data as metadata
|
||||
)
|
||||
except Exception as e:
|
||||
# Handle errors
|
||||
logging.error(f"Search error: {e}")
|
||||
# Initialize TextContent with additional parameters
|
||||
return TextContent(
|
||||
type="text",
|
||||
text="", # Empty string instead of None
|
||||
**{"metadata": {}} # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
async def search_image_single(query: str, num: int = 5) -> Optional[Dict[str, Any]]:
|
||||
"""Execute a single search query, returns None on error"""
|
||||
try:
|
||||
url = os.getenv('PIC_SEARCH_URL')
|
||||
searchMode = os.getenv('PIC_SEARCH_SEARCHMODE')
|
||||
source = os.getenv('PIC_SEARCH_SOURCE')
|
||||
domain = os.getenv('PIC_SEARCH_DOMAIN')
|
||||
uid = os.getenv('PIC_SEARCH_UID')
|
||||
if not url or not searchMode or not source or not domain:
|
||||
return None
|
||||
|
||||
headers = {
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
data = {
|
||||
"domain": domain,
|
||||
"extParams": {
|
||||
"contentType": "llmWholeImage"
|
||||
},
|
||||
"page": 0,
|
||||
"pageSize": num,
|
||||
"query": query,
|
||||
"searchMode": searchMode,
|
||||
"source": source,
|
||||
"userId": uid
|
||||
}
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
try:
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
if response.status != 200:
|
||||
return None
|
||||
|
||||
result = await response.json()
|
||||
return result
|
||||
except aiohttp.ClientError:
|
||||
return None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
load_dotenv(override=True)
|
||||
logger.info("Starting readweb-server MCP server!")
|
||||
mcp.run(transport="stdio")
|
||||
Vendored
+17
@@ -0,0 +1,17 @@
|
||||
[project]
|
||||
name = "terminal-server"
|
||||
version = "0.1.0"
|
||||
description = "Add your description here"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
aiohttp= ">=3.12.15"
|
||||
dependencies = [
|
||||
"mcp",
|
||||
"aiohttp~=3.12.15",
|
||||
"requests~=2.32.4",
|
||||
"fastmcp~=2.11.3",
|
||||
"python-magic~=0.4.27",
|
||||
"chardet~=3.0.4",
|
||||
"pandas~=2.3.0",
|
||||
|
||||
]
|
||||
Vendored
+226
@@ -0,0 +1,226 @@
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import magic
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import requests
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import Any, Literal
|
||||
|
||||
|
||||
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)"
|
||||
)
|
||||
|
||||
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")
|
||||
|
||||
|
||||
def _validate_file_path(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)
|
||||
if not path.is_absolute():
|
||||
path = path.expanduser().resolve()
|
||||
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"File not found: {path}")
|
||||
|
||||
return path
|
||||
|
||||
def is_url(path_or_url: str) -> bool:
|
||||
"""
|
||||
Check if the given string is a URL.
|
||||
|
||||
Args:
|
||||
path_or_url: String to check
|
||||
|
||||
Returns:
|
||||
bool: True if the string is a URL, False otherwise
|
||||
"""
|
||||
parsed = urlparse(path_or_url)
|
||||
return bool(parsed.scheme and parsed.netloc)
|
||||
|
||||
|
||||
def get_mime_type(file_path: str, default_mime: str | None = None) -> str:
|
||||
"""
|
||||
Detect MIME type of a file using python-magic if available,
|
||||
otherwise fallback to extension-based detection.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
default_mime: Default MIME type to return if detection fails
|
||||
|
||||
Returns:
|
||||
str: Detected MIME type
|
||||
"""
|
||||
# Try using python-magic for accurate MIME type detection
|
||||
try:
|
||||
mime = magic.Magic(mime=True)
|
||||
return mime.from_file(file_path)
|
||||
except (AttributeError, IOError):
|
||||
# Fallback to extension-based detection
|
||||
extension_mime_map = {
|
||||
# Audio formats
|
||||
".mp3": "audio/mpeg",
|
||||
".wav": "audio/wav",
|
||||
".ogg": "audio/ogg",
|
||||
".m4a": "audio/mp4",
|
||||
".flac": "audio/flac",
|
||||
# Image formats
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".png": "image/png",
|
||||
".gif": "image/gif",
|
||||
".webp": "image/webp",
|
||||
".bmp": "image/bmp",
|
||||
".tiff": "image/tiff",
|
||||
# Video formats
|
||||
".mp4": "video/mp4",
|
||||
".avi": "video/x-msvideo",
|
||||
".mov": "video/quicktime",
|
||||
".mkv": "video/x-matroska",
|
||||
".webm": "video/webm",
|
||||
}
|
||||
|
||||
ext = Path(file_path).suffix.lower()
|
||||
return extension_mime_map.get(ext, default_mime or "application/octet-stream")
|
||||
|
||||
|
||||
def get_file_from_source(
|
||||
source: str,
|
||||
max_size_mb: float = 100.0,
|
||||
timeout: int = 60,
|
||||
) -> tuple[str, str, bytes]:
|
||||
"""
|
||||
Unified function to get file content from a URL or local path with validation.
|
||||
|
||||
Args:
|
||||
source: URL or local file path
|
||||
max_size_mb: Maximum allowed file size in MB
|
||||
timeout: Timeout for URL requests in seconds
|
||||
|
||||
Returns:
|
||||
Tuple[str, str, bytes]: (file_path, mime_type, file_content)
|
||||
- For URLs, file_path will be a temporary file path
|
||||
- For local files, file_path will be the original path
|
||||
|
||||
Raises:
|
||||
ValueError: When file doesn't exist, exceeds size limit, or has invalid MIME type
|
||||
IOError: When file cannot be read
|
||||
requests.RequestException: When URL request fails
|
||||
"""
|
||||
max_size_bytes = max_size_mb * 1024 * 1024
|
||||
|
||||
if is_url(source):
|
||||
# Handle URL source
|
||||
try:
|
||||
# Make a HEAD request first to check content length
|
||||
head_response = requests.head(source, timeout=timeout, allow_redirects=True)
|
||||
head_response.raise_for_status()
|
||||
|
||||
# Check content length if available
|
||||
content_length = head_response.headers.get("content-length")
|
||||
if content_length and int(content_length) > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({int(content_length) / (1024 * 1024):.2f} MB) "
|
||||
f"exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Download the file
|
||||
response = requests.get(source, timeout=timeout, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
# Read content with size checking
|
||||
content = b""
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
if len(content) + len(chunk) > max_size_bytes:
|
||||
raise ValueError(f"File size exceeds maximum allowed size ({max_size_mb} MB)")
|
||||
content += chunk
|
||||
|
||||
# Create temporary file
|
||||
parsed_url = urlparse(source)
|
||||
filename = os.path.basename(parsed_url.path) or "downloaded_file"
|
||||
|
||||
# Create temporary file with proper extension
|
||||
suffix = Path(filename).suffix or ".tmp"
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
|
||||
temp_file.write(content)
|
||||
temp_path = temp_file.name
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(temp_path)
|
||||
|
||||
return temp_path, mime_type, content
|
||||
|
||||
except requests.RequestException as e:
|
||||
raise requests.RequestException(f"Failed to download file from URL: {e}: {traceback.format_exc()}")
|
||||
except Exception as e:
|
||||
raise IOError(f"Error processing URL: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
else:
|
||||
# Handle local file path
|
||||
file_path = Path(source)
|
||||
|
||||
# Check if file exists
|
||||
if not file_path.exists():
|
||||
raise ValueError(f"File does not exist: {source}")
|
||||
|
||||
if not file_path.is_file():
|
||||
raise ValueError(f"Path is not a file: {source}")
|
||||
|
||||
# Check file size
|
||||
file_size = file_path.stat().st_size
|
||||
if file_size > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({file_size / (1024 * 1024):.2f} MB) exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Read file content
|
||||
try:
|
||||
with open(file_path, "rb") as f:
|
||||
content = f.read()
|
||||
except Exception as e:
|
||||
raise IOError(f"Cannot read file {source}: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(str(file_path))
|
||||
|
||||
return str(file_path), mime_type, content
|
||||
Vendored
+572
@@ -0,0 +1,572 @@
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import platform
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Union
|
||||
|
||||
import chardet
|
||||
import pandas as pd
|
||||
from dotenv import load_dotenv
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.types import TextContent
|
||||
from pydantic import Field, BaseModel
|
||||
|
||||
from base import ActionResponse
|
||||
|
||||
load_dotenv()
|
||||
workspace = Path.home()
|
||||
|
||||
command_history: list[dict] = []
|
||||
max_history_size = 50
|
||||
|
||||
# Define dangerous commands for safety
|
||||
dangerous_commands = [
|
||||
"rm -rf /",
|
||||
"mkfs",
|
||||
"dd if=",
|
||||
":(){ :|:& };:", # Unix
|
||||
"del /f /s /q",
|
||||
"format",
|
||||
"diskpart", # Windows
|
||||
"sudo rm",
|
||||
"sudo dd",
|
||||
"sudo mkfs", # Sudo variants
|
||||
]
|
||||
|
||||
# Get current platform info
|
||||
platform_info = {
|
||||
"system": platform.system(),
|
||||
"platform": platform.platform(),
|
||||
"architecture": platform.architecture()[0],
|
||||
}
|
||||
|
||||
|
||||
class CommandResult(BaseModel):
|
||||
"""Individual command execution result with structured data."""
|
||||
|
||||
command: str
|
||||
success: bool
|
||||
stdout: str
|
||||
stderr: str
|
||||
return_code: int
|
||||
duration: str
|
||||
timestamp: str
|
||||
|
||||
|
||||
class TerminalMetadata(BaseModel):
|
||||
"""Metadata for terminal operation results."""
|
||||
|
||||
command: str
|
||||
platform: str
|
||||
working_directory: str
|
||||
timeout_seconds: int
|
||||
execution_time: float | None = None
|
||||
return_code: int | None = None
|
||||
safety_check_passed: bool = True
|
||||
error_type: str | None = None
|
||||
history_count: int | None = None
|
||||
|
||||
|
||||
mcp = FastMCP(
|
||||
"terminal-server",
|
||||
instructions="""
|
||||
Terminal MCP Server
|
||||
|
||||
This module provides MCP server functionality for executing terminal commands safely.
|
||||
It supports command execution with timeout controls and returns LLM-friendly formatted results.
|
||||
|
||||
Key features:
|
||||
- Execute terminal commands with configurable timeouts
|
||||
- Cross-platform command execution support
|
||||
- Command history tracking and retrieval
|
||||
- Safety checks for dangerous commands
|
||||
- LLM-optimized output formatting
|
||||
|
||||
Main functions:
|
||||
- mcp_execute_command: Execute terminal commands with safety checks
|
||||
- mcp_get_command_history: Retrieve recent command execution history
|
||||
- mcp_get_terminal_capabilities: Get terminal service capabilities
|
||||
""",
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Execute a terminal command with safety checks and timeout controls.
|
||||
|
||||
This tool provides secure command execution with:
|
||||
- Cross-platform compatibility (Windows, macOS, Linux)
|
||||
- Configurable timeout controls
|
||||
- Safety checks for dangerous commands
|
||||
- LLM-optimized result formatting
|
||||
- Command history tracking
|
||||
|
||||
Specialized Feature:
|
||||
- Execute Python code and output the result to stdout
|
||||
- Example (Directly execute simple Python code): `python -c "nums = [1, 2, 3, 4]\nsum_of_nums = sum(nums)\nprint(f'{sum_of_nums=}')"`
|
||||
- Example (Execute code from a file): `python my_script.py`
|
||||
"""
|
||||
)
|
||||
async def execute_command(
|
||||
command: str = Field(description="Terminal command to execute"),
|
||||
timeout: int = Field(
|
||||
default=30, description="Command timeout in seconds (default: 30)"
|
||||
),
|
||||
output_format: str = Field(
|
||||
default="markdown", description="Output format: 'markdown', 'json', or 'text'"
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
if isinstance(command, FieldInfo):
|
||||
command = command.default
|
||||
if isinstance(timeout, FieldInfo):
|
||||
timeout = timeout.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
try:
|
||||
# Safety check
|
||||
is_safe, safety_reason = _check_command_safety(command)
|
||||
if not is_safe:
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"Command rejected for security reasons: {safety_reason}",
|
||||
metadata=TerminalMetadata(
|
||||
command=command,
|
||||
platform=platform_info["system"],
|
||||
working_directory=str(workspace),
|
||||
timeout_seconds=timeout,
|
||||
safety_check_passed=False,
|
||||
error_type="security_violation",
|
||||
).model_dump(),
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
logging.info(f"🔧 Executing command: {command}")
|
||||
|
||||
# Execute command
|
||||
start_time = time.time()
|
||||
result = await _execute_command_async(command, timeout)
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
# Format output
|
||||
formatted_output = _format_command_output(result, output_format)
|
||||
|
||||
# Create metadata
|
||||
metadata = TerminalMetadata(
|
||||
command=command,
|
||||
platform=platform_info["system"],
|
||||
working_directory=str(workspace),
|
||||
timeout_seconds=timeout,
|
||||
execution_time=execution_time,
|
||||
return_code=result.return_code,
|
||||
safety_check_passed=True,
|
||||
)
|
||||
|
||||
if result.success:
|
||||
logging.info(
|
||||
"✅ Command completed successfully",
|
||||
)
|
||||
else:
|
||||
logging.info(f"❌ Command failed with return code {result.return_code}")
|
||||
metadata.error_type = "execution_failure"
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=result.success,
|
||||
message=formatted_output,
|
||||
metadata=metadata.model_dump(),
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
)
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to execute command: {str(e)}"
|
||||
logging.error(f"Command execution error: {traceback.format_exc()}")
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=error_msg,
|
||||
metadata=TerminalMetadata(
|
||||
command=command,
|
||||
platform=platform_info["system"],
|
||||
working_directory=str(workspace),
|
||||
timeout_seconds=timeout,
|
||||
safety_check_passed=True,
|
||||
error_type="internal_error",
|
||||
).model_dump(),
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Retrieve recent command execution history.
|
||||
"""
|
||||
)
|
||||
async def get_command_history(
|
||||
count: int = Field(
|
||||
default=10, description="Number of recent commands to return (default: 10)"
|
||||
),
|
||||
output_format: str = Field(
|
||||
default="markdown", description="Output format: 'markdown', 'json', or 'text'"
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
if isinstance(count, FieldInfo):
|
||||
count = count.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
try:
|
||||
# Get recent history
|
||||
recent_history = command_history[-count:] if command_history else []
|
||||
|
||||
if not recent_history:
|
||||
message = "No command history available."
|
||||
else:
|
||||
if output_format == "json":
|
||||
message = json.dumps(recent_history, indent=2)
|
||||
elif output_format == "text":
|
||||
history_lines = []
|
||||
for i, entry in enumerate(recent_history, 1):
|
||||
status = "SUCCESS" if entry["success"] else "FAILED"
|
||||
history_lines.append(
|
||||
f"{i}. [{entry['timestamp']}] {entry['command']} - {status}"
|
||||
f" ({entry.get('duration', 'N/A')})"
|
||||
)
|
||||
message = "\n".join(history_lines)
|
||||
else: # markdown
|
||||
history_lines = [
|
||||
"# Command History",
|
||||
f"Showing {len(recent_history)} recent commands:\n",
|
||||
]
|
||||
|
||||
for i, entry in enumerate(recent_history, 1):
|
||||
status_emoji = "✅" if entry["success"] else "❌"
|
||||
history_lines.extend(
|
||||
[
|
||||
f"## {i}. {status_emoji} `{entry['command']}`",
|
||||
f"- **Timestamp:** {entry['timestamp']}",
|
||||
f"- **Duration:** {entry.get('duration', 'N/A')}",
|
||||
"",
|
||||
]
|
||||
)
|
||||
|
||||
message = "\n".join(history_lines)
|
||||
|
||||
metadata = TerminalMetadata(
|
||||
command="get_command_history",
|
||||
platform=platform_info["system"],
|
||||
working_directory=str(workspace),
|
||||
timeout_seconds=0,
|
||||
history_count=len(recent_history),
|
||||
)
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=message, metadata=metadata.model_dump()
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
)
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to retrieve command history: {str(e)}"
|
||||
logging.error(f"History retrieval error: {traceback.format_exc()}")
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=error_msg,
|
||||
metadata=TerminalMetadata(
|
||||
command="get_command_history",
|
||||
platform=platform_info["system"],
|
||||
working_directory=str(workspace),
|
||||
timeout_seconds=0,
|
||||
error_type="internal_error",
|
||||
).model_dump(),
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Get information about terminal service capabilities and configuration.
|
||||
"""
|
||||
)
|
||||
async def get_terminal_capabilities() -> Union[str, TextContent]:
|
||||
capabilities = {
|
||||
"platform_info": platform_info,
|
||||
"supported_features": [
|
||||
"Cross-platform command execution",
|
||||
"Configurable timeout controls",
|
||||
"Command history tracking",
|
||||
"Safety checks for dangerous commands",
|
||||
"Multiple output formats (markdown, json, text)",
|
||||
"LLM-optimized result formatting",
|
||||
"Async command execution",
|
||||
],
|
||||
"supported_formats": ["markdown", "json", "text"],
|
||||
"configuration": {
|
||||
"max_history_size": max_history_size,
|
||||
"current_history_count": len(command_history),
|
||||
"working_directory": str(workspace),
|
||||
"dangerous_commands_count": len(dangerous_commands),
|
||||
},
|
||||
"safety_features": [
|
||||
"Dangerous command detection",
|
||||
"Timeout controls",
|
||||
"Error handling and logging",
|
||||
"Command validation",
|
||||
],
|
||||
}
|
||||
|
||||
formatted_info = f"""# Terminal Service Capabilities
|
||||
|
||||
## Platform Information
|
||||
- **System:** {platform_info["system"]}
|
||||
- **Platform:** {platform_info["platform"]}
|
||||
- **Architecture:** {platform_info["architecture"]}
|
||||
|
||||
## Features
|
||||
{chr(10).join(f"- {feature}" for feature in capabilities["supported_features"])}
|
||||
|
||||
## Supported Output Formats
|
||||
{chr(10).join(f"- {fmt}" for fmt in capabilities["supported_formats"])}
|
||||
|
||||
## Current Configuration
|
||||
- **Max History Size:** {capabilities["configuration"]["max_history_size"]}
|
||||
- **Current History Count:** {capabilities["configuration"]["current_history_count"]}
|
||||
- **Working Directory:** {capabilities["configuration"]["working_directory"]}
|
||||
- **Dangerous Commands Monitored:** {capabilities["configuration"]["dangerous_commands_count"]}
|
||||
|
||||
## Safety Features
|
||||
{chr(10).join(f"- {feature}" for feature in capabilities["safety_features"])}
|
||||
"""
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=formatted_info, metadata=capabilities
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
)
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
def _check_command_safety(command: str) -> tuple[bool, str | None]:
|
||||
"""Check if command is safe to execute.
|
||||
|
||||
Args:
|
||||
command: Command string to check
|
||||
|
||||
Returns:
|
||||
Tuple of (is_safe, reason_if_unsafe)
|
||||
"""
|
||||
command_lower = command.lower().strip()
|
||||
|
||||
for dangerous_cmd in dangerous_commands:
|
||||
if dangerous_cmd.lower() in command_lower:
|
||||
return False, f"Command contains dangerous pattern: {dangerous_cmd}"
|
||||
|
||||
return True, None
|
||||
|
||||
|
||||
def _format_command_output(
|
||||
result: CommandResult, output_format: str = "markdown"
|
||||
) -> str:
|
||||
"""Format command execution results for LLM consumption.
|
||||
|
||||
Args:
|
||||
result: Command execution result
|
||||
output_format: Format type ('markdown', 'json', 'text')
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if output_format == "json":
|
||||
return json.dumps(result.model_dump(), indent=2)
|
||||
|
||||
elif output_format == "text":
|
||||
output_parts = [
|
||||
f"Command: {result.command}",
|
||||
f"Status: {'SUCCESS' if result.success else 'FAILED'}",
|
||||
f"Duration: {result.duration}",
|
||||
f"Return Code: {result.return_code}",
|
||||
]
|
||||
|
||||
if result.stdout:
|
||||
output_parts.extend(["\nOutput:", result.stdout])
|
||||
|
||||
if result.stderr:
|
||||
output_parts.extend(["\nErrors/Warnings:", result.stderr])
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
else: # markdown (default)
|
||||
status_emoji = "✅" if result.success else "❌"
|
||||
|
||||
output_parts = [
|
||||
f"# Terminal Command Execution {status_emoji}",
|
||||
f"**Command:** `{result.command}`",
|
||||
f"**Status:** {'SUCCESS' if result.success else 'FAILED'}",
|
||||
f"**Duration:** {result.duration}",
|
||||
f"**Return Code:** {result.return_code}",
|
||||
f"**Timestamp:** {result.timestamp}",
|
||||
]
|
||||
|
||||
if result.stdout:
|
||||
output_parts.extend(["\n## Output", "```", result.stdout.strip(), "```"])
|
||||
|
||||
if result.stderr:
|
||||
output_parts.extend(
|
||||
["\n## Errors/Warnings", "```", result.stderr.strip(), "```"]
|
||||
)
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
|
||||
async def _execute_command_async(command: str, timeout: int) -> CommandResult:
|
||||
"""Execute command asynchronously with timeout.
|
||||
|
||||
Args:
|
||||
command: Command to execute
|
||||
timeout: Timeout in seconds
|
||||
|
||||
Returns:
|
||||
CommandResult with execution details
|
||||
"""
|
||||
start_time = datetime.now()
|
||||
|
||||
try:
|
||||
# Create appropriate subprocess for platform
|
||||
if platform_info["system"] == "Windows":
|
||||
process = await asyncio.create_subprocess_shell(
|
||||
command,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
shell=True,
|
||||
)
|
||||
else:
|
||||
process = await asyncio.create_subprocess_shell(
|
||||
command,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
shell=True,
|
||||
executable="/bin/bash",
|
||||
)
|
||||
|
||||
try:
|
||||
stdout, stderr = await asyncio.wait_for(process.communicate(), timeout)
|
||||
stdout = stdout.decode("utf-8", errors="replace")
|
||||
stderr = stderr.decode("utf-8", errors="replace")
|
||||
return_code = process.returncode
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
try:
|
||||
process.kill()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
duration = str(datetime.now() - start_time)
|
||||
return CommandResult(
|
||||
command=command,
|
||||
success=False,
|
||||
stdout="",
|
||||
stderr=f"Command timed out after {timeout} seconds",
|
||||
return_code=-1,
|
||||
duration=duration,
|
||||
timestamp=start_time.isoformat(),
|
||||
)
|
||||
|
||||
duration = str(datetime.now() - start_time)
|
||||
result = CommandResult(
|
||||
command=command,
|
||||
success=return_code == 0,
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
return_code=return_code,
|
||||
duration=duration,
|
||||
timestamp=start_time.isoformat(),
|
||||
)
|
||||
|
||||
# Add to history
|
||||
command_history.append(
|
||||
{
|
||||
"timestamp": start_time.isoformat(),
|
||||
"command": command,
|
||||
"success": return_code == 0,
|
||||
"duration": duration,
|
||||
}
|
||||
)
|
||||
|
||||
# Maintain history size limit
|
||||
if len(command_history) > max_history_size:
|
||||
command_history.pop(0)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
duration = str(datetime.now() - start_time)
|
||||
return CommandResult(
|
||||
command=command,
|
||||
success=False,
|
||||
stdout="",
|
||||
stderr=f"Error executing command: {str(e)}",
|
||||
return_code=-1,
|
||||
duration=duration,
|
||||
timestamp=start_time.isoformat(),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
load_dotenv(override=True)
|
||||
logging.info("Starting terminal-server MCP server!")
|
||||
mcp.run(transport="stdio")
|
||||
Vendored
+19
@@ -0,0 +1,19 @@
|
||||
[project]
|
||||
name = "wayback-server"
|
||||
version = "0.1.0"
|
||||
description = "Add your description here"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
aiohttp= ">=3.12.15"
|
||||
dependencies = [
|
||||
"mcp",
|
||||
"aiohttp~=3.12.15",
|
||||
"requests~=2.32.4",
|
||||
"fastmcp~=2.11.3",
|
||||
"python-magic~=0.4.27",
|
||||
"chardet~=3.0.4",
|
||||
"pandas~=2.3.0",
|
||||
"beautifulsoup4~=4.12.3",
|
||||
"waybackpy~=3.0.6",
|
||||
|
||||
]
|
||||
+226
@@ -0,0 +1,226 @@
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import magic
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import requests
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import Any, Literal
|
||||
|
||||
|
||||
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)"
|
||||
)
|
||||
|
||||
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")
|
||||
|
||||
|
||||
def _validate_file_path(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)
|
||||
if not path.is_absolute():
|
||||
path = path.expanduser().resolve()
|
||||
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"File not found: {path}")
|
||||
|
||||
return path
|
||||
|
||||
def is_url(path_or_url: str) -> bool:
|
||||
"""
|
||||
Check if the given string is a URL.
|
||||
|
||||
Args:
|
||||
path_or_url: String to check
|
||||
|
||||
Returns:
|
||||
bool: True if the string is a URL, False otherwise
|
||||
"""
|
||||
parsed = urlparse(path_or_url)
|
||||
return bool(parsed.scheme and parsed.netloc)
|
||||
|
||||
|
||||
def get_mime_type(file_path: str, default_mime: str | None = None) -> str:
|
||||
"""
|
||||
Detect MIME type of a file using python-magic if available,
|
||||
otherwise fallback to extension-based detection.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
default_mime: Default MIME type to return if detection fails
|
||||
|
||||
Returns:
|
||||
str: Detected MIME type
|
||||
"""
|
||||
# Try using python-magic for accurate MIME type detection
|
||||
try:
|
||||
mime = magic.Magic(mime=True)
|
||||
return mime.from_file(file_path)
|
||||
except (AttributeError, IOError):
|
||||
# Fallback to extension-based detection
|
||||
extension_mime_map = {
|
||||
# Audio formats
|
||||
".mp3": "audio/mpeg",
|
||||
".wav": "audio/wav",
|
||||
".ogg": "audio/ogg",
|
||||
".m4a": "audio/mp4",
|
||||
".flac": "audio/flac",
|
||||
# Image formats
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".png": "image/png",
|
||||
".gif": "image/gif",
|
||||
".webp": "image/webp",
|
||||
".bmp": "image/bmp",
|
||||
".tiff": "image/tiff",
|
||||
# Video formats
|
||||
".mp4": "video/mp4",
|
||||
".avi": "video/x-msvideo",
|
||||
".mov": "video/quicktime",
|
||||
".mkv": "video/x-matroska",
|
||||
".webm": "video/webm",
|
||||
}
|
||||
|
||||
ext = Path(file_path).suffix.lower()
|
||||
return extension_mime_map.get(ext, default_mime or "application/octet-stream")
|
||||
|
||||
|
||||
def get_file_from_source(
|
||||
source: str,
|
||||
max_size_mb: float = 100.0,
|
||||
timeout: int = 60,
|
||||
) -> tuple[str, str, bytes]:
|
||||
"""
|
||||
Unified function to get file content from a URL or local path with validation.
|
||||
|
||||
Args:
|
||||
source: URL or local file path
|
||||
max_size_mb: Maximum allowed file size in MB
|
||||
timeout: Timeout for URL requests in seconds
|
||||
|
||||
Returns:
|
||||
Tuple[str, str, bytes]: (file_path, mime_type, file_content)
|
||||
- For URLs, file_path will be a temporary file path
|
||||
- For local files, file_path will be the original path
|
||||
|
||||
Raises:
|
||||
ValueError: When file doesn't exist, exceeds size limit, or has invalid MIME type
|
||||
IOError: When file cannot be read
|
||||
requests.RequestException: When URL request fails
|
||||
"""
|
||||
max_size_bytes = max_size_mb * 1024 * 1024
|
||||
|
||||
if is_url(source):
|
||||
# Handle URL source
|
||||
try:
|
||||
# Make a HEAD request first to check content length
|
||||
head_response = requests.head(source, timeout=timeout, allow_redirects=True)
|
||||
head_response.raise_for_status()
|
||||
|
||||
# Check content length if available
|
||||
content_length = head_response.headers.get("content-length")
|
||||
if content_length and int(content_length) > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({int(content_length) / (1024 * 1024):.2f} MB) "
|
||||
f"exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Download the file
|
||||
response = requests.get(source, timeout=timeout, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
# Read content with size checking
|
||||
content = b""
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
if len(content) + len(chunk) > max_size_bytes:
|
||||
raise ValueError(f"File size exceeds maximum allowed size ({max_size_mb} MB)")
|
||||
content += chunk
|
||||
|
||||
# Create temporary file
|
||||
parsed_url = urlparse(source)
|
||||
filename = os.path.basename(parsed_url.path) or "downloaded_file"
|
||||
|
||||
# Create temporary file with proper extension
|
||||
suffix = Path(filename).suffix or ".tmp"
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
|
||||
temp_file.write(content)
|
||||
temp_path = temp_file.name
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(temp_path)
|
||||
|
||||
return temp_path, mime_type, content
|
||||
|
||||
except requests.RequestException as e:
|
||||
raise requests.RequestException(f"Failed to download file from URL: {e}: {traceback.format_exc()}")
|
||||
except Exception as e:
|
||||
raise IOError(f"Error processing URL: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
else:
|
||||
# Handle local file path
|
||||
file_path = Path(source)
|
||||
|
||||
# Check if file exists
|
||||
if not file_path.exists():
|
||||
raise ValueError(f"File does not exist: {source}")
|
||||
|
||||
if not file_path.is_file():
|
||||
raise ValueError(f"Path is not a file: {source}")
|
||||
|
||||
# Check file size
|
||||
file_size = file_path.stat().st_size
|
||||
if file_size > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({file_size / (1024 * 1024):.2f} MB) exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Read file content
|
||||
try:
|
||||
with open(file_path, "rb") as f:
|
||||
content = f.read()
|
||||
except Exception as e:
|
||||
raise IOError(f"Cannot read file {source}: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(str(file_path))
|
||||
|
||||
return str(file_path), mime_type, content
|
||||
Vendored
+615
@@ -0,0 +1,615 @@
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any, Union, Literal
|
||||
|
||||
import chardet
|
||||
import pandas as pd
|
||||
import requests
|
||||
from bs4 import BeautifulSoup
|
||||
from dotenv import load_dotenv
|
||||
from pydantic.fields import FieldInfo
|
||||
from waybackpy import WaybackMachineCDXServerAPI
|
||||
|
||||
from mcp.server import FastMCP
|
||||
from mcp.types import TextContent
|
||||
from pydantic import Field, BaseModel
|
||||
|
||||
from base import ActionResponse
|
||||
|
||||
load_dotenv()
|
||||
workspace = Path.home()
|
||||
user_agent = "AWorld/1.0"
|
||||
default_timeout = 30
|
||||
max_content_length = 8192
|
||||
|
||||
|
||||
class ArchivedVersion(BaseModel):
|
||||
"""Individual archived version with structured data."""
|
||||
|
||||
timestamp: str
|
||||
url: str
|
||||
status_code: str
|
||||
digest: str
|
||||
length: str
|
||||
mime_type: str
|
||||
|
||||
|
||||
class WaybackMetadata(BaseModel):
|
||||
"""Metadata for Wayback Machine operation results."""
|
||||
|
||||
url: str
|
||||
operation: str # 'list_versions', 'get_content', 'save_page'
|
||||
timestamp: str | None = None
|
||||
total_versions: int | None = None
|
||||
date_range: dict[str, str | None] | None = None
|
||||
content_length: int | None = None
|
||||
text_extracted: bool = False
|
||||
truncated: bool = False
|
||||
execution_time: float | None = None
|
||||
error_type: str | None = None
|
||||
user_agent: str = (
|
||||
"AWorld/1.0 (https://github.com/inclusionAI/AWorld; qintong.wqt@antgroup.com)"
|
||||
)
|
||||
|
||||
|
||||
mcp = FastMCP(
|
||||
"wayback-server",
|
||||
instructions="""
|
||||
MCP service for Wayback Machine operations.
|
||||
|
||||
Provides comprehensive Wayback Machine functionality including:
|
||||
- Listing archived versions of URLs with flexible filtering
|
||||
- Fetching content from specific archived snapshots
|
||||
- LLM-friendly content formatting and text extraction
|
||||
- Robust error handling and detailed logging
|
||||
- Multiple output formats (markdown, JSON, text)
|
||||
""",
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
List available archived versions of a URL from the Wayback Machine.
|
||||
|
||||
This function queries the Wayback Machine CDX API to retrieve all available
|
||||
archived snapshots for a given URL, with optional date range filtering.
|
||||
"""
|
||||
)
|
||||
async def list_archived_versions(
|
||||
url: str = Field(
|
||||
description="The URL of the website to check for archived versions"
|
||||
),
|
||||
limit: int = Field(
|
||||
default=10, description="Maximum number of versions to return (0 for all)"
|
||||
),
|
||||
from_date: str | None = Field(
|
||||
default=None, description="Start date filter (YYYYMMDDhhmmss)"
|
||||
),
|
||||
to_date: str | None = Field(
|
||||
default=None, description="End date filter (YYYYMMDDhhmmss)"
|
||||
),
|
||||
output_format: str = Field(
|
||||
default="markdown", description="Output format: 'markdown', 'json', or 'text'"
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
start_time = time.time()
|
||||
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(url, FieldInfo):
|
||||
url = url.default
|
||||
if isinstance(limit, FieldInfo):
|
||||
limit = limit.default
|
||||
if isinstance(from_date, FieldInfo):
|
||||
from_date = from_date.default
|
||||
if isinstance(to_date, FieldInfo):
|
||||
to_date = to_date.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
try:
|
||||
# Validate parameters
|
||||
url, _ = _validate_wayback_parameters(url)
|
||||
|
||||
logging.info(f"Listing archived versions for: {url}")
|
||||
|
||||
# Query Wayback Machine CDX API
|
||||
cdx_api = WaybackMachineCDXServerAPI(url, user_agent=user_agent)
|
||||
all_snapshots = list(cdx_api.snapshots())
|
||||
|
||||
# Apply date filtering
|
||||
if from_date or to_date:
|
||||
snapshots = [
|
||||
s
|
||||
for s in all_snapshots
|
||||
if (not from_date or s.timestamp >= from_date)
|
||||
and (not to_date or s.timestamp <= to_date)
|
||||
]
|
||||
else:
|
||||
snapshots = all_snapshots
|
||||
|
||||
if not snapshots:
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message="No archived versions found for the specified URL and date range.",
|
||||
metadata=WaybackMetadata(
|
||||
url=url,
|
||||
operation="list_versions",
|
||||
total_versions=0,
|
||||
date_range={"from_date": from_date, "to_date": to_date},
|
||||
execution_time=time.time() - start_time,
|
||||
error_type="no_results",
|
||||
).model_dump(),
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
# Convert to structured format
|
||||
versions = [
|
||||
ArchivedVersion(
|
||||
timestamp=snapshot.timestamp,
|
||||
url=snapshot.archive_url,
|
||||
status_code=snapshot.statuscode,
|
||||
digest=snapshot.digest,
|
||||
length=snapshot.length,
|
||||
mime_type=snapshot.mimetype,
|
||||
)
|
||||
for snapshot in snapshots
|
||||
]
|
||||
|
||||
# Apply limit
|
||||
if limit > 0 and len(versions) > limit:
|
||||
versions = versions[:limit]
|
||||
|
||||
# Format output
|
||||
query_info = {
|
||||
"url": url,
|
||||
"from_date": from_date,
|
||||
"to_date": to_date,
|
||||
"total_found": len(snapshots),
|
||||
}
|
||||
|
||||
if output_format == "json":
|
||||
message = [version.model_dump() for version in versions]
|
||||
else:
|
||||
message = _format_versions_for_llm(versions, query_info)
|
||||
|
||||
execution_time = time.time() - start_time
|
||||
logging.info(
|
||||
f"Found {len(versions)} archived versions in {execution_time:.2f}s"
|
||||
)
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True,
|
||||
message=message,
|
||||
metadata=WaybackMetadata(
|
||||
url=url,
|
||||
operation="list_versions",
|
||||
total_versions=len(snapshots),
|
||||
date_range={"from_date": from_date, "to_date": to_date},
|
||||
execution_time=execution_time,
|
||||
).model_dump(),
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
)
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to list archived versions: {str(e)}"
|
||||
logging.info(error_msg)
|
||||
logging.error(f"Error in mcp_list_archived_versions: {traceback.format_exc()}")
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=error_msg,
|
||||
metadata=WaybackMetadata(
|
||||
url=url or "unknown",
|
||||
operation="list_versions",
|
||||
execution_time=time.time() - start_time,
|
||||
error_type=type(e).__name__,
|
||||
).model_dump(),
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Fetch content from a specific archived page version.
|
||||
|
||||
This function retrieves the content of a specific archived snapshot from the
|
||||
Wayback Machine, with options for text extraction and content truncation.
|
||||
"""
|
||||
)
|
||||
async def get_archived_content(
|
||||
url: str = Field(
|
||||
description="The URL of the website to fetch archived content from"
|
||||
),
|
||||
timestamp: str = Field(
|
||||
description="The timestamp of the desired version (YYYYMMDDhhmmss)"
|
||||
),
|
||||
extract_text_only: bool = Field(
|
||||
default=True, description="Extract only text content, removing HTML tags"
|
||||
),
|
||||
truncate_content: bool = Field(
|
||||
default=False, description="Truncate content to manageable length for LLMs"
|
||||
),
|
||||
output_format: str = Field(
|
||||
default="markdown", description="Output format: 'markdown', 'json', or 'text'"
|
||||
),
|
||||
) -> Union[str, TextContent]:
|
||||
start_time = time.time()
|
||||
|
||||
# Handle FieldInfo objects
|
||||
if isinstance(url, FieldInfo):
|
||||
url = url.default
|
||||
if isinstance(timestamp, FieldInfo):
|
||||
timestamp = timestamp.default
|
||||
if isinstance(extract_text_only, FieldInfo):
|
||||
extract_text_only = extract_text_only.default
|
||||
if isinstance(truncate_content, FieldInfo):
|
||||
truncate_content = truncate_content.default
|
||||
if isinstance(output_format, FieldInfo):
|
||||
output_format = output_format.default
|
||||
|
||||
try:
|
||||
# Validate parameters
|
||||
url, timestamp = _validate_wayback_parameters(url, timestamp)
|
||||
|
||||
logging.info(f"Fetching archived content: {url} at {timestamp}")
|
||||
|
||||
# Query Wayback Machine for closest snapshot
|
||||
cdx_api = WaybackMachineCDXServerAPI(url, user_agent=user_agent)
|
||||
snapshot = cdx_api.near(wayback_machine_timestamp=timestamp)
|
||||
|
||||
if not snapshot or not snapshot.archive_url:
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=f"No archived version found for {url} at timestamp {timestamp}",
|
||||
metadata=WaybackMetadata(
|
||||
url=url,
|
||||
operation="get_content",
|
||||
timestamp=timestamp,
|
||||
execution_time=time.time() - start_time,
|
||||
error_type="no_snapshot",
|
||||
).model_dump(),
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
# Fetch content
|
||||
response = requests.get(snapshot.archive_url, timeout=default_timeout)
|
||||
response.raise_for_status()
|
||||
content = response.text
|
||||
original_length = len(content)
|
||||
|
||||
# Extract text if requested
|
||||
if extract_text_only:
|
||||
soup = BeautifulSoup(content, "html.parser")
|
||||
content = soup.get_text(separator=" ", strip=True)
|
||||
|
||||
# Truncate if requested
|
||||
truncated = False
|
||||
if truncate_content and len(content) > max_content_length:
|
||||
content = content[:max_content_length] + "..."
|
||||
truncated = True
|
||||
|
||||
# Prepare content data
|
||||
content_data = {
|
||||
"url": url,
|
||||
"timestamp": timestamp,
|
||||
"fetched_timestamp": snapshot.timestamp,
|
||||
"content": content,
|
||||
"original_content_length": original_length,
|
||||
"truncated": truncated,
|
||||
"extract_text_only": extract_text_only,
|
||||
}
|
||||
|
||||
# Format output
|
||||
if output_format == "json":
|
||||
message = content_data
|
||||
elif output_format == "text":
|
||||
message = content
|
||||
else: # markdown
|
||||
message = _format_content_for_llm(content_data, output_format)
|
||||
|
||||
execution_time = time.time() - start_time
|
||||
logging.info(f"Retrieved {len(content):,} characters in {execution_time:.2f}s")
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True,
|
||||
message=message,
|
||||
metadata=WaybackMetadata(
|
||||
url=url,
|
||||
operation="get_content",
|
||||
timestamp=snapshot.timestamp,
|
||||
content_length=len(content),
|
||||
text_extracted=extract_text_only,
|
||||
truncated=truncated,
|
||||
execution_time=execution_time,
|
||||
).model_dump(),
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
)
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to fetch archived content: {str(e)}"
|
||||
logging.info(error_msg)
|
||||
logging.error(f"Error in mcp_get_archived_content: {traceback.format_exc()}")
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=False,
|
||||
message=error_msg,
|
||||
metadata=WaybackMetadata(
|
||||
url=url or "unknown",
|
||||
operation="get_content",
|
||||
timestamp=timestamp or "unknown",
|
||||
execution_time=time.time() - start_time,
|
||||
error_type=type(e).__name__,
|
||||
).model_dump(),
|
||||
)
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(
|
||||
action_response.model_dump()
|
||||
), # Empty string instead of None
|
||||
**{"metadata": {}}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
description="""
|
||||
Get Wayback Machine service capabilities and configuration.
|
||||
"""
|
||||
)
|
||||
async def get_wayback_capabilities() -> Union[str, TextContent]:
|
||||
capabilities = {
|
||||
"service": "Wayback Machine MCP Server",
|
||||
"version": "1.0.0",
|
||||
"description": "Interact with the Internet Archive's Wayback Machine",
|
||||
"operations": {
|
||||
"list_versions": "List archived versions of URLs with date filtering",
|
||||
"get_content": "Fetch content from specific archived snapshots",
|
||||
},
|
||||
"features": {
|
||||
"date_filtering": True,
|
||||
"text_extraction": True,
|
||||
"content_truncation": True,
|
||||
"multiple_formats": ["markdown", "json", "text"],
|
||||
"error_handling": True,
|
||||
"logging": True,
|
||||
},
|
||||
"configuration": {
|
||||
"user_agent": user_agent,
|
||||
"default_timeout": default_timeout,
|
||||
"max_content_length": max_content_length,
|
||||
},
|
||||
"limits": {
|
||||
"max_content_length": max_content_length,
|
||||
"request_timeout": default_timeout,
|
||||
},
|
||||
}
|
||||
|
||||
message = f"""# Wayback Machine Service Capabilities
|
||||
|
||||
## Service Information
|
||||
- **Service:** {capabilities["service"]}
|
||||
- **Version:** {capabilities["version"]}
|
||||
- **Description:** {capabilities["description"]}
|
||||
|
||||
## Available Operations
|
||||
- **List Versions:** {capabilities["operations"]["list_versions"]}
|
||||
- **Get Content:** {capabilities["operations"]["get_content"]}
|
||||
- **Save Page:** {capabilities["operations"]["save_page"]}
|
||||
|
||||
## Features
|
||||
- **Date Filtering:** {capabilities["features"]["date_filtering"]}
|
||||
- **Text Extraction:** {capabilities["features"]["text_extraction"]}
|
||||
- **Content Truncation:** {capabilities["features"]["content_truncation"]}
|
||||
- **Output Formats:** {", ".join(capabilities["features"]["multiple_formats"])}
|
||||
- **Error Handling:** {capabilities["features"]["error_handling"]}
|
||||
- **Logging:** {capabilities["features"]["logging"]}
|
||||
|
||||
## Configuration
|
||||
- **User Agent:** {capabilities["configuration"]["user_agent"]}
|
||||
- **Default Timeout:** {capabilities["configuration"]["default_timeout"]} seconds
|
||||
- **Max Content Length:** {capabilities["configuration"]["max_content_length"]:,} characters
|
||||
|
||||
## Limits
|
||||
- **Max Content Length:** {capabilities["limits"]["max_content_length"]:,} characters
|
||||
- **Request Timeout:** {capabilities["limits"]["request_timeout"]} seconds
|
||||
"""
|
||||
|
||||
action_response = ActionResponse(
|
||||
success=True, message=message, metadata=capabilities
|
||||
)
|
||||
output_dict = {
|
||||
"artifact_type": "MARKDOWN",
|
||||
"artifact_data": json.dumps(
|
||||
action_response.model_dump()
|
||||
)
|
||||
}
|
||||
return TextContent(
|
||||
type="text",
|
||||
text=json.dumps(action_response.model_dump()), # Empty string instead of None
|
||||
**{"metadata": output_dict}, # Pass as additional fields
|
||||
)
|
||||
|
||||
|
||||
def _format_versions_for_llm(versions: list[ArchivedVersion], query_info: dict) -> str:
|
||||
"""Format archived versions list for LLM consumption.
|
||||
|
||||
Args:
|
||||
versions: List of archived versions
|
||||
query_info: Query information including URL and filters
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if not versions:
|
||||
return f"No archived versions found for URL: {query_info.get('url', 'Unknown')}"
|
||||
|
||||
output_parts = [
|
||||
f"# Wayback Machine Archives for {query_info.get('url', 'Unknown')}",
|
||||
f"\nFound **{len(versions)}** archived versions",
|
||||
]
|
||||
|
||||
if query_info.get("from_date") or query_info.get("to_date"):
|
||||
date_filter = []
|
||||
if query_info.get("from_date"):
|
||||
date_filter.append(f"From: {query_info['from_date']}")
|
||||
if query_info.get("to_date"):
|
||||
date_filter.append(f"To: {query_info['to_date']}")
|
||||
output_parts.append(f"\n**Date Filter:** {' | '.join(date_filter)}")
|
||||
|
||||
output_parts.append("\n## Available Versions:")
|
||||
|
||||
for i, version in enumerate(versions[:10], 1): # Show first 10
|
||||
timestamp_formatted = _format_timestamp(version.timestamp)
|
||||
output_parts.append(
|
||||
f"\n{i}. **{timestamp_formatted}**\n"
|
||||
f" - Archive URL: {version.url}\n"
|
||||
f" - Status: {version.status_code} | Size: {version.length} bytes\n"
|
||||
f" - Type: {version.mime_type}"
|
||||
)
|
||||
|
||||
if len(versions) > 10:
|
||||
output_parts.append(f"\n... and {len(versions) - 10} more versions")
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
|
||||
def _format_content_for_llm(content_data: dict, output_format: str = "markdown") -> str:
|
||||
"""Format archived content for LLM consumption.
|
||||
|
||||
Args:
|
||||
content_data: Content data dictionary
|
||||
output_format: Format type ('markdown', 'json', 'text')
|
||||
|
||||
Returns:
|
||||
Formatted string suitable for LLM consumption
|
||||
"""
|
||||
if output_format == "json":
|
||||
return json.dumps(content_data, indent=2)
|
||||
|
||||
elif output_format == "text":
|
||||
return content_data.get("content", "")
|
||||
|
||||
else: # markdown (default)
|
||||
output_parts = [
|
||||
f"# Archived Content from {content_data.get('url', 'Unknown')}",
|
||||
f"\n**Requested Timestamp:** {content_data.get('timestamp', 'Unknown')}",
|
||||
f"**Actual Timestamp:** {_format_timestamp(content_data.get('fetched_timestamp', ''))}",
|
||||
f"**Content Length:** {content_data.get('original_content_length', 0):,} characters",
|
||||
]
|
||||
|
||||
if content_data.get("truncated"):
|
||||
output_parts.append(
|
||||
f"**Note:** Content truncated to {max_content_length:,} characters"
|
||||
)
|
||||
|
||||
if content_data.get("extract_text_only"):
|
||||
output_parts.append("**Note:** Text-only extraction applied")
|
||||
|
||||
output_parts.extend(
|
||||
["\n## Content:", "\n---\n", content_data.get("content", ""), "\n---"]
|
||||
)
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
|
||||
def _format_timestamp(timestamp: str) -> str:
|
||||
"""Format Wayback Machine timestamp to human-readable format.
|
||||
|
||||
Args:
|
||||
timestamp: Wayback timestamp (YYYYMMDDhhmmss)
|
||||
|
||||
Returns:
|
||||
Human-readable timestamp
|
||||
"""
|
||||
try:
|
||||
if len(timestamp) >= 14:
|
||||
dt = datetime.strptime(timestamp[:14], "%Y%m%d%H%M%S")
|
||||
return dt.strftime("%Y-%m-%d %H:%M:%S UTC")
|
||||
return timestamp
|
||||
except (ValueError, TypeError):
|
||||
return timestamp or "Unknown"
|
||||
|
||||
|
||||
def _validate_wayback_parameters(
|
||||
url: str, timestamp: str = None
|
||||
) -> tuple[str, str | None]:
|
||||
"""Validate and normalize Wayback Machine parameters.
|
||||
|
||||
Args:
|
||||
url: URL to validate
|
||||
timestamp: Optional timestamp to validate
|
||||
|
||||
Returns:
|
||||
Tuple of (validated_url, validated_timestamp)
|
||||
|
||||
Raises:
|
||||
ValueError: If parameters are invalid
|
||||
"""
|
||||
if not url or not url.strip():
|
||||
raise ValueError("URL cannot be empty")
|
||||
|
||||
url = url.strip()
|
||||
if not url.startswith(("http://", "https://")):
|
||||
url = "https://" + url
|
||||
|
||||
validated_timestamp = None
|
||||
if timestamp:
|
||||
timestamp = timestamp.strip()
|
||||
if len(timestamp) < 8:
|
||||
raise ValueError("Timestamp must be at least 8 characters (YYYYMMDD)")
|
||||
validated_timestamp = timestamp
|
||||
|
||||
return url, validated_timestamp
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv(override=True)
|
||||
logging.info("Starting wayback-server MCP server!")
|
||||
mcp.run(transport="stdio")
|
||||
+19
@@ -0,0 +1,19 @@
|
||||
[project]
|
||||
name = "wiki-server"
|
||||
version = "0.1.0"
|
||||
description = "Add your description here"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
aiohttp= ">=3.12.15"
|
||||
dependencies = [
|
||||
"mcp",
|
||||
"aiohttp~=3.12.15",
|
||||
"requests~=2.32.4",
|
||||
"fastmcp~=2.11.3",
|
||||
"python-magic~=0.4.27",
|
||||
"beautifulsoup4~=4.12.3",
|
||||
"waybackpy~=3.0.6",
|
||||
"wikipedia~=1.4.0",
|
||||
"pandas~=2.3.0",
|
||||
"chardet"
|
||||
]
|
||||
+226
@@ -0,0 +1,226 @@
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import magic
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import requests
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import Any, Literal
|
||||
|
||||
|
||||
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)"
|
||||
)
|
||||
|
||||
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")
|
||||
|
||||
|
||||
def _validate_file_path(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)
|
||||
if not path.is_absolute():
|
||||
path = path.expanduser().resolve()
|
||||
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"File not found: {path}")
|
||||
|
||||
return path
|
||||
|
||||
def is_url(path_or_url: str) -> bool:
|
||||
"""
|
||||
Check if the given string is a URL.
|
||||
|
||||
Args:
|
||||
path_or_url: String to check
|
||||
|
||||
Returns:
|
||||
bool: True if the string is a URL, False otherwise
|
||||
"""
|
||||
parsed = urlparse(path_or_url)
|
||||
return bool(parsed.scheme and parsed.netloc)
|
||||
|
||||
|
||||
def get_mime_type(file_path: str, default_mime: str | None = None) -> str:
|
||||
"""
|
||||
Detect MIME type of a file using python-magic if available,
|
||||
otherwise fallback to extension-based detection.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
default_mime: Default MIME type to return if detection fails
|
||||
|
||||
Returns:
|
||||
str: Detected MIME type
|
||||
"""
|
||||
# Try using python-magic for accurate MIME type detection
|
||||
try:
|
||||
mime = magic.Magic(mime=True)
|
||||
return mime.from_file(file_path)
|
||||
except (AttributeError, IOError):
|
||||
# Fallback to extension-based detection
|
||||
extension_mime_map = {
|
||||
# Audio formats
|
||||
".mp3": "audio/mpeg",
|
||||
".wav": "audio/wav",
|
||||
".ogg": "audio/ogg",
|
||||
".m4a": "audio/mp4",
|
||||
".flac": "audio/flac",
|
||||
# Image formats
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".png": "image/png",
|
||||
".gif": "image/gif",
|
||||
".webp": "image/webp",
|
||||
".bmp": "image/bmp",
|
||||
".tiff": "image/tiff",
|
||||
# Video formats
|
||||
".mp4": "video/mp4",
|
||||
".avi": "video/x-msvideo",
|
||||
".mov": "video/quicktime",
|
||||
".mkv": "video/x-matroska",
|
||||
".webm": "video/webm",
|
||||
}
|
||||
|
||||
ext = Path(file_path).suffix.lower()
|
||||
return extension_mime_map.get(ext, default_mime or "application/octet-stream")
|
||||
|
||||
|
||||
def get_file_from_source(
|
||||
source: str,
|
||||
max_size_mb: float = 100.0,
|
||||
timeout: int = 60,
|
||||
) -> tuple[str, str, bytes]:
|
||||
"""
|
||||
Unified function to get file content from a URL or local path with validation.
|
||||
|
||||
Args:
|
||||
source: URL or local file path
|
||||
max_size_mb: Maximum allowed file size in MB
|
||||
timeout: Timeout for URL requests in seconds
|
||||
|
||||
Returns:
|
||||
Tuple[str, str, bytes]: (file_path, mime_type, file_content)
|
||||
- For URLs, file_path will be a temporary file path
|
||||
- For local files, file_path will be the original path
|
||||
|
||||
Raises:
|
||||
ValueError: When file doesn't exist, exceeds size limit, or has invalid MIME type
|
||||
IOError: When file cannot be read
|
||||
requests.RequestException: When URL request fails
|
||||
"""
|
||||
max_size_bytes = max_size_mb * 1024 * 1024
|
||||
|
||||
if is_url(source):
|
||||
# Handle URL source
|
||||
try:
|
||||
# Make a HEAD request first to check content length
|
||||
head_response = requests.head(source, timeout=timeout, allow_redirects=True)
|
||||
head_response.raise_for_status()
|
||||
|
||||
# Check content length if available
|
||||
content_length = head_response.headers.get("content-length")
|
||||
if content_length and int(content_length) > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({int(content_length) / (1024 * 1024):.2f} MB) "
|
||||
f"exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Download the file
|
||||
response = requests.get(source, timeout=timeout, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
# Read content with size checking
|
||||
content = b""
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
if len(content) + len(chunk) > max_size_bytes:
|
||||
raise ValueError(f"File size exceeds maximum allowed size ({max_size_mb} MB)")
|
||||
content += chunk
|
||||
|
||||
# Create temporary file
|
||||
parsed_url = urlparse(source)
|
||||
filename = os.path.basename(parsed_url.path) or "downloaded_file"
|
||||
|
||||
# Create temporary file with proper extension
|
||||
suffix = Path(filename).suffix or ".tmp"
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
|
||||
temp_file.write(content)
|
||||
temp_path = temp_file.name
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(temp_path)
|
||||
|
||||
return temp_path, mime_type, content
|
||||
|
||||
except requests.RequestException as e:
|
||||
raise requests.RequestException(f"Failed to download file from URL: {e}: {traceback.format_exc()}")
|
||||
except Exception as e:
|
||||
raise IOError(f"Error processing URL: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
else:
|
||||
# Handle local file path
|
||||
file_path = Path(source)
|
||||
|
||||
# Check if file exists
|
||||
if not file_path.exists():
|
||||
raise ValueError(f"File does not exist: {source}")
|
||||
|
||||
if not file_path.is_file():
|
||||
raise ValueError(f"Path is not a file: {source}")
|
||||
|
||||
# Check file size
|
||||
file_size = file_path.stat().st_size
|
||||
if file_size > max_size_bytes:
|
||||
raise ValueError(
|
||||
f"File size ({file_size / (1024 * 1024):.2f} MB) exceeds maximum allowed size ({max_size_mb} MB)"
|
||||
)
|
||||
|
||||
# Read file content
|
||||
try:
|
||||
with open(file_path, "rb") as f:
|
||||
content = f.read()
|
||||
except Exception as e:
|
||||
raise IOError(f"Cannot read file {source}: {e}: {traceback.format_exc()}") from e
|
||||
|
||||
# Get MIME type
|
||||
mime_type = get_mime_type(str(file_path))
|
||||
|
||||
return str(file_path), mime_type, content
|
||||
+1254
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,8 @@
|
||||
#!/bin/sh
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
sh ../virtualpc-mcp/mcp_server/build-image.sh && \
|
||||
|
||||
sh ../gaia-mcp-server/build-image.sh && \
|
||||
|
||||
sh -c "docker run -p 4242:4242 -p 5901:5901 --rm gaia-mcp-server"
|
||||
@@ -0,0 +1,8 @@
|
||||
#!/bin/sh
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
export MCP_SERVERS_PATH="$(pwd)/mcp_servers"
|
||||
|
||||
sh mcp_servers/init_env.sh && \
|
||||
|
||||
sh -c "cd ../virtualpc-mcp/mcp_server/mcp_server_proxy && uv run -m mcp_server_proxy.main"
|
||||
@@ -0,0 +1,8 @@
|
||||
#!/bin/sh
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
sh virtualpc-mcp/mcp_server/build-image.sh && \
|
||||
|
||||
sh gaia-mcp-server/build-image.sh && \
|
||||
|
||||
sh virtualpc-mcp/run-local.sh
|
||||
+201
@@ -0,0 +1,201 @@
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
import time
|
||||
import traceback
|
||||
from dotenv import dotenv_values, set_key
|
||||
import httpx
|
||||
from mcp import ClientSession
|
||||
from mcp.types import (
|
||||
LoggingMessageNotificationParams,
|
||||
ElicitResult,
|
||||
ElicitRequestParams,
|
||||
)
|
||||
from mcp.client.streamable_http import streamablehttp_client
|
||||
from mcp.shared.context import RequestContext
|
||||
|
||||
from aworld.utils.common import get_local_ip
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
LOCAL_MCP_TOKEN_SECRET = "123321"
|
||||
|
||||
|
||||
def _jwt_part(value: dict) -> str:
|
||||
raw = json.dumps(value, separators=(",", ":")).encode()
|
||||
return base64.urlsafe_b64encode(raw).rstrip(b"=").decode()
|
||||
|
||||
|
||||
def gen_local_mcp_token(app: str = "local_debug") -> str:
|
||||
secret = os.getenv("MCP_GATEWAY_TOKEN_SECRET", LOCAL_MCP_TOKEN_SECRET)
|
||||
header = {"alg": "HS256", "typ": "JWT"}
|
||||
payload = {"app": app, "version": 1, "time": time.time()}
|
||||
signing_input = f"{_jwt_part(header)}.{_jwt_part(payload)}"
|
||||
signature = hmac.new(
|
||||
secret.encode(),
|
||||
signing_input.encode(),
|
||||
hashlib.sha256,
|
||||
).digest()
|
||||
encoded_signature = base64.urlsafe_b64encode(signature).rstrip(b"=").decode()
|
||||
return f"{signing_input}.{encoded_signature}"
|
||||
|
||||
|
||||
class TranEnv:
|
||||
|
||||
def __init__(self):
|
||||
self.base_dir = Path(__file__).parent.parent
|
||||
self.env_dir = self.base_dir / "env"
|
||||
self.mcp_config = None
|
||||
self.mcp_variables = None
|
||||
|
||||
def get_env_config(self):
|
||||
if self.mcp_variables:
|
||||
url = f"http://{self.mcp_variables['ip']}:{self.mcp_variables['port']}/mcp"
|
||||
self.mcp_config = {
|
||||
"mcpServers": {
|
||||
"virtualpc-mcp": {
|
||||
"type": "streamable-http",
|
||||
"url": url,
|
||||
"headers": {
|
||||
"Authorization": f"Bearer {self.mcp_variables['token']}",
|
||||
},
|
||||
"timeout": 600,
|
||||
"sse_read_timeout": 600,
|
||||
"client_session_timeout_seconds": 600,
|
||||
}
|
||||
}
|
||||
}
|
||||
return self.mcp_config
|
||||
return None
|
||||
|
||||
async def create_env(self, mode: str = "local", docker_dir: str = None) -> bool:
|
||||
if mode == "local":
|
||||
if not docker_dir:
|
||||
logger.error("You must provide --docker_dir to specify the Docker directory to build (relative to env).")
|
||||
return False
|
||||
|
||||
image_ready = await self._build_image(docker_dir)
|
||||
assert image_ready, "Image is not ready!"
|
||||
|
||||
service_ready = await self._start_service()
|
||||
assert service_ready, "Service config is not ready!"
|
||||
|
||||
service_ready = await self._check_service_ready()
|
||||
assert service_ready, "Service is not ready!"
|
||||
|
||||
self.mcp_variables = {
|
||||
"ip": get_local_ip(),
|
||||
"port": 8000,
|
||||
"token": gen_local_mcp_token(),
|
||||
}
|
||||
|
||||
logger.info("✅ Service is ready!")
|
||||
return True
|
||||
else:
|
||||
logger.warning(f"Mode {mode} is not supported!")
|
||||
return False
|
||||
|
||||
async def _build_image(self, docker_dir: str):
|
||||
try:
|
||||
# Use asyncio.create_subprocess_exec for async subprocess execution
|
||||
logger.info(f"Building {docker_dir} image...")
|
||||
process1 = await asyncio.create_subprocess_exec(
|
||||
"sh",
|
||||
"build-image.sh",
|
||||
cwd=self.env_dir / "virtualpc-mcp" / "mcp_server",
|
||||
)
|
||||
await process1.wait()
|
||||
|
||||
if process1.returncode != 0:
|
||||
logger.error("Failed to build virtualpc-mcp image")
|
||||
return False
|
||||
|
||||
target_dir = self.env_dir / docker_dir
|
||||
if not target_dir.exists():
|
||||
logger.error(f"Specified Docker directory does not exist: {target_dir}")
|
||||
return False
|
||||
|
||||
process2 = await asyncio.create_subprocess_exec(
|
||||
"sh",
|
||||
"build-image.sh",
|
||||
cwd=target_dir,
|
||||
)
|
||||
await process2.wait()
|
||||
|
||||
if process2.returncode != 0:
|
||||
logger.error(f"Failed to build {docker_dir} image")
|
||||
return False
|
||||
|
||||
logger.info("All images built successfully")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to build image: {traceback.format_exc()}")
|
||||
return False
|
||||
|
||||
async def _start_service(self):
|
||||
try:
|
||||
logger.info("Starting virtualpc-mcp service...")
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
"sh",
|
||||
"run-local.sh",
|
||||
cwd=self.env_dir / "virtualpc-mcp",
|
||||
)
|
||||
|
||||
# Wait a bit for the service to start
|
||||
await asyncio.sleep(2)
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to start service: {traceback.format_exc()}")
|
||||
return None
|
||||
|
||||
async def _check_service_ready(self) -> bool:
|
||||
url = "http://localhost:8000/health"
|
||||
|
||||
max_retries = 180
|
||||
for i in range(max_retries):
|
||||
try:
|
||||
# Try to establish MCP connection to check if service is ready
|
||||
async with httpx.AsyncClient() as client:
|
||||
resp = await client.get(url)
|
||||
resp.raise_for_status()
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Waiting for service ready: {(i+1)}/{max_retries} attempts"
|
||||
)
|
||||
await asyncio.sleep(10)
|
||||
else:
|
||||
logger.error(
|
||||
f"Service at {url} is not reachable after {max_retries} attempts."
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
parser = argparse.ArgumentParser(description="Env construction")
|
||||
parser.add_argument("--docker_dir", help="Docker directory to build (relative to env, e.g., gaia-mcp-server)")
|
||||
|
||||
async def main():
|
||||
try:
|
||||
args = parser.parse_args()
|
||||
if not args.docker_dir:
|
||||
parser.error("You must use --docker_dir to specify the Docker directory to build (e.g., gaia-mcp-server)")
|
||||
train_env = TranEnv()
|
||||
env_started = await train_env.create_env(docker_dir=args.docker_dir)
|
||||
if env_started:
|
||||
mcp_variables = json.dumps(train_env.mcp_variables, ensure_ascii=False, indent=4)
|
||||
print(mcp_variables)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to start env: {traceback.format_exc()}")
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
Vendored
+29
@@ -0,0 +1,29 @@
|
||||
FROM docker:dind
|
||||
|
||||
# Configure mirror
|
||||
ARG ENABLE_MIRROR=false
|
||||
ENV PIP_INDEX_URL=${ENABLE_MIRROR:+https://mirrors.aliyun.com/pypi/simple}
|
||||
|
||||
RUN apk add --no-cache python3 py3-pip curl tini
|
||||
|
||||
COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
|
||||
|
||||
WORKDIR /app/container_server
|
||||
|
||||
COPY container_server/pyproject.toml /app/container_server/pyproject.toml
|
||||
COPY container_server/src /app/container_server/src
|
||||
|
||||
RUN uv sync --python-preference=only-system
|
||||
|
||||
EXPOSE 8000
|
||||
|
||||
HEALTHCHECK --interval=30s --timeout=10s --start-period=60s --retries=5 \
|
||||
CMD curl -f http://localhost:8000/health || exit 1
|
||||
|
||||
ENV DOCKER_MODE=dind
|
||||
|
||||
ENTRYPOINT ["tini", "--"]
|
||||
|
||||
CMD ["sh", "-c", "(dockerd-entrypoint.sh) & (sleep 5 && uv run --no-sync -m container_server.main) & wait -n"]
|
||||
|
||||
|
||||
+20
@@ -0,0 +1,20 @@
|
||||
[project]
|
||||
name = "container-server"
|
||||
version = "0.1.0"
|
||||
description = "Container Server"
|
||||
requires-python = ">=3.12"
|
||||
dependencies = [
|
||||
"docker",
|
||||
"aiohttp",
|
||||
"pydantic",
|
||||
"pydantic-settings",
|
||||
"fastapi[standard]",
|
||||
"python-dotenv",
|
||||
"websockets",
|
||||
"pyyaml",
|
||||
"mcp>=1.12.2",
|
||||
]
|
||||
|
||||
[build-system]
|
||||
requires = ["setuptools"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
Vendored
Vendored
+17
@@ -0,0 +1,17 @@
|
||||
import os
|
||||
|
||||
debug_mode = os.getenv("DEBUG_MODE", "false").lower() == "true"
|
||||
container_server_port = int(os.getenv("CONTAINER_SERVER_PORT", "9000"))
|
||||
|
||||
docker_registry_url = os.getenv("DOCKER_REGISTRY_URL")
|
||||
docker_registry_user_name = os.getenv("DOCKER_REGISTRY_USER_NAME")
|
||||
docker_registry_password = os.getenv("DOCKER_REGISTRY_PASSWORD")
|
||||
|
||||
gateway_server_addr = os.getenv("GATEWAY_SERVER_ADDR", "http://mcp-gateway:8000")
|
||||
|
||||
mcp_server_image_id = os.getenv(
|
||||
"VIRTUALPC_MCP_SERVER_IMAGE_ID",
|
||||
"aworld-registry-registry-vpc.ap-southeast-1.cr.aliyuncs.com/aworld/mcp-server",
|
||||
)
|
||||
|
||||
docker_mode = os.getenv("DOCKER_MODE", "dind")
|
||||
+231
@@ -0,0 +1,231 @@
|
||||
from functools import cache
|
||||
import socket
|
||||
import logging
|
||||
import traceback
|
||||
import asyncio
|
||||
import uuid
|
||||
import httpx
|
||||
import threading
|
||||
|
||||
from .dockers import docker_helper
|
||||
|
||||
from .configs import (
|
||||
container_server_port,
|
||||
docker_registry_url,
|
||||
docker_registry_user_name,
|
||||
docker_registry_password,
|
||||
mcp_server_image_id,
|
||||
gateway_server_addr,
|
||||
docker_mode,
|
||||
debug_mode,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
token = str(uuid.uuid4())
|
||||
|
||||
|
||||
async def wait_docker_ready(timeout: int = 30):
|
||||
for i in range(timeout):
|
||||
try:
|
||||
cmd = ["docker", "ps"]
|
||||
p = await asyncio.subprocess.create_subprocess_exec(
|
||||
*cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT
|
||||
)
|
||||
stdout, _ = await p.communicate()
|
||||
if p.returncode == 0:
|
||||
logger.info(f"Docker daemon is ready! \n{stdout.decode()}")
|
||||
return
|
||||
else:
|
||||
logger.warning(f"Docker daemon is not ready! {stdout.decode()}")
|
||||
except:
|
||||
logger.error(f"Docker daemon is not ready! {traceback.format_exc()}")
|
||||
await asyncio.sleep(1)
|
||||
else:
|
||||
logger.error(f"Docker daemon is not ready after {timeout} seconds!")
|
||||
raise Exception(f"Docker daemon is not ready after {timeout} seconds!")
|
||||
|
||||
|
||||
async def start_container_server_register_task():
|
||||
register_url = f"{gateway_server_addr}/api/container_server/register"
|
||||
local_ip_addr = get_local_ip()
|
||||
|
||||
async def register():
|
||||
try:
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.post(
|
||||
register_url,
|
||||
json={
|
||||
"token": token,
|
||||
"ip_addr": local_ip_addr,
|
||||
"port": container_server_port,
|
||||
"cpu_load": [],
|
||||
"memory_usage": [],
|
||||
},
|
||||
timeout=10.0,
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Register container server failed: {register_url}, {e}")
|
||||
raise
|
||||
|
||||
async def init_register():
|
||||
for _ in range(20):
|
||||
try:
|
||||
await register()
|
||||
logger.info(f"Init register success")
|
||||
break
|
||||
except:
|
||||
await asyncio.sleep(3)
|
||||
else:
|
||||
logger.error(f"Init register failed after 20 times!")
|
||||
raise Exception(f"Init register failed after 20 times!")
|
||||
|
||||
await init_register()
|
||||
|
||||
async def update_register():
|
||||
while True:
|
||||
try:
|
||||
await register()
|
||||
except:
|
||||
logger.error(f"Update register failed: {traceback.format_exc()}")
|
||||
await asyncio.sleep(10)
|
||||
|
||||
asyncio.create_task(update_register())
|
||||
|
||||
|
||||
async def load_mcp_server_image():
|
||||
if docker_registry_url and docker_registry_user_name and docker_registry_password:
|
||||
await docker_helper.login_async(
|
||||
registry_url=docker_registry_url,
|
||||
username=docker_registry_user_name,
|
||||
password=docker_registry_password,
|
||||
)
|
||||
if not debug_mode:
|
||||
await docker_helper.pull_async(mcp_server_image_id)
|
||||
|
||||
|
||||
async def start_mcp_server_life_cycle_manager():
|
||||
pass
|
||||
|
||||
|
||||
async def clean_mcp_server_container():
|
||||
pass
|
||||
|
||||
|
||||
async def create_mcp_server_container():
|
||||
mcp_port = docker_helper.get_available_port()
|
||||
novnc_port = docker_helper.get_available_port()
|
||||
ip_addr = get_local_ip()
|
||||
|
||||
container_name = f"mcp_server_{str(uuid.uuid4()).replace('-', '')}"
|
||||
logger.info(
|
||||
f"Create mcp server container: {container_name}, image_id: {mcp_server_image_id}, mcp_port: {mcp_port}, novnc_port: {novnc_port}"
|
||||
)
|
||||
try:
|
||||
ports = {4242: f"{mcp_port}", 5901: f"{novnc_port}"}
|
||||
|
||||
network = ""
|
||||
if docker_mode == "host":
|
||||
network = "visualvirtualpc_virtualpc-network"
|
||||
|
||||
container = await docker_helper.run_async(
|
||||
image_id=mcp_server_image_id,
|
||||
container_name=container_name,
|
||||
ports=ports,
|
||||
network=network,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"Create mcp server container success, waiting for Ready: {container.id}, ip_addr: {ip_addr}, mcp_port: {mcp_port}, novnc_port: {novnc_port}"
|
||||
)
|
||||
|
||||
def tail_logs():
|
||||
logs = container.logs(stream=True, tail=100, follow=True)
|
||||
try:
|
||||
buffer = []
|
||||
for line in logs:
|
||||
buffer.append(line.decode())
|
||||
if len(buffer) >= 20:
|
||||
logger.info(f"VPC[{container.name}] >>> \n{'> '.join(buffer)}\n")
|
||||
buffer.clear()
|
||||
if buffer:
|
||||
logger.info(f"VPC[{container.name}] >>> \n{'> '.join(buffer)}\n")
|
||||
logger.info(f"VPC [{container.name}] logs end!")
|
||||
except Exception as e:
|
||||
logger.error(f"Error in tail_logs: {e}")
|
||||
|
||||
# Start log tailing in background thread
|
||||
log_thread = threading.Thread(
|
||||
target=tail_logs, name=f"VPC_{container.name}_logs", daemon=True
|
||||
)
|
||||
log_thread.start()
|
||||
|
||||
async def health_check(timeout: float = 3.0):
|
||||
try:
|
||||
# async with httpx.AsyncClient() as client:
|
||||
# response = await client.get(
|
||||
# f"http://{ip_addr}:{mcp_port}/health",
|
||||
# timeout=httpx.Timeout(timeout),
|
||||
# )
|
||||
# response.raise_for_status()
|
||||
# return True
|
||||
return await docker_helper.check_health(container.id)
|
||||
except Exception as e:
|
||||
logger.error(f"Check mcp server health error! {e}")
|
||||
return False
|
||||
|
||||
max_check = 30
|
||||
for i in range(max_check):
|
||||
if await health_check():
|
||||
logger.info(
|
||||
f"MCP server {ip_addr}:{mcp_port} is ready: {i+1}/{max_check}"
|
||||
)
|
||||
break
|
||||
else:
|
||||
logger.warning(
|
||||
f"MCP server {ip_addr}:{mcp_port} is not ready: {i+1}/{max_check}"
|
||||
)
|
||||
await asyncio.sleep(3)
|
||||
else:
|
||||
logger.error(
|
||||
f"MCP server {ip_addr}:{mcp_port} is not ready after {max_check} times!"
|
||||
)
|
||||
raise Exception(
|
||||
f"MCP server {ip_addr}:{mcp_port} is not ready after {max_check} times!"
|
||||
)
|
||||
|
||||
if docker_mode == "host":
|
||||
ip_addr = await docker_helper.get_container_ip(container.id)
|
||||
mcp_port = 4242
|
||||
novnc_port = 5901
|
||||
|
||||
return container.id, ip_addr, mcp_port, novnc_port
|
||||
except:
|
||||
logger.error(f"Create mcp server container failed: {traceback.format_exc()}")
|
||||
raise
|
||||
|
||||
|
||||
async def shutdown_mcp_server_container(container_id: str):
|
||||
try:
|
||||
await docker_helper.stop_async(container_id)
|
||||
except:
|
||||
logger.error(f"Shutdown mcp server container failed: {traceback.format_exc()}")
|
||||
raise
|
||||
|
||||
|
||||
@cache
|
||||
def get_local_ip() -> str | None:
|
||||
try:
|
||||
host_name = socket.gethostname()
|
||||
_, _, ip_list = socket.gethostbyname_ex(host_name)
|
||||
for ip in ip_list:
|
||||
if not ip.startswith("127."):
|
||||
return ip
|
||||
except Exception as e:
|
||||
logger.error(f"Get local ip failed: {traceback.format_exc()}")
|
||||
|
||||
raise RuntimeError("Get local ip failed")
|
||||
+229
@@ -0,0 +1,229 @@
|
||||
import asyncio
|
||||
import traceback
|
||||
from typing import Tuple
|
||||
import docker
|
||||
import logging
|
||||
import socket
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
client = docker.from_env(timeout=600)
|
||||
|
||||
|
||||
async def login_async(registry_url: str, username: str, password: str):
|
||||
try:
|
||||
return await asyncio.to_thread(login, registry_url, username, password)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in login_async: {e}")
|
||||
raise
|
||||
|
||||
|
||||
async def pull_async(image_id: str):
|
||||
try:
|
||||
return await asyncio.to_thread(pull, image_id)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in pull_async: {e}")
|
||||
raise
|
||||
|
||||
|
||||
async def run_async(
|
||||
image_id: str,
|
||||
container_name: str,
|
||||
ports: dict[int, str] = {},
|
||||
network: str = "",
|
||||
volumes: dict[str, str] = {},
|
||||
environments: dict[str, str] = {},
|
||||
):
|
||||
try:
|
||||
return await asyncio.to_thread(
|
||||
run, image_id, container_name, ports, network, volumes, environments
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in run_async: {e}")
|
||||
raise
|
||||
|
||||
|
||||
async def exec_async(container_id: str, cmd: list[str]) -> Tuple[int, str]:
|
||||
try:
|
||||
return await asyncio.to_thread(exec, container_id, cmd)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in exec_async: {e}")
|
||||
raise
|
||||
|
||||
|
||||
async def stop_async(container_id: str):
|
||||
try:
|
||||
return await asyncio.to_thread(stop, container_id)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in stop_async: {e}")
|
||||
raise
|
||||
|
||||
|
||||
async def check_health(container_id: str):
|
||||
try:
|
||||
container = client.containers.get(container_id)
|
||||
container.reload()
|
||||
return container.health == "healthy"
|
||||
except Exception as e:
|
||||
logger.error(f"Error in check_health: {e}")
|
||||
return False
|
||||
|
||||
|
||||
async def get_container_ip(container_id: str):
|
||||
try:
|
||||
container = client.containers.get(container_id)
|
||||
container.reload()
|
||||
nets = container.attrs["NetworkSettings"]["Networks"]
|
||||
net = list(nets.values())[0]
|
||||
return net["IPAddress"]
|
||||
except Exception as e:
|
||||
logger.error(f"Error in get_container_ip: {e}")
|
||||
return None
|
||||
|
||||
|
||||
async def build_image_async(image_id: str, dockerfile: str, context_path: str):
|
||||
async def build():
|
||||
try:
|
||||
cmd = [
|
||||
"docker",
|
||||
"build",
|
||||
"--platform",
|
||||
"linux/amd64",
|
||||
"-t",
|
||||
image_id,
|
||||
"-f",
|
||||
dockerfile,
|
||||
context_path,
|
||||
]
|
||||
p = await asyncio.subprocess.create_subprocess_exec(
|
||||
*cmd,
|
||||
cwd=context_path,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.STDOUT,
|
||||
)
|
||||
|
||||
assert p.stdout is not None
|
||||
while True:
|
||||
line = await p.stdout.readline()
|
||||
if not line:
|
||||
break
|
||||
logger.info(line.decode(errors="ignore").rstrip())
|
||||
|
||||
rc = await p.wait()
|
||||
if rc == 0:
|
||||
logger.info("Build mcp server image success!")
|
||||
return
|
||||
else:
|
||||
raise Exception(f"Build mcp server image error! return code: {rc}")
|
||||
except Exception as e:
|
||||
logger.error(f"Build mcp server image error! {e}")
|
||||
raise
|
||||
|
||||
for _ in range(3):
|
||||
try:
|
||||
await build()
|
||||
break
|
||||
except Exception as e:
|
||||
await asyncio.sleep(1)
|
||||
else:
|
||||
logger.info(f"Build mcp server image failed after 3 times!")
|
||||
raise Exception("Build mcp server image failed after 3 times!")
|
||||
|
||||
|
||||
def pull(image_id: str):
|
||||
logger.info(f"Pulling image {image_id}")
|
||||
try:
|
||||
if client.images.get(image_id):
|
||||
logger.info(f"Image {image_id} already exists, skipping pull")
|
||||
return
|
||||
img = client.images.pull(image_id)
|
||||
logger.info(f"Pulled image {image_id}, {img}")
|
||||
except:
|
||||
logger.error(f"Failed to pull image {image_id}\n{traceback.format_exc()}")
|
||||
raise
|
||||
|
||||
|
||||
def run(
|
||||
image_id: str,
|
||||
container_name: str,
|
||||
ports: dict[int, str] = {},
|
||||
network: str = "",
|
||||
volumes: dict[str, str] = {},
|
||||
environments: dict[str, str] = {},
|
||||
):
|
||||
logger.info(
|
||||
f"Creating container {container_name} with args: {{'name': {container_name}, 'image': {image_id}, 'ports': {ports}, 'volumes': {volumes}, 'environments': {environments}}}"
|
||||
)
|
||||
try:
|
||||
container = client.containers.run(
|
||||
name=container_name,
|
||||
image=image_id,
|
||||
detach=True,
|
||||
auto_remove=True,
|
||||
ports=ports,
|
||||
network=network,
|
||||
volumes=volumes,
|
||||
environment=environments,
|
||||
cpu_period=100000,
|
||||
cpu_quota=90000,
|
||||
mem_limit="2G",
|
||||
)
|
||||
logger.info(f"Created container {container_name} response: {container}")
|
||||
return container
|
||||
except:
|
||||
logger.error(
|
||||
f"Failed to create container {container_name} with image {image_id}\n{traceback.format_exc()}"
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
def exec(container_id: str, cmd: list[str]) -> Tuple[int, str]:
|
||||
logger.info(f"Executing command {cmd} on container {container_id}")
|
||||
try:
|
||||
container = client.containers.get(container_id)
|
||||
exit_code, output = container.exec_run(cmd)
|
||||
logger.info(
|
||||
f"Command {cmd} executed on container {container_id} with result: {exit_code} {output}"
|
||||
)
|
||||
return exit_code, output
|
||||
except:
|
||||
logger.error(
|
||||
f"Failed to execute command {cmd} on container {container_id}\n{traceback.format_exc()}"
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
def stop(container_id: str):
|
||||
logger.info(f"Stop container {container_id}")
|
||||
try:
|
||||
container = client.containers.get(container_id)
|
||||
container.stop()
|
||||
logger.info(f"Stopped container {container_id}")
|
||||
except:
|
||||
logger.error(
|
||||
f"Failed to stop container {container_id}\n{traceback.format_exc()}"
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
def login(registry_url: str, username: str, password: str):
|
||||
logger.info(f"Logging in to {registry_url} with username {username}")
|
||||
try:
|
||||
result = client.login(
|
||||
registry=registry_url, username=username, password=password
|
||||
)
|
||||
logger.info(
|
||||
f"Logged in to {registry_url} with username {username} result: {result}"
|
||||
)
|
||||
except:
|
||||
logger.error(
|
||||
f"Failed to login to {registry_url} with username {username}\n{traceback.format_exc()}"
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
def get_available_port() -> int:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
port = s.getsockname()[1]
|
||||
return port
|
||||
Vendored
+58
@@ -0,0 +1,58 @@
|
||||
from contextlib import asynccontextmanager
|
||||
from fastapi import FastAPI, Request
|
||||
import logging
|
||||
|
||||
from .configs import container_server_port
|
||||
from .routers import api_server
|
||||
from . import container_server_manager
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
||||
)
|
||||
|
||||
# Suppress httpx INFO logs
|
||||
logging.getLogger("httpx").setLevel(logging.WARNING)
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
"""FastAPI lifespan context manager for startup and shutdown events"""
|
||||
# Startup
|
||||
await container_server_manager.wait_docker_ready()
|
||||
await container_server_manager.load_mcp_server_image()
|
||||
await container_server_manager.start_mcp_server_life_cycle_manager()
|
||||
await container_server_manager.start_container_server_register_task()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
# Shutdown
|
||||
# await mcp_server_manager.clean_mcp_server_container()
|
||||
pass
|
||||
|
||||
|
||||
# FastAPI application setup with lifespan
|
||||
app = FastAPI(
|
||||
title="MCP Container Server",
|
||||
description="MCP Container Server",
|
||||
version="1.0.0",
|
||||
lifespan=lifespan,
|
||||
)
|
||||
|
||||
|
||||
app.include_router(api_server.router, prefix="/api")
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
async def health(request: Request):
|
||||
return {"status": "success", "message": "Container server is healthy"}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
|
||||
# Run the server
|
||||
uvicorn.run(app, host="0.0.0.0", port=container_server_port)
|
||||
+45
@@ -0,0 +1,45 @@
|
||||
import logging
|
||||
from fastapi import APIRouter, Request
|
||||
|
||||
from .. import container_server_manager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.post("/container/create")
|
||||
async def create_container(request: Request, body: dict):
|
||||
token = body.get("token")
|
||||
logger.info(f"Create container: token={token}")
|
||||
container_id, ip_addr, mcp_port, novnc_port = (
|
||||
await container_server_manager.create_mcp_server_container()
|
||||
)
|
||||
logger.info(
|
||||
f"Container created: container_id={container_id}, ip_addr={ip_addr}, mcp_port={mcp_port}, novnc_port={novnc_port}"
|
||||
)
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"message": f"MCP server created: {ip_addr}:{mcp_port}",
|
||||
"data": {
|
||||
"ip_addr": ip_addr,
|
||||
"mcp_port": mcp_port,
|
||||
"novnc_port": novnc_port,
|
||||
"container_id": container_id,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@router.post("/container/shutdown")
|
||||
async def shutdown_container(request: Request, body: dict):
|
||||
token = body.get("token")
|
||||
container_id = body.get("container_id")
|
||||
logger.info(f"Shutdown container: token={token}, container_id={container_id}")
|
||||
await container_server_manager.shutdown_mcp_server_container(container_id)
|
||||
logger.info(f"Container shutdown: container_id={container_id}")
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"message": f"MCP server shutdown: {container_id}",
|
||||
}
|
||||
+1320
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,50 @@
|
||||
services:
|
||||
mcp-gateway:
|
||||
build:
|
||||
context: mcp_gateway
|
||||
dockerfile: Dockerfile_mcp_gateway
|
||||
platform: linux/amd64
|
||||
ports:
|
||||
- 8000:8000
|
||||
environment:
|
||||
MCP_GATEWAY_TOKEN_SECRET: "${MCP_GATEWAY_TOKEN_SECRET:-123321}"
|
||||
VNC_AUTH: false
|
||||
MCP_GATEWAY_REDIS_URL: redis://redis-server:6379/0
|
||||
restart: on-failure:3
|
||||
depends_on:
|
||||
- redis-server
|
||||
deploy:
|
||||
resources:
|
||||
limits:
|
||||
memory: 2G
|
||||
|
||||
container-server:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: container_server/Dockerfile_container_server
|
||||
platform: linux/amd64
|
||||
restart: on-failure:3
|
||||
privileged: true
|
||||
volumes:
|
||||
- /var/run/docker.sock:/var/run/docker.sock
|
||||
command: ["uv", "run", "--no-sync", "-m", "container_server.main"]
|
||||
environment:
|
||||
VIRTUALPC_MCP_SERVER_IMAGE_ID: gaia-mcp-server
|
||||
GATEWAY_SERVER_ADDR: http://mcp-gateway:8000
|
||||
deploy:
|
||||
resources:
|
||||
limits:
|
||||
memory: 8G
|
||||
|
||||
redis-server:
|
||||
image: redis:7.2-alpine
|
||||
platform: linux/amd64
|
||||
ports:
|
||||
- 6379:6379
|
||||
volumes:
|
||||
- ./.data/redis-data:/data
|
||||
restart: on-failure:3
|
||||
deploy:
|
||||
resources:
|
||||
limits:
|
||||
memory: 2G
|
||||
+39
@@ -0,0 +1,39 @@
|
||||
FROM python:3.13-slim-bookworm
|
||||
|
||||
# Avoid prompts from apt
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
# Configure mirror
|
||||
ARG ENABLE_MIRROR=false
|
||||
ENV PIP_INDEX_URL=${ENABLE_MIRROR:+https://mirrors.aliyun.com/pypi/simple}
|
||||
RUN if [ "$ENABLE_MIRROR" = "true" ]; then \
|
||||
rm -rfv /etc/apt/sources.list.d/* && \
|
||||
echo "deb http://ftp.cn.debian.org/debian bookworm main contrib non-free non-free-firmware" > /etc/apt/sources.list && \
|
||||
echo "deb http://ftp.cn.debian.org/debian bookworm-updates main contrib non-free non-free-firmware" >> /etc/apt/sources.list && \
|
||||
echo "deb http://ftp.cn.debian.org/debian bookworm-backports main contrib non-free non-free-firmware" >> /etc/apt/sources.list && \
|
||||
echo "deb http://ftp.cn.debian.org/debian-security bookworm-security main contrib non-free non-free-firmware" >> /etc/apt/sources.list && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*; \
|
||||
fi
|
||||
|
||||
# Install dependencies
|
||||
RUN apt-get update && apt-get install -y \
|
||||
curl \
|
||||
--no-install-recommends \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
RUN pip install -U --break-system-packages uv
|
||||
|
||||
COPY pyproject.toml /app/mcp_gateway/pyproject.toml
|
||||
COPY src /app/mcp_gateway/src
|
||||
|
||||
WORKDIR /app/mcp_gateway
|
||||
|
||||
RUN uv sync --python-preference=only-system
|
||||
|
||||
EXPOSE 8000
|
||||
|
||||
HEALTHCHECK --interval=30s --timeout=10s --start-period=10s --retries=5 \
|
||||
CMD curl -f http://localhost:8000/health || exit 1
|
||||
|
||||
CMD ["uv", "run", "--no-sync", "-m", "mcp_gateway.main"]
|
||||
@@ -0,0 +1,28 @@
|
||||
[project]
|
||||
name = "mcp-gateway-server"
|
||||
version = "0.1.0"
|
||||
description = "MCP Gateway Server"
|
||||
requires-python = ">=3.12"
|
||||
dependencies = [
|
||||
"docker",
|
||||
"aiohttp",
|
||||
"pydantic",
|
||||
"pydantic-settings",
|
||||
"fastapi[standard]",
|
||||
"typer",
|
||||
"aiofiles",
|
||||
"python-dotenv",
|
||||
"websockets",
|
||||
"pyyaml",
|
||||
"mcp>=1.12.2",
|
||||
"httpx[http2]",
|
||||
"pyjwt>=2.10.1",
|
||||
"redis>=6.4.0",
|
||||
]
|
||||
|
||||
[build-system]
|
||||
requires = ["setuptools"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project.scripts]
|
||||
mcp-gateway = "mcp_gateway.main:app"
|
||||
Vendored
Vendored
+43
@@ -0,0 +1,43 @@
|
||||
import logging
|
||||
import traceback
|
||||
from fastapi import Request
|
||||
import jwt
|
||||
|
||||
from ..utils.common_utils import get_remote_addr
|
||||
|
||||
from ..configs import token_secret
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def check_auth(request: Request) -> bool:
|
||||
payload = get_auth_payload(dict(request.headers))
|
||||
logger.info(
|
||||
f"Gateway auth: remote.addr={get_remote_addr(request)}, payload={payload}"
|
||||
)
|
||||
return payload is not None
|
||||
|
||||
|
||||
def get_auth_payload(headers: dict) -> str | None:
|
||||
"""Check if the request is authorized"""
|
||||
try:
|
||||
token = headers.get("Authorization")
|
||||
if not token:
|
||||
return None
|
||||
|
||||
token = token[len("Bearer ") :]
|
||||
if not token:
|
||||
return None
|
||||
payload = decode_token(token)
|
||||
return payload
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Failed to check auth, remote.addr={get_remote_addr(request)}, request.headers={request.headers} \n{traceback.format_exc()}"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def decode_token(token: str) -> str:
|
||||
"""Decode the token"""
|
||||
payload = jwt.decode(token, token_secret, algorithms=["HS256"])
|
||||
return payload
|
||||
+12
@@ -0,0 +1,12 @@
|
||||
import os
|
||||
|
||||
cluster_name = os.getenv("CLUSTER_NAME", "mcp_gateway")
|
||||
|
||||
debug_mode = os.getenv("DEBUG_MODE", "false").lower() == "true"
|
||||
|
||||
vnc_auth = os.getenv("VNC_AUTH", "true").lower() == "true"
|
||||
|
||||
redis_url = os.getenv("MCP_GATEWAY_REDIS_URL", "redis://localhost:6379/0")
|
||||
|
||||
token_secret = os.getenv("MCP_GATEWAY_TOKEN_SECRET")
|
||||
assert token_secret, "MCP_GATEWAY_TOKEN_SECRET is not set"
|
||||
chapter9/gaia-experience/AWorld/env/virtualpc-mcp/mcp_gateway/src/mcp_gateway/containers/__init__.py
Vendored
+1
@@ -0,0 +1 @@
|
||||
from .container_server_manager import ContainerServer, ContainerServerManager, container_server_manager_builder
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user