623 lines
20 KiB
Python
623 lines
20 KiB
Python
"""
|
|
Azure Blob Storage AI Agent - MCP (Model Context Protocol) 版本
|
|
使用 MCP 协议实现智能文件操作功能
|
|
"""
|
|
import os
|
|
import logging
|
|
import json
|
|
from typing import Optional, Dict, Any, List
|
|
from datetime import datetime
|
|
from fastapi import FastAPI, HTTPException
|
|
from pydantic import BaseModel, Field
|
|
from azure.storage.blob import BlobServiceClient, ContainerClient
|
|
import uvicorn
|
|
import asyncio
|
|
|
|
# 配置日志
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# 环境变量配置
|
|
SERVICE_HOST = os.getenv("SERVICE_HOST", "0.0.0.0")
|
|
SERVICE_PORT = int(os.getenv("SERVICE_PORT", "8080"))
|
|
POD_NAME = os.getenv("POD_NAME", "azure-blob-agent-mcp")
|
|
TEMPLATE_TYPE = os.getenv("TEMPLATE_TYPE", "azure_blob_agent_mcp")
|
|
AGENT_FRAMEWORK = os.getenv("AGENT_FRAMEWORK", "mcp")
|
|
|
|
# 工具配置 (从环境变量传入的 JSON)
|
|
TOOLS_CONFIG = json.loads(os.getenv("TOOLS_CONFIG", "{}"))
|
|
TOOL_ENDPOINT = os.getenv("TOOL_ENDPOINT", "")
|
|
TOOL_API_KEY = os.getenv("TOOL_API_KEY", "")
|
|
|
|
# 模型配置
|
|
MODEL_PROVIDER = os.getenv("MODEL_PROVIDER", "openai")
|
|
MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4")
|
|
MODEL_API_KEY = os.getenv("MODEL_API_KEY", "")
|
|
MODEL_ENDPOINT = os.getenv("MODEL_ENDPOINT", "https://api.openai.com/v1")
|
|
|
|
# 存储配置
|
|
AZURE_STORAGE_CONNECTION_STRING = os.getenv("AZURE_STORAGE_CONNECTION_STRING", "")
|
|
STORAGE_ACCOUNT_NAME = os.getenv("STORAGE_ACCOUNT_NAME", "")
|
|
|
|
# 用户标识
|
|
USER_ID = os.getenv("USER_ID", "")
|
|
TENANT_ID = os.getenv("TENANT_ID", "")
|
|
NAMESPACE = os.getenv("NAMESPACE", "ai-agents")
|
|
|
|
# 全局存储客户端
|
|
blob_service_client: Optional[BlobServiceClient] = None
|
|
connection_string: Optional[str] = None
|
|
|
|
# MCP 工具注册表
|
|
mcp_tools: Dict[str, Any] = {}
|
|
|
|
# FastAPI应用
|
|
app = FastAPI(
|
|
title="Azure Blob Storage AI Agent (MCP)",
|
|
description="基于 MCP 协议的智能 Azure Blob 存储管理代理",
|
|
version="1.0.0"
|
|
)
|
|
|
|
|
|
# ==================== 请求/响应模型 ====================
|
|
|
|
class ConnectRequest(BaseModel):
|
|
"""连接请求"""
|
|
connection_string: str = Field(..., description="Azure Storage连接字符串")
|
|
|
|
|
|
class MCPToolRequest(BaseModel):
|
|
"""MCP 工具调用请求"""
|
|
tool_name: str = Field(..., description="工具名称")
|
|
parameters: Dict[str, Any] = Field(default_factory=dict, description="工具参数")
|
|
|
|
|
|
class MCPQueryRequest(BaseModel):
|
|
"""MCP 查询请求"""
|
|
query: str = Field(..., description="自然语言查询或操作指令")
|
|
container_name: Optional[str] = Field(None, description="指定容器名称")
|
|
context: Optional[Dict] = Field(default_factory=dict, description="上下文信息")
|
|
|
|
|
|
class HealthResponse(BaseModel):
|
|
"""健康检查响应"""
|
|
status: str
|
|
connected: bool
|
|
framework: str
|
|
user_id: Optional[str] = None
|
|
namespace: Optional[str] = None
|
|
connection_info: Optional[Dict] = None
|
|
|
|
|
|
# ==================== MCP 工具定义 ====================
|
|
|
|
class MCPTool:
|
|
"""MCP 工具基类"""
|
|
|
|
def __init__(self, name: str, description: str, parameters_schema: Dict):
|
|
self.name = name
|
|
self.description = description
|
|
self.parameters_schema = parameters_schema
|
|
|
|
async def execute(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""执行工具"""
|
|
raise NotImplementedError
|
|
|
|
def to_mcp_spec(self) -> Dict:
|
|
"""转换为 MCP 工具规范"""
|
|
return {
|
|
"name": self.name,
|
|
"description": self.description,
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": self.parameters_schema,
|
|
"required": list(self.parameters_schema.keys())
|
|
}
|
|
}
|
|
|
|
|
|
class ListContainersTool(MCPTool):
|
|
"""列出所有容器工具"""
|
|
|
|
def __init__(self):
|
|
super().__init__(
|
|
name="list_containers",
|
|
description="列出 Azure Blob Storage 中的所有容器",
|
|
parameters_schema={}
|
|
)
|
|
|
|
async def execute(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
|
|
global blob_service_client
|
|
|
|
if not blob_service_client:
|
|
return {"error": "未连接到 Azure Blob Storage"}
|
|
|
|
try:
|
|
containers = blob_service_client.list_containers()
|
|
container_list = []
|
|
for container in containers:
|
|
container_list.append({
|
|
"name": container.name,
|
|
"last_modified": str(container.last_modified)
|
|
})
|
|
|
|
return {
|
|
"success": True,
|
|
"containers": container_list,
|
|
"count": len(container_list)
|
|
}
|
|
except Exception as e:
|
|
logger.error(f"列出容器失败: {str(e)}")
|
|
return {"error": str(e)}
|
|
|
|
|
|
class ListBlobsTool(MCPTool):
|
|
"""列出容器中的 blob 工具"""
|
|
|
|
def __init__(self):
|
|
super().__init__(
|
|
name="list_blobs",
|
|
description="列出指定容器中的所有文件",
|
|
parameters_schema={
|
|
"container_name": {
|
|
"type": "string",
|
|
"description": "容器名称"
|
|
}
|
|
}
|
|
)
|
|
|
|
async def execute(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
|
|
global blob_service_client
|
|
|
|
if not blob_service_client:
|
|
return {"error": "未连接到 Azure Blob Storage"}
|
|
|
|
container_name = parameters.get("container_name")
|
|
if not container_name:
|
|
return {"error": "缺少参数: container_name"}
|
|
|
|
try:
|
|
container_client = blob_service_client.get_container_client(container_name)
|
|
blobs = container_client.list_blobs()
|
|
|
|
blob_list = []
|
|
total_size = 0
|
|
for blob in blobs:
|
|
blob_info = {
|
|
"name": blob.name,
|
|
"size": blob.size,
|
|
"size_mb": round(blob.size / (1024 * 1024), 2),
|
|
"content_type": blob.content_settings.content_type if blob.content_settings else "unknown",
|
|
"last_modified": str(blob.last_modified)
|
|
}
|
|
blob_list.append(blob_info)
|
|
total_size += blob.size
|
|
|
|
return {
|
|
"success": True,
|
|
"container": container_name,
|
|
"blobs": blob_list,
|
|
"count": len(blob_list),
|
|
"total_size_mb": round(total_size / (1024 * 1024), 2)
|
|
}
|
|
except Exception as e:
|
|
logger.error(f"列出 blob 失败: {str(e)}")
|
|
return {"error": str(e)}
|
|
|
|
|
|
class GetBlobInfoTool(MCPTool):
|
|
"""获取 blob 信息工具"""
|
|
|
|
def __init__(self):
|
|
super().__init__(
|
|
name="get_blob_info",
|
|
description="获取特定文件的详细信息",
|
|
parameters_schema={
|
|
"container_name": {
|
|
"type": "string",
|
|
"description": "容器名称"
|
|
},
|
|
"blob_name": {
|
|
"type": "string",
|
|
"description": "文件名称"
|
|
}
|
|
}
|
|
)
|
|
|
|
async def execute(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
|
|
global blob_service_client
|
|
|
|
if not blob_service_client:
|
|
return {"error": "未连接到 Azure Blob Storage"}
|
|
|
|
container_name = parameters.get("container_name")
|
|
blob_name = parameters.get("blob_name")
|
|
|
|
if not container_name or not blob_name:
|
|
return {"error": "缺少参数: container_name 或 blob_name"}
|
|
|
|
try:
|
|
blob_client = blob_service_client.get_blob_client(container_name, blob_name)
|
|
properties = blob_client.get_blob_properties()
|
|
|
|
return {
|
|
"success": True,
|
|
"blob_name": blob_name,
|
|
"container": container_name,
|
|
"size": properties.size,
|
|
"size_mb": round(properties.size / (1024 * 1024), 2),
|
|
"content_type": properties.content_settings.content_type if properties.content_settings else "unknown",
|
|
"creation_time": str(properties.creation_time),
|
|
"last_modified": str(properties.last_modified),
|
|
"etag": properties.etag,
|
|
"metadata": properties.metadata if properties.metadata else {}
|
|
}
|
|
except Exception as e:
|
|
logger.error(f"获取 blob 信息失败: {str(e)}")
|
|
return {"error": str(e)}
|
|
|
|
|
|
class SearchBlobsTool(MCPTool):
|
|
"""搜索 blob 工具"""
|
|
|
|
def __init__(self):
|
|
super().__init__(
|
|
name="search_blobs",
|
|
description="在容器中搜索包含关键字的文件",
|
|
parameters_schema={
|
|
"container_name": {
|
|
"type": "string",
|
|
"description": "容器名称"
|
|
},
|
|
"keyword": {
|
|
"type": "string",
|
|
"description": "搜索关键字"
|
|
}
|
|
}
|
|
)
|
|
|
|
async def execute(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
|
|
global blob_service_client
|
|
|
|
if not blob_service_client:
|
|
return {"error": "未连接到 Azure Blob Storage"}
|
|
|
|
container_name = parameters.get("container_name")
|
|
keyword = parameters.get("keyword")
|
|
|
|
if not container_name or not keyword:
|
|
return {"error": "缺少参数: container_name 或 keyword"}
|
|
|
|
try:
|
|
container_client = blob_service_client.get_container_client(container_name)
|
|
blobs = container_client.list_blobs()
|
|
|
|
matched_blobs = []
|
|
for blob in blobs:
|
|
if keyword.lower() in blob.name.lower():
|
|
matched_blobs.append({
|
|
"name": blob.name,
|
|
"size": blob.size,
|
|
"size_kb": round(blob.size / 1024, 2),
|
|
"last_modified": str(blob.last_modified)
|
|
})
|
|
|
|
return {
|
|
"success": True,
|
|
"container": container_name,
|
|
"keyword": keyword,
|
|
"results": matched_blobs,
|
|
"count": len(matched_blobs)
|
|
}
|
|
except Exception as e:
|
|
logger.error(f"搜索 blob 失败: {str(e)}")
|
|
return {"error": str(e)}
|
|
|
|
|
|
class GetStorageStatsTool(MCPTool):
|
|
"""获取存储统计工具"""
|
|
|
|
def __init__(self):
|
|
super().__init__(
|
|
name="get_storage_stats",
|
|
description="获取存储的统计信息，包括容器数量、文件数量、总大小等",
|
|
parameters_schema={}
|
|
)
|
|
|
|
async def execute(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
|
|
global blob_service_client
|
|
|
|
if not blob_service_client:
|
|
return {"error": "未连接到 Azure Blob Storage"}
|
|
|
|
try:
|
|
containers = list(blob_service_client.list_containers())
|
|
total_containers = len(containers)
|
|
total_blobs = 0
|
|
total_size = 0
|
|
|
|
container_stats = []
|
|
for container in containers:
|
|
container_client = blob_service_client.get_container_client(container.name)
|
|
blobs = list(container_client.list_blobs())
|
|
blob_count = len(blobs)
|
|
container_size = sum(blob.size for blob in blobs)
|
|
|
|
total_blobs += blob_count
|
|
total_size += container_size
|
|
|
|
container_stats.append({
|
|
"name": container.name,
|
|
"blobs": blob_count,
|
|
"size_mb": round(container_size / (1024 * 1024), 2)
|
|
})
|
|
|
|
return {
|
|
"success": True,
|
|
"total_containers": total_containers,
|
|
"total_blobs": total_blobs,
|
|
"total_size_mb": round(total_size / (1024 * 1024), 2),
|
|
"container_stats": container_stats
|
|
}
|
|
except Exception as e:
|
|
logger.error(f"获取统计信息失败: {str(e)}")
|
|
return {"error": str(e)}
|
|
|
|
|
|
# ==================== MCP 工具注册 ====================
|
|
|
|
def register_tools():
|
|
"""注册所有 MCP 工具"""
|
|
global mcp_tools
|
|
|
|
tools = [
|
|
ListContainersTool(),
|
|
ListBlobsTool(),
|
|
GetBlobInfoTool(),
|
|
SearchBlobsTool(),
|
|
GetStorageStatsTool()
|
|
]
|
|
|
|
for tool in tools:
|
|
mcp_tools[tool.name] = tool
|
|
|
|
logger.info(f"✅ 注册了 {len(mcp_tools)} 个 MCP 工具")
|
|
|
|
|
|
# ==================== API 端点 ====================
|
|
|
|
@app.get("/health", response_model=HealthResponse)
|
|
async def health_check():
|
|
"""健康检查"""
|
|
global blob_service_client, connection_string
|
|
|
|
connected = blob_service_client is not None
|
|
|
|
connection_info = None
|
|
if connected:
|
|
try:
|
|
account_info = blob_service_client.get_account_information()
|
|
connection_info = {
|
|
"account_kind": account_info.get('account_kind', 'unknown'),
|
|
"sku_name": account_info.get('sku_name', 'unknown'),
|
|
"connected_at": str(datetime.now())
|
|
}
|
|
except Exception as e:
|
|
logger.error(f"获取账户信息失败: {str(e)}")
|
|
|
|
return HealthResponse(
|
|
status="healthy" if connected else "not_connected",
|
|
connected=connected,
|
|
framework=AGENT_FRAMEWORK,
|
|
user_id=USER_ID,
|
|
namespace=NAMESPACE,
|
|
connection_info=connection_info
|
|
)
|
|
|
|
|
|
@app.post("/connect")
|
|
async def connect_to_storage(request: ConnectRequest):
|
|
"""连接到 Azure Blob Storage"""
|
|
global blob_service_client, connection_string
|
|
|
|
try:
|
|
blob_service_client = BlobServiceClient.from_connection_string(
|
|
request.connection_string
|
|
)
|
|
|
|
account_info = blob_service_client.get_account_information()
|
|
connection_string = request.connection_string
|
|
|
|
logger.info(f"✅ 成功连接到 Azure Blob Storage (User: {USER_ID})")
|
|
|
|
return {
|
|
"status": "connected",
|
|
"message": "成功连接到 Azure Blob Storage",
|
|
"framework": AGENT_FRAMEWORK,
|
|
"user_id": USER_ID,
|
|
"account_info": {
|
|
"account_kind": account_info.get('account_kind'),
|
|
"sku_name": account_info.get('sku_name')
|
|
}
|
|
}
|
|
except Exception as e:
|
|
logger.error(f"❌ 连接失败: {str(e)}")
|
|
blob_service_client = None
|
|
connection_string = None
|
|
raise HTTPException(status_code=400, detail=f"连接失败: {str(e)}")
|
|
|
|
|
|
@app.get("/mcp/tools")
|
|
async def list_mcp_tools():
|
|
"""列出所有可用的 MCP 工具"""
|
|
if not blob_service_client:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail="未连接到 Azure Blob Storage，请先调用 /connect"
|
|
)
|
|
|
|
tools_spec = [tool.to_mcp_spec() for tool in mcp_tools.values()]
|
|
|
|
return {
|
|
"tools": tools_spec,
|
|
"count": len(tools_spec),
|
|
"framework": AGENT_FRAMEWORK
|
|
}
|
|
|
|
|
|
@app.post("/mcp/call")
|
|
async def call_mcp_tool(request: MCPToolRequest):
|
|
"""调用 MCP 工具"""
|
|
if not blob_service_client:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail="未连接到 Azure Blob Storage，请先调用 /connect"
|
|
)
|
|
|
|
tool_name = request.tool_name
|
|
if tool_name not in mcp_tools:
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail=f"工具 '{tool_name}' 不存在"
|
|
)
|
|
|
|
try:
|
|
tool = mcp_tools[tool_name]
|
|
result = await tool.execute(request.parameters)
|
|
|
|
return {
|
|
"tool": tool_name,
|
|
"result": result,
|
|
"timestamp": str(datetime.now())
|
|
}
|
|
except Exception as e:
|
|
logger.error(f"工具调用失败: {str(e)}")
|
|
raise HTTPException(status_code=500, detail=f"工具调用失败: {str(e)}")
|
|
|
|
|
|
@app.post("/query")
|
|
async def query_storage(request: MCPQueryRequest):
|
|
"""使用自然语言查询存储 (简化版 - 实际应集成 LLM)"""
|
|
if not blob_service_client:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail="未连接到 Azure Blob Storage，请先调用 /connect"
|
|
)
|
|
|
|
try:
|
|
query = request.query.lower()
|
|
result = None
|
|
|
|
# 简单的规则匹配 (实际应使用 LLM 进行意图识别)
|
|
if "容器" in query and ("列出" in query or "显示" in query or "有哪些" in query):
|
|
tool = mcp_tools["list_containers"]
|
|
result = await tool.execute({})
|
|
elif "统计" in query or "有多少" in query or "占用" in query:
|
|
tool = mcp_tools["get_storage_stats"]
|
|
result = await tool.execute({})
|
|
elif request.container_name:
|
|
if "文件" in query or "blob" in query.lower():
|
|
tool = mcp_tools["list_blobs"]
|
|
result = await tool.execute({"container_name": request.container_name})
|
|
|
|
if result:
|
|
return {
|
|
"status": "success",
|
|
"query": request.query,
|
|
"result": result,
|
|
"framework": AGENT_FRAMEWORK
|
|
}
|
|
else:
|
|
return {
|
|
"status": "info",
|
|
"query": request.query,
|
|
"message": "未能匹配到合适的工具，请使用 /mcp/tools 查看可用工具",
|
|
"available_tools": list(mcp_tools.keys())
|
|
}
|
|
except Exception as e:
|
|
logger.error(f"查询执行失败: {str(e)}")
|
|
raise HTTPException(status_code=500, detail=f"查询失败: {str(e)}")
|
|
|
|
|
|
@app.get("/")
|
|
async def root():
|
|
"""根端点"""
|
|
return {
|
|
"service": "Azure Blob Storage AI Agent",
|
|
"version": "1.0.0",
|
|
"framework": AGENT_FRAMEWORK,
|
|
"pod_name": POD_NAME,
|
|
"template": TEMPLATE_TYPE,
|
|
"user_id": USER_ID,
|
|
"namespace": NAMESPACE,
|
|
"connected": blob_service_client is not None,
|
|
"tools_count": len(mcp_tools),
|
|
"endpoints": {
|
|
"health": "/health",
|
|
"connect": "POST /connect",
|
|
"list_tools": "GET /mcp/tools",
|
|
"call_tool": "POST /mcp/call",
|
|
"query": "POST /query"
|
|
}
|
|
}
|
|
|
|
|
|
# ==================== 主函数 ====================
|
|
|
|
def init_storage_connection():
|
|
"""启动时初始化存储连接"""
|
|
global blob_service_client, connection_string
|
|
|
|
if AZURE_STORAGE_CONNECTION_STRING:
|
|
try:
|
|
logger.info("检测到环境变量中的连接字符串，尝试连接...")
|
|
blob_service_client = BlobServiceClient.from_connection_string(
|
|
AZURE_STORAGE_CONNECTION_STRING
|
|
)
|
|
|
|
account_info = blob_service_client.get_account_information()
|
|
connection_string = AZURE_STORAGE_CONNECTION_STRING
|
|
|
|
logger.info(f"✅ 成功连接到 Azure Blob Storage")
|
|
logger.info(f" - Account Kind: {account_info.get('account_kind')}")
|
|
logger.info(f" - SKU: {account_info.get('sku_name')}")
|
|
except Exception as e:
|
|
logger.error(f"❌ 启动时连接失败: {str(e)}")
|
|
logger.info("💡 提示: 可以稍后通过 /connect API 手动连接")
|
|
blob_service_client = None
|
|
connection_string = None
|
|
else:
|
|
logger.info("💡 未设置 AZURE_STORAGE_CONNECTION_STRING，需通过 /connect API 手动连接")
|
|
|
|
|
|
def main():
|
|
"""启动服务"""
|
|
logger.info(f"🚀 启动 Azure Blob Storage AI Agent (MCP)")
|
|
logger.info(f" - Framework: {AGENT_FRAMEWORK}")
|
|
logger.info(f" - Pod名称: {POD_NAME}")
|
|
logger.info(f" - 模板类型: {TEMPLATE_TYPE}")
|
|
logger.info(f" - User ID: {USER_ID}")
|
|
logger.info(f" - Namespace: {NAMESPACE}")
|
|
logger.info(f" - 模型: {MODEL_NAME} @ {MODEL_PROVIDER}")
|
|
logger.info(f" - 服务地址: http://{SERVICE_HOST}:{SERVICE_PORT}")
|
|
|
|
# 注册 MCP 工具
|
|
register_tools()
|
|
|
|
# 初始化存储连接
|
|
init_storage_connection()
|
|
|
|
uvicorn.run(
|
|
app,
|
|
host=SERVICE_HOST,
|
|
port=SERVICE_PORT,
|
|
log_level="info"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|