537 lines
18 KiB
Python
537 lines
18 KiB
Python
"""
|
||
Azure Blob Storage AI Agent - 使用LangChain + LiteLLM实现
|
||
通过HTTP API接收连接字符串，并提供智能文件操作功能
|
||
"""
|
||
import os
|
||
import logging
|
||
import time
|
||
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
|
||
from agent_callback_utils import AgentCallbackHandler, CallbackContextManager
|
||
|
||
# 配置日志
|
||
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")
|
||
MODEL_NAME = os.getenv("MODEL_NAME") or os.getenv("LITELLM_MODEL", "gpt-3.5-turbo")
|
||
|
||
# Azure Storage 连接字符串（从环境变量获取）
|
||
AZURE_STORAGE_CONNECTION_STRING = os.getenv("AZURE_STORAGE_CONNECTION_STRING", "")
|
||
|
||
# 全局变量
|
||
blob_service_client: Optional[BlobServiceClient] = None
|
||
connection_string: Optional[str] = None
|
||
callback_handler: Optional[AgentCallbackHandler] = 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="自然语言查询或操作指令")
|
||
litellm_api_key: str = Field(..., description="LiteLLM API密钥")
|
||
user_id: Optional[str] = Field(None, description="用户ID（用于计费回调）")
|
||
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(litellm_api_key: str) -> 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=MODEL_NAME,
|
||
api_base=LITELLM_API_BASE,
|
||
api_key=litellm_api_key,
|
||
temperature=0
|
||
)
|
||
logger.info(f"✅ LiteLLM初始化成功: {MODEL_NAME} @ {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, callback_handler
|
||
|
||
if not blob_service_client:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail="未连接到Azure Blob Storage，请先调用 /connect"
|
||
)
|
||
|
||
# 初始化回调处理器
|
||
if not callback_handler:
|
||
callback_handler = AgentCallbackHandler()
|
||
|
||
# 使用上下文管理器自动处理回调
|
||
with CallbackContextManager(
|
||
handler=callback_handler,
|
||
user_id=request.user_id,
|
||
request_id=f"blob-{int(time.time())}"
|
||
) as ctx:
|
||
try:
|
||
ctx.add_tool("azure_blob_storage")
|
||
|
||
# 创建Agent
|
||
agent = create_blob_agent(request.litellm_api_key)
|
||
|
||
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():
|
||
"""启动服务"""
|
||
global callback_handler
|
||
|
||
logger.info(f"🚀 启动 Azure Blob Storage AI Agent")
|
||
logger.info(f" - Pod名称: {POD_NAME}")
|
||
logger.info(f" - 模板类型: {TEMPLATE_TYPE}")
|
||
logger.info(f" - LiteLLM: {MODEL_NAME} @ {LITELLM_API_BASE}")
|
||
logger.info(f" - 服务地址: http://{SERVICE_HOST}:{SERVICE_PORT}")
|
||
logger.info(f" ℹ️ API key 将从请求中获取")
|
||
|
||
# 初始化存储连接
|
||
init_storage_connection()
|
||
|
||
# 初始化回调处理器
|
||
callback_handler = AgentCallbackHandler()
|
||
|
||
uvicorn.run(
|
||
app,
|
||
host=SERVICE_HOST,
|
||
port=SERVICE_PORT,
|
||
log_level="info"
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|