Files
agent_management/agent_templates/azure_blob_agent_a2a.py
T
2026-01-12 15:04:11 +00:00

653 lines
22 KiB
Python

"""
Azure Blob Storage AI Agent - A2A (Agent-to-Agent) 版本
支持 Agent 之间的协作和通信
"""
import os
import logging
import json
import httpx
from typing import Optional, Dict, Any, List
from datetime import datetime
from fastapi import FastAPI, HTTPException, Header
from pydantic import BaseModel, Field
from azure.storage.blob import BlobServiceClient, ContainerClient
import uvicorn
# 配置日志
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-a2a")
TEMPLATE_TYPE = os.getenv("TEMPLATE_TYPE", "azure_blob_agent_a2a")
AGENT_FRAMEWORK = os.getenv("AGENT_FRAMEWORK", "a2a")
# 工具配置
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")
# A2A Agent 配置
AGENT_ID = os.getenv("AGENT_ID", POD_NAME)
AGENT_ROLE = os.getenv("AGENT_ROLE", "storage_manager")
AGENT_CAPABILITIES = json.loads(os.getenv("AGENT_CAPABILITIES", '["blob_storage", "file_operations"]'))
# 全局存储客户端
blob_service_client: Optional[BlobServiceClient] = None
connection_string: Optional[str] = None
# A2A Agent 注册表 (其他可协作的 Agent)
registered_agents: Dict[str, Dict] = {}
# FastAPI应用
app = FastAPI(
title="Azure Blob Storage AI Agent (A2A)",
description="支持 Agent-to-Agent 协作的智能 Azure Blob 存储管理代理",
version="1.0.0"
)
# ==================== 请求/响应模型 ====================
class ConnectRequest(BaseModel):
"""连接请求"""
connection_string: str = Field(..., description="Azure Storage连接字符串")
class A2AMessage(BaseModel):
"""A2A 消息格式"""
message_id: str = Field(..., description="消息ID")
from_agent: str = Field(..., description="发送者 Agent ID")
to_agent: str = Field(..., description="接收者 Agent ID")
message_type: str = Field(..., description="消息类型: request/response/notification")
action: str = Field(..., description="请求的动作")
parameters: Dict[str, Any] = Field(default_factory=dict, description="参数")
context: Optional[Dict] = Field(default_factory=dict, description="上下文")
timestamp: Optional[str] = None
class A2AQueryRequest(BaseModel):
"""A2A 查询请求"""
query: str = Field(..., description="自然语言查询")
container_name: Optional[str] = None
requester_agent: Optional[str] = Field(None, description="请求者 Agent ID")
context: Optional[Dict] = Field(default_factory=dict)
class A2ARegisterRequest(BaseModel):
"""A2A Agent 注册请求"""
agent_id: str
agent_role: str
capabilities: List[str]
endpoint: str
class HealthResponse(BaseModel):
"""健康检查响应"""
status: str
connected: bool
framework: str
agent_id: str
agent_role: str
capabilities: List[str]
user_id: Optional[str] = None
namespace: Optional[str] = None
registered_agents_count: int = 0
connection_info: Optional[Dict] = None
# ==================== A2A 操作处理器 ====================
class A2AActionHandler:
"""A2A 动作处理器"""
@staticmethod
async def handle_list_containers(parameters: Dict) -> Dict:
"""处理列出容器请求"""
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)}
@staticmethod
async def handle_list_blobs(parameters: Dict) -> Dict:
"""处理列出 blob 请求"""
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)}
@staticmethod
async def handle_get_blob_info(parameters: Dict) -> Dict:
"""处理获取 blob 信息请求"""
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)}
@staticmethod
async def handle_search_blobs(parameters: Dict) -> Dict:
"""处理搜索 blob 请求"""
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)}
@staticmethod
async def handle_get_stats(parameters: Dict) -> Dict:
"""处理获取统计信息请求"""
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)}
# 动作路由表
ACTION_HANDLERS = {
"list_containers": A2AActionHandler.handle_list_containers,
"list_blobs": A2AActionHandler.handle_list_blobs,
"get_blob_info": A2AActionHandler.handle_get_blob_info,
"search_blobs": A2AActionHandler.handle_search_blobs,
"get_stats": A2AActionHandler.handle_get_stats,
}
# ==================== 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,
agent_id=AGENT_ID,
agent_role=AGENT_ROLE,
capabilities=AGENT_CAPABILITIES,
user_id=USER_ID,
namespace=NAMESPACE,
registered_agents_count=len(registered_agents),
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 (Agent: {AGENT_ID}, User: {USER_ID})")
return {
"status": "connected",
"message": "成功连接到 Azure Blob Storage",
"framework": AGENT_FRAMEWORK,
"agent_id": AGENT_ID,
"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("/a2a/capabilities")
async def get_capabilities():
"""获取 Agent 能力"""
return {
"agent_id": AGENT_ID,
"agent_role": AGENT_ROLE,
"capabilities": AGENT_CAPABILITIES,
"supported_actions": list(ACTION_HANDLERS.keys()),
"framework": AGENT_FRAMEWORK
}
@app.post("/a2a/register")
async def register_agent(request: A2ARegisterRequest):
"""注册其他 Agent"""
global registered_agents
registered_agents[request.agent_id] = {
"agent_id": request.agent_id,
"agent_role": request.agent_role,
"capabilities": request.capabilities,
"endpoint": request.endpoint,
"registered_at": str(datetime.now())
}
logger.info(f"✅ Agent '{request.agent_id}' 注册成功")
return {
"status": "registered",
"agent_id": request.agent_id,
"message": f"Agent '{request.agent_id}' 已注册"
}
@app.get("/a2a/agents")
async def list_registered_agents():
"""列出已注册的 Agent"""
return {
"agents": list(registered_agents.values()),
"count": len(registered_agents)
}
@app.post("/a2a/message")
async def handle_a2a_message(message: A2AMessage):
"""处理 A2A 消息"""
if not blob_service_client:
raise HTTPException(
status_code=400,
detail="未连接到 Azure Blob Storage,请先调用 /connect"
)
# 验证消息目标
if message.to_agent != AGENT_ID:
raise HTTPException(
status_code=400,
detail=f"消息目标不匹配: 期望 {AGENT_ID}, 收到 {message.to_agent}"
)
# 处理消息
if message.message_type == "request":
action = message.action
if action not in ACTION_HANDLERS:
return {
"message_id": message.message_id,
"status": "error",
"error": f"不支持的动作: {action}",
"supported_actions": list(ACTION_HANDLERS.keys())
}
try:
handler = ACTION_HANDLERS[action]
result = await handler(message.parameters)
return {
"message_id": message.message_id,
"from_agent": AGENT_ID,
"to_agent": message.from_agent,
"message_type": "response",
"action": action,
"result": result,
"timestamp": str(datetime.now())
}
except Exception as e:
logger.error(f"处理 A2A 消息失败: {str(e)}")
return {
"message_id": message.message_id,
"status": "error",
"error": str(e)
}
return {
"message_id": message.message_id,
"status": "info",
"message": f"收到消息类型: {message.message_type}"
}
@app.post("/query")
async def query_storage(request: A2AQueryRequest):
"""查询存储(支持 A2A 上下文)"""
if not blob_service_client:
raise HTTPException(
status_code=400,
detail="未连接到 Azure Blob Storage,请先调用 /connect"
)
try:
query = request.query.lower()
result = None
action_used = None
# 简单的规则匹配
if "容器" in query and ("列出" in query or "显示" in query or "有哪些" in query):
result = await A2AActionHandler.handle_list_containers({})
action_used = "list_containers"
elif "统计" in query or "有多少" in query or "占用" in query:
result = await A2AActionHandler.handle_get_stats({})
action_used = "get_stats"
elif request.container_name:
if "文件" in query or "blob" in query.lower():
result = await A2AActionHandler.handle_list_blobs({"container_name": request.container_name})
action_used = "list_blobs"
return {
"status": "success" if result else "info",
"query": request.query,
"action": action_used,
"result": result,
"agent_id": AGENT_ID,
"requester": request.requester_agent,
"framework": AGENT_FRAMEWORK
}
except Exception as e:
logger.error(f"查询执行失败: {str(e)}")
raise HTTPException(status_code=500, detail=f"查询失败: {str(e)}")
@app.post("/a2a/collaborate")
async def collaborate_with_agent(
target_agent_id: str,
action: str,
parameters: Dict[str, Any]
):
"""与其他 Agent 协作"""
if target_agent_id not in registered_agents:
raise HTTPException(
status_code=404,
detail=f"Agent '{target_agent_id}' 未注册"
)
target_agent = registered_agents[target_agent_id]
# 创建 A2A 消息
message = A2AMessage(
message_id=f"{AGENT_ID}_{datetime.now().timestamp()}",
from_agent=AGENT_ID,
to_agent=target_agent_id,
message_type="request",
action=action,
parameters=parameters,
timestamp=str(datetime.now())
)
try:
# 发送请求到目标 Agent
async with httpx.AsyncClient() as client:
response = await client.post(
f"{target_agent['endpoint']}/a2a/message",
json=message.dict(),
timeout=30.0
)
response.raise_for_status()
return {
"status": "success",
"target_agent": target_agent_id,
"action": action,
"response": response.json()
}
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,
"agent_id": AGENT_ID,
"agent_role": AGENT_ROLE,
"capabilities": AGENT_CAPABILITIES,
"pod_name": POD_NAME,
"template": TEMPLATE_TYPE,
"user_id": USER_ID,
"namespace": NAMESPACE,
"connected": blob_service_client is not None,
"registered_agents": len(registered_agents),
"endpoints": {
"health": "/health",
"connect": "POST /connect",
"capabilities": "GET /a2a/capabilities",
"register_agent": "POST /a2a/register",
"list_agents": "GET /a2a/agents",
"handle_message": "POST /a2a/message",
"collaborate": "POST /a2a/collaborate",
"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 (A2A)")
logger.info(f" - Framework: {AGENT_FRAMEWORK}")
logger.info(f" - Agent ID: {AGENT_ID}")
logger.info(f" - Agent Role: {AGENT_ROLE}")
logger.info(f" - Capabilities: {AGENT_CAPABILITIES}")
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}")
# 初始化存储连接
init_storage_connection()
uvicorn.run(
app,
host=SERVICE_HOST,
port=SERVICE_PORT,
log_level="info"
)
if __name__ == "__main__":
main()