forked from zhanggangyong/agent_management
515 lines
17 KiB
Python
515 lines
17 KiB
Python
"""
|
|
Azure Blob Storage AI Agent - 使用LangChain + LiteLLM实现
|
|
通过HTTP API接收连接字符串，并提供智能文件操作功能
|
|
"""
|
|
import os
|
|
import logging
|
|
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
|
|
from langchain.agents import Tool, AgentExecutor, create_react_agent
|
|
from langchain.prompts import PromptTemplate
|
|
from langchain_community.chat_models import ChatLiteLLM
|
|
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")
|
|
TEMPLATE_TYPE = os.getenv("TEMPLATE_TYPE", "azure_blob_agent")
|
|
|
|
# LiteLLM配置
|
|
LITELLM_API_BASE = os.getenv("LITELLM_API_BASE", "http://localhost:4000")
|
|
LITELLM_MODEL = os.getenv("LITELLM_MODEL", "gpt-3.5-turbo")
|
|
LITELLM_API_KEY = os.getenv("LITELLM_API_KEY", "sk-1234")
|
|
|
|
# Azure Storage 连接字符串（可选，也可通过API动态传入）
|
|
AZURE_STORAGE_CONNECTION_STRING = os.getenv("AZURE_STORAGE_CONNECTION_STRING", "")
|
|
|
|
# 全局存储客户端
|
|
blob_service_client: Optional[BlobServiceClient] = None
|
|
connection_string: Optional[str] = None
|
|
|
|
# FastAPI应用
|
|
app = FastAPI(
|
|
title="Azure Blob Storage AI Agent",
|
|
description="智能Azure Blob存储管理代理",
|
|
version="1.0.0"
|
|
)
|
|
|
|
|
|
# ==================== 请求/响应模型 ====================
|
|
|
|
class ConnectRequest(BaseModel):
|
|
"""连接请求"""
|
|
connection_string: str = Field(..., description="Azure Storage连接字符串")
|
|
|
|
|
|
class QueryRequest(BaseModel):
|
|
"""查询请求"""
|
|
query: str = Field(..., description="自然语言查询或操作指令")
|
|
container_name: Optional[str] = Field(None, description="指定容器名称")
|
|
|
|
|
|
class HealthResponse(BaseModel):
|
|
"""健康检查响应"""
|
|
status: str
|
|
connected: bool
|
|
connection_info: Optional[Dict] = None
|
|
|
|
|
|
# ==================== Azure Blob Storage 工具函数 ====================
|
|
|
|
def list_containers_tool() -> str:
|
|
"""列出所有容器"""
|
|
global blob_service_client
|
|
|
|
if not blob_service_client:
|
|
return "错误: 未连接到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)
|
|
})
|
|
|
|
if not container_list:
|
|
return "当前没有容器"
|
|
|
|
result = "容器列表:\n"
|
|
for i, c in enumerate(container_list, 1):
|
|
result += f"{i}. {c['name']} (最后修改: {c['last_modified']})\n"
|
|
|
|
return result
|
|
except Exception as e:
|
|
logger.error(f"列出容器失败: {str(e)}")
|
|
return f"错误: {str(e)}"
|
|
|
|
|
|
def list_blobs_in_container(container_name: str) -> str:
|
|
"""列出指定容器中的所有blob"""
|
|
global blob_service_client
|
|
|
|
if not blob_service_client:
|
|
return "错误: 未连接到Azure Blob Storage"
|
|
|
|
try:
|
|
container_client = blob_service_client.get_container_client(container_name)
|
|
blobs = container_client.list_blobs()
|
|
|
|
blob_list = []
|
|
for blob in blobs:
|
|
blob_list.append({
|
|
"name": blob.name,
|
|
"size": blob.size,
|
|
"content_type": blob.content_settings.content_type if blob.content_settings else "unknown",
|
|
"last_modified": str(blob.last_modified)
|
|
})
|
|
|
|
if not blob_list:
|
|
return f"容器 '{container_name}' 中没有文件"
|
|
|
|
result = f"容器 '{container_name}' 中的文件列表:\n"
|
|
total_size = 0
|
|
for i, b in enumerate(blob_list, 1):
|
|
size_mb = b['size'] / (1024 * 1024)
|
|
result += f"{i}. {b['name']} ({size_mb:.2f}MB, {b['content_type']})\n"
|
|
total_size += b['size']
|
|
|
|
result += f"\n总计: {len(blob_list)} 个文件, {total_size / (1024 * 1024):.2f}MB"
|
|
|
|
return result
|
|
except Exception as e:
|
|
logger.error(f"列出blob失败: {str(e)}")
|
|
return f"错误: {str(e)}"
|
|
|
|
|
|
def get_blob_info(container_name: str, blob_name: str) -> str:
|
|
"""获取blob的详细信息"""
|
|
global blob_service_client
|
|
|
|
if not blob_service_client:
|
|
return "错误: 未连接到Azure Blob Storage"
|
|
|
|
try:
|
|
blob_client = blob_service_client.get_blob_client(container_name, blob_name)
|
|
properties = blob_client.get_blob_properties()
|
|
|
|
info = f"文件信息: {blob_name}\n"
|
|
info += f"- 容器: {container_name}\n"
|
|
info += f"- 大小: {properties.size / (1024 * 1024):.2f}MB\n"
|
|
info += f"- 类型: {properties.content_settings.content_type if properties.content_settings else 'unknown'}\n"
|
|
info += f"- 创建时间: {properties.creation_time}\n"
|
|
info += f"- 最后修改: {properties.last_modified}\n"
|
|
info += f"- ETag: {properties.etag}\n"
|
|
|
|
if properties.metadata:
|
|
info += f"- 元数据: {properties.metadata}\n"
|
|
|
|
return info
|
|
except Exception as e:
|
|
logger.error(f"获取blob信息失败: {str(e)}")
|
|
return f"错误: {str(e)}"
|
|
|
|
|
|
def search_blobs(container_name: str, keyword: str) -> str:
|
|
"""在容器中搜索包含关键字的blob"""
|
|
global blob_service_client
|
|
|
|
if not blob_service_client:
|
|
return "错误: 未连接到Azure Blob Storage"
|
|
|
|
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,
|
|
"last_modified": str(blob.last_modified)
|
|
})
|
|
|
|
if not matched_blobs:
|
|
return f"在容器 '{container_name}' 中没有找到包含 '{keyword}' 的文件"
|
|
|
|
result = f"搜索结果 (关键字: '{keyword}'):\n"
|
|
for i, b in enumerate(matched_blobs, 1):
|
|
result += f"{i}. {b['name']} ({b['size'] / 1024:.2f}KB)\n"
|
|
|
|
return result
|
|
except Exception as e:
|
|
logger.error(f"搜索blob失败: {str(e)}")
|
|
return f"错误: {str(e)}"
|
|
|
|
|
|
def get_storage_stats() -> str:
|
|
"""获取存储统计信息"""
|
|
global blob_service_client
|
|
|
|
if not blob_service_client:
|
|
return "错误: 未连接到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": container_size / (1024 * 1024)
|
|
})
|
|
|
|
result = "存储统计信息:\n"
|
|
result += f"- 总容器数: {total_containers}\n"
|
|
result += f"- 总文件数: {total_blobs}\n"
|
|
result += f"- 总大小: {total_size / (1024 * 1024):.2f}MB\n\n"
|
|
|
|
if container_stats:
|
|
result += "各容器详情:\n"
|
|
for stat in container_stats:
|
|
result += f" • {stat['name']}: {stat['blobs']} 个文件, {stat['size_mb']:.2f}MB\n"
|
|
|
|
return result
|
|
except Exception as e:
|
|
logger.error(f"获取统计信息失败: {str(e)}")
|
|
return f"错误: {str(e)}"
|
|
|
|
|
|
# ==================== 创建LangChain Agent ====================
|
|
|
|
def create_blob_agent() -> Optional[AgentExecutor]:
|
|
"""创建Azure Blob Storage Agent"""
|
|
global blob_service_client
|
|
|
|
if not blob_service_client:
|
|
logger.warning("尚未连接到Azure Blob Storage")
|
|
return None
|
|
|
|
# 初始化LiteLLM
|
|
try:
|
|
llm = ChatLiteLLM(
|
|
model=LITELLM_MODEL,
|
|
api_base=LITELLM_API_BASE,
|
|
api_key=LITELLM_API_KEY,
|
|
temperature=0
|
|
)
|
|
logger.info(f"✅ LiteLLM初始化成功: {LITELLM_MODEL} @ {LITELLM_API_BASE}")
|
|
except Exception as e:
|
|
logger.error(f"❌ LiteLLM初始化失败: {str(e)}")
|
|
return None
|
|
|
|
# 定义工具
|
|
tools = [
|
|
Tool(
|
|
name="list_containers",
|
|
func=list_containers_tool,
|
|
description="列出所有Azure Blob Storage容器。当用户询问'有哪些容器'、'显示容器列表'时使用此工具。"
|
|
),
|
|
Tool(
|
|
name="list_blobs",
|
|
func=lambda input_str: list_blobs_in_container(input_str),
|
|
description="列出指定容器中的所有文件。输入参数是容器名称。当用户询问'容器X中有什么文件'、'列出XXX容器的文件'时使用此工具。"
|
|
),
|
|
Tool(
|
|
name="get_blob_info",
|
|
func=lambda input_str: get_blob_info(*input_str.split(",")),
|
|
description="获取特定文件的详细信息。输入格式: '容器名,文件名'。当用户询问'文件XXX的详细信息'、'XXX文件的属性'时使用此工具。"
|
|
),
|
|
Tool(
|
|
name="search_blobs",
|
|
func=lambda input_str: search_blobs(*input_str.split(",", 1)),
|
|
description="在容器中搜索文件。输入格式: '容器名,关键字'。当用户询问'搜索包含XXX的文件'、'查找XXX'时使用此工具。"
|
|
),
|
|
Tool(
|
|
name="get_storage_stats",
|
|
func=get_storage_stats,
|
|
description="获取存储的统计信息，包括容器数量、文件数量、总大小等。当用户询问'存储统计'、'有多少文件'、'占用多少空间'时使用此工具。"
|
|
),
|
|
]
|
|
|
|
# 定义Agent Prompt
|
|
template = """你是一个Azure Blob Storage管理助手。你可以帮助用户管理和查询Azure存储中的文件。
|
|
|
|
可用工具:
|
|
{tools}
|
|
|
|
工具名称: {tool_names}
|
|
|
|
回答问题时请使用以下格式:
|
|
|
|
Question: 用户的输入问题
|
|
Thought: 你应该思考如何回答这个问题
|
|
Action: 要使用的工具名称，必须是以下之一: [{tool_names}]
|
|
Action Input: 传递给工具的输入
|
|
Observation: 工具返回的结果
|
|
... (这个 Thought/Action/Action Input/Observation 可以重复N次)
|
|
Thought: 我现在知道最终答案了
|
|
Final Answer: 对用户问题的最终回答
|
|
|
|
重要提示:
|
|
- 如果用户只是说"列出容器"或"显示容器"，使用 list_containers 工具
|
|
- 如果用户说"显示XXX容器的文件"，使用 list_blobs 工具，传入容器名
|
|
- 搜索时需要同时提供容器名和关键字
|
|
- 获取文件信息时需要提供容器名和文件名，用逗号分隔
|
|
- 始终用中文回答
|
|
|
|
开始!
|
|
|
|
Question: {input}
|
|
Thought: {agent_scratchpad}"""
|
|
|
|
prompt = PromptTemplate(
|
|
template=template,
|
|
input_variables=["input", "agent_scratchpad"],
|
|
partial_variables={
|
|
"tools": "\n".join([f"- {tool.name}: {tool.description}" for tool in tools]),
|
|
"tool_names": ", ".join([tool.name for tool in tools])
|
|
}
|
|
)
|
|
|
|
# 创建Agent
|
|
agent = create_react_agent(llm, tools, prompt)
|
|
|
|
# 创建Agent执行器
|
|
agent_executor = AgentExecutor(
|
|
agent=agent,
|
|
tools=tools,
|
|
verbose=True,
|
|
handle_parsing_errors=True,
|
|
max_iterations=5
|
|
)
|
|
|
|
logger.info("✅ Azure Blob Storage Agent创建成功")
|
|
return agent_executor
|
|
|
|
|
|
# ==================== 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,
|
|
connection_info=connection_info
|
|
)
|
|
|
|
|
|
@app.post("/connect")
|
|
async def connect_to_storage(request: ConnectRequest):
|
|
"""连接到Azure Blob Storage"""
|
|
global blob_service_client, connection_string
|
|
|
|
try:
|
|
# 创建BlobServiceClient
|
|
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")
|
|
|
|
return {
|
|
"status": "connected",
|
|
"message": "成功连接到Azure Blob Storage",
|
|
"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.post("/query")
|
|
async def query_storage(request: QueryRequest):
|
|
"""使用自然语言查询存储"""
|
|
global blob_service_client
|
|
|
|
if not blob_service_client:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail="未连接到Azure Blob Storage，请先调用 /connect"
|
|
)
|
|
|
|
try:
|
|
# 创建Agent
|
|
agent = create_blob_agent()
|
|
|
|
if not agent:
|
|
raise HTTPException(status_code=500, detail="Agent创建失败")
|
|
|
|
# 执行查询
|
|
logger.info(f"收到查询: {request.query}")
|
|
result = agent.invoke({"input": request.query})
|
|
|
|
return {
|
|
"status": "success",
|
|
"query": request.query,
|
|
"answer": result.get("output", "无法生成答案"),
|
|
"intermediate_steps": str(result.get("intermediate_steps", []))
|
|
}
|
|
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",
|
|
"pod_name": POD_NAME,
|
|
"template": TEMPLATE_TYPE,
|
|
"connected": blob_service_client is not None,
|
|
"endpoints": {
|
|
"health": "/health",
|
|
"connect": "POST /connect",
|
|
"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")
|
|
logger.info(f" - Pod名称: {POD_NAME}")
|
|
logger.info(f" - 模板类型: {TEMPLATE_TYPE}")
|
|
logger.info(f" - LiteLLM: {LITELLM_MODEL} @ {LITELLM_API_BASE}")
|
|
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()
|