commit b2a4b8fce6eba32dfb005c87ab4e8edee729c22d Author: Ubuntu Date: Sun Jan 25 08:03:43 2026 +0000 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 全球分区和中国分区 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..9460c99 --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +__pycache__/ +*.pyc +*.pyo +.env diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..5252a14 --- /dev/null +++ b/Dockerfile @@ -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"] diff --git a/README.md b/README.md new file mode 100644 index 0000000..088757a --- /dev/null +++ b/README.md @@ -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"`。 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..811555a --- /dev/null +++ b/requirements.txt @@ -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 diff --git a/run_api_server.py b/run_api_server.py new file mode 100644 index 0000000..fdd7da9 --- /dev/null +++ b/run_api_server.py @@ -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") diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..1982123 --- /dev/null +++ b/src/__init__.py @@ -0,0 +1 @@ +"""Agent 源代码包""" diff --git a/src/server/__init__.py b/src/server/__init__.py new file mode 100644 index 0000000..fc4a3a8 --- /dev/null +++ b/src/server/__init__.py @@ -0,0 +1 @@ +"""服务器模块""" diff --git a/src/server/api_server.py b/src/server/api_server.py new file mode 100644 index 0000000..e1b896f --- /dev/null +++ b/src/server/api_server.py @@ -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) diff --git a/src/server/mcp_server.py b/src/server/mcp_server.py new file mode 100644 index 0000000..9abe2d1 --- /dev/null +++ b/src/server/mcp_server.py @@ -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'(.*?)', html, re.IGNORECASE | re.DOTALL) + title = title_match.group(1).strip() if title_match else "AWS Documentation" + + # 提取主要内容（通常在
或
中） + main_match = re.search(r']*>(.*?)
', html, re.IGNORECASE | re.DOTALL) + if not main_match: + main_match = re.search(r']*id=["\']main-content["\'][^>]*>(.*?)', html, re.IGNORECASE | re.DOTALL) + + content = main_match.group(1) if main_match else html + + # 简单的 HTML 标签清理 + content = re.sub(r']*>.*?', '', content, flags=re.IGNORECASE | re.DOTALL) + content = re.sub(r']*>.*?', '', 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()