AWS Documentation Agent - AWS 文档查询工具
功能: - search_aws_documentation: 搜索 AWS 文档 - get_aws_doc: 获取完整文档内容(Markdown) - recommend_aws_content: 获取推荐内容 - get_aws_services_list: 获取服务列表 特性: - 支持 AWS 文档网站直接访问 - 集成 Pydantic AI 进行结果整理和解释 - 提供 REST API 和 MCP 协议端点 - 支持 API Key 验证 - 支持 AWS 全球分区和中国分区
This commit is contained in:
@@ -0,0 +1,4 @@
|
||||
__pycache__/
|
||||
*.pyc
|
||||
*.pyo
|
||||
.env
|
||||
+20
@@ -0,0 +1,20 @@
|
||||
FROM python:3.12-slim
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
ENV PYTHONDONTWRITEBYTECODE=1
|
||||
|
||||
RUN apt-get update && apt-get install -y gcc curl && rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY requirements.txt .
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
COPY . .
|
||||
|
||||
EXPOSE 8000
|
||||
|
||||
HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \
|
||||
CMD curl -f http://localhost:8000/health || exit 1
|
||||
|
||||
CMD ["python", "run_api_server.py"]
|
||||
@@ -0,0 +1,208 @@
|
||||
# AWS Documentation Agent
|
||||
|
||||
基于 **Pydantic AI** 和 **FastMCP** 的 AWS 文档查询 Agent。
|
||||
|
||||
## 功能概述
|
||||
|
||||
AWS Documentation Agent 提供以下功能:
|
||||
|
||||
- **搜索 AWS 文档** - 在 AWS 官方文档中搜索服务指南、API 参考、教程等
|
||||
- **获取文档内容** - 根据 URL 获取完整的文档内容(Markdown 格式)
|
||||
- **获取推荐内容** - 根据文档获取相关推荐内容
|
||||
- **获取服务列表** - 查看 AWS 服务列表
|
||||
|
||||
## 快速开始
|
||||
|
||||
### 1. 环境变量配置
|
||||
|
||||
```bash
|
||||
# LiteLLM Gateway 配置
|
||||
export LITELLM_GATEWAY_URL="https://litellm.graystone-fb459c5d.southeastasia.azurecontainerapps.io/v1"
|
||||
export OPENAI_API_KEY="your-api-key"
|
||||
|
||||
# AWS 文档分区(可选,默认 aws)
|
||||
export AWS_DOCUMENTATION_PARTITION="aws" # 或 "aws-cn" 用于中国区域
|
||||
|
||||
# 服务端口(可选,默认 8000)
|
||||
export API_PORT=8000
|
||||
```
|
||||
|
||||
### 2. 本地测试
|
||||
|
||||
```bash
|
||||
cd aws-documentation
|
||||
pip install -r requirements.txt
|
||||
python run_api_server.py
|
||||
```
|
||||
|
||||
### 3. 构建镜像
|
||||
|
||||
```bash
|
||||
docker build -t aws-documentation:latest .
|
||||
```
|
||||
|
||||
### 4. 运行容器
|
||||
|
||||
```bash
|
||||
docker run -d -p 8000:8000 \
|
||||
-e LITELLM_GATEWAY_URL="https://..." \
|
||||
-e OPENAI_API_KEY="your-key" \
|
||||
-e AWS_DOCUMENTATION_PARTITION="aws" \
|
||||
aws-documentation:latest
|
||||
```
|
||||
|
||||
## 项目结构
|
||||
|
||||
```
|
||||
aws-documentation/
|
||||
├── Dockerfile
|
||||
├── requirements.txt
|
||||
├── run_api_server.py # 启动脚本
|
||||
└── src/
|
||||
├── __init__.py
|
||||
└── server/
|
||||
├── __init__.py
|
||||
├── api_server.py # FastAPI + MCP HTTP
|
||||
└── mcp_server.py # MCP 工具定义
|
||||
```
|
||||
|
||||
## API 端点
|
||||
|
||||
### REST API
|
||||
|
||||
- `GET /` - 服务信息
|
||||
- `GET /health` - 健康检查
|
||||
- `POST /api/v1/search` - 搜索文档(需要 API Key)
|
||||
- `POST /api/v1/doc` - 获取文档内容(需要 API Key)
|
||||
- `POST /api/v1/recommend` - 获取推荐内容(需要 API Key)
|
||||
- `GET /api/v1/services` - 获取服务列表(需要 API Key)
|
||||
|
||||
### MCP 端点
|
||||
|
||||
- `POST /mcp` - MCP HTTP 端点
|
||||
- `GET /mcp/sse` - MCP SSE 端点
|
||||
- `POST /mcp/sse` - MCP SSE POST 端点
|
||||
|
||||
## 使用示例
|
||||
|
||||
### 搜索文档
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/api/v1/search \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "api-key: your-api-key" \
|
||||
-d '{
|
||||
"query": "S3 bucket naming",
|
||||
"limit": 5
|
||||
}'
|
||||
```
|
||||
|
||||
### 获取文档内容
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/api/v1/doc \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "api-key: your-api-key" \
|
||||
-d '{
|
||||
"doc_url": "https://docs.aws.amazon.com/AmazonS3/latest/userguide/bucketnamingrules.html"
|
||||
}'
|
||||
```
|
||||
|
||||
### 获取推荐内容
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/api/v1/recommend \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "api-key: your-api-key" \
|
||||
-d '{
|
||||
"doc_url": "https://docs.aws.amazon.com/AmazonS3/latest/userguide/bucketnamingrules.html"
|
||||
}'
|
||||
```
|
||||
|
||||
### MCP 调用
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/mcp \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "api-key: your-api-key" \
|
||||
-d '{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "tools/call",
|
||||
"params": {
|
||||
"name": "search_aws_documentation",
|
||||
"arguments": {
|
||||
"query": "Lambda Python",
|
||||
"limit": 10
|
||||
}
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
## 环境变量
|
||||
|
||||
| 变量 | 必需 | 说明 |
|
||||
|------|------|------|
|
||||
| LITELLM_GATEWAY_URL | 是 | LiteLLM Gateway URL |
|
||||
| OPENAI_API_KEY | 是 | LLM API Key |
|
||||
| AWS_DOCUMENTATION_PARTITION | 否 | AWS 分区,默认 aws(可选:aws-cn 用于中国区域) |
|
||||
| API_PORT | 否 | 服务端口,默认 8000 |
|
||||
|
||||
## 工具说明
|
||||
|
||||
### search_aws_documentation
|
||||
|
||||
搜索 AWS 文档。
|
||||
|
||||
**参数:**
|
||||
- `query` (string, 必需): 搜索关键词
|
||||
- `limit` (integer, 可选): 最大返回结果数,默认 10
|
||||
|
||||
**示例:**
|
||||
```json
|
||||
{
|
||||
"query": "S3 bucket",
|
||||
"limit": 5
|
||||
}
|
||||
```
|
||||
|
||||
### get_aws_doc
|
||||
|
||||
获取 AWS 文档内容。
|
||||
|
||||
**参数:**
|
||||
- `doc_url` (string, 必需): AWS 文档 URL
|
||||
|
||||
**示例:**
|
||||
```json
|
||||
{
|
||||
"doc_url": "https://docs.aws.amazon.com/AmazonS3/latest/userguide/..."
|
||||
}
|
||||
```
|
||||
|
||||
### recommend_aws_content
|
||||
|
||||
获取 AWS 文档推荐内容。
|
||||
|
||||
**参数:**
|
||||
- `doc_url` (string, 必需): AWS 文档 URL
|
||||
|
||||
### get_aws_services_list
|
||||
|
||||
获取 AWS 服务列表。
|
||||
|
||||
**参数:** 无
|
||||
|
||||
## 注册到 Agent Manager
|
||||
|
||||
在 `k8s_manager.py` 中添加:
|
||||
|
||||
```python
|
||||
# TEMPLATE_PORTS
|
||||
"aws-documentation": 8000,
|
||||
|
||||
# image_map
|
||||
"aws-documentation": "agnettaiji.azurecr.io/ai-agents/aws-documentation:latest",
|
||||
```
|
||||
|
||||
在 `app.py` 的 `valid_templates` 中添加 `"aws-documentation"`。
|
||||
@@ -0,0 +1,13 @@
|
||||
# Pydantic AI
|
||||
pydantic-ai>=0.0.14
|
||||
|
||||
# MCP
|
||||
mcp>=0.9.0
|
||||
fastmcp>=0.1.0
|
||||
|
||||
# FastAPI
|
||||
fastapi>=0.109.0
|
||||
uvicorn[standard]>=0.27.0
|
||||
|
||||
# HTTP Client
|
||||
aiohttp>=3.9.0
|
||||
@@ -0,0 +1,17 @@
|
||||
#!/usr/bin/env python
|
||||
"""启动 API 服务器"""
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).parent))
|
||||
|
||||
if __name__ == '__main__':
|
||||
from src.server.api_server import app
|
||||
import uvicorn
|
||||
import os
|
||||
|
||||
host = os.getenv('API_HOST', '0.0.0.0')
|
||||
port = int(os.getenv('API_PORT', '8000'))
|
||||
|
||||
print(f"🚀 启动 Agent API: http://{host}:{port}")
|
||||
uvicorn.run(app, host=host, port=port, log_level="info")
|
||||
@@ -0,0 +1 @@
|
||||
"""Agent 源代码包"""
|
||||
@@ -0,0 +1 @@
|
||||
"""服务器模块"""
|
||||
@@ -0,0 +1,295 @@
|
||||
"""
|
||||
HTTP API 服务器
|
||||
|
||||
提供 REST API 和 MCP HTTP/SSE 端点。
|
||||
"""
|
||||
import json
|
||||
import uuid
|
||||
import os
|
||||
from typing import Optional, Dict, Any, AsyncGenerator
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from fastapi import FastAPI, HTTPException, Request, Header, Depends
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import StreamingResponse, JSONResponse
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from .mcp_server import TOOL_MAP, TOOL_LIST
|
||||
|
||||
# ==================== 配置 ====================
|
||||
|
||||
SERVER_NAME = "AWS Documentation API"
|
||||
|
||||
|
||||
# ==================== FastAPI 应用 ====================
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
print(f"🚀 {SERVER_NAME} 启动")
|
||||
yield
|
||||
print(f"🛑 {SERVER_NAME} 关闭")
|
||||
|
||||
app = FastAPI(
|
||||
title=SERVER_NAME,
|
||||
version="1.0.0",
|
||||
lifespan=lifespan
|
||||
)
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
|
||||
# ==================== API Key 验证 ====================
|
||||
|
||||
async def verify_api_key(
|
||||
api_key: Optional[str] = Header(None, alias="api-key"),
|
||||
authorization: Optional[str] = Header(None)
|
||||
) -> str:
|
||||
"""验证 API Key"""
|
||||
if api_key and api_key.strip() and api_key.strip() != "sk":
|
||||
return api_key.strip()
|
||||
|
||||
if authorization:
|
||||
key = authorization[7:].strip() if authorization.startswith("Bearer ") else authorization.strip()
|
||||
if key and key != "sk":
|
||||
return key
|
||||
|
||||
raise HTTPException(status_code=401, detail="缺少 API Key")
|
||||
|
||||
|
||||
def get_api_key_from_request(request: Request) -> Optional[str]:
|
||||
"""从请求头提取 API Key(不验证)"""
|
||||
api_key = request.headers.get("api-key") or request.headers.get("api_key")
|
||||
if not api_key:
|
||||
auth = request.headers.get("Authorization")
|
||||
if auth:
|
||||
api_key = auth[7:] if auth.startswith("Bearer ") else auth
|
||||
return api_key
|
||||
|
||||
|
||||
# ==================== 健康检查 ====================
|
||||
|
||||
@app.get("/")
|
||||
async def root():
|
||||
return {
|
||||
"service": SERVER_NAME,
|
||||
"status": "running",
|
||||
"tools": list(TOOL_MAP.keys())
|
||||
}
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
async def health():
|
||||
return {"status": "healthy", "service": SERVER_NAME}
|
||||
|
||||
|
||||
# ==================== MCP 端点 ====================
|
||||
|
||||
sessions: Dict[str, Dict] = {}
|
||||
|
||||
|
||||
async def handle_mcp_request(data: Dict, session_id: str = None, api_key: str = None) -> Dict:
|
||||
"""处理 MCP JSON-RPC 请求"""
|
||||
method = data.get("method")
|
||||
params = data.get("params", {})
|
||||
req_id = data.get("id")
|
||||
|
||||
# tools/call 需要验证 API Key
|
||||
if method == "tools/call" and (not api_key or api_key == "sk"):
|
||||
return {"jsonrpc": "2.0", "id": req_id, "error": {"code": -32001, "message": "缺少 API Key"}}
|
||||
|
||||
try:
|
||||
if method == "initialize":
|
||||
session_id = session_id or str(uuid.uuid4())
|
||||
sessions[session_id] = {"initialized": True}
|
||||
return {
|
||||
"jsonrpc": "2.0", "id": req_id,
|
||||
"result": {
|
||||
"protocolVersion": "2024-11-05",
|
||||
"capabilities": {"tools": {}},
|
||||
"serverInfo": {"name": SERVER_NAME, "version": "1.0.0"}
|
||||
}
|
||||
}
|
||||
|
||||
elif method == "tools/list":
|
||||
return {"jsonrpc": "2.0", "id": req_id, "result": {"tools": TOOL_LIST}}
|
||||
|
||||
elif method == "tools/call":
|
||||
tool_name = params.get("name")
|
||||
args = params.get("arguments", {})
|
||||
|
||||
if tool_name not in TOOL_MAP:
|
||||
raise ValueError(f"Unknown tool: {tool_name}")
|
||||
|
||||
# 设置 API Key 到环境变量
|
||||
old_key = os.environ.get('OPENAI_API_KEY')
|
||||
if api_key:
|
||||
os.environ['OPENAI_API_KEY'] = api_key
|
||||
|
||||
try:
|
||||
result = await TOOL_MAP[tool_name](**args)
|
||||
finally:
|
||||
if old_key:
|
||||
os.environ['OPENAI_API_KEY'] = old_key
|
||||
|
||||
return {
|
||||
"jsonrpc": "2.0", "id": req_id,
|
||||
"result": {"content": [{"type": "text", "text": str(result)}]}
|
||||
}
|
||||
|
||||
elif method == "ping":
|
||||
return {"jsonrpc": "2.0", "id": req_id, "result": {}}
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unknown method: {method}")
|
||||
|
||||
except Exception as e:
|
||||
return {"jsonrpc": "2.0", "id": req_id, "error": {"code": -32603, "message": str(e)}}
|
||||
|
||||
|
||||
@app.post("/mcp")
|
||||
async def mcp_endpoint(request: Request):
|
||||
"""MCP HTTP 端点"""
|
||||
try:
|
||||
body = await request.json()
|
||||
session_id = request.headers.get("x-mcp-session-id")
|
||||
api_key = get_api_key_from_request(request)
|
||||
response = await handle_mcp_request(body, session_id, api_key)
|
||||
return JSONResponse(content=response, headers={"x-mcp-session-id": session_id or ""})
|
||||
except Exception as e:
|
||||
return JSONResponse(status_code=400, content={"jsonrpc": "2.0", "error": {"code": -32700, "message": str(e)}})
|
||||
|
||||
|
||||
@app.get("/mcp/sse")
|
||||
async def mcp_sse(request: Request):
|
||||
"""MCP SSE 端点"""
|
||||
session_id = request.headers.get("x-mcp-session-id") or str(uuid.uuid4())
|
||||
|
||||
async def stream() -> AsyncGenerator[str, None]:
|
||||
yield f"data: {json.dumps({'type': 'connection', 'sessionId': session_id})}\n\n"
|
||||
import asyncio
|
||||
while True:
|
||||
await asyncio.sleep(30)
|
||||
yield f"data: {json.dumps({'type': 'ping'})}\n\n"
|
||||
|
||||
return StreamingResponse(stream(), media_type="text/event-stream",
|
||||
headers={"Cache-Control": "no-cache", "x-mcp-session-id": session_id})
|
||||
|
||||
|
||||
@app.post("/mcp/sse")
|
||||
async def mcp_sse_post(request: Request):
|
||||
"""MCP SSE POST 端点"""
|
||||
try:
|
||||
body = await request.json()
|
||||
session_id = request.headers.get("x-mcp-session-id") or str(uuid.uuid4())
|
||||
api_key = get_api_key_from_request(request)
|
||||
|
||||
async def stream() -> AsyncGenerator[str, None]:
|
||||
response = await handle_mcp_request(body, session_id, api_key)
|
||||
yield f"data: {json.dumps(response)}\n\n"
|
||||
|
||||
return StreamingResponse(stream(), media_type="text/event-stream",
|
||||
headers={"Cache-Control": "no-cache", "x-mcp-session-id": session_id})
|
||||
except Exception as e:
|
||||
return JSONResponse(status_code=400, content={"jsonrpc": "2.0", "error": {"code": -32700, "message": str(e)}})
|
||||
|
||||
|
||||
# ==================== 业务 API ====================
|
||||
|
||||
class SearchRequest(BaseModel):
|
||||
"""搜索请求模型"""
|
||||
query: str = Field(..., description="搜索关键词")
|
||||
limit: Optional[int] = Field(10, description="最大返回结果数")
|
||||
|
||||
|
||||
class DocRequest(BaseModel):
|
||||
"""文档请求模型"""
|
||||
doc_url: str = Field(..., description="AWS 文档 URL")
|
||||
|
||||
|
||||
class QueryResponse(BaseModel):
|
||||
"""响应模型"""
|
||||
success: bool
|
||||
result: Optional[str] = None
|
||||
error: Optional[str] = None
|
||||
|
||||
|
||||
@app.post("/api/v1/search", response_model=QueryResponse)
|
||||
async def api_search(request: SearchRequest, api_key: str = Depends(verify_api_key)):
|
||||
"""搜索 AWS 文档"""
|
||||
try:
|
||||
old_key = os.environ.get('OPENAI_API_KEY')
|
||||
os.environ['OPENAI_API_KEY'] = api_key
|
||||
|
||||
try:
|
||||
result = await TOOL_MAP['search_aws_documentation'](
|
||||
query=request.query,
|
||||
limit=request.limit
|
||||
)
|
||||
return QueryResponse(success=True, result=result)
|
||||
finally:
|
||||
if old_key:
|
||||
os.environ['OPENAI_API_KEY'] = old_key
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@app.post("/api/v1/doc", response_model=QueryResponse)
|
||||
async def api_get_doc(request: DocRequest, api_key: str = Depends(verify_api_key)):
|
||||
"""获取 AWS 文档内容"""
|
||||
try:
|
||||
old_key = os.environ.get('OPENAI_API_KEY')
|
||||
os.environ['OPENAI_API_KEY'] = api_key
|
||||
|
||||
try:
|
||||
result = await TOOL_MAP['get_aws_doc'](doc_url=request.doc_url)
|
||||
return QueryResponse(success=True, result=result)
|
||||
finally:
|
||||
if old_key:
|
||||
os.environ['OPENAI_API_KEY'] = old_key
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@app.post("/api/v1/recommend", response_model=QueryResponse)
|
||||
async def api_recommend(request: DocRequest, api_key: str = Depends(verify_api_key)):
|
||||
"""获取 AWS 文档推荐内容"""
|
||||
try:
|
||||
old_key = os.environ.get('OPENAI_API_KEY')
|
||||
os.environ['OPENAI_API_KEY'] = api_key
|
||||
|
||||
try:
|
||||
result = await TOOL_MAP['recommend_aws_content'](doc_url=request.doc_url)
|
||||
return QueryResponse(success=True, result=result)
|
||||
finally:
|
||||
if old_key:
|
||||
os.environ['OPENAI_API_KEY'] = old_key
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@app.get("/api/v1/services", response_model=QueryResponse)
|
||||
async def api_list_services(api_key: str = Depends(verify_api_key)):
|
||||
"""获取 AWS 服务列表"""
|
||||
try:
|
||||
old_key = os.environ.get('OPENAI_API_KEY')
|
||||
os.environ['OPENAI_API_KEY'] = api_key
|
||||
|
||||
try:
|
||||
result = await TOOL_MAP['get_aws_services_list']()
|
||||
return QueryResponse(success=True, result=result)
|
||||
finally:
|
||||
if old_key:
|
||||
os.environ['OPENAI_API_KEY'] = old_key
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
import uvicorn
|
||||
uvicorn.run(app, host="0.0.0.0", port=8000)
|
||||
@@ -0,0 +1,436 @@
|
||||
"""
|
||||
AWS Documentation MCP 服务器 - AWS 文档查询工具
|
||||
|
||||
使用 Pydantic AI 和 FastMCP 框架,通过 AWS 文档网站查询文档。
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import aiohttp
|
||||
import re
|
||||
from typing import Optional, List
|
||||
from urllib.parse import quote, urljoin, urlparse
|
||||
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
from pydantic_ai import Agent
|
||||
|
||||
# ==================== 配置 ====================
|
||||
|
||||
# LiteLLM Gateway 配置
|
||||
_BASE_URL = os.getenv('OPENAI_BASE_URL',
|
||||
os.getenv('LLM_BASE_URL', 'https://litellm.graystone-fb459c5d.southeastasia.azurecontainerapps.io/v1'))
|
||||
_API_KEY = os.getenv('OPENAI_API_KEY', 'sk')
|
||||
|
||||
os.environ.setdefault('OPENAI_API_KEY', _API_KEY)
|
||||
os.environ.setdefault('OPENAI_BASE_URL', _BASE_URL)
|
||||
|
||||
# 模型名称(pydantic_ai 需要 openai: 前缀)
|
||||
def _get_model_name() -> str:
|
||||
model = os.getenv('MODEL_NAME', os.getenv('LITELLM_MODEL', 'taiji/gpt-4o-mini'))
|
||||
return model if ':' in model else f'openai:{model}'
|
||||
|
||||
MODEL_NAME = _get_model_name()
|
||||
|
||||
# AWS 文档配置
|
||||
AWS_DOCUMENTATION_PARTITION = os.getenv('AWS_DOCUMENTATION_PARTITION', 'aws') # aws 或 aws-cn
|
||||
AWS_DOC_BASE_URL = 'https://docs.aws.amazon.com' if AWS_DOCUMENTATION_PARTITION == 'aws' else 'https://docs.amazonaws.cn'
|
||||
AWS_DOC_SEARCH_URL = f'{AWS_DOC_BASE_URL}/search/api.html'
|
||||
|
||||
# ==================== MCP 服务器 ====================
|
||||
|
||||
server = FastMCP('AWS Documentation')
|
||||
|
||||
# 系统提示词
|
||||
SYSTEM_PROMPT = '''你是一个专业的 AWS 文档查询助手。
|
||||
你可以帮助用户查询 AWS 官方文档、服务指南、API 参考等内容。
|
||||
请根据用户的查询提供准确、详细的 AWS 文档信息。'''
|
||||
|
||||
|
||||
def get_agent() -> Agent:
|
||||
"""创建 Agent 实例(每次调用使用最新的 API Key)"""
|
||||
return Agent(MODEL_NAME, system_prompt=SYSTEM_PROMPT)
|
||||
|
||||
|
||||
# ==================== AWS 文档工具函数 ====================
|
||||
|
||||
async def search_aws_docs_api(query: str, limit: int = 10) -> dict:
|
||||
"""
|
||||
通过 AWS 文档搜索 API 搜索文档
|
||||
|
||||
AWS 文档搜索使用特定的搜索端点
|
||||
"""
|
||||
try:
|
||||
# AWS 文档搜索 API 端点
|
||||
search_url = f"{AWS_DOC_BASE_URL}/search/api.html"
|
||||
|
||||
# 构建搜索参数
|
||||
params = {
|
||||
'q': query,
|
||||
'limit': limit
|
||||
}
|
||||
|
||||
headers = {
|
||||
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36',
|
||||
'Accept': 'application/json'
|
||||
}
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.get(
|
||||
search_url,
|
||||
params=params,
|
||||
headers=headers,
|
||||
timeout=aiohttp.ClientTimeout(total=30)
|
||||
) as resp:
|
||||
if resp.status == 200:
|
||||
try:
|
||||
return await resp.json()
|
||||
except:
|
||||
# 如果不是 JSON,尝试解析 HTML
|
||||
text = await resp.text()
|
||||
return {"html_response": text[:1000]}
|
||||
else:
|
||||
return {"error": f"HTTP {resp.status}", "status": resp.status}
|
||||
|
||||
except Exception as e:
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
async def fetch_aws_doc_page(url: str) -> str:
|
||||
"""
|
||||
获取 AWS 文档页面内容并转换为 Markdown
|
||||
|
||||
从 AWS 文档 URL 获取 HTML 并提取主要内容
|
||||
"""
|
||||
try:
|
||||
# 验证 URL 是 AWS 文档 URL
|
||||
if 'docs.aws.amazon.com' not in url and 'docs.amazonaws.cn' not in url:
|
||||
raise Exception(f"URL 必须是 AWS 文档 URL (docs.aws.amazon.com 或 docs.amazonaws.cn)")
|
||||
|
||||
headers = {
|
||||
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36',
|
||||
'Accept': 'text/html,application/xhtml+xml'
|
||||
}
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.get(url, headers=headers, timeout=aiohttp.ClientTimeout(total=30)) as resp:
|
||||
if resp.status != 200:
|
||||
raise Exception(f"无法获取文档: HTTP {resp.status}")
|
||||
|
||||
html = await resp.text()
|
||||
|
||||
# 简单的 HTML 到 Markdown 转换
|
||||
# 提取主要内容区域
|
||||
import re
|
||||
|
||||
# 提取标题
|
||||
title_match = re.search(r'<title>(.*?)</title>', html, re.IGNORECASE | re.DOTALL)
|
||||
title = title_match.group(1).strip() if title_match else "AWS Documentation"
|
||||
|
||||
# 提取主要内容(通常在 <main> 或 <div id="main-content"> 中)
|
||||
main_match = re.search(r'<main[^>]*>(.*?)</main>', html, re.IGNORECASE | re.DOTALL)
|
||||
if not main_match:
|
||||
main_match = re.search(r'<div[^>]*id=["\']main-content["\'][^>]*>(.*?)</div>', html, re.IGNORECASE | re.DOTALL)
|
||||
|
||||
content = main_match.group(1) if main_match else html
|
||||
|
||||
# 简单的 HTML 标签清理
|
||||
content = re.sub(r'<script[^>]*>.*?</script>', '', content, flags=re.IGNORECASE | re.DOTALL)
|
||||
content = re.sub(r'<style[^>]*>.*?</style>', '', content, flags=re.IGNORECASE | re.DOTALL)
|
||||
content = re.sub(r'<[^>]+>', '', content) # 移除所有 HTML 标签
|
||||
content = re.sub(r'\s+', ' ', content) # 合并空白字符
|
||||
content = content.strip()
|
||||
|
||||
return f"# {title}\n\n{content[:5000]}" # 限制长度
|
||||
|
||||
except Exception as e:
|
||||
raise Exception(f"获取文档失败: {str(e)}")
|
||||
|
||||
|
||||
# ==================== MCP 工具定义 ====================
|
||||
|
||||
@server.tool()
|
||||
async def search_aws_documentation(
|
||||
query: str,
|
||||
limit: Optional[int] = 10
|
||||
) -> str:
|
||||
"""
|
||||
搜索 AWS 文档
|
||||
|
||||
在 AWS 官方文档中搜索服务指南、API 参考、教程等内容。
|
||||
|
||||
Args:
|
||||
query: 搜索关键词(例如:"S3 bucket", "Lambda Python", "EC2 instance")
|
||||
limit: 最大返回结果数,默认 10
|
||||
|
||||
Returns:
|
||||
搜索结果(JSON 格式)
|
||||
"""
|
||||
try:
|
||||
# 尝试使用 AWS 文档搜索
|
||||
search_result = await search_aws_docs_api(query, limit)
|
||||
|
||||
# 如果搜索 API 返回错误,使用 AI Agent 提供建议
|
||||
if "error" in search_result:
|
||||
agent = get_agent()
|
||||
suggestion = await agent.run(f"""
|
||||
用户想要搜索 AWS 文档:{query}
|
||||
|
||||
由于 AWS 文档搜索 API 暂时不可用,请提供:
|
||||
1. 最相关的 AWS 服务名称
|
||||
2. 建议的文档页面 URL(docs.aws.amazon.com 格式)
|
||||
3. 相关的 AWS 文档主题
|
||||
|
||||
搜索关键词:{query}
|
||||
""")
|
||||
|
||||
return json.dumps({
|
||||
"success": False,
|
||||
"query": query,
|
||||
"error": search_result.get("error"),
|
||||
"suggestion": suggestion.output if hasattr(suggestion, 'output') else str(suggestion),
|
||||
"note": "建议直接使用 get_aws_doc 工具获取特定文档,或访问 https://docs.aws.amazon.com 手动搜索"
|
||||
}, ensure_ascii=False, indent=2)
|
||||
|
||||
# 使用 AI Agent 整理搜索结果
|
||||
agent = get_agent()
|
||||
summary = await agent.run(f"""
|
||||
请整理以下 AWS 文档搜索结果:
|
||||
|
||||
用户查询:{query}
|
||||
|
||||
搜索结果:
|
||||
{json.dumps(search_result, ensure_ascii=False, indent=2)[:2000]}
|
||||
|
||||
请提供:
|
||||
1. 搜索结果摘要
|
||||
2. 最相关的文档链接和标题
|
||||
3. 简要说明这些文档的内容
|
||||
""")
|
||||
|
||||
return json.dumps({
|
||||
"success": True,
|
||||
"query": query,
|
||||
"partition": AWS_DOCUMENTATION_PARTITION,
|
||||
"raw_results": search_result,
|
||||
"ai_explanation": summary.output if hasattr(summary, 'output') else str(summary)
|
||||
}, ensure_ascii=False, indent=2)
|
||||
|
||||
except Exception as e:
|
||||
return json.dumps({
|
||||
"success": False,
|
||||
"error": str(e),
|
||||
"query": query
|
||||
}, ensure_ascii=False, indent=2)
|
||||
|
||||
|
||||
@server.tool()
|
||||
async def get_aws_doc(
|
||||
doc_url: str
|
||||
) -> str:
|
||||
"""
|
||||
获取 AWS 文档内容
|
||||
|
||||
根据文档 URL 获取完整的文档内容并转换为 Markdown 格式。
|
||||
|
||||
Args:
|
||||
doc_url: AWS 文档 URL(必须是 docs.aws.amazon.com 或 docs.amazonaws.cn 域名的有效链接)
|
||||
|
||||
Returns:
|
||||
文档内容(JSON 格式,包含 Markdown 格式的完整内容)
|
||||
"""
|
||||
try:
|
||||
# 获取文档内容
|
||||
markdown_content = await fetch_aws_doc_page(doc_url)
|
||||
|
||||
# 使用 AI Agent 总结文档内容
|
||||
agent = get_agent()
|
||||
summary = await agent.run(f"""
|
||||
请总结以下 AWS 文档的主要内容:
|
||||
|
||||
文档 URL: {doc_url}
|
||||
|
||||
文档内容(Markdown):
|
||||
{markdown_content[:4000]}
|
||||
|
||||
请提供:
|
||||
1. 文档主题和核心内容
|
||||
2. 关键概念和要点
|
||||
3. 适用场景
|
||||
4. 主要章节概述
|
||||
""")
|
||||
|
||||
return json.dumps({
|
||||
"success": True,
|
||||
"doc_url": doc_url,
|
||||
"markdown_content": markdown_content,
|
||||
"ai_summary": summary.output if hasattr(summary, 'output') else str(summary)
|
||||
}, ensure_ascii=False, indent=2)
|
||||
|
||||
except Exception as e:
|
||||
return json.dumps({
|
||||
"success": False,
|
||||
"error": str(e),
|
||||
"doc_url": doc_url
|
||||
}, ensure_ascii=False, indent=2)
|
||||
|
||||
|
||||
@server.tool()
|
||||
async def recommend_aws_content(
|
||||
doc_url: str
|
||||
) -> str:
|
||||
"""
|
||||
获取 AWS 文档推荐内容
|
||||
|
||||
根据文档 URL 获取相关推荐内容(类似 AWS 文档页面底部的"相关主题")。
|
||||
|
||||
Args:
|
||||
doc_url: AWS 文档 URL
|
||||
|
||||
Returns:
|
||||
推荐内容列表(JSON 格式)
|
||||
"""
|
||||
try:
|
||||
# 获取文档页面
|
||||
markdown_content = await fetch_aws_doc_page(doc_url)
|
||||
|
||||
# 使用 AI Agent 分析并推荐相关内容
|
||||
agent = get_agent()
|
||||
recommendations = await agent.run(f"""
|
||||
基于以下 AWS 文档,推荐相关的文档主题:
|
||||
|
||||
文档 URL: {doc_url}
|
||||
文档内容:
|
||||
{markdown_content[:2000]}
|
||||
|
||||
请推荐:
|
||||
1. 相关的 AWS 服务文档
|
||||
2. 相关的教程或指南
|
||||
3. 相关的 API 参考
|
||||
4. 相关的概念说明
|
||||
|
||||
每个推荐应包含:
|
||||
- 主题名称
|
||||
- 建议的文档 URL(docs.aws.amazon.com 格式)
|
||||
- 推荐理由
|
||||
""")
|
||||
|
||||
return json.dumps({
|
||||
"success": True,
|
||||
"doc_url": doc_url,
|
||||
"recommendations": recommendations.output if hasattr(recommendations, 'output') else str(recommendations)
|
||||
}, ensure_ascii=False, indent=2)
|
||||
|
||||
except Exception as e:
|
||||
return json.dumps({
|
||||
"success": False,
|
||||
"error": str(e),
|
||||
"doc_url": doc_url
|
||||
}, ensure_ascii=False, indent=2)
|
||||
|
||||
|
||||
@server.tool()
|
||||
async def get_aws_services_list() -> str:
|
||||
"""
|
||||
获取 AWS 服务列表
|
||||
|
||||
获取 AWS 提供的服务列表(特别适用于中国区域)。
|
||||
|
||||
Returns:
|
||||
AWS 服务列表(JSON 格式)
|
||||
"""
|
||||
try:
|
||||
# 使用 AI Agent 生成 AWS 服务列表
|
||||
agent = get_agent()
|
||||
services = await agent.run(f"""
|
||||
请列出主要的 AWS 服务类别和服务名称。
|
||||
|
||||
分区:{AWS_DOCUMENTATION_PARTITION}
|
||||
|
||||
请按类别组织:
|
||||
1. 计算服务(Compute)
|
||||
2. 存储服务(Storage)
|
||||
3. 数据库服务(Database)
|
||||
4. 网络服务(Networking)
|
||||
5. 安全服务(Security)
|
||||
6. 机器学习服务(Machine Learning)
|
||||
7. 分析服务(Analytics)
|
||||
8. 开发工具(Developer Tools)
|
||||
9. 管理工具(Management Tools)
|
||||
10. 其他服务
|
||||
|
||||
每个服务应包含:
|
||||
- 服务名称
|
||||
- 简要描述
|
||||
- 文档链接(docs.aws.amazon.com 格式)
|
||||
""")
|
||||
|
||||
return json.dumps({
|
||||
"success": True,
|
||||
"partition": AWS_DOCUMENTATION_PARTITION,
|
||||
"services": services.output if hasattr(services, 'output') else str(services)
|
||||
}, ensure_ascii=False, indent=2)
|
||||
|
||||
except Exception as e:
|
||||
return json.dumps({
|
||||
"success": False,
|
||||
"error": str(e)
|
||||
}, ensure_ascii=False, indent=2)
|
||||
|
||||
|
||||
# ==================== 工具映射(供 API 使用)====================
|
||||
|
||||
TOOL_MAP = {
|
||||
'search_aws_documentation': search_aws_documentation,
|
||||
'get_aws_doc': get_aws_doc,
|
||||
'recommend_aws_content': recommend_aws_content,
|
||||
'get_aws_services_list': get_aws_services_list,
|
||||
}
|
||||
|
||||
TOOL_LIST = [
|
||||
{
|
||||
"name": "search_aws_documentation",
|
||||
"description": "搜索 AWS 文档",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {"type": "string", "description": "搜索关键词(例如:S3 bucket, Lambda Python, EC2 instance)"},
|
||||
"limit": {"type": "integer", "description": "最大返回结果数", "default": 10}
|
||||
},
|
||||
"required": ["query"]
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "get_aws_doc",
|
||||
"description": "获取 AWS 文档的完整内容(Markdown 格式)",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"doc_url": {"type": "string", "description": "AWS 文档 URL(必须是 docs.aws.amazon.com 或 docs.amazonaws.cn 域名的有效链接)"}
|
||||
},
|
||||
"required": ["doc_url"]
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "recommend_aws_content",
|
||||
"description": "获取 AWS 文档推荐内容",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"doc_url": {"type": "string", "description": "AWS 文档 URL"}
|
||||
},
|
||||
"required": ["doc_url"]
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "get_aws_services_list",
|
||||
"description": "获取 AWS 服务列表",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": []
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
server.run()
|
||||
Reference in New Issue
Block a user