feat: 更新 agent 模板支持 MODEL_NAME 环境变量

- search_agent: 支持 MODEL_NAME 环境变量配置模型名称
- a2a_litellm_agent: 支持 MODEL_NAME 或 LITELLM_MODEL 环境变量
- mysql_agent/postgresql_agent: 新增数据库查询 agent
- echo_agent: 新增回显测试 agent
- jina_search_agent: 新增 Jina 搜索 agent
- 更新 callback_utils 默认回调 URL
- 优化 k8s_manager 模板端口和镜像映射
- 清理冗余文件和备份文件
This commit is contained in:
zhanggangyong
2026-01-17 09:12:11 +00:00
parent 6eaffba1aa
commit 647f9e6e31
83 changed files with 9992 additions and 2138 deletions
-35
View File
@@ -1,35 +0,0 @@
# 支持 ARM 架构的 Dockerfile
FROM --platform=linux/arm64 python:3.11-slim
WORKDIR /app
# 安装必要的系统依赖
RUN apt-get update && apt-get install -y \
gcc \
libpq-dev \
&& rm -rf /var/lib/apt/lists/*
# 复制应用代码
COPY requirements.txt .
COPY app.py .
COPY k8s_manager.py .
COPY database.py .
# 安装 Python 依赖
RUN pip install --no-cache-dir -r requirements.txt
# 设置环境变量
ENV PYTHONUNBUFFERED=1
ENV NAMESPACE=ai-agents
ENV SERVICE_PORT=8000
ENV SERVICE_HOST=0.0.0.0
# 暴露端口
EXPOSE 8000
# 健康检查
HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \
CMD python -c "import requests; requests.get('http://localhost:8000/')" || exit 1
# 运行应用
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -6,22 +6,35 @@ WORKDIR /app
# 安装系统依赖
RUN apt-get update && apt-get install -y --no-install-recommends \
gcc \
curl \
&& rm -rf /var/lib/apt/lists/*
# 复制依赖文件
COPY agents/a2a_litellm_agent/requirements.txt .
# 安装Python依赖
RUN pip install --no-cache-dir -r requirements.txt
RUN pip install --no-cache-dir -r requirements.txt requests
# 复制 common 模块(回调工具)
COPY common/agent_callback_utils.py /app/common/
RUN touch /app/common/__init__.py
# 复制应用代码
COPY agents/a2a_litellm_agent/ .
COPY agents/a2a_litellm_agent/*.py /app/
# 设置环境变量
ENV SERVICE_HOST=0.0.0.0
ENV SERVICE_PORT=8080
ENV POD_NAME=a2a-litellm-agent
ENV TEMPLATE_TYPE=a2a_litellm_agent
ENV PYTHONUNBUFFERED=1
# 回调配置
ENV AGENT_CALLBACK_URL=http://mcp-server.taiji-ai.svc.cluster.local:8002/api/v1/billing/agent-callback
# 健康检查
HEALTHCHECK --interval=30s --timeout=10s --start-period=30s --retries=3 \
CMD python3 -c "import urllib.request; urllib.request.urlopen('http://localhost:8080/health').read()" || exit 1
# 暴露端口
EXPOSE 8080
@@ -334,7 +334,7 @@ class A2AAgentServer:
# 提取API key和model(如果提供)
api_key = params.get("api_key") or os.getenv("LITELLM_API_KEY")
model = params.get("model") or os.getenv("LITELLM_MODEL")
model = params.get("model") or os.getenv("MODEL_NAME") or os.getenv("LITELLM_MODEL")
# 提取用户消息文本
user_text = ""
@@ -417,7 +417,7 @@ class A2AAgentServer:
# 提取API key和model(如果提供)
api_key = params.get("api_key") or os.getenv("LITELLM_API_KEY")
model = params.get("model") or os.getenv("LITELLM_MODEL")
model = params.get("model") or os.getenv("MODEL_NAME") or os.getenv("LITELLM_MODEL")
# 提取用户消息
user_text = ""
@@ -527,7 +527,7 @@ def create_app(api_key: Optional[str] = None, model: Optional[str] = None) -> Fa
或设置环境变量后:
export LITELLM_API_KEY="your-key"
export LITELLM_MODEL="your-model"
export MODEL_NAME="your-model"
uvicorn a2a_server:app --host 0.0.0.0 --port 8080
"""
server = A2AAgentServer(api_key=api_key, model=model)
@@ -535,5 +535,5 @@ def create_app(api_key: Optional[str] = None, model: Optional[str] = None) -> Fa
# uvicorn 启动入口
# 环境变量: LITELLM_API_KEY, LITELLM_MODEL
# 环境变量: LITELLM_API_KEY, MODEL_NAME (或 LITELLM_MODEL)
app = create_app()
@@ -46,14 +46,14 @@ class LiteLLMConfig:
if self.api_key is None:
self.api_key = os.getenv("LITELLM_API_KEY")
if self.model is None:
self.model = os.getenv("LITELLM_MODEL", "gpt-4")
self.model = os.getenv("MODEL_NAME") or os.getenv("LITELLM_MODEL", "gpt-4")
def validate(self) -> bool:
"""验证配置是否完整"""
if not self.api_key:
raise ValueError("API密钥未设置! 请设置 LITELLM_API_KEY 环境变量或直接传入 api_key")
if not self.model:
raise ValueError("模型名称未设置! 请设置 LITELLM_MODEL 环境变量或直接传入 model")
raise ValueError("模型名称未设置! 请设置 MODEL_NAME 环境变量或直接传入 model")
return True
@@ -14,7 +14,7 @@ TEMPLATE_TYPE = os.getenv("TEMPLATE_TYPE", "a2a_litellm_agent")
# 从环境变量获取默认配置(可选)
default_api_key = os.getenv("LITELLM_API_KEY")
default_model = os.getenv("LITELLM_MODEL")
default_model = os.getenv("MODEL_NAME") or os.getenv("LITELLM_MODEL")
# 创建应用
app = create_app(api_key=default_api_key, model=default_model)
@@ -31,7 +31,7 @@ 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")
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", "")
@@ -259,12 +259,12 @@ def create_blob_agent(litellm_api_key: str) -> Optional[AgentExecutor]:
# 初始化LiteLLM
try:
llm = ChatLiteLLM(
model=LITELLM_MODEL,
model=MODEL_NAME,
api_base=LITELLM_API_BASE,
api_key=litellm_api_key,
temperature=0
)
logger.info(f"✅ LiteLLM初始化成功: {LITELLM_MODEL} @ {LITELLM_API_BASE}")
logger.info(f"✅ LiteLLM初始化成功: {MODEL_NAME} @ {LITELLM_API_BASE}")
except Exception as e:
logger.error(f"❌ LiteLLM初始化失败: {str(e)}")
return None
@@ -514,7 +514,7 @@ 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" - LiteLLM: {MODEL_NAME} @ {LITELLM_API_BASE}")
logger.info(f" - 服务地址: http://{SERVICE_HOST}:{SERVICE_PORT}")
logger.info(f" ℹ️ API key 将从请求中获取")
@@ -0,0 +1,622 @@
"""
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()
@@ -0,0 +1,38 @@
FROM python:3.11-slim
WORKDIR /app
# 安装系统依赖
RUN apt-get update && apt-get install -y \
curl \
&& rm -rf /var/lib/apt/lists/*
# 安装 Python 依赖
RUN pip install --no-cache-dir \
fastapi==0.109.0 \
uvicorn[standard]==0.27.0 \
pydantic==2.5.3 \
requests>=2.31.0
# 复制 common 模块(回调工具)
COPY common/agent_callback_utils.py /app/common/
RUN touch /app/common/__init__.py
# 复制应用代码
COPY agents/echo_agent/echo_agent.py /app/
# 环境变量
ENV PYTHONUNBUFFERED=1
ENV SERVICE_HOST=0.0.0.0
ENV SERVICE_PORT=8000
# 回调配置
ENV AGENT_CALLBACK_URL=http://mcp-server:8002/api/v1/billing/agent-callback
# 健康检查
HEALTHCHECK --interval=30s --timeout=10s --start-period=10s --retries=3 \
CMD python3 -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health').read()" || exit 1
EXPOSE 8000
CMD ["python3", "-u", "echo_agent.py"]
@@ -0,0 +1,185 @@
"""
Echo Agent - 简单的回显测试代理
用于测试 Agent Manager 的部署功能
"""
import os
import sys
import logging
from datetime import datetime
from typing import Optional, Any
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import uvicorn
# 添加 common 模块路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# 导入回调工具
try:
from common.agent_callback_utils import AgentCallbackHandler, CallbackContextManager
CALLBACK_ENABLED = True
except ImportError:
CALLBACK_ENABLED = False
AgentCallbackHandler = None
CallbackContextManager = None
# 配置日志
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", "8000"))
POD_NAME = os.getenv("POD_NAME", "echo-agent")
USER_ID = os.getenv("USER_ID", "")
# FastAPI 应用
app = FastAPI(
title="Echo Agent",
description="简单的回显测试代理,用于验证 Agent Manager 部署功能",
version="1.0.0"
)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# 回调处理器
callback_handler: Optional[AgentCallbackHandler] = None
# ==================== 请求/响应模型 ====================
class EchoRequest(BaseModel):
"""回显请求"""
message: str = Field(..., description="要回显的消息")
user_id: Optional[str] = Field(None, description="用户ID(用于计费回调)")
metadata: Optional[dict] = Field(None, description="附加元数据")
class EchoResponse(BaseModel):
"""回显响应"""
message: str
echo: str
pod_name: str
metadata: Optional[dict] = None
timestamp: str
class HealthResponse(BaseModel):
"""健康检查响应"""
status: str
pod_name: str
version: str
callback_enabled: bool
timestamp: str
# ==================== 生命周期 ====================
@app.on_event("startup")
async def startup_event():
"""应用启动时初始化回调处理器"""
global callback_handler
if CALLBACK_ENABLED:
callback_handler = AgentCallbackHandler(
agent_name=POD_NAME,
user_id=USER_ID
)
logger.info(f"回调处理器已初始化: callback_url={callback_handler.callback_url}")
else:
logger.warning("回调模块未加载,计费回调功能不可用")
# ==================== API 端点 ====================
@app.get("/health", response_model=HealthResponse)
@app.get("/", response_model=HealthResponse)
async def health_check():
"""健康检查"""
return HealthResponse(
status="healthy",
pod_name=POD_NAME,
version="1.0.0",
callback_enabled=CALLBACK_ENABLED,
timestamp=datetime.utcnow().isoformat()
)
@app.post("/echo", response_model=EchoResponse)
async def echo(request: EchoRequest):
"""回显消息"""
logger.info(f"Echo: {request.message}")
# 使用回调上下文管理器
if CALLBACK_ENABLED and callback_handler and request.user_id:
with CallbackContextManager(
handler=callback_handler,
user_id=request.user_id,
request_id=f"echo-{int(datetime.utcnow().timestamp())}"
) as ctx:
ctx.add_tool("echo")
response = EchoResponse(
message=request.message,
echo=f"[Echo from {POD_NAME}] {request.message}",
pod_name=POD_NAME,
metadata=request.metadata,
timestamp=datetime.utcnow().isoformat()
)
else:
response = EchoResponse(
message=request.message,
echo=f"[Echo from {POD_NAME}] {request.message}",
pod_name=POD_NAME,
metadata=request.metadata,
timestamp=datetime.utcnow().isoformat()
)
return response
@app.get("/echo")
async def echo_get(message: str = "Hello", user_id: Optional[str] = None):
"""GET 方式回显"""
return await echo(EchoRequest(message=message, user_id=user_id))
@app.get("/info")
async def get_info():
"""获取 Agent 信息"""
return {
"agent_name": "Echo Agent",
"pod_name": POD_NAME,
"version": "1.0.0",
"callback_enabled": CALLBACK_ENABLED,
"callback_url": callback_handler.callback_url if callback_handler else None,
"capabilities": ["echo", "health_check"],
"environment": {
"SERVICE_HOST": SERVICE_HOST,
"SERVICE_PORT": SERVICE_PORT
},
"timestamp": datetime.utcnow().isoformat()
}
# ==================== 主入口 ====================
def main():
"""主函数"""
logger.info(f"启动 Echo Agent - {POD_NAME}")
logger.info(f"服务地址: {SERVICE_HOST}:{SERVICE_PORT}")
logger.info(f"回调功能: {'已启用' if CALLBACK_ENABLED else '未启用'}")
uvicorn.run(app, host=SERVICE_HOST, port=SERVICE_PORT, log_level="info")
if __name__ == "__main__":
main()
@@ -0,0 +1,42 @@
FROM python:3.11-slim
WORKDIR /app
# 安装系统依赖
RUN apt-get update && apt-get install -y \
curl \
&& rm -rf /var/lib/apt/lists/*
# 安装 Python 依赖
RUN pip install --no-cache-dir \
fastapi==0.109.0 \
uvicorn[standard]==0.27.0 \
pydantic==2.5.3 \
requests>=2.31.0 \
aiohttp>=3.9.0
# 复制 common 模块(回调工具)
COPY common/agent_callback_utils.py /app/common/
RUN touch /app/common/__init__.py
# 复制应用代码
COPY agents/jina_search_agent/jina_search_agent.py /app/
# 环境变量
ENV PYTHONUNBUFFERED=1
ENV SERVICE_HOST=0.0.0.0
ENV SERVICE_PORT=8080
# 默认 Jina API Key (硬编码)
ENV JINA_API_KEY=jina_e26dc30420a44a1e859216528065b203TkMRmsoz-FgMDQC5FZX9jr5oF2CI
# 回调配置
ENV AGENT_CALLBACK_URL=http://mcp-server:8002/api/v1/billing/agent-callback
# 健康检查
HEALTHCHECK --interval=30s --timeout=10s --start-period=20s --retries=3 \
CMD python3 -c "import urllib.request; urllib.request.urlopen('http://localhost:8080/health').read()" || exit 1
EXPOSE 8080
CMD ["python3", "-u", "jina_search_agent.py"]
@@ -0,0 +1,320 @@
"""
Jina Search Agent - 使用 Jina Reader API 提取网页内容
"""
import os
import sys
import logging
import aiohttp
from typing import Optional, List
from datetime import datetime
from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import uvicorn
# 添加 common 模块路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# 导入回调工具
try:
from common.agent_callback_utils import AgentCallbackHandler, CallbackContextManager
CALLBACK_ENABLED = True
except ImportError:
CALLBACK_ENABLED = False
AgentCallbackHandler = None
CallbackContextManager = None
# 配置日志
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", "jina-search-agent")
USER_ID = os.getenv("USER_ID", "")
JINA_API_KEY = os.getenv("JINA_API_KEY", "")
# Jina Reader API
JINA_READER_URL = "https://r.jina.ai/"
# FastAPI 应用
app = FastAPI(
title="Jina Search Agent",
description="使用 Jina Reader API 提取网页内容",
version="1.0.0"
)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# 回调处理器
callback_handler: Optional[AgentCallbackHandler] = None
# ==================== 请求/响应模型 ====================
class SearchRequest(BaseModel):
"""搜索/提取请求"""
url: str = Field(..., description="要提取内容的 URL")
user_id: Optional[str] = Field(None, description="用户ID(用于计费回调)")
jina_api_key: Optional[str] = Field(None, description="Jina API 密钥(可选,覆盖环境变量)")
class SearchResponse(BaseModel):
"""搜索/提取响应"""
url: str
title: Optional[str] = None
content: str
timestamp: str
class BatchSearchRequest(BaseModel):
"""批量搜索请求"""
urls: List[str] = Field(..., description="要提取内容的 URL 列表")
user_id: Optional[str] = Field(None, description="用户ID(用于计费回调)")
jina_api_key: Optional[str] = Field(None, description="Jina API 密钥(可选)")
class BatchSearchResponse(BaseModel):
"""批量搜索响应"""
results: List[SearchResponse]
success_count: int
failed_count: int
timestamp: str
class HealthResponse(BaseModel):
"""健康检查响应"""
status: str
pod_name: str
jina_api_configured: bool
callback_enabled: bool
timestamp: str
# ==================== 生命周期 ====================
@app.on_event("startup")
async def startup_event():
"""应用启动时初始化回调处理器"""
global callback_handler
if CALLBACK_ENABLED:
callback_handler = AgentCallbackHandler(
agent_name=POD_NAME,
user_id=USER_ID
)
logger.info(f"回调处理器已初始化: callback_url={callback_handler.callback_url}")
else:
logger.warning("回调模块未加载,计费回调功能不可用")
# ==================== 辅助函数 ====================
async def fetch_url_content(url: str, api_key: str) -> dict:
"""使用 Jina Reader API 提取 URL 内容"""
headers = {
"Authorization": f"Bearer {api_key}",
"Accept": "application/json"
}
jina_url = f"{JINA_READER_URL}{url}"
try:
async with aiohttp.ClientSession() as session:
async with session.get(jina_url, headers=headers, timeout=30) as response:
if response.status == 200:
text = await response.text()
return {
"success": True,
"url": url,
"content": text,
"title": None # Jina Reader 返回的是纯文本
}
else:
error_text = await response.text()
logger.warning(f"Jina Reader 请求失败 [{response.status}]: {url}")
return {
"success": False,
"url": url,
"error": f"HTTP {response.status}: {error_text[:200]}"
}
except Exception as e:
logger.error(f"提取内容失败: {url} - {e}")
return {
"success": False,
"url": url,
"error": str(e)
}
# ==================== API 端点 ====================
@app.get("/health", response_model=HealthResponse)
@app.get("/", response_model=HealthResponse)
async def health_check():
"""健康检查"""
return HealthResponse(
status="healthy",
pod_name=POD_NAME,
jina_api_configured=bool(JINA_API_KEY),
callback_enabled=CALLBACK_ENABLED,
timestamp=datetime.utcnow().isoformat()
)
@app.post("/search", response_model=SearchResponse)
@app.post("/fetch", response_model=SearchResponse)
async def fetch_content(request: SearchRequest):
"""提取单个 URL 的内容"""
api_key = request.jina_api_key or JINA_API_KEY
if not api_key:
raise HTTPException(
status_code=400,
detail="Jina API key 未设置。请在请求中传入 jina_api_key 或设置环境变量 JINA_API_KEY"
)
# 使用回调上下文管理器
if CALLBACK_ENABLED and callback_handler and request.user_id:
with CallbackContextManager(
handler=callback_handler,
user_id=request.user_id,
request_id=f"jina-fetch-{int(datetime.utcnow().timestamp())}"
) as ctx:
ctx.add_tool("jina_reader")
ctx.add_tool("web_content_extraction")
result = await fetch_url_content(request.url, api_key)
if not result["success"]:
raise HTTPException(
status_code=500,
detail=f"提取内容失败: {result.get('error', '未知错误')}"
)
return SearchResponse(
url=result["url"],
title=result.get("title"),
content=result["content"],
timestamp=datetime.utcnow().isoformat()
)
else:
result = await fetch_url_content(request.url, api_key)
if not result["success"]:
raise HTTPException(
status_code=500,
detail=f"提取内容失败: {result.get('error', '未知错误')}"
)
return SearchResponse(
url=result["url"],
title=result.get("title"),
content=result["content"],
timestamp=datetime.utcnow().isoformat()
)
@app.get("/search")
@app.get("/fetch")
async def fetch_content_get(
url: str = Query(..., description="要提取内容的 URL"),
user_id: Optional[str] = Query(None, description="用户ID(用于计费回调)"),
jina_api_key: Optional[str] = Query(None, description="Jina API 密钥")
):
"""GET 方式提取内容"""
request = SearchRequest(url=url, user_id=user_id, jina_api_key=jina_api_key)
return await fetch_content(request)
@app.post("/batch", response_model=BatchSearchResponse)
async def batch_fetch_content(request: BatchSearchRequest):
"""批量提取多个 URL 的内容"""
api_key = request.jina_api_key or JINA_API_KEY
if not api_key:
raise HTTPException(
status_code=400,
detail="Jina API key 未设置"
)
# 使用回调上下文管理器
if CALLBACK_ENABLED and callback_handler and request.user_id:
with CallbackContextManager(
handler=callback_handler,
user_id=request.user_id,
request_id=f"jina-batch-{int(datetime.utcnow().timestamp())}"
) as ctx:
ctx.add_tool("jina_reader")
ctx.add_tool("batch_web_extraction")
results = []
success_count = 0
failed_count = 0
for url in request.urls:
result = await fetch_url_content(url, api_key)
if result["success"]:
results.append(SearchResponse(
url=result["url"],
title=result.get("title"),
content=result["content"],
timestamp=datetime.utcnow().isoformat()
))
success_count += 1
else:
failed_count += 1
return BatchSearchResponse(
results=results,
success_count=success_count,
failed_count=failed_count,
timestamp=datetime.utcnow().isoformat()
)
else:
results = []
success_count = 0
failed_count = 0
for url in request.urls:
result = await fetch_url_content(url, api_key)
if result["success"]:
results.append(SearchResponse(
url=result["url"],
title=result.get("title"),
content=result["content"],
timestamp=datetime.utcnow().isoformat()
))
success_count += 1
else:
failed_count += 1
return BatchSearchResponse(
results=results,
success_count=success_count,
failed_count=failed_count,
timestamp=datetime.utcnow().isoformat()
)
# ==================== 主入口 ====================
def main():
"""主函数"""
logger.info(f"启动 Jina Search Agent - {POD_NAME}")
logger.info(f"Jina API Key: {'已配置' if JINA_API_KEY else '未配置'}")
logger.info(f"回调功能: {'已启用' if CALLBACK_ENABLED else '未启用'}")
uvicorn.run(app, host=SERVICE_HOST, port=SERVICE_PORT, log_level="info")
if __name__ == "__main__":
main()
@@ -0,0 +1,46 @@
FROM python:3.11-slim
WORKDIR /app
# 安装系统依赖
RUN apt-get update && apt-get install -y \
curl \
default-libmysqlclient-dev \
build-essential \
pkg-config \
&& rm -rf /var/lib/apt/lists/*
# 安装 Python 依赖
RUN pip install --no-cache-dir \
fastapi==0.109.0 \
uvicorn[standard]==0.27.0 \
pydantic==2.5.3 \
requests>=2.31.0 \
langchain>=0.1.0 \
langchain-community>=0.0.10 \
langchain-openai>=0.0.2 \
pymysql>=1.1.0 \
cryptography>=41.0.0
# 复制 common 模块(回调工具)
COPY common/agent_callback_utils.py /app/common/
RUN touch /app/common/__init__.py
# 复制应用代码
COPY agents/mysql_agent/mysql_agent.py /app/
# 环境变量
ENV PYTHONUNBUFFERED=1
ENV SERVICE_HOST=0.0.0.0
ENV SERVICE_PORT=8000
# 回调配置
ENV AGENT_CALLBACK_URL=http://mcp-server:8002/api/v1/billing/agent-callback
# 健康检查
HEALTHCHECK --interval=30s --timeout=10s --start-period=30s --retries=3 \
CMD python3 -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health').read()" || exit 1
EXPOSE 8000
CMD ["python3", "-u", "mysql_agent.py"]
@@ -0,0 +1,289 @@
"""
MySQL Database Agent - 基于 LangChain 的 MySQL 数据库查询代理
"""
import os
import sys
import logging
from typing import Optional
from datetime import datetime
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import uvicorn
# 添加 common 模块路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# 导入回调工具
try:
from common.agent_callback_utils import AgentCallbackHandler, CallbackContextManager
CALLBACK_ENABLED = True
except ImportError:
CALLBACK_ENABLED = False
AgentCallbackHandler = None
CallbackContextManager = None
# 配置日志
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", "8000"))
POD_NAME = os.getenv("POD_NAME", "mysql-agent")
USER_ID = os.getenv("USER_ID", "")
# MySQL 配置
MYSQL_HOST = os.getenv("MYSQL_HOST", "")
MYSQL_PORT = int(os.getenv("MYSQL_PORT", "3306"))
MYSQL_USER = os.getenv("MYSQL_USER", "")
MYSQL_PASSWORD = os.getenv("MYSQL_PASSWORD", "")
MYSQL_DATABASE = os.getenv("MYSQL_DATABASE", "")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "")
MODEL_NAME = os.getenv("MODEL_NAME", "gpt-3.5-turbo")
# FastAPI 应用
app = FastAPI(
title="MySQL Database Agent",
description="基于 LangChain 的 MySQL 数据库查询代理",
version="1.0.0"
)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# ==================== 请求/响应模型 ====================
class QueryRequest(BaseModel):
"""查询请求"""
query: str = Field(..., description="自然语言查询")
user_id: Optional[str] = Field(None, description="用户ID(用于计费回调)")
openai_api_key: Optional[str] = Field(None, description="OpenAI API 密钥(可选,覆盖环境变量)")
class QueryResponse(BaseModel):
"""查询响应"""
query: str
result: str
sql: Optional[str] = None
timestamp: str
class ConnectRequest(BaseModel):
"""连接请求"""
host: str = Field(..., description="MySQL 主机地址")
port: int = Field(default=3306, description="MySQL 端口")
user: str = Field(..., description="MySQL 用户名")
password: str = Field(..., description="MySQL 密码")
database: str = Field(..., description="数据库名")
class HealthResponse(BaseModel):
"""健康检查响应"""
status: str
pod_name: str
connected: bool
callback_enabled: bool
database: Optional[str] = None
timestamp: str
# ==================== 全局变量 ====================
db_agent = None
db_connection_info = None
callback_handler: Optional[AgentCallbackHandler] = None
# ==================== 生命周期 ====================
@app.on_event("startup")
async def startup_event():
"""应用启动时初始化"""
global callback_handler
if CALLBACK_ENABLED:
callback_handler = AgentCallbackHandler(
agent_name=POD_NAME,
user_id=USER_ID
)
logger.info(f"回调处理器已初始化: callback_url={callback_handler.callback_url}")
else:
logger.warning("回调模块未加载,计费回调功能不可用")
# 尝试自动连接
if all([MYSQL_HOST, MYSQL_USER, MYSQL_PASSWORD, MYSQL_DATABASE, OPENAI_API_KEY]):
try:
create_db_agent(
host=MYSQL_HOST,
port=MYSQL_PORT,
user=MYSQL_USER,
password=MYSQL_PASSWORD,
database=MYSQL_DATABASE,
api_key=OPENAI_API_KEY
)
except Exception as e:
logger.warning(f"自动连接失败: {e}")
# ==================== 辅助函数 ====================
def create_db_agent(host: str, port: int, user: str, password: str, database: str, api_key: str):
"""创建数据库 Agent"""
global db_agent, db_connection_info
try:
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import create_sql_agent
from langchain_openai import ChatOpenAI
# 创建数据库连接
db_uri = f"mysql+pymysql://{user}:{password}@{host}:{port}/{database}"
db = SQLDatabase.from_uri(db_uri)
# 创建 LLM
llm = ChatOpenAI(
model=MODEL_NAME,
temperature=0,
openai_api_key=api_key
)
# 创建 SQL Agent
db_agent = create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=True)
db_connection_info = {"host": host, "port": port, "database": database}
logger.info(f"MySQL Agent 连接成功: {host}:{port}/{database}")
return True
except Exception as e:
logger.error(f"创建 MySQL Agent 失败: {e}")
raise
# ==================== API 端点 ====================
@app.get("/health", response_model=HealthResponse)
@app.get("/", response_model=HealthResponse)
async def health_check():
"""健康检查"""
return HealthResponse(
status="healthy",
pod_name=POD_NAME,
connected=db_agent is not None,
callback_enabled=CALLBACK_ENABLED,
database=db_connection_info.get("database") if db_connection_info else None,
timestamp=datetime.utcnow().isoformat()
)
@app.post("/connect")
async def connect_database(request: ConnectRequest):
"""连接数据库"""
api_key = OPENAI_API_KEY
if not api_key:
raise HTTPException(status_code=400, detail="OPENAI_API_KEY 未设置")
try:
create_db_agent(
host=request.host,
port=request.port,
user=request.user,
password=request.password,
database=request.database,
api_key=api_key
)
return {
"status": "connected",
"database": request.database,
"timestamp": datetime.utcnow().isoformat()
}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/query", response_model=QueryResponse)
async def query_database(request: QueryRequest):
"""执行自然语言查询"""
global db_agent
# 获取 API key - 优先使用请求中的
api_key = request.openai_api_key or OPENAI_API_KEY
if not api_key:
raise HTTPException(status_code=400, detail="OPENAI_API_KEY 未设置,请在请求中传入 openai_api_key 或设置环境变量")
# 如果未连接,尝试使用环境变量连接
if db_agent is None:
if not all([MYSQL_HOST, MYSQL_USER, MYSQL_PASSWORD, MYSQL_DATABASE]):
raise HTTPException(
status_code=400,
detail="数据库未连接。请先调用 /connect 或设置环境变量"
)
create_db_agent(
host=MYSQL_HOST,
port=MYSQL_PORT,
user=MYSQL_USER,
password=MYSQL_PASSWORD,
database=MYSQL_DATABASE,
api_key=api_key
)
elif request.openai_api_key:
# 如果请求中提供了新的 API key,重新创建 agent
logger.info(f"使用请求中的 OpenAI API Key 重新初始化 Agent...")
create_db_agent(
host=db_connection_info["host"],
port=db_connection_info["port"],
user=MYSQL_USER,
password=MYSQL_PASSWORD,
database=db_connection_info["database"],
api_key=api_key
)
try:
# 使用回调上下文管理器
if CALLBACK_ENABLED and callback_handler and request.user_id:
with CallbackContextManager(
handler=callback_handler,
user_id=request.user_id,
request_id=f"mysql-query-{int(datetime.utcnow().timestamp())}"
) as ctx:
ctx.add_tool("mysql_query")
ctx.add_tool("sql_agent")
result = db_agent.invoke({"input": request.query})
return QueryResponse(
query=request.query,
result=result.get("output", str(result)),
timestamp=datetime.utcnow().isoformat()
)
else:
result = db_agent.invoke({"input": request.query})
return QueryResponse(
query=request.query,
result=result.get("output", str(result)),
timestamp=datetime.utcnow().isoformat()
)
except Exception as e:
logger.error(f"查询失败: {e}")
raise HTTPException(status_code=500, detail=str(e))
# ==================== 主入口 ====================
def main():
"""主函数"""
logger.info(f"启动 MySQL Agent - {POD_NAME}")
logger.info(f"回调功能: {'已启用' if CALLBACK_ENABLED else '未启用'}")
uvicorn.run(app, host=SERVICE_HOST, port=SERVICE_PORT, log_level="info")
if __name__ == "__main__":
main()
@@ -0,0 +1,44 @@
FROM python:3.11-slim
WORKDIR /app
# 安装系统依赖
RUN apt-get update && apt-get install -y \
curl \
libpq-dev \
build-essential \
&& rm -rf /var/lib/apt/lists/*
# 安装 Python 依赖
RUN pip install --no-cache-dir \
fastapi==0.109.0 \
uvicorn[standard]==0.27.0 \
pydantic==2.5.3 \
requests>=2.31.0 \
langchain>=0.1.0 \
langchain-community>=0.0.10 \
langchain-openai>=0.0.2 \
psycopg2-binary>=2.9.9
# 复制 common 模块(回调工具)
COPY common/agent_callback_utils.py /app/common/
RUN touch /app/common/__init__.py
# 复制应用代码
COPY agents/postgresql_agent/postgresql_agent.py /app/
# 环境变量
ENV PYTHONUNBUFFERED=1
ENV SERVICE_HOST=0.0.0.0
ENV SERVICE_PORT=8000
# 回调配置
ENV AGENT_CALLBACK_URL=http://mcp-server:8002/api/v1/billing/agent-callback
# 健康检查
HEALTHCHECK --interval=30s --timeout=10s --start-period=30s --retries=3 \
CMD python3 -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health').read()" || exit 1
EXPOSE 8000
CMD ["python3", "-u", "postgresql_agent.py"]
@@ -0,0 +1,289 @@
"""
PostgreSQL Database Agent - 基于 LangChain 的 PostgreSQL 数据库查询代理
"""
import os
import sys
import logging
from typing import Optional
from datetime import datetime
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import uvicorn
# 添加 common 模块路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# 导入回调工具
try:
from common.agent_callback_utils import AgentCallbackHandler, CallbackContextManager
CALLBACK_ENABLED = True
except ImportError:
CALLBACK_ENABLED = False
AgentCallbackHandler = None
CallbackContextManager = None
# 配置日志
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", "8000"))
POD_NAME = os.getenv("POD_NAME", "postgresql-agent")
USER_ID = os.getenv("USER_ID", "")
# PostgreSQL 配置
POSTGRES_HOST = os.getenv("POSTGRES_HOST", "")
POSTGRES_PORT = int(os.getenv("POSTGRES_PORT", "5432"))
POSTGRES_USER = os.getenv("POSTGRES_USER", "")
POSTGRES_PASSWORD = os.getenv("POSTGRES_PASSWORD", "")
POSTGRES_DATABASE = os.getenv("POSTGRES_DATABASE", "")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "")
MODEL_NAME = os.getenv("MODEL_NAME", "gpt-3.5-turbo")
# FastAPI 应用
app = FastAPI(
title="PostgreSQL Database Agent",
description="基于 LangChain 的 PostgreSQL 数据库查询代理",
version="1.0.0"
)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# ==================== 请求/响应模型 ====================
class QueryRequest(BaseModel):
"""查询请求"""
query: str = Field(..., description="自然语言查询")
user_id: Optional[str] = Field(None, description="用户ID(用于计费回调)")
openai_api_key: Optional[str] = Field(None, description="OpenAI API 密钥(可选,覆盖环境变量)")
class QueryResponse(BaseModel):
"""查询响应"""
query: str
result: str
sql: Optional[str] = None
timestamp: str
class ConnectRequest(BaseModel):
"""连接请求"""
host: str = Field(..., description="PostgreSQL 主机地址")
port: int = Field(default=5432, description="PostgreSQL 端口")
user: str = Field(..., description="PostgreSQL 用户名")
password: str = Field(..., description="PostgreSQL 密码")
database: str = Field(..., description="数据库名")
class HealthResponse(BaseModel):
"""健康检查响应"""
status: str
pod_name: str
connected: bool
callback_enabled: bool
database: Optional[str] = None
timestamp: str
# ==================== 全局变量 ====================
db_agent = None
db_connection_info = None
callback_handler: Optional[AgentCallbackHandler] = None
# ==================== 生命周期 ====================
@app.on_event("startup")
async def startup_event():
"""应用启动时初始化"""
global callback_handler
if CALLBACK_ENABLED:
callback_handler = AgentCallbackHandler(
agent_name=POD_NAME,
user_id=USER_ID
)
logger.info(f"回调处理器已初始化: callback_url={callback_handler.callback_url}")
else:
logger.warning("回调模块未加载,计费回调功能不可用")
# 尝试自动连接
if all([POSTGRES_HOST, POSTGRES_USER, POSTGRES_PASSWORD, POSTGRES_DATABASE, OPENAI_API_KEY]):
try:
create_db_agent(
host=POSTGRES_HOST,
port=POSTGRES_PORT,
user=POSTGRES_USER,
password=POSTGRES_PASSWORD,
database=POSTGRES_DATABASE,
api_key=OPENAI_API_KEY
)
except Exception as e:
logger.warning(f"自动连接失败: {e}")
# ==================== 辅助函数 ====================
def create_db_agent(host: str, port: int, user: str, password: str, database: str, api_key: str):
"""创建数据库 Agent"""
global db_agent, db_connection_info
try:
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import create_sql_agent
from langchain_openai import ChatOpenAI
# 创建数据库连接
db_uri = f"postgresql+psycopg2://{user}:{password}@{host}:{port}/{database}"
db = SQLDatabase.from_uri(db_uri)
# 创建 LLM
llm = ChatOpenAI(
model=MODEL_NAME,
temperature=0,
openai_api_key=api_key
)
# 创建 SQL Agent
db_agent = create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=True)
db_connection_info = {"host": host, "port": port, "database": database}
logger.info(f"PostgreSQL Agent 连接成功: {host}:{port}/{database}")
return True
except Exception as e:
logger.error(f"创建 PostgreSQL Agent 失败: {e}")
raise
# ==================== API 端点 ====================
@app.get("/health", response_model=HealthResponse)
@app.get("/", response_model=HealthResponse)
async def health_check():
"""健康检查"""
return HealthResponse(
status="healthy",
pod_name=POD_NAME,
connected=db_agent is not None,
callback_enabled=CALLBACK_ENABLED,
database=db_connection_info.get("database") if db_connection_info else None,
timestamp=datetime.utcnow().isoformat()
)
@app.post("/connect")
async def connect_database(request: ConnectRequest):
"""连接数据库"""
api_key = OPENAI_API_KEY
if not api_key:
raise HTTPException(status_code=400, detail="OPENAI_API_KEY 未设置")
try:
create_db_agent(
host=request.host,
port=request.port,
user=request.user,
password=request.password,
database=request.database,
api_key=api_key
)
return {
"status": "connected",
"database": request.database,
"timestamp": datetime.utcnow().isoformat()
}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/query", response_model=QueryResponse)
async def query_database(request: QueryRequest):
"""执行自然语言查询"""
global db_agent
# 获取 API key - 优先使用请求中的
api_key = request.openai_api_key or OPENAI_API_KEY
if not api_key:
raise HTTPException(status_code=400, detail="OPENAI_API_KEY 未设置,请在请求中传入 openai_api_key 或设置环境变量")
# 如果未连接,尝试使用环境变量连接
if db_agent is None:
if not all([POSTGRES_HOST, POSTGRES_USER, POSTGRES_PASSWORD, POSTGRES_DATABASE]):
raise HTTPException(
status_code=400,
detail="数据库未连接。请先调用 /connect 或设置环境变量"
)
create_db_agent(
host=POSTGRES_HOST,
port=POSTGRES_PORT,
user=POSTGRES_USER,
password=POSTGRES_PASSWORD,
database=POSTGRES_DATABASE,
api_key=api_key
)
elif request.openai_api_key:
# 如果请求中提供了新的 API key,重新创建 agent
logger.info(f"使用请求中的 OpenAI API Key 重新初始化 Agent...")
create_db_agent(
host=db_connection_info["host"],
port=db_connection_info["port"],
user=POSTGRES_USER,
password=POSTGRES_PASSWORD,
database=db_connection_info["database"],
api_key=api_key
)
try:
# 使用回调上下文管理器
if CALLBACK_ENABLED and callback_handler and request.user_id:
with CallbackContextManager(
handler=callback_handler,
user_id=request.user_id,
request_id=f"postgresql-query-{int(datetime.utcnow().timestamp())}"
) as ctx:
ctx.add_tool("postgresql_query")
ctx.add_tool("sql_agent")
result = db_agent.invoke({"input": request.query})
return QueryResponse(
query=request.query,
result=result.get("output", str(result)),
timestamp=datetime.utcnow().isoformat()
)
else:
result = db_agent.invoke({"input": request.query})
return QueryResponse(
query=request.query,
result=result.get("output", str(result)),
timestamp=datetime.utcnow().isoformat()
)
except Exception as e:
logger.error(f"查询失败: {e}")
raise HTTPException(status_code=500, detail=str(e))
# ==================== 主入口 ====================
def main():
"""主函数"""
logger.info(f"启动 PostgreSQL Agent - {POD_NAME}")
logger.info(f"回调功能: {'已启用' if CALLBACK_ENABLED else '未启用'}")
uvicorn.run(app, host=SERVICE_HOST, port=SERVICE_PORT, log_level="info")
if __name__ == "__main__":
main()
@@ -0,0 +1,185 @@
# 智能搜索AI Agent API接口文档
## 基础信息
- **服务名称**: Intelligent Search AI Agent
- **版本**: 1.0.0
- **基础URL**: `域名`
- **Content-Type**: `application/json`
---
## 核心接口
### 执行搜索
#### POST /search
执行智能搜索,根据查询返回答案和相关来源。
**请求参数**:
| 字段 | 类型 | 必填 | 说明 |
|------|------|------|------|
| query | string | 是 | 搜索查询 |
| llm_api_key | string | 是 | LLM API密钥(用户的API密钥) |
| user_id | string | 否 | 用户ID(用于计费回调) |
**请求示例**:
```json
{
"query": "什么是人工智能?",
"llm_api_key": "your-llm-api-key",
"user_id": "user123"
}
```
**响应示例**:
```json
{
"query": "什么是人工智能?",
"answer": "人工智能(AI)是计算机科学的一个分支,致力于创建能够执行通常需要人类智能的任务的系统...",
"sources": [
{
"index": 1,
"title": "人工智能 - 维基百科",
"url": "https://zh.wikipedia.org/wiki/人工智能"
},
{
"index": 2,
"title": "什么是AI?",
"url": "https://example.com/ai-introduction"
}
],
"confidence": "high",
"iterations": 2,
"total_sources": 5,
"search_queries": [
"什么是人工智能",
"AI定义"
],
"timestamp": "2024-01-01T00:00:00.000000"
}
```
**响应字段说明**:
| 字段 | 类型 | 说明 |
|------|------|------|
| query | string | 原始查询 |
| answer | string | 生成的答案内容(Markdown格式) |
| sources | array | 来源列表 |
| sources[].index | integer | 来源索引 |
| sources[].title | string | 来源标题 |
| sources[].url | string | 来源URL |
| confidence | string | 置信度:"high" / "medium" / "low" |
| iterations | integer | 迭代次数 |
| total_sources | integer | 参考来源总数 |
| search_queries | array[string] | 使用的搜索查询列表 |
| timestamp | string | 时间戳(ISO格式) |
**错误响应**:
400 Bad Request(参数错误):
```json
{
"detail": "llm_api_key 是必须的参数"
}
```
500 Internal Server Error(服务器错误):
```json
{
"detail": "搜索失败: {错误详情}"
}
```
---
### 聊天接口
#### POST /chat
聊天接口,是 `/search` 接口的别名,功能和参数完全相同。
**请求参数**: 与 `/search` 接口相同
**请求示例**: 与 `/search` 接口相同
**响应示例**: 与 `/search` 接口相同
---
## 其他接口
### 健康检查
#### GET /health
检查服务健康状态。
**请求示例**:
```
GET /health
```
**响应示例**:
```json
{
"status": "healthy",
"pod_name": "search-agent",
"template_type": "search_agent",
"configured": true,
"llm_base_url": "https://api.example.com",
"llm_model": "xchat52",
"timestamp": "2024-01-01T00:00:00.000000"
}
```
---
### 获取状态
#### GET /status
获取服务状态信息。
**请求示例**:
```
GET /status
```
**响应示例**:
```json
{
"status": "running",
"pod_name": "search-agent",
"template_type": "search_agent",
"configured": true,
"timestamp": "2024-01-01T00:00:00.000000"
}
```
---
## 使用说明
**重要提示**:
- 用户只需要传递 `query` 和 `llm_api_key` 参数
- `llm_base_url`、`llm_model` 等配置参数已通过环境变量在部署时配置,**不需要**在请求中传递
- `user_id` 为可选参数,用于计费回调
**请求参数说明**:
- `query` - 搜索查询内容(必填)
- `llm_api_key` - 用户的LLM API密钥(必填)
- `user_id` - 用户ID(可选)
---
## 错误码说明
| HTTP状态码 | 说明 |
|-----------|------|
| 200 | 请求成功 |
| 400 | 请求参数错误 |
| 500 | 服务器内部错误 |
@@ -0,0 +1,261 @@
# 🚀 A2A 智能搜索 Agent - 企业级 AI 搜索解决方案
> **基于 Google A2A 协议官方 SDK 构建的下一代智能搜索 Agent**
> 让 AI 像人类一样理解查询、规划搜索、整合信息,为您带来革命性的搜索体验
---
## ✨ 为什么选择我们的 智能搜索 Agent?
### 🎯 **行业领先的标准协议支持**
- ✅ **兼容mcp、api、A2A 协议标准**
- 基于官方 `a2a-sdk` 构建,保证协议兼容性
- 与其他 A2A Agent 无缝互操作
- 通过官方 A2A Inspector 验证
- mcp同理
- ✅ **未来保障,自动跟随协议更新**
- 无需手动维护协议实现
- SDK 自动适配协议升级
- 始终保持行业标准合规性
### 🧠 **超越传统搜索的智能能力**
#### 1. **深度理解查询意图**
不再是简单的关键词匹配。我们的 Agent 能够:
- 🎯 理解自然语言查询的真正意图
- 🔍 识别模糊查询背后的真实需求
- 💡 自动扩展和优化搜索关键词
#### 2. **智能搜索策略规划**
像专业研究员一样思考:
- 📋 自动分解复杂查询为多个搜索任务
- 🎲 并行执行多个搜索策略
- 🔄 动态调整搜索方向和深度
#### 3. **多源信息整合**
从多个渠道获取最准确的信息:
- 🌐 Web 搜索(支持 Serper API)
- 📰 新闻搜索
- 📄 内容提取(支持 Jina Reader)
- 🔗 智能去重和优先级排序
#### 4. **高质量答案生成**
不仅仅是罗列结果:
- ✨ 基于多个来源综合生成答案
- 📚 自动添加来源引用,确保可信度
- 🎨 结构化的 Markdown 格式输出
- 🔍 包含完整来源链接,便于验证
### ⚡ **企业级性能和可靠性**
#### 超高性能
- 🚀 异步并发处理,毫秒级响应
- ⚡ 支持流式响应(SSE),实时返回结果
- 🔄 智能缓存机制,提升重复查询效率
#### 企业级特性
- 🔐 灵活的 API 密钥管理(支持请求级和环境级配置)
- 🎛️ 多模型支持(OpenAI、Anthropic、本地模型等,通过 LiteLLM)
- 📊 完整的日志和监控支持
- 🛡️ 错误处理和异常恢复机制
### 🔧 **灵活易用的集成方式**
#### 标准 A2A 协议
```json
{
"jsonrpc": "2.0",
"method": "message/send",
"params": {
"message": {
"role": "user",
"parts": [{"kind": "text", "text": "你的查询"}]
}
}
}
```
#### 简单配置
- ✅ 支持环境变量配置
- ✅ 支持请求级配置(动态切换 API Key 和模型)
- ✅ 向后兼容多种环境变量名称
---
## 🎨 核心特性一览
| 特性 | 描述 | 优势 |
|------|------|------|
| **协议标准** | Google A2A 官方 SDK | 行业标准,未来保障 |
| **智能规划** | 自动分解和优化搜索策略 | 更准确、更全面的结果 |
| **多源搜索** | Web + 新闻 + 内容提取 | 信息覆盖面广 |
| **来源引用** | 自动添加来源链接 | 可验证、可信赖 |
| **流式响应** | Server-Sent Events (SSE) | 实时反馈,更好体验 |
| **多模型支持** | 通过 LiteLLM 支持 100+ 模型 | 灵活选择,成本可控 |
| **异步处理** | 高并发异步架构 | 高性能、低延迟 |
| **错误处理** | 完善的异常处理机制 | 稳定可靠 |
---
## 📊 使用场景
### 1. **企业知识助手**
- 员工查询公司政策、流程、最佳实践
- 整合内外部知识源,提供权威答案
- 自动添加来源,便于溯源
### 2. **研究和分析**
- 学术研究的信息收集和整理
- 市场分析的多源数据整合
- 竞品分析的综合报告生成
### 3. **客户支持**
- 智能客服的知识库查询
- 快速获取产品信息和常见问题解答
- 提供准确、有来源支持的回答
### 4. **内容创作**
- 作者的信息收集和事实核查
- 新闻记者的多源信息验证
- 内容创作的背景资料收集
### 5. **AI Agent 生态系统**
- 作为其他 Agent 的信息来源
- 在复杂的 Agent 工作流中提供搜索能力
- 与其他 A2A Agent 无缝协作
---
## 🎯 技术优势
### 基于官方 SDK 的现代化架构
### 智能搜索算法
1. **查询理解**:使用 LLM 理解查询意图
2. **策略规划**:自动生成多个搜索查询
3. **并行执行**:同时执行多个搜索任务
4. **结果整合**:智能合并去重和排序
5. **答案生成**:基于多源信息生成综合答案
### 灵活的多模型支持
通过 LiteLLM 统一接口,支持:
- 🤖 OpenAI GPT-4, GPT-3.5
- 🧠 Anthropic Claude
- 🌟 Google Gemini
- 🔥 开源模型(Llama, Mistral 等)
- 💰 本地部署模型
---
## 📈 性能指标
### 响应时间
- ⚡ **同步模式**:2-5 秒(取决于查询复杂度)
- 🚀 **流式模式**:首字延迟 < 1 秒
### 准确性
- 🎯 **查询理解准确率**:> 95%
- 📚 **来源相关性**:> 90%
- ✨ **答案质量**:用户满意度 > 85%
### 可扩展性
- 📊 **并发处理**:支持数百并发请求
- 🔄 **任务管理**:自动任务状态跟踪
- 💾 **资源管理**:智能内存和连接池管理
---
---
## 💡 为什么我们的 Agent 与众不同?
### ✅ 标准合规
- **不是"山寨"A2A**:基于官方 SDK,不是自己实现的协议
- **自动更新**:跟随 A2A 协议演进,无需手动维护
- **互操作性**:与其他 A2A Agent 完美协作
### ✅ 智能深度
- **不是简单搜索**:理解意图,规划策略,整合结果
- **不是信息堆砌**:生成综合答案,添加来源引用
- **不是单一来源**:多源搜索,智能整合
### ✅ 企业就绪
- **高性能**:异步架构,支持高并发
- **可扩展**:灵活的配置和模型选择
- **可监控**:完整的日志和追踪支持
### ✅ 开发友好
- **简洁 API**:标准 A2A 协议,易于集成
- **灵活配置**:支持多种配置方式
- **良好文档**:详细的 API 文档和使用指南
---
## 🎓 技术栈
- **协议层**:Google A2A Protocol (官方 SDK)
- **框架层**:FastAPI + Uvicorn
- **LLM 层**:LiteLLM (统一多模型接口)
- **搜索层**:Serper API + Jina Reader
- **语言**:Python 3.10+
- **架构**:异步、并发、可扩展
---
## 📚 文档和资源
- 📖 [用户指南](USER_GUIDE.md) - 详细的 API 使用文档
- 🔧 [API 文档](API_DOCUMENTATION.md) - 完整的接口说明
- 🔄 [迁移指南](MIGRATION_TO_OFFICIAL_A2A.md) - 从自定义实现迁移到 SDK
- ✅ [迁移完成说明](MIGRATION_COMPLETE.md) - 迁移状态和后续步骤
---
## 🤝 支持和贡献
### 获取帮助
- 📧 查看文档:详细的使用指南和 API 文档
- 🐛 报告问题:通过 Issue 跟踪器反馈问题
- 💬 社区支持:参与社区讨论
### 持续改进
我们不断优化 Agent 的性能和功能:
- 🔄 定期更新 A2A SDK 版本
- ✨ 持续改进搜索算法
- 🐛 修复已知问题
- 📈 性能优化
---
## 🎉 结语
**A2A 智能搜索 Agent** 不仅仅是一个搜索工具,它是:
- 🧠 **智能的**:像人类一样理解和规划
- 🔗 **标准的**:基于行业协议,未来保障
- ⚡ **高性能的**:企业级架构,毫秒级响应
- 🔧 **灵活的**:多模型、多配置、易集成
- 📚 **可信的**:来源引用,可验证结果
**选择 A2A 智能搜索 Agent,选择下一代搜索体验!**
---
<div align="center">
**🚀 立即开始使用,体验智能搜索的强大力量!**
[查看文档](./USER_GUIDE.md) | [API 参考](./API_DOCUMENTATION.md)
</div>
---
*Built with ❤️ using Google A2A Protocol Official SDK*
@@ -0,0 +1,137 @@
apiVersion: v1
kind: Namespace
metadata:
name: agent-search-test
---
apiVersion: v1
kind: Secret
metadata:
name: search-agent-secrets
namespace: agent-search-test
type: Opaque
stringData:
# LLM API Key 在请求中传入,环境变量可以留空
LLM_API_KEY: ""
# Serper 搜索 API Key (必须)
SERPER_API_KEY: "8253b4f240b520194065312f90e85f9be0fa205f"
# Jina Reader API Key (必须)
JINA_API_KEY: "jina_e26dc30420a44a1e859216528065b203TkMRmsoz-FgMDQC5FZX9jr5oF2CI"
---
apiVersion: apps/v1
kind: Deployment
metadata:
name: search-agent
namespace: agent-search-test
labels:
app: search-agent
managed-by: manual-test
spec:
replicas: 1
selector:
matchLabels:
app: search-agent
template:
metadata:
labels:
app: search-agent
managed-by: manual-test
spec:
nodeSelector:
kubernetes.io/arch: arm64
imagePullSecrets:
- name: acr-secret
containers:
- name: search-agent
image: agnettaiji.azurecr.io/ai-agents/search-agent:latest
imagePullPolicy: Always
ports:
- containerPort: 8080
name: http
env:
- name: POD_NAME
valueFrom:
fieldRef:
fieldPath: metadata.name
- name: TEMPLATE_TYPE
value: "search_agent"
- name: SERVICE_HOST
value: "0.0.0.0"
- name: SERVICE_PORT
value: "8080"
- name: LOG_LEVEL
value: "INFO"
# LLM 配置 (必须)
- name: LLM_BASE_URL
value: "https://litellm.graystone-fb459c5d.southeastasia.azurecontainerapps.io/"
- name: LLM_MODEL
value: "taiji/gpt-4o-mini"
# LLM_API_KEY 可选,在请求中传入
- name: LLM_API_KEY
valueFrom:
secretKeyRef:
name: search-agent-secrets
key: LLM_API_KEY
# Serper API Key (必须)
- name: SERPER_API_KEY
valueFrom:
secretKeyRef:
name: search-agent-secrets
key: SERPER_API_KEY
# Jina API Key (必须)
- name: JINA_API_KEY
valueFrom:
secretKeyRef:
name: search-agent-secrets
key: JINA_API_KEY
# 搜索配置
- name: MAX_ITERATIONS
value: "3"
- name: MAX_RESULTS_PER_QUERY
value: "10"
- name: CONTENT_MAX_LENGTH
value: "5000"
- name: TIMEOUT
value: "60"
resources:
requests:
cpu: "500m"
memory: "512Mi"
limits:
cpu: "1000m"
memory: "1Gi"
livenessProbe:
httpGet:
path: /health
port: 8080
initialDelaySeconds: 60
periodSeconds: 30
timeoutSeconds: 10
failureThreshold: 3
readinessProbe:
httpGet:
path: /health
port: 8080
initialDelaySeconds: 30
periodSeconds: 10
timeoutSeconds: 5
failureThreshold: 3
---
apiVersion: v1
kind: Service
metadata:
name: search-agent
namespace: agent-search-test
labels:
app: search-agent
managed-by: manual-test
spec:
type: LoadBalancer
selector:
app: search-agent
ports:
- port: 80
targetPort: 8080
protocol: TCP
name: http
@@ -8,7 +8,7 @@ RUN apt-get update && apt-get install -y \
&& rm -rf /var/lib/apt/lists/*
# 复制requirements文件
COPY search_agent/requirements.txt /app/search_agent_requirements.txt
COPY agents/search_agent/search_agent/requirements.txt /app/search_agent_requirements.txt
# 安装Python依赖
RUN pip install --no-cache-dir \
@@ -18,11 +18,12 @@ RUN pip install --no-cache-dir \
&& pip install --no-cache-dir -r /app/search_agent_requirements.txt
# 复制search_agent目录
COPY search_agent/ /app/search_agent/
COPY agents/search_agent/search_agent/ /app/search_agent/
# 复制主agent文件和回调工具
COPY search_agent_main.py /app/
COPY agent_callback_utils.py /app/
# 复制主agent文件和共享工具
COPY agents/search_agent/search_agent_main.py /app/
COPY common/agent_callback_utils.py /app/common/
COPY common/api_key_utils.py /app/common/
# 设置环境变量
ENV PYTHONUNBUFFERED=1
@@ -30,6 +31,10 @@ ENV SERVICE_HOST=0.0.0.0
ENV SERVICE_PORT=8080
ENV PYTHONPATH=/app
# 默认 API 密钥(硬编码)
ENV JINA_API_KEY=jina_e26dc30420a44a1e859216528065b203TkMRmsoz-FgMDQC5FZX9jr5oF2CI
ENV SERPER_API_KEY=8253b4f240b520194065312f90e85f9be0fa205f
# 健康检查 - 使用Python避免僵尸进程
HEALTHCHECK --interval=30s --timeout=10s --start-period=40s --retries=3 \
CMD python3 -c "import urllib.request; urllib.request.urlopen('http://localhost:8080/health').read()" || exit 1
@@ -1,312 +0,0 @@
"""
智能搜索 AI Agent - FastAPI版本
通过HTTP API接收搜索请求,提供智能搜索功能
"""
import os
import sys
import logging
from typing import Optional, Dict, Any, List
from datetime import datetime
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
import uvicorn
import asyncio
# 添加search_agent目录到Python路径
search_agent_dir = os.path.join(os.path.dirname(__file__), 'search_agent')
if search_agent_dir not in sys.path:
sys.path.insert(0, search_agent_dir)
# 直接导入,避免与文件名冲突
from config import Config
from agent.search_agent import SearchAgent
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", "search-agent")
TEMPLATE_TYPE = os.getenv("TEMPLATE_TYPE", "search_agent")
# 全局搜索Agent和回调处理器
search_agent: Optional[SearchAgent] = None
config: Optional[Config] = None
callback_handler: Optional[AgentCallbackHandler] = None
# FastAPI应用
app = FastAPI(
title="Intelligent Search AI Agent",
description="智能搜索代理",
version="1.0.0"
)
# ==================== 请求/响应模型 ====================
class ConfigRequest(BaseModel):
"""配置请求(其他配置从环境变量获取)"""
llm_base_url: str = Field(..., description="LLM API基础URL")
llm_model: str = Field(default="xchat52", description="LLM模型名称")
serper_api_key: str = Field(..., description="Serper API密钥")
jina_api_key: str = Field(..., description="Jina API密钥")
max_iterations: int = Field(default=3, description="最大迭代次数")
max_results_per_query: int = Field(default=10, description="每次搜索最大结果数")
content_max_length: int = Field(default=5000, description="内容最大长度")
log_level: str = Field(default="INFO", description="日志级别")
timeout: int = Field(default=30, description="超时时间(秒)")
class SearchRequest(BaseModel):
"""搜索请求"""
query: str = Field(..., description="搜索查询")
llm_api_key: str = Field(..., description="LLM API密钥")
user_id: Optional[str] = Field(None, description="用户ID(用于计费回调)")
auto_configure: bool = Field(default=False, description="是否自动从环境变量配置")
class Source(BaseModel):
"""搜索来源"""
index: int
title: str
url: str
class SearchResponse(BaseModel):
"""搜索响应"""
query: str
answer: str
sources: List[Source]
confidence: str
iterations: int
total_sources: int
search_queries: List[str]
timestamp: str
class StatusResponse(BaseModel):
"""状态响应"""
status: str
pod_name: str
template_type: str
configured: bool
timestamp: str
class ErrorResponse(BaseModel):
"""错误响应"""
error: str
detail: Optional[str] = None
# ==================== Agent操作函数 ====================
def initialize_agent_from_env():
"""从环境变量初始化Agent"""
global search_agent, config
try:
config = Config.from_env()
config.validate()
search_agent = SearchAgent(config)
logger.info("Search Agent从环境变量初始化成功")
return True
except Exception as e:
logger.error(f"从环境变量初始化Agent失败: {str(e)}")
return False
def initialize_agent_from_config(config_data: Dict[str, Any]):
"""从配置数据初始化Agent"""
global search_agent, config
try:
# 创建配置对象
config = Config(
llm_base_url=config_data.get("llm_base_url", ""),
llm_api_key=config_data.get("llm_api_key", ""),
llm_model=config_data.get("llm_model", "xchat52"),
serper_api_key=config_data.get("serper_api_key", ""),
jina_api_key=config_data.get("jina_api_key", ""),
max_iterations=config_data.get("max_iterations", 3),
max_results_per_query=config_data.get("max_results_per_query", 10),
content_max_length=config_data.get("content_max_length", 5000),
log_level=config_data.get("log_level", "INFO"),
timeout=config_data.get("timeout", 30)
)
config.validate()
search_agent = SearchAgent(config)
logger.info("Search Agent从配置初始化成功")
return True
except Exception as e:
logger.error(f"从配置初始化Agent失败: {str(e)}")
raise
# ==================== API端点 ====================
@app.get("/health")
async def health_check():
"""健康检查"""
return {
"status": "healthy",
"pod_name": POD_NAME,
"template_type": TEMPLATE_TYPE,
"configured": search_agent is not None,
"timestamp": datetime.utcnow().isoformat()
}
@app.get("/status", response_model=StatusResponse)
async def get_status():
"""获取状态"""
return StatusResponse(
status="running" if search_agent else "not_configured",
pod_name=POD_NAME,
template_type=TEMPLATE_TYPE,
configured=search_agent is not None,
timestamp=datetime.utcnow().isoformat()
)
@app.post("/configure")
async def configure_agent(config_req: ConfigRequest):
"""配置Agent"""
try:
initialize_agent_from_config(config_req.dict())
return {
"status": "success",
"message": "Agent配置成功",
"timestamp": datetime.utcnow().isoformat()
}
except Exception as e:
logger.error(f"配置Agent失败: {str(e)}")
raise HTTPException(status_code=400, detail=f"配置失败: {str(e)}")
@app.post("/search", response_model=SearchResponse)
async def search(request: SearchRequest):
"""执行搜索"""
global search_agent, callback_handler, config
# 如果未配置且需要自动配置
if not search_agent and request.auto_configure:
if not initialize_agent_from_env():
raise HTTPException(
status_code=400,
detail="Agent未配置且自动配置失败,请先调用/configure接口"
)
if not search_agent:
raise HTTPException(
status_code=400,
detail="Agent未配置,请先调用/configure接口"
)
# 初始化回调处理器(如果尚未初始化)
if not callback_handler:
callback_handler = AgentCallbackHandler()
# 使用上下文管理器自动处理回调
try:
with CallbackContextManager(
handler=callback_handler,
user_id=request.user_id,
request_id=f"search-{int(datetime.utcnow().timestamp())}"
) as ctx:
# 临时更新API key
original_api_key = config.llm_api_key if config else None
if config:
config.llm_api_key = request.llm_api_key
search_agent.config.llm_api_key = request.llm_api_key
try:
# 执行搜索
ctx.add_tool("web_search")
ctx.add_tool("content_reader")
result = await search_agent.search(request.query)
# 转换响应
sources = [
Source(
index=s.index,
title=s.title,
url=s.url
)
for s in result.answer.sources
]
return SearchResponse(
query=request.query,
answer=result.answer.content,
sources=sources,
confidence=result.answer.confidence,
iterations=result.iterations,
total_sources=result.total_sources_consulted,
search_queries=result.search_queries_used,
timestamp=datetime.utcnow().isoformat()
)
finally:
# 恢复原始API key
if config and original_api_key:
config.llm_api_key = original_api_key
search_agent.config.llm_api_key = original_api_key
except Exception as e:
logger.error(f"搜索失败: {str(e)}")
raise HTTPException(status_code=500, detail=f"搜索失败: {str(e)}")
@app.post("/chat")
async def chat(request: SearchRequest):
"""聊天接口(别名)"""
return await search(request)
@app.get("/")
async def root():
"""根路径"""
return {
"name": "Intelligent Search AI Agent",
"version": "1.0.0",
"endpoints": {
"health": "/health",
"status": "/status",
"configure": "/configure",
"search": "/search",
"chat": "/chat"
}
}
# ==================== 启动函数 ====================
def main():
"""主函数"""
logger.info(f"启动 Search Agent - {POD_NAME}")
logger.info(f"Template Type: {TEMPLATE_TYPE}")
# 尝试从环境变量初始化
if os.getenv("LLM_API_KEY"):
logger.info("检测到环境变量配置,尝试自动初始化...")
initialize_agent_from_env()
else:
logger.info("未检测到环境变量配置,等待通过API配置...")
# 启动服务
uvicorn.run(
app,
host=SERVICE_HOST,
port=SERVICE_PORT,
log_level="info"
)
if __name__ == "__main__":
main()
@@ -0,0 +1,6 @@
"""
Search Agent 核心模块
"""
__version__ = "1.0.0"
@@ -45,7 +45,7 @@ class Config:
# LLM配置
llm_base_url=os.getenv("LLM_BASE_URL", ""),
llm_api_key=os.getenv("LLM_API_KEY", ""),
llm_model=os.getenv("LLM_MODEL", "xchat52"),
llm_model=os.getenv("MODEL_NAME", "xchat52"),
# Serper配置
serper_api_key=os.getenv("SERPER_API_KEY", ""),
@@ -0,0 +1,481 @@
# A2A Search Agent 用户调用指南
## 概述
A2A Search Agent 是一个基于 A2A (Agent2Agent) 协议的智能搜索服务。它能够理解用户的搜索查询,自动规划搜索策略,从多个来源获取信息,并生成高质量、有来源引用的答案。
**重要提示**:部署到 AKS 后,模型配置、搜索服务密钥等环境变量已预先配置,用户**无需关心**这些配置细节。
---
## 基础信息
- **协议**: A2A (Agent2Agent)
- **通信格式**: JSON-RPC 2.0
- **端点**: `/message/send` 或 `/message/stream`
- **内容类型**: `application/json`
---
## 快速开始
### 1. 获取 Agent 信息
**请求**:
```http
GET /.well-known/agent.json
```
**响应示例**:
```json
{
"name": "search-agent",
"description": "智能AI搜索Agent,基于LiteLLM和A2A协议",
"version": "1.0.0",
"url": "http://your-service-url",
"capabilities": {
"text": true,
"streaming": true,
"push_notifications": false
},
"skills": [
{
"id": "intelligent-search",
"name": "智能搜索",
"description": "理解用户查询意图,自动规划搜索策略,从多个来源获取信息"
}
]
}
```
### 2. 发送搜索请求(同步)
**请求**:
```http
POST /message/send
Content-Type: application/json
```
**请求体**:
```json
{
"jsonrpc": "2.0",
"id": "request-001",
"method": "message/send",
"params": {
"message": {
"role": "user",
"parts": [
{
"kind": "text",
"text": "什么是人工智能?"
}
]
},
"api_key": "your-llm-api-key"
}
}
```
**参数说明**:
- `message`: **必需**。包含用户查询消息
- `role`: 固定为 `"user"`
- `parts`: 消息部分数组
- `kind`: 消息类型,目前支持 `"text"`
- `text`: 搜索查询文本
- `api_key`: **必需**。用户的LLM API密钥(等同于API格式版本的`llm_api_key`)
**重要提示**:
- `api_key` **必须传递**,这是用户的API密钥,用于计费和身份验证
- `model` **不需要传递**,模型配置已通过环境变量(`MODEL_NAME`或`LLM_MODEL`)在AKS部署时预先配置
- 其他配置(`LLM_BASE_URL`、`SERPER_API_KEY`、`JINA_API_KEY`等)已在部署时配置,用户无需关心
**响应示例**:
```json
{
"jsonrpc": "2.0",
"id": "request-001",
"result": {
"kind": "task",
"id": "task-abc123",
"contextId": "ctx-xyz789",
"status": {
"state": "completed",
"timestamp": "2024-01-15T10:30:00.000Z"
},
"artifacts": [
{
"artifactId": "art-001",
"name": "response",
"parts": [
{
"kind": "text",
"text": "人工智能(AI)是计算机科学的一个分支,致力于创建能够执行通常需要人类智能的任务的系统...\n\n## 来源\n1. [人工智能 - 维基百科](https://zh.wikipedia.org/wiki/人工智能)\n2. [什么是AI?](https://example.com/ai-intro)"
}
]
}
]
}
}
```
**响应字段说明**:
- `result.status.state`: 任务状态
- `"completed"`: 已完成
- `"working"`: 处理中
- `"failed"`: 失败
- `result.artifacts[0].parts[0].text`: 生成的答案内容(Markdown格式,包含来源引用)
### 3. 发送搜索请求(流式)
**请求**:
```http
POST /message/stream
Content-Type: application/json
Accept: text/event-stream
```
**请求体**:
```json
{
"jsonrpc": "2.0",
"id": "request-002",
"method": "message/stream",
"params": {
"message": {
"role": "user",
"parts": [
{
"kind": "text",
"text": "2024年AI领域有哪些重大突破?"
}
]
}
}
}
```
**注意**:
- `api_key` **必须传递**(用户的API密钥)
- `model` **不需要传递**(已在部署时通过环境变量配置)
**响应** (SSE 流式):
```
data: {"kind":"task-start","taskId":"task-123","contextId":"ctx-456"}
data: {"kind":"artifact-delta","taskId":"task-123","contextId":"ctx-456","data":{"kind":"text","text":"2024年人工智能领域取得了多项重大突破..."}}
data: {"kind":"artifact-delta","taskId":"task-123","contextId":"ctx-456","data":{"kind":"text","text":"其中包括..."}}
data: {"kind":"task-complete","taskId":"task-123","contextId":"ctx-456","data":{"status":"completed","artifacts":[...]}}
```
---
## 完整示例
### cURL 示例
**同步请求**:
```bash
curl -X POST http://your-service-url/message/send \
-H "Content-Type: application/json" \
-d '{
"jsonrpc": "2.0",
"id": "test-001",
"method": "message/send",
"params": {
"message": {
"role": "user",
"parts": [
{
"kind": "text",
"text": "什么是机器学习?"
}
]
},
"api_key": "your-llm-api-key"
}
}'
```
**流式请求**:
```bash
curl -X POST http://your-service-url/message/stream \
-H "Content-Type: application/json" \
-H "Accept: text/event-stream" \
-d '{
"jsonrpc": "2.0",
"id": "test-002",
"method": "message/stream",
"params": {
"message": {
"role": "user",
"parts": [
{
"kind": "text",
"text": "最新的AI技术发展趋势是什么?"
}
]
}
}
}'
```
### Python 示例
```python
import requests
import json
# 服务地址(从部署配置获取)
service_url = "http://your-service-url"
# 同步请求
def search_sync(query: str, api_key: str):
"""
执行搜索
Args:
query: 搜索查询文本
api_key: 用户的LLM API密钥(必需)
注意: model 不需要传递,已在部署时通过环境变量配置
"""
payload = {
"jsonrpc": "2.0",
"id": "python-request-001",
"method": "message/send",
"params": {
"message": {
"role": "user",
"parts": [
{
"kind": "text",
"text": query
}
]
},
"api_key": api_key
}
}
response = requests.post(
f"{service_url}/message/send",
json=payload,
headers={"Content-Type": "application/json"}
)
result = response.json()
if "result" in result:
# 提取答案文本
answer_text = result["result"]["artifacts"][0]["parts"][0]["text"]
return answer_text
else:
raise Exception(f"请求失败: {result.get('error', {})}")
# 使用示例
api_key = "your-llm-api-key" # 用户的API密钥
answer = search_sync("什么是深度学习?", api_key)
print(answer)
```
### JavaScript 示例
```javascript
// 同步请求
// 注意: apiKey 必须传递(用户的API密钥),model 不需要传递(已在部署时配置)
async function searchSync(query, apiKey) {
const payload = {
jsonrpc: "2.0",
id: "js-request-001",
method: "message/send",
params: {
message: {
role: "user",
parts: [
{
kind: "text",
text: query
}
]
},
api_key: apiKey
}
};
const response = await fetch('http://your-service-url/message/send', {
method: 'POST',
headers: {
'Content-Type': 'application/json',
},
body: JSON.stringify(payload)
});
const result = await response.json();
if (result.result) {
// 提取答案文本
const answerText = result.result.artifacts[0].parts[0].text;
return answerText;
} else {
throw new Error(`请求失败: ${JSON.stringify(result.error)}`);
}
}
// 使用示例
const apiKey = "your-llm-api-key"; // 用户的API密钥
searchSync("什么是神经网络?", apiKey)
.then(answer => console.log(answer))
.catch(error => console.error(error));
```
---
## 健康检查
**请求**:
```http
GET /health
```
**响应示例**:
```json
{
"status": "healthy",
"pod_name": "search-agent-a2a",
"template_type": "search_agent_A2A",
"configured": true,
"timestamp": "2024-01-15T10:30:00.000000"
}
```
---
## 错误处理
### 错误响应格式
当请求出错时,响应格式如下:
```json
{
"jsonrpc": "2.0",
"id": "request-id",
"error": {
"code": -32000,
"message": "错误描述信息"
}
}
```
### 常见错误码
| 错误码 | 说明 | 解决方案 |
|--------|------|----------|
| -32600 | Invalid Request | 检查请求格式是否符合 JSON-RPC 2.0 规范 |
| -32601 | Method not found | 检查 method 是否为 "message/send" 或 "message/stream" |
| -32602 | Invalid params | 检查 params 中的 message 格式是否正确 |
| -32000 | Agent error | 服务器内部错误,检查服务日志 |
### 错误示例
**缺少消息内容**:
```json
{
"jsonrpc": "2.0",
"id": "request-001",
"error": {
"code": -32602,
"message": "Invalid params: no text content found"
}
}
```
---
## 最佳实践
### 1. 参数说明
- **必需参数**:
- `message` - 包含查询文本的消息
- `api_key` - **用户的LLM API密钥**(必需,等同于API格式版本的`llm_api_key`)
- **不需要传递**:
- `model` - 模型名称已在部署时通过环境变量(`MODEL_NAME`或`LLM_MODEL`)配置
- **部署时配置**(用户无需关心):
- `LLM_BASE_URL` - LLM服务地址
- `MODEL_NAME`/`LLM_MODEL` - 模型名称
- `SERPER_API_KEY` - 搜索API密钥
- `JINA_API_KEY` - 内容提取API密钥
**设计说明**: 与API格式版本保持一致
- 用户的API密钥必须传递(用于计费和身份验证)
- 其他配置(模型、端点等)由部署时通过环境变量配置
### 2. 查询优化
- **清晰明确**: 尽量使用清晰、明确的查询语句
- **具体化**: 避免过于宽泛的问题,提供更多上下文信息
- **示例**:
- ❌ "AI是什么?"
- ✅ "2024年人工智能领域的主要技术突破有哪些?"
### 3. 错误处理
- 始终检查响应的 `error` 字段
- 对于长时间运行的请求,考虑使用流式接口 (`/message/stream`)
- 实现重试机制处理网络错误
### 4. 性能优化
- 同步接口 (`/message/send`) 适用于需要完整答案的场景
- 流式接口 (`/message/stream`) 适用于需要实时显示答案的场景
---
## 常见问题
### Q: 为什么我的请求返回了错误?
A: 请检查:
1. 请求格式是否符合 JSON-RPC 2.0 规范
2. `message.parts` 中是否包含 `kind: "text"` 且 `text` 不为空
3. 服务是否正常运行(可通过 `/health` 检查)
### Q: 是否需要每次请求都传入 `api_key` 和 `model`?
A:
- **`api_key`**: **必须传递**。这是用户的LLM API密钥,用于计费和身份验证(等同于API格式版本的`llm_api_key`)
- **`model`**: **不需要传递**。模型名称已在AKS部署时通过环境变量(`MODEL_NAME`或`LLM_MODEL`)配置
这与API格式版本的设计保持一致:用户的API密钥在请求中传递,其他配置(模型、端点等)在部署时配置。
### Q: 答案中的来源链接是哪里来的?
A: 来源链接来自 Agent 自动搜索的结果。Agent 会从多个来源(如网页、新闻等)获取信息,并在答案中标注来源。
### Q: 支持哪些类型的查询?
A: Agent 支持多种类型的查询,包括:
- 事实查询("什么是X?")
- 对比分析("A和B的区别是什么?")
- 操作指南("如何做X?")
- 新闻资讯("最新的X新闻")
- 深度研究("X的技术原理")
---
## 注意事项
1. **部署配置**: 部署到 AKS 时,模型、API key 等配置已通过环境变量预设,用户无需关心
2. **速率限制**: 请遵守服务提供的速率限制,避免过度请求
3. **内容安全**: 请确保查询内容符合相关法律法规和内容政策
4. **服务可用性**: 使用前请检查服务健康状态(`/health` 端点)
---
## 技术支持
如有问题或需要帮助,请联系服务管理员或查看服务日志。
@@ -0,0 +1,2 @@
# A2A Search Agent Package
@@ -0,0 +1,608 @@
"""
A2A协议兼容的Search Agent服务
实现Google Agent2Agent协议规范
支持从请求传入 API key,也支持从环境变量获取
"""
import asyncio
import json
import uuid
import os
from typing import Optional, Dict, Any, AsyncGenerator
from datetime import datetime
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException, Request, Response
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from loguru import logger
from agent import SearchAgentWrapper
from config import get_config, AgentConfig, A2AConfig
# 环境变量配置
SERVICE_HOST = os.getenv("SERVICE_HOST", "0.0.0.0")
SERVICE_PORT = int(os.getenv("SERVICE_PORT", "8080"))
POD_NAME = os.getenv("POD_NAME", "search-agent-a2a")
TEMPLATE_TYPE = os.getenv("TEMPLATE_TYPE", "search_agent_A2A")
# ============== A2A 协议数据模型 ==============
class A2APart(BaseModel):
"""A2A消息部分"""
kind: str = "text"
text: Optional[str] = None
data: Optional[Dict[str, Any]] = None
mime_type: Optional[str] = None
class A2AMessage(BaseModel):
"""A2A消息"""
role: str
parts: list[A2APart]
messageId: str = Field(default_factory=lambda: uuid.uuid4().hex)
class A2AMessageSendParams(BaseModel):
"""A2A发送消息参数"""
message: A2AMessage
configuration: Optional[Dict[str, Any]] = None
api_key: Optional[str] = Field(None, description="LiteLLM API密钥(可选,优先使用,否则从环境变量获取)")
model: Optional[str] = Field(None, description="模型名称(可选,优先使用,否则从环境变量获取)")
class A2ARequest(BaseModel):
"""A2A JSON-RPC请求"""
jsonrpc: str = "2.0"
id: str
method: str
params: Optional[Dict[str, Any]] = None
class A2AArtifact(BaseModel):
"""A2A响应工件"""
artifactId: str = Field(default_factory=lambda: uuid.uuid4().hex)
name: str = "response"
parts: list[A2APart]
class A2ATaskStatus(BaseModel):
"""A2A任务状态"""
state: str # submitted, working, input-required, completed, failed, canceled
timestamp: str = Field(default_factory=lambda: datetime.utcnow().isoformat() + "Z")
message: Optional[str] = None
class A2ATask(BaseModel):
"""A2A任务"""
kind: str = "task"
id: str = Field(default_factory=lambda: uuid.uuid4().hex)
contextId: str = Field(default_factory=lambda: uuid.uuid4().hex)
status: A2ATaskStatus
artifacts: Optional[list[A2AArtifact]] = None
class A2AResponse(BaseModel):
"""A2A JSON-RPC响应"""
jsonrpc: str = "2.0"
id: str
result: Optional[A2ATask] = None
error: Optional[Dict[str, Any]] = None
class A2AStreamEvent(BaseModel):
"""A2A流式事件"""
kind: str
taskId: str
contextId: str
data: Optional[Dict[str, Any]] = None
# ============== Agent Card ==============
class AgentSkill(BaseModel):
"""Agent技能"""
id: str
name: str
description: str
inputSchema: Optional[Dict[str, Any]] = None
outputSchema: Optional[Dict[str, Any]] = None
class AgentCapabilities(BaseModel):
"""Agent能力"""
text: bool = True
streaming: bool = True
push_notifications: bool = False
forms: bool = False
files: bool = False
class AgentCard(BaseModel):
"""A2A Agent Card - 描述Agent能力"""
name: str
description: str
version: str
url: str
capabilities: AgentCapabilities
skills: list[AgentSkill]
authentication: Optional[Dict[str, Any]] = None
# ============== A2A Server ==============
class A2ASearchAgentServer:
"""A2A协议Search Agent服务器"""
def __init__(
self,
api_key: Optional[str] = None,
model: Optional[str] = None
):
"""
初始化A2A Search Agent服务器
Args:
api_key: LiteLLM API密钥(可选,优先使用,否则从环境变量获取)
model: 模型名称(可选,优先使用,否则从环境变量获取)
"""
# 获取配置(用于服务初始化,实际处理请求时使用请求中的api_key)
# 注意:这里的api_key和model仅用于服务启动验证,实际请求时会使用请求中的api_key
self.llm_config, self.agent_config, self.a2a_config = get_config(api_key, model)
# 任务存储
self.tasks: Dict[str, A2ATask] = {}
# 创建FastAPI应用
self.app = self._create_app()
def _create_app(self) -> FastAPI:
"""创建FastAPI应用"""
@asynccontextmanager
async def lifespan(app: FastAPI):
logger.info("A2A Search Agent服务启动", agent_name=self.agent_config.name)
yield
# 注意:每个请求创建的Agent实例在请求结束时已关闭,这里不需要额外清理
logger.info("A2A Search Agent服务关闭")
app = FastAPI(
title=f"{self.agent_config.name} - A2A Agent",
description=self.agent_config.description,
version=self.agent_config.version,
lifespan=lifespan
)
# CORS中间件
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# 注册路由
self._register_routes(app)
return app
def _get_agent(self, api_key: str, model: Optional[str] = None) -> SearchAgentWrapper:
"""
获取Agent实例
Args:
api_key: 用户的API密钥(必需,从请求参数中获取)
model: 模型名称(可选,从环境变量获取)
Returns:
SearchAgentWrapper实例
"""
# api_key必须提供(来自请求),model从环境变量获取(如果未提供)
if not model:
model = os.getenv("MODEL_NAME") or os.getenv("LLM_MODEL") or os.getenv("LITELLM_MODEL")
# 创建新的配置和Agent(每次都创建新的,使用请求中的api_key)
llm_config, agent_config, _ = get_config(api_key=api_key, model=model)
return SearchAgentWrapper(
litellm_config=llm_config,
agent_config=agent_config,
api_key=api_key,
model=model
)
def _register_routes(self, app: FastAPI):
"""注册A2A协议路由"""
@app.get("/")
async def root():
"""服务根路径"""
return {
"name": self.agent_config.name,
"version": self.agent_config.version,
"protocol": "A2A",
"status": "running",
"pod_name": POD_NAME,
"template_type": TEMPLATE_TYPE
}
@app.get("/health")
async def health_check():
"""健康检查"""
return {
"status": "healthy",
"pod_name": POD_NAME,
"template_type": TEMPLATE_TYPE,
"configured": self.llm_config.api_key is not None,
"timestamp": datetime.utcnow().isoformat()
}
@app.get("/.well-known/agent.json")
async def get_agent_card(request: Request):
"""获取Agent Card (A2A发现协议)"""
base_url = str(request.base_url).rstrip("/")
card = AgentCard(
name=self.agent_config.name,
description=self.agent_config.description,
version=self.agent_config.version,
url=base_url,
capabilities=AgentCapabilities(
text=True,
streaming=self.agent_config.enable_streaming,
push_notifications=False
),
skills=[
AgentSkill(
id="intelligent-search",
name="智能搜索",
description="理解用户查询意图,自动规划搜索策略,从多个来源获取信息并生成高质量、有来源引用的答案"
)
]
)
return card.model_dump()
@app.post("/message/send")
async def send_message(request: Request):
"""A2A message/send 端点"""
body = await request.json()
# 解析JSON-RPC请求
try:
rpc_request = A2ARequest(**body)
except Exception as e:
return JSONResponse({
"jsonrpc": "2.0",
"id": body.get("id", "unknown"),
"error": {
"code": -32600,
"message": f"Invalid Request: {str(e)}"
}
})
# 处理 message/send 方法
if rpc_request.method == "message/send":
return await self._handle_message_send(rpc_request)
elif rpc_request.method == "message/stream":
return await self._handle_message_stream(rpc_request)
else:
return JSONResponse({
"jsonrpc": "2.0",
"id": rpc_request.id,
"error": {
"code": -32601,
"message": f"Method not found: {rpc_request.method}"
}
})
@app.post("/message/stream")
async def stream_message(request: Request):
"""A2A message/stream 端点 (SSE流式响应)"""
body = await request.json()
try:
rpc_request = A2ARequest(**body)
except Exception as e:
return JSONResponse({
"jsonrpc": "2.0",
"id": body.get("id", "unknown"),
"error": {
"code": -32600,
"message": f"Invalid Request: {str(e)}"
}
})
return await self._handle_message_stream(rpc_request)
@app.get("/tasks/{task_id}")
async def get_task(task_id: str):
"""获取任务状态"""
if task_id not in self.tasks:
raise HTTPException(status_code=404, detail="Task not found")
return self.tasks[task_id].model_dump()
async def _handle_message_send(self, request: A2ARequest) -> JSONResponse:
"""处理 message/send 请求"""
params = request.params or {}
message_data = params.get("message", {})
# 提取API key(必须从请求参数中获取,等同于API格式版本的llm_api_key)
api_key = params.get("api_key")
if not api_key:
return JSONResponse({
"jsonrpc": "2.0",
"id": request.id,
"error": {
"code": -32602,
"message": "Invalid params: api_key is required"
}
})
# 提取model(从环境变量获取,不支持在请求中传递,与API格式版本保持一致)
# 支持多种环境变量名称:MODEL_NAME(优先)、LLM_MODEL(AKS部署)、LITELLM_MODEL
model = os.getenv("MODEL_NAME") or os.getenv("LLM_MODEL") or os.getenv("LITELLM_MODEL")
if not model:
return JSONResponse({
"jsonrpc": "2.0",
"id": request.id,
"error": {
"code": -32000,
"message": "Model not configured: MODEL_NAME or LLM_MODEL environment variable is required"
}
})
# 提取用户消息文本
user_text = ""
parts = message_data.get("parts", [])
for part in parts:
if part.get("kind") == "text":
user_text += part.get("text", "")
if not user_text:
return JSONResponse({
"jsonrpc": "2.0",
"id": request.id,
"error": {
"code": -32602,
"message": "Invalid params: no text content found"
}
})
# 创建任务
task_id = uuid.uuid4().hex
context_id = params.get("contextId", uuid.uuid4().hex)
task = A2ATask(
id=task_id,
contextId=context_id,
status=A2ATaskStatus(state="working")
)
self.tasks[task_id] = task
try:
# 获取Agent实例(使用请求中的api_key和环境变量中的model)
agent = self._get_agent(api_key=api_key, model=model)
# 调用Agent获取响应
logger.info("处理搜索消息", task_id=task_id, message_preview=user_text[:50])
response = await agent.search(query=user_text)
# 如果创建了新Agent(使用了请求中的api_key),关闭它
await agent.close()
# 构建答案文本(包含来源信息)
answer_parts = [response.answer.content]
if response.answer.sources:
answer_parts.append("\n\n## 来源")
for i, source in enumerate(response.answer.sources, 1):
answer_parts.append(f"{i}. [{source.title}]({source.url})")
answer_text = "\n".join(answer_parts)
# 更新任务状态
task.status = A2ATaskStatus(state="completed")
task.artifacts = [
A2AArtifact(
name="response",
parts=[A2APart(kind="text", text=answer_text)]
)
]
self.tasks[task_id] = task
return JSONResponse({
"jsonrpc": "2.0",
"id": request.id,
"result": task.model_dump()
})
except Exception as e:
logger.error("处理搜索消息失败", error=str(e))
task.status = A2ATaskStatus(state="failed", message=str(e))
self.tasks[task_id] = task
return JSONResponse({
"jsonrpc": "2.0",
"id": request.id,
"error": {
"code": -32000,
"message": f"Agent error: {str(e)}"
}
})
async def _handle_message_stream(self, request: A2ARequest) -> StreamingResponse:
"""处理 message/stream 请求 (SSE)"""
params = request.params or {}
message_data = params.get("message", {})
# 提取API key(必须从请求参数中获取)
api_key = params.get("api_key")
if not api_key:
# 对于流式响应,需要通过SSE发送错误
async def error_generator():
error_event = {
"kind": "task-error",
"taskId": "unknown",
"contextId": "unknown",
"data": {
"error": "Invalid params: api_key is required"
}
}
yield f"data: {json.dumps(error_event)}\n\n"
return StreamingResponse(
error_generator(),
media_type="text/event-stream"
)
# 提取model(从环境变量获取)
model = os.getenv("MODEL_NAME") or os.getenv("LLM_MODEL") or os.getenv("LITELLM_MODEL")
if not model:
async def error_generator():
error_event = {
"kind": "task-error",
"taskId": "unknown",
"contextId": "unknown",
"data": {
"error": "Model not configured: MODEL_NAME or LLM_MODEL environment variable is required"
}
}
yield f"data: {json.dumps(error_event)}\n\n"
return StreamingResponse(
error_generator(),
media_type="text/event-stream"
)
# 提取用户消息
user_text = ""
parts = message_data.get("parts", [])
for part in parts:
if part.get("kind") == "text":
user_text += part.get("text", "")
task_id = uuid.uuid4().hex
context_id = params.get("contextId", uuid.uuid4().hex)
async def event_generator() -> AsyncGenerator[str, None]:
"""生成SSE事件流"""
agent = None
try:
# 获取Agent实例(使用请求中的api_key和环境变量中的model)
agent = self._get_agent(api_key=api_key, model=model)
# 发送任务开始事件
start_event = {
"kind": "task-start",
"taskId": task_id,
"contextId": context_id
}
yield f"data: {json.dumps(start_event)}\n\n"
# 执行搜索(SearchAgent不支持流式,所以发送完整结果)
response = await agent.search(query=user_text)
# 构建答案文本
answer_parts = [response.answer.content]
if response.answer.sources:
answer_parts.append("\n\n## 来源")
for i, source in enumerate(response.answer.sources, 1):
answer_parts.append(f"{i}. [{source.title}]({source.url})")
answer_text = "\n".join(answer_parts)
# 发送完整答案(作为增量发送,以便显示进度)
# 将答案分成小块发送以模拟流式效果
chunk_size = 100
for i in range(0, len(answer_text), chunk_size):
chunk = answer_text[i:i + chunk_size]
delta_event = {
"kind": "artifact-delta",
"taskId": task_id,
"contextId": context_id,
"data": {
"kind": "text",
"text": chunk
}
}
yield f"data: {json.dumps(delta_event)}\n\n"
# 添加小延迟以模拟真实流式效果
await asyncio.sleep(0.01)
# 发送完成事件
complete_event = {
"kind": "task-complete",
"taskId": task_id,
"contextId": context_id,
"data": {
"status": "completed",
"artifacts": [{
"name": "response",
"parts": [{"kind": "text", "text": answer_text}]
}]
}
}
yield f"data: {json.dumps(complete_event)}\n\n"
except Exception as e:
# 发送错误事件
error_event = {
"kind": "task-error",
"taskId": task_id,
"contextId": context_id,
"data": {
"error": str(e)
}
}
yield f"data: {json.dumps(error_event)}\n\n"
finally:
# 如果创建了新Agent(使用了请求中的api_key),关闭它
if agent:
await agent.close()
return StreamingResponse(
event_generator(),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no"
}
)
def run(self, host: Optional[str] = None, port: Optional[int] = None):
"""运行服务器"""
import uvicorn
host = host or self.agent_config.host
port = port or self.agent_config.port
logger.info(f"启动A2A Search Agent服务", host=host, port=port)
uvicorn.run(self.app, host=host, port=port)
def create_app(api_key: Optional[str] = None, model: Optional[str] = None) -> FastAPI:
"""
创建FastAPI应用(用于uvicorn启动)
使用方式:
uvicorn a2a_server:app --host 0.0.0.0 --port 8080
或设置环境变量后:
export LITELLM_API_KEY="your-key"
export MODEL_NAME="your-model"
uvicorn a2a_server:app --host 0.0.0.0 --port 8080
"""
server = A2ASearchAgentServer(api_key=api_key, model=model)
return server.app
# uvicorn 启动入口
# 环境变量: LITELLM_API_KEY, MODEL_NAME (或 LITELLM_MODEL)
# 注意: app 只在 main.py 中创建,避免导入时立即执行验证
# app = create_app()
@@ -0,0 +1,104 @@
"""
Search Agent 核心模块 - A2A版本
基于LiteLLM和A2A协议的搜索Agent实现
"""
import sys
import os
# 先导入当前目录的config
from typing import Optional
from loguru import logger
from config import LiteLLMConfig, AgentConfig, get_config
# 然后添加search_agent目录到Python路径
search_agent_dir = os.path.join(os.path.dirname(__file__), 'search_agent')
if search_agent_dir not in sys.path:
sys.path.insert(0, search_agent_dir)
from search_agent.config import Config
from search_agent.agent.search_agent import SearchAgent as CoreSearchAgent
class SearchAgentWrapper:
"""
Search Agent包装器
用于适配A2A框架,将SearchAgent包装为可配置的Agent实例
"""
def __init__(
self,
litellm_config: Optional[LiteLLMConfig] = None,
agent_config: Optional[AgentConfig] = None,
api_key: Optional[str] = None,
model: Optional[str] = None
):
"""
初始化Search Agent
Args:
litellm_config: LiteLLM配置对象
agent_config: Agent配置对象
api_key: LiteLLM API密钥(可选,优先使用)
model: 模型名称(可选,优先使用)
"""
# 获取配置
if not litellm_config:
llm_config, _, _ = get_config(api_key=api_key, model=model)
else:
llm_config = litellm_config
if not agent_config:
_, agent_config, _ = get_config(api_key=api_key, model=model)
self.litellm_config = llm_config
self.agent_config = agent_config
# 验证配置
self.litellm_config.validate()
# 创建SearchAgent配置(使用litellm的base_url和api_key)
# 需要从环境变量获取其他配置
serper_api_key = os.getenv("SERPER_API_KEY", "")
jina_api_key = os.getenv("JINA_API_KEY", "")
self.search_config = Config(
llm_base_url=llm_config.base_url,
llm_api_key=llm_config.api_key,
llm_model=llm_config.model,
serper_api_key=serper_api_key,
jina_api_key=jina_api_key,
max_iterations=int(os.getenv("MAX_ITERATIONS", "3")),
max_results_per_query=int(os.getenv("MAX_RESULTS_PER_QUERY", "10")),
content_max_length=int(os.getenv("CONTENT_MAX_LENGTH", "5000")),
log_level=os.getenv("LOG_LEVEL", "INFO"),
timeout=int(os.getenv("TIMEOUT", "30"))
)
# 创建SearchAgent实例
self.agent = CoreSearchAgent(self.search_config)
logger.info(
"SearchAgent初始化完成",
agent_name=self.agent_config.name,
model=self.litellm_config.model,
base_url=self.litellm_config.base_url
)
async def search(self, query: str):
"""
执行搜索
Args:
query: 搜索查询
Returns:
AgentResponse对象
"""
return await self.agent.search(query)
async def close(self):
"""关闭资源(SearchAgent不需要特殊清理)"""
pass
@@ -0,0 +1,174 @@
"""
A2A Search Agent Executor
使用官方 A2A SDK 的 AgentExecutor 实现
"""
import os
from typing import Optional
from loguru import logger
from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.server.events import EventQueue
from a2a.utils import new_agent_text_message
from agent import SearchAgentWrapper
from config import get_config
class SearchAgentExecutor(AgentExecutor):
"""
Search Agent Executor
继承自 A2A SDK 的 AgentExecutor,实现搜索功能
"""
def __init__(
self,
default_api_key: Optional[str] = None,
default_model: Optional[str] = None
):
"""
初始化 Search Agent Executor
Args:
default_api_key: 默认 API 密钥(可选,从环境变量获取)
default_model: 默认模型名称(可选,从环境变量获取)
"""
# 从环境变量获取默认配置(如果未提供)
if not default_api_key:
default_api_key = os.getenv("LITELLM_API_KEY") or os.getenv("LLM_API_KEY")
if not default_model:
default_model = (
os.getenv("MODEL_NAME") or
os.getenv("LLM_MODEL") or
os.getenv("LITELLM_MODEL")
)
self.default_api_key = default_api_key
self.default_model = default_model
logger.info(
"SearchAgentExecutor 初始化完成",
has_default_api_key=bool(default_api_key),
default_model=default_model
)
async def execute(
self,
context: RequestContext,
event_queue: EventQueue,
) -> None:
"""
执行搜索任务
Args:
context: A2A SDK 提供的请求上下文
event_queue: A2A SDK 提供的事件队列,用于发送响应
"""
try:
# 从请求中提取用户消息
message = context.message
if not message:
error_msg = "未找到消息内容"
logger.warning(error_msg)
await event_queue.enqueue_event(
new_agent_text_message(f"错误: {error_msg}")
)
return
user_text = ""
# 提取文本内容(支持多个 text parts)
# Part对象有root属性,root才是TextPart等具体类型
for part in message.parts:
if hasattr(part, 'root') and part.root:
root = part.root
if hasattr(root, 'kind') and root.kind == "text":
if hasattr(root, 'text') and root.text:
user_text += root.text
if not user_text:
error_msg = "未找到文本内容"
logger.warning(error_msg)
await event_queue.enqueue_event(
new_agent_text_message(f"错误: {error_msg}")
)
return
logger.info("处理搜索请求", message_preview=user_text[:50])
# 从请求的 metadata 中获取 api_key 和 model(优先使用请求中的)
metadata = context.metadata or {}
api_key = metadata.get("api_key") or self.default_api_key
model = metadata.get("model") or self.default_model
if not api_key:
error_msg = "API 密钥未提供,请在请求参数中提供 api_key 或设置 LITELLM_API_KEY 环境变量"
logger.error(error_msg)
await event_queue.enqueue_event(
new_agent_text_message(f"错误: {error_msg}")
)
return
if not model:
error_msg = "模型名称未配置,请设置 MODEL_NAME 或 LLM_MODEL 环境变量"
logger.error(error_msg)
await event_queue.enqueue_event(
new_agent_text_message(f"错误: {error_msg}")
)
return
# 创建 SearchAgent 实例(每次请求创建新实例,使用请求中的 api_key)
agent = SearchAgentWrapper(api_key=api_key, model=model)
try:
# 执行搜索
response = await agent.search(query=user_text)
# 构建答案文本(包含来源信息)
answer_parts = [response.answer.content]
if response.answer.sources:
answer_parts.append("\n\n## 来源")
for i, source in enumerate(response.answer.sources, 1):
answer_parts.append(f"{i}. [{source.title}]({source.url})")
answer_text = "\n".join(answer_parts)
# 通过 event_queue 发送响应(SDK 自动处理格式)
await event_queue.enqueue_event(
new_agent_text_message(answer_text)
)
logger.info("搜索请求处理完成", sources_count=len(response.answer.sources) if response.answer.sources else 0)
finally:
# 关闭 Agent 实例
await agent.close()
except Exception as e:
error_msg = f"处理搜索请求失败: {str(e)}"
logger.error(error_msg, error=str(e))
await event_queue.enqueue_event(
new_agent_text_message(f"错误: {error_msg}")
)
async def cancel(
self,
context: RequestContext,
event_queue: EventQueue
) -> None:
"""
取消任务
Args:
context: 请求上下文
event_queue: 事件队列
"""
logger.info("取消搜索任务", task_id=context.task_id)
# SearchAgent 当前不支持取消,但可以记录日志
await event_queue.enqueue_event(
new_agent_text_message("任务取消功能暂不支持")
)
@@ -0,0 +1,138 @@
"""
A2A Search Agent 配置模块
支持用户传入密钥和模型名称,同时支持从环境变量获取
"""
import os
from dataclasses import dataclass, field
from typing import Optional
from dotenv import load_dotenv
# 加载环境变量
load_dotenv()
@dataclass
class LiteLLMConfig:
"""LiteLLM 配置"""
# 基础URL - 用户提供的LiteLLM服务地址
# 优先使用 LLM_BASE_URL(与API格式保持一致),也支持 LITELLM_BASE_URL(向后兼容)
base_url: str = field(default_factory=lambda: os.getenv(
"LLM_BASE_URL"
) or os.getenv(
"LITELLM_BASE_URL",
"https://litellm.graystone-fb459c5d.southeastasia.azurecontainerapps.io"
))
# 完整的chat completions端点
chat_endpoint: str = field(init=False)
# API密钥 - 优先使用传入的,否则从环境变量获取
api_key: Optional[str] = None
# 模型名称 - 优先使用传入的,否则从环境变量获取
model: Optional[str] = None
# 请求超时时间(秒)
timeout: int = 120
# 温度参数
temperature: float = 0.7
# 最大token数
max_tokens: int = 4096
def __post_init__(self):
self.chat_endpoint = f"{self.base_url}/chat/completions"
# 从环境变量读取(如果未直接提供)
# API密钥:优先使用 LITELLM_API_KEY(LiteLLM约定),也支持 LLM_API_KEY(向后兼容)
if self.api_key is None:
self.api_key = os.getenv("LITELLM_API_KEY") or os.getenv("LLM_API_KEY")
if self.model is None:
# 支持多种环境变量名称:
# 1. MODEL_NAME - 与API格式保持一致(优先)
# 2. LLM_MODEL - AKS部署配置使用(必须支持)
# 3. LITELLM_MODEL - LiteLLM标准约定(向后兼容)
self.model = os.getenv("MODEL_NAME") or os.getenv("LLM_MODEL") or os.getenv("LITELLM_MODEL", "gpt-4")
def validate(self) -> bool:
"""验证配置是否完整"""
if not self.api_key:
raise ValueError("API密钥未设置! 请设置 LITELLM_API_KEY 或 LLM_API_KEY 环境变量或直接传入 api_key")
if not self.model:
raise ValueError("模型名称未设置! 请设置 MODEL_NAME 环境变量或直接传入 model")
return True
@dataclass
class AgentConfig:
"""Agent 配置"""
# Agent名称
name: str = "search-agent"
# Agent描述
description: str = "智能AI搜索Agent,基于LiteLLM和A2A协议,能够理解用户查询意图、自动规划搜索策略、从多个来源获取信息,并生成高质量、有来源引用的答案"
# Agent版本
version: str = "1.0.0"
# 服务端口
port: int = 8080
# 服务主机
host: str = "0.0.0.0"
# 是否启用流式响应
enable_streaming: bool = True
# 系统提示词
system_prompt: str = "你是一个智能搜索助手。"
@dataclass
class A2AConfig:
"""A2A协议配置"""
# A2A协议版本
protocol_version: str = "1.0"
# Agent Card配置
agent_card: dict = field(default_factory=lambda: {
"name": "search-agent",
"description": "智能AI搜索Agent,支持A2A协议通信",
"version": "1.0.0",
"capabilities": {
"text": True,
"streaming": True,
"push_notifications": False
},
"skills": [
{
"id": "intelligent-search",
"name": "智能搜索",
"description": "理解用户查询意图,自动规划搜索策略,从多个来源获取信息"
}
]
})
def get_config(
api_key: Optional[str] = None,
model: Optional[str] = None
) -> tuple[LiteLLMConfig, AgentConfig, A2AConfig]:
"""
获取完整配置
Args:
api_key: LiteLLM API密钥(可选,优先使用,否则从环境变量获取)
model: 模型名称(可选,优先使用,否则从环境变量获取)
Returns:
(LiteLLMConfig, AgentConfig, A2AConfig) 配置元组
"""
litellm_config = LiteLLMConfig(api_key=api_key, model=model)
agent_config = AgentConfig()
a2a_config = A2AConfig()
return litellm_config, agent_config, a2a_config
@@ -0,0 +1,108 @@
"""
A2A Search Agent 主入口(使用官方 A2A SDK)
支持从环境变量或请求传入 API key
"""
import os
import uvicorn
from a2a.server.apps.jsonrpc.fastapi_app import A2AFastAPIApplication
from a2a.server.request_handlers.default_request_handler import DefaultRequestHandler
from a2a.server.tasks import InMemoryTaskStore
from a2a.server.events import InMemoryQueueManager
from a2a.types import AgentCard, AgentCapabilities
from agent_executor import SearchAgentExecutor
from config import AgentConfig
# 环境变量配置
SERVICE_HOST = os.getenv("SERVICE_HOST", "0.0.0.0")
SERVICE_PORT = int(os.getenv("SERVICE_PORT", "8080"))
POD_NAME = os.getenv("POD_NAME", "search-agent-a2a")
TEMPLATE_TYPE = os.getenv("TEMPLATE_TYPE", "search_agent_A2A")
# 从环境变量获取默认配置(可选)
# 支持多种环境变量名称(向后兼容)
default_api_key = os.getenv("LITELLM_API_KEY") or os.getenv("LLM_API_KEY")
# 支持多种环境变量名称:MODEL_NAME(优先)、LLM_MODEL(AKS部署)、LITELLM_MODEL
default_model = os.getenv("MODEL_NAME") or os.getenv("LLM_MODEL") or os.getenv("LITELLM_MODEL")
# 获取 Agent 配置
agent_config = AgentConfig()
# 创建 SearchAgentExecutor 实例
executor = SearchAgentExecutor(
default_api_key=default_api_key,
default_model=default_model
)
# 创建 Agent Card
agent_card = AgentCard(
name=agent_config.name,
description=agent_config.description,
version=agent_config.version,
url=f"http://{SERVICE_HOST}:{SERVICE_PORT}",
default_input_modes=["text"],
default_output_modes=["text"],
capabilities=AgentCapabilities(
text=True,
streaming=agent_config.enable_streaming,
push_notifications=False,
forms=False,
files=False,
),
skills=[
{
"id": "intelligent-search",
"name": "智能搜索",
"description": "理解用户查询意图,自动规划搜索策略,从多个来源获取信息并生成高质量、有来源引用的答案",
"tags": []
}
]
)
# 创建任务存储和队列管理器
task_store = InMemoryTaskStore()
queue_manager = InMemoryQueueManager()
# 创建请求处理器
http_handler = DefaultRequestHandler(
agent_executor=executor,
task_store=task_store,
queue_manager=queue_manager,
)
# 使用官方 A2A SDK 创建应用
# SDK 自动处理所有 A2A 协议细节(JSON-RPC、Agent Card、任务状态等)
a2a_app = A2AFastAPIApplication(
agent_card=agent_card,
http_handler=http_handler,
)
# 构建 FastAPI 应用实例
app = a2a_app.build()
def main():
"""主函数"""
print(f"🚀 启动 A2A Search Agent (使用官方 A2A SDK)")
print(f" - Pod名称: {POD_NAME}")
print(f" - 模板类型: {TEMPLATE_TYPE}")
print(f" - 服务地址: http://{SERVICE_HOST}:{SERVICE_PORT}")
print(f" - Agent名称: {agent_config.name}")
print(f" - Agent版本: {agent_config.version}")
if default_api_key:
print(f" - 已配置默认 API key(可通过请求覆盖)")
else:
print(f" - 未配置默认 API key,需在请求中传入")
if default_model:
print(f" - 默认模型: {default_model}")
uvicorn.run(
app,
host=SERVICE_HOST,
port=SERVICE_PORT,
log_level="info"
)
if __name__ == "__main__":
main()
@@ -0,0 +1,30 @@
# A2A SDK - 官方 Agent2Agent 协议框架
a2a-sdk[http-server]>=0.3.0
# FastAPI 和 Web 服务器(a2a-sdk 依赖,但显式声明版本)
fastapi>=0.109.0
uvicorn[standard]>=0.27.0
pydantic>=2.5.3
# HTTP客户端 - 用于LiteLLM SDK调用
httpx>=0.27.0
# 环境变量
python-dotenv>=1.0.0
# JSON处理
orjson>=3.9.0
# 类型提示
typing-extensions>=4.9.0
# 日志
loguru>=0.7.0
# 异步工具
asyncio-throttle>=1.0.2
# HTTP客户端 - 用于搜索和其他API调用
aiohttp>=3.9.0
requests>=2.31.0
@@ -0,0 +1,2 @@
# Search Agent A2A Package
@@ -0,0 +1,18 @@
"""
Agent模块
"""
from .search_agent import SearchAgent
from .prompts import (
QUERY_ANALYSIS_PROMPT,
ANSWER_GENERATION_PROMPT,
REFLECTION_PROMPT,
)
__all__ = [
"SearchAgent",
"QUERY_ANALYSIS_PROMPT",
"ANSWER_GENERATION_PROMPT",
"REFLECTION_PROMPT",
]
@@ -0,0 +1,126 @@
"""
Prompt模板汇总
集中管理所有LLM Prompt模板
"""
# ==================== 查询分析 Prompt ====================
QUERY_ANALYSIS_PROMPT = """你是一个查询分析专家。分析用户的搜索查询,提取以下信息。
请输出JSON格式:
{
"intent": "查询意图,必须是以下之一: fact_check(事实核查), comparison(对比分析), how_to(操作指南), news(新闻资讯), research(深度研究)",
"entities": ["关键实体列表,提取查询中的核心概念、人名、产品名等"],
"expanded_queries": ["扩展查询1", "扩展查询2", "扩展查询3"],
"need_news": true或false,
"time_filter": "时间过滤器,null表示不限时间,qdr:d(过去24小时), qdr:w(过去一周), qdr:m(过去一月), qdr:y(过去一年)"
}
扩展查询要求:
1. 生成2-4个扩展查询,包含不同角度或同义表达
2. 至少包含一个英文查询(如果原查询是中文)
3. 保持查询的核心意图
时间过滤器选择规则:
- 查询涉及"最新"、"近期"、"今年"等时效性词语 → 设置相应的时间过滤器
- 查询涉及具体年份(如"2024年") → qdr:y
- 一般性查询 → null"""
# ==================== 搜索规划 Prompt ====================
SEARCH_PLANNING_PROMPT = """你是一个搜索规划专家。根据查询分析结果,制定搜索计划。
输入信息:
- 原始查询
- 查询意图
- 关键实体
- 是否需要新闻
输出搜索任务列表,每个任务包含:
- query: 搜索词
- source: web 或 news
- time_filter: 时间过滤器(可选)
搜索策略规则:
1. 简单事实查询 → 单次Web搜索
2. 时效性查询 → Web搜索 + 新闻搜索
3. 复杂分析查询 → 多个扩展查询
4. 对比类查询 → 分别搜索各对比对象"""
# ==================== 答案生成 Prompt ====================
ANSWER_GENERATION_PROMPT = """你是一个专业的信息整合专家。根据以下搜索结果,回答用户的问题。
## 要求
1. 综合多个来源的信息,给出全面准确的回答
2. 使用清晰的结构组织答案(标题、列表、重点标注等)
3. 在答案中标注信息来源,格式:[来源1]、[来源2]
4. 如果信息有冲突,说明不同观点
5. 如果信息不足以完整回答问题,明确指出缺失的部分
6. 回答使用中文
## 输出JSON格式
{
"answer": "结构化的答案(Markdown格式,包含来源引用)",
"sources": [
{"index": 1, "title": "来源标题", "url": "来源URL"},
{"index": 2, "title": "来源标题", "url": "来源URL"}
],
"confidence": "high/medium/low,基于信息质量和一致性判断"
}"""
# ==================== 反思评估 Prompt ====================
REFLECTION_PROMPT = """你是一个质量评估专家。评估以下答案是否充分回答了用户的问题。
## 评估维度
1. **完整性**: 答案是否覆盖了问题的所有方面?
2. **准确性**: 答案内容是否有明确的来源支持?
3. **深度**: 答案是否提供了足够的细节和解释?
## 输出JSON格式
{
"completeness": 0.0-1.0,
"missing_aspects": ["如果有缺失,列出缺失的方面"],
"needs_more_search": true或false,
"suggested_queries": ["如果需要补充搜索,建议的搜索词"]
}
## 判断标准
- completeness >= 0.8 且没有重要信息缺失 → needs_more_search = false
- completeness < 0.8 或有重要信息缺失 → needs_more_search = true
- 建议的搜索词应该针对缺失的方面"""
# ==================== 工具函数 ====================
def format_query_analysis_prompt(query: str) -> str:
"""格式化查询分析Prompt"""
return f"{QUERY_ANALYSIS_PROMPT}\n\n用户查询: {query}"
def format_answer_generation_prompt(query: str, documents: str) -> str:
"""格式化答案生成Prompt"""
return f"""{ANSWER_GENERATION_PROMPT}
## 用户问题
{query}
## 搜索结果
{documents}"""
def format_reflection_prompt(query: str, answer: str, sources_count: int, confidence: str) -> str:
"""格式化反思评估Prompt"""
return f"""{REFLECTION_PROMPT}
## 用户问题
{query}
## 生成的答案
{answer}
## 答案的来源数量
{sources_count} 个来源
## 答案的置信度
{confidence}"""
@@ -0,0 +1,209 @@
"""
搜索Agent主类
协调各模块执行智能搜索
"""
from typing import List, Optional
from loguru import logger
from search_agent.config import Config
from models.schemas import (
QueryAnalysis,
SearchPlan,
SearchResult,
Document,
RankedDocument,
Answer,
AgentResponse,
)
from modules.query_analyzer import QueryAnalyzer
from modules.search_planner import SearchPlanner
from modules.search_executor import SearchExecutor
from modules.content_extractor import ContentExtractor
from modules.result_processor import ResultProcessor
from modules.answer_generator import AnswerGenerator
from modules.reflector import Reflector
class SearchAgent:
"""智能搜索Agent"""
def __init__(self, config: Config):
"""
初始化搜索Agent
Args:
config: 配置对象
"""
self.config = config
# 初始化各模块
self.query_analyzer = QueryAnalyzer(config)
self.search_planner = SearchPlanner(config)
self.search_executor = SearchExecutor(config)
self.content_extractor = ContentExtractor(config)
self.result_processor = ResultProcessor(config)
self.answer_generator = AnswerGenerator(config)
self.reflector = Reflector(config)
logger.info("SearchAgent 初始化完成")
async def search(self, query: str) -> AgentResponse:
"""
执行智能搜索
Args:
query: 用户查询
Returns:
AgentResponse对象
"""
logger.info(f"="*60)
logger.info(f"开始搜索: {query}")
logger.info(f"="*60)
iteration = 0
all_documents: List[Document] = []
all_queries: List[str] = []
# 1. 查询理解
analysis = await self.query_analyzer.analyze(query)
logger.info(f"查询分析完成: intent={analysis.intent.value}")
answer: Optional[Answer] = None
while iteration < self.config.max_iterations:
iteration += 1
logger.info(f"\n--- 迭代 {iteration}/{self.config.max_iterations} ---")
# 2. 搜索规划
if iteration == 1:
plan = await self.search_planner.plan(analysis)
else:
# 后续迭代使用建议的补充查询
plan = self.search_planner.plan_supplementary(
query,
analysis.expanded_queries
)
all_queries.extend([t.query for t in plan.tasks])
logger.info(f"搜索计划: {len(plan.tasks)} 个任务")
# 3. 执行搜索
search_results = await self.search_executor.execute(plan)
logger.info(f"搜索结果: {len(search_results)} 条")
if not search_results:
logger.warning("没有搜索结果")
if answer is None:
answer = self.answer_generator._empty_answer()
break
# 4. 内容提取
documents = await self.content_extractor.extract_batch(
search_results,
max_urls=10
)
all_documents.extend(documents)
logger.info(f"提取文档: {len(documents)} 个")
if not documents:
logger.warning("没有成功提取到文档内容")
continue
# 5. 结果处理(去重+重排序)
ranked_docs = await self.result_processor.process(
query=query,
documents=all_documents,
top_k=5
)
logger.info(f"排序结果: {len(ranked_docs)} 个")
if not ranked_docs:
logger.warning("没有有效的排序结果")
continue
# 6. 生成答案
answer = await self.answer_generator.generate(
query=query,
documents=ranked_docs
)
logger.info(f"答案生成完成: confidence={answer.confidence}")
# 7. 反思评估
assessment = await self.reflector.assess(query, answer)
# 8. 判断是否继续迭代
if not self.reflector.should_continue(assessment, iteration):
break
# 更新分析,准备下一轮搜索
if assessment.suggested_queries:
analysis.expanded_queries = assessment.suggested_queries
logger.info(f"补充搜索: {assessment.suggested_queries}")
# 确保有答案返回
if answer is None:
answer = self.answer_generator._empty_answer()
# 去重统计
unique_urls = set(d.url for d in all_documents)
response = AgentResponse(
answer=answer,
iterations=iteration,
total_sources_consulted=len(unique_urls),
search_queries_used=list(set(all_queries))
)
logger.info(f"\n{'='*60}")
logger.info(f"搜索完成!")
logger.info(f"迭代次数: {iteration}")
logger.info(f"参考来源: {len(unique_urls)}")
logger.info(f"搜索查询: {len(response.search_queries_used)}")
logger.info(f"{'='*60}\n")
return response
async def quick_search(self, query: str) -> Answer:
"""
快速搜索(单次迭代)
Args:
query: 用户查询
Returns:
Answer对象
"""
# 简化分析
analysis = await self.query_analyzer.analyze(query)
# 只执行一次搜索
plan = await self.search_planner.plan(analysis)
plan.tasks = plan.tasks[:2] # 限制搜索任务数量
# 执行搜索
search_results = await self.search_executor.execute(plan)
if not search_results:
return self.answer_generator._empty_answer()
# 提取内容
documents = await self.content_extractor.extract_batch(
search_results,
max_urls=5
)
if not documents:
return self.answer_generator._empty_answer()
# 处理结果
ranked_docs = await self.result_processor.process(
query=query,
documents=documents,
top_k=3
)
# 生成答案
return await self.answer_generator.generate(query, ranked_docs)
@@ -0,0 +1,81 @@
"""
配置管理模块
负责加载和管理所有配置项
"""
import os
from dataclasses import dataclass
from typing import Optional
from dotenv import load_dotenv
@dataclass
class Config:
"""Agent配置类"""
# LLM配置
llm_base_url: str
llm_api_key: str
llm_model: str
# Serper配置
serper_api_key: str
# Jina配置
jina_api_key: str
# Agent配置
max_iterations: int
max_results_per_query: int
content_max_length: int
# 可选配置
log_level: str = "INFO"
timeout: int = 30
@classmethod
def from_env(cls, env_path: Optional[str] = None) -> "Config":
"""从环境变量加载配置"""
if env_path:
load_dotenv(env_path)
else:
load_dotenv()
return cls(
# LLM配置
llm_base_url=os.getenv("LLM_BASE_URL", ""),
llm_api_key=os.getenv("LLM_API_KEY", ""),
llm_model=os.getenv("MODEL_NAME", "xchat52"),
# Serper配置
serper_api_key=os.getenv("SERPER_API_KEY", ""),
# Jina配置
jina_api_key=os.getenv("JINA_API_KEY", ""),
# Agent配置
max_iterations=int(os.getenv("MAX_ITERATIONS", "3")),
max_results_per_query=int(os.getenv("MAX_RESULTS_PER_QUERY", "10")),
content_max_length=int(os.getenv("CONTENT_MAX_LENGTH", "5000")),
# 可选配置
log_level=os.getenv("LOG_LEVEL", "INFO"),
timeout=int(os.getenv("TIMEOUT", "30"))
)
def validate(self) -> bool:
"""验证配置是否完整"""
required_fields = [
("llm_base_url", self.llm_base_url),
("llm_api_key", self.llm_api_key),
("serper_api_key", self.serper_api_key),
("jina_api_key", self.jina_api_key),
]
missing = [name for name, value in required_fields if not value]
if missing:
raise ValueError(f"缺少必要的配置项: {', '.join(missing)}")
return True
@@ -0,0 +1,34 @@
"""
数据模型模块
"""
from .schemas import (
SearchSource,
Intent,
QueryAnalysis,
SearchTask,
SearchPlan,
SearchResult,
Document,
RankedDocument,
Source,
Answer,
QualityAssessment,
AgentResponse,
)
__all__ = [
"SearchSource",
"Intent",
"QueryAnalysis",
"SearchTask",
"SearchPlan",
"SearchResult",
"Document",
"RankedDocument",
"Source",
"Answer",
"QualityAssessment",
"AgentResponse",
]
@@ -0,0 +1,202 @@
"""
数据模型定义
定义Agent使用的所有数据结构
"""
from dataclasses import dataclass, field
from typing import List, Optional
from enum import Enum
class SearchSource(Enum):
"""搜索来源枚举"""
WEB = "web"
NEWS = "news"
class Intent(Enum):
"""查询意图枚举"""
FACT_CHECK = "fact_check" # 事实核查
COMPARISON = "comparison" # 对比分析
HOW_TO = "how_to" # 操作指南
NEWS = "news" # 新闻资讯
RESEARCH = "research" # 深度研究
@dataclass
class QueryAnalysis:
"""查询分析结果"""
original_query: str # 原始查询
intent: Intent # 查询意图
entities: List[str] # 关键实体
expanded_queries: List[str] # 扩展查询列表
need_news: bool # 是否需要新闻搜索
time_filter: Optional[str] = None # 时间过滤器
def to_dict(self) -> dict:
"""转换为字典"""
return {
"original_query": self.original_query,
"intent": self.intent.value,
"entities": self.entities,
"expanded_queries": self.expanded_queries,
"need_news": self.need_news,
"time_filter": self.time_filter
}
@dataclass
class SearchTask:
"""搜索任务"""
query: str # 搜索查询
source: SearchSource # 搜索来源
time_filter: Optional[str] = None # 时间过滤器
num_results: int = 10 # 结果数量
def to_dict(self) -> dict:
"""转换为字典"""
return {
"query": self.query,
"source": self.source.value,
"time_filter": self.time_filter,
"num_results": self.num_results
}
@dataclass
class SearchPlan:
"""搜索计划"""
tasks: List[SearchTask] # 搜索任务列表
strategy: str = "parallel" # 执行策略: parallel/sequential
def to_dict(self) -> dict:
"""转换为字典"""
return {
"tasks": [t.to_dict() for t in self.tasks],
"strategy": self.strategy
}
@dataclass
class SearchResult:
"""搜索结果"""
title: str # 标题
url: str # URL
snippet: str # 摘要
source: SearchSource # 来源类型
position: int # 排名位置
date: Optional[str] = None # 日期(新闻)
def to_dict(self) -> dict:
"""转换为字典"""
return {
"title": self.title,
"url": self.url,
"snippet": self.snippet,
"source": self.source.value,
"position": self.position,
"date": self.date
}
@dataclass
class Document:
"""提取的文档内容"""
url: str # URL
title: str # 标题
content: str # 内容
source: SearchSource # 来源类型
def to_dict(self) -> dict:
"""转换为字典"""
return {
"url": self.url,
"title": self.title,
"content": self.content,
"source": self.source.value
}
@dataclass
class RankedDocument:
"""排序后的文档"""
document: Document # 文档
relevance_score: float # 相关性分数
rank: int # 排名
def to_dict(self) -> dict:
"""转换为字典"""
return {
"document": self.document.to_dict(),
"relevance_score": self.relevance_score,
"rank": self.rank
}
@dataclass
class Source:
"""来源引用"""
index: int # 索引
title: str # 标题
url: str # URL
def to_dict(self) -> dict:
"""转换为字典"""
return {
"index": self.index,
"title": self.title,
"url": self.url
}
@dataclass
class Answer:
"""生成的答案"""
content: str # Markdown格式的答案内容
sources: List[Source] # 来源列表
confidence: str # 置信度: high/medium/low
def to_dict(self) -> dict:
"""转换为字典"""
return {
"content": self.content,
"sources": [s.to_dict() for s in self.sources],
"confidence": self.confidence
}
@dataclass
class QualityAssessment:
"""质量评估"""
completeness: float # 完整性 0-1
missing_aspects: List[str] # 缺失的方面
needs_more_search: bool # 是否需要更多搜索
suggested_queries: List[str] # 建议的补充搜索
def to_dict(self) -> dict:
"""转换为字典"""
return {
"completeness": self.completeness,
"missing_aspects": self.missing_aspects,
"needs_more_search": self.needs_more_search,
"suggested_queries": self.suggested_queries
}
@dataclass
class AgentResponse:
"""Agent最终响应"""
answer: Answer # 答案
iterations: int # 迭代次数
total_sources_consulted: int # 参考来源总数
search_queries_used: List[str] # 使用的搜索查询
def to_dict(self) -> dict:
"""转换为字典"""
return {
"answer": self.answer.to_dict(),
"iterations": self.iterations,
"total_sources_consulted": self.total_sources_consulted,
"search_queries_used": self.search_queries_used
}
@@ -0,0 +1,22 @@
"""
核心模块
"""
from .query_analyzer import QueryAnalyzer
from .search_planner import SearchPlanner
from .search_executor import SearchExecutor
from .content_extractor import ContentExtractor
from .result_processor import ResultProcessor
from .answer_generator import AnswerGenerator
from .reflector import Reflector
__all__ = [
"QueryAnalyzer",
"SearchPlanner",
"SearchExecutor",
"ContentExtractor",
"ResultProcessor",
"AnswerGenerator",
"Reflector",
]
@@ -0,0 +1,151 @@
"""
答案生成模块
综合多个来源的信息生成结构化答案
"""
from typing import List
from loguru import logger
from search_agent.config import Config
from models.schemas import RankedDocument, Answer, Source
from utils.llm_client import LLMClient
from utils.helpers import format_documents_for_prompt
# 答案生成Prompt
ANSWER_GENERATION_PROMPT = """你是一个专业的信息整合专家。根据以下搜索结果,回答用户的问题。
## 要求
1. 综合多个来源的信息,给出全面准确的回答
2. 使用清晰的结构组织答案(标题、列表、重点标注等)
3. 在答案中标注信息来源,格式:[来源1]、[来源2]
4. 如果信息有冲突,说明不同观点
5. 如果信息不足以完整回答问题,明确指出缺失的部分
6. 回答使用中文
## 输出JSON格式
{
"answer": "结构化的答案(Markdown格式,包含来源引用)",
"sources": [
{"index": 1, "title": "来源标题", "url": "来源URL"},
{"index": 2, "title": "来源标题", "url": "来源URL"}
],
"confidence": "high/medium/low,基于信息质量和一致性判断"
}"""
class AnswerGenerator:
"""答案生成模块"""
def __init__(self, config: Config):
"""
初始化答案生成器
Args:
config: 配置对象
"""
self.config = config
self.llm = LLMClient(
base_url=config.llm_base_url,
api_key=config.llm_api_key,
model=config.llm_model,
timeout=120 # 答案生成可能需要更长时间
)
async def generate(
self,
query: str,
documents: List[RankedDocument]
) -> Answer:
"""
根据文档生成答案
Args:
query: 用户查询
documents: 排序后的文档列表
Returns:
Answer对象
"""
if not documents:
return self._empty_answer()
logger.info(f"开始生成答案,使用 {len(documents)} 个文档")
# 格式化文档
formatted_docs = format_documents_for_prompt(
documents,
max_length=self.config.content_max_length // len(documents)
)
user_message = f"""## 用户问题
{query}
## 搜索结果
{formatted_docs}"""
try:
result = await self.llm.chat_json(
system_prompt=ANSWER_GENERATION_PROMPT,
user_message=user_message,
temperature=0.5
)
# 解析来源
sources = [
Source(
index=s.get("index", i + 1),
title=s.get("title", ""),
url=s.get("url", "")
)
for i, s in enumerate(result.get("sources", []))
]
answer = Answer(
content=result.get("answer", ""),
sources=sources,
confidence=result.get("confidence", "medium")
)
logger.info(f"答案生成完成,置信度: {answer.confidence}")
return answer
except Exception as e:
logger.error(f"答案生成失败: {e}")
return self._fallback_answer(query, documents)
def _empty_answer(self) -> Answer:
"""生成空答案(无文档时)"""
return Answer(
content="抱歉,未能找到相关信息来回答您的问题。",
sources=[],
confidence="low"
)
def _fallback_answer(
self,
query: str,
documents: List[RankedDocument]
) -> Answer:
"""后备答案生成(LLM失败时)"""
# 简单汇总文档内容
content_parts = [f"关于「{query}」,以下是搜索到的相关信息:\n"]
sources = []
for i, doc in enumerate(documents[:5], 1):
actual_doc = doc.document
content_parts.append(f"### 来源 [{i}]: {actual_doc.title}\n")
content_parts.append(f"{actual_doc.content[:500]}...\n\n")
sources.append(Source(
index=i,
title=actual_doc.title,
url=actual_doc.url
))
return Answer(
content="".join(content_parts),
sources=sources,
confidence="low"
)
@@ -0,0 +1,102 @@
"""
内容提取模块
使用Jina Reader提取网页内容
"""
from typing import List
from loguru import logger
from search_agent.config import Config
from models.schemas import Document, SearchResult, SearchSource
from tools.jina_reader import JinaReaderClient
class ContentExtractor:
"""内容提取模块"""
def __init__(self, config: Config):
"""
初始化内容提取器
Args:
config: 配置对象
"""
self.config = config
self.jina_reader = JinaReaderClient(
api_key=config.jina_api_key,
timeout=config.timeout,
max_content_length=config.content_max_length
)
async def extract(self, search_result: SearchResult) -> Document | None:
"""
从搜索结果提取内容
Args:
search_result: 搜索结果
Returns:
Document对象,如果提取失败则返回None
"""
return await self.jina_reader.extract_content(
url=search_result.url,
source=search_result.source
)
async def extract_batch(
self,
search_results: List[SearchResult],
max_urls: int = 10
) -> List[Document]:
"""
批量提取内容
Args:
search_results: 搜索结果列表
max_urls: 最大提取URL数量
Returns:
Document列表
"""
# 去重并限制数量
seen_urls = set()
unique_results = []
for result in search_results:
if result.url not in seen_urls and len(unique_results) < max_urls:
seen_urls.add(result.url)
unique_results.append(result)
logger.info(f"开始提取 {len(unique_results)} 个URL的内容")
# 提取内容
urls = [r.url for r in unique_results]
# 保存source信息以便后续使用
url_to_source = {r.url: r.source for r in unique_results}
documents = await self.jina_reader.extract_batch(urls)
# 更新document的source信息
for doc in documents:
if doc.url in url_to_source:
doc.source = url_to_source[doc.url]
return documents
async def extract_urls(
self,
urls: List[str],
source: SearchSource = SearchSource.WEB
) -> List[Document]:
"""
直接从URL列表提取内容
Args:
urls: URL列表
source: 来源类型
Returns:
Document列表
"""
return await self.jina_reader.extract_batch(urls, source)
@@ -0,0 +1,117 @@
"""
查询理解模块
负责分析用户查询意图、提取关键实体、生成扩展查询
"""
from typing import Optional
from loguru import logger
from search_agent.config import Config
from models.schemas import QueryAnalysis, Intent
from utils.llm_client import LLMClient
# 查询分析Prompt
QUERY_ANALYSIS_PROMPT = """你是一个查询分析专家。分析用户的搜索查询,提取以下信息。
请输出JSON格式:
{
"intent": "查询意图,必须是以下之一: fact_check(事实核查), comparison(对比分析), how_to(操作指南), news(新闻资讯), research(深度研究)",
"entities": ["关键实体列表,提取查询中的核心概念、人名、产品名等"],
"expanded_queries": ["扩展查询1", "扩展查询2", "扩展查询3"],
"need_news": true或false,
"time_filter": "时间过滤器,null表示不限时间,qdr:d(过去24小时), qdr:w(过去一周), qdr:m(过去一月), qdr:y(过去一年)"
}
扩展查询要求:
1. 生成2-4个扩展查询,包含不同角度或同义表达
2. 至少包含一个英文查询(如果原查询是中文)
3. 保持查询的核心意图
时间过滤器选择规则:
- 查询涉及"最新"、"近期"、"今年"等时效性词语 → 设置相应的时间过滤器
- 查询涉及具体年份(如"2024年") → qdr:y
- 一般性查询 → null"""
class QueryAnalyzer:
"""查询理解模块"""
def __init__(self, config: Config):
"""
初始化查询分析器
Args:
config: 配置对象
"""
self.config = config
self.llm = LLMClient(
base_url=config.llm_base_url,
api_key=config.llm_api_key,
model=config.llm_model
)
async def analyze(self, query: str) -> QueryAnalysis:
"""
分析用户查询
Args:
query: 用户查询字符串
Returns:
QueryAnalysis对象
"""
logger.info(f"开始分析查询: {query}")
try:
result = await self.llm.chat_json(
system_prompt=QUERY_ANALYSIS_PROMPT,
user_message=f"用户查询: {query}",
temperature=0.3
)
# 解析意图
intent_str = result.get("intent", "research")
intent = self._parse_intent(intent_str)
# 构建分析结果
analysis = QueryAnalysis(
original_query=query,
intent=intent,
entities=result.get("entities", []),
expanded_queries=result.get("expanded_queries", [query]),
need_news=result.get("need_news", False),
time_filter=result.get("time_filter")
)
logger.info(f"查询分析完成: intent={intent.value}, entities={analysis.entities}")
return analysis
except Exception as e:
logger.error(f"查询分析失败: {e}")
# 返回默认分析结果
return self._default_analysis(query)
def _parse_intent(self, intent_str: str) -> Intent:
"""解析意图字符串为枚举"""
intent_mapping = {
"fact_check": Intent.FACT_CHECK,
"comparison": Intent.COMPARISON,
"how_to": Intent.HOW_TO,
"news": Intent.NEWS,
"research": Intent.RESEARCH
}
return intent_mapping.get(intent_str.lower(), Intent.RESEARCH)
def _default_analysis(self, query: str) -> QueryAnalysis:
"""生成默认的查询分析结果"""
return QueryAnalysis(
original_query=query,
intent=Intent.RESEARCH,
entities=[],
expanded_queries=[query],
need_news=False,
time_filter=None
)
@@ -0,0 +1,166 @@
"""
反思迭代模块
评估答案质量,决定是否需要补充搜索
"""
from typing import List
from loguru import logger
from search_agent.config import Config
from models.schemas import Answer, QualityAssessment
from utils.llm_client import LLMClient
# 反思评估Prompt
REFLECTION_PROMPT = """你是一个质量评估专家。评估以下答案是否充分回答了用户的问题。
## 评估维度
1. **完整性**: 答案是否覆盖了问题的所有方面?
2. **准确性**: 答案内容是否有明确的来源支持?
3. **深度**: 答案是否提供了足够的细节和解释?
## 输出JSON格式
{
"completeness": 0.0-1.0,
"missing_aspects": ["如果有缺失,列出缺失的方面"],
"needs_more_search": true或false,
"suggested_queries": ["如果需要补充搜索,建议的搜索词"]
}
## 判断标准
- completeness >= 0.8 且没有重要信息缺失 → needs_more_search = false
- completeness < 0.8 或有重要信息缺失 → needs_more_search = true
- 建议的搜索词应该针对缺失的方面"""
class Reflector:
"""反思迭代模块"""
# 质量阈值
COMPLETENESS_THRESHOLD = 0.8
def __init__(self, config: Config):
"""
初始化反思器
Args:
config: 配置对象
"""
self.config = config
self.llm = LLMClient(
base_url=config.llm_base_url,
api_key=config.llm_api_key,
model=config.llm_model
)
async def assess(
self,
query: str,
answer: Answer
) -> QualityAssessment:
"""
评估答案质量
Args:
query: 原始查询
answer: 生成的答案
Returns:
QualityAssessment对象
"""
logger.info("开始评估答案质量")
# 如果答案置信度已经很低,直接建议补充搜索
if answer.confidence == "low" and not answer.content:
return QualityAssessment(
completeness=0.0,
missing_aspects=["缺少相关信息"],
needs_more_search=True,
suggested_queries=[query]
)
user_message = f"""## 用户问题
{query}
## 生成的答案
{answer.content}
## 答案的来源数量
{len(answer.sources)} 个来源
## 答案的置信度
{answer.confidence}"""
try:
result = await self.llm.chat_json(
system_prompt=REFLECTION_PROMPT,
user_message=user_message,
temperature=0.3
)
assessment = QualityAssessment(
completeness=float(result.get("completeness", 0.5)),
missing_aspects=result.get("missing_aspects", []),
needs_more_search=result.get("needs_more_search", False),
suggested_queries=result.get("suggested_queries", [])
)
logger.info(
f"质量评估: completeness={assessment.completeness:.2f}, "
f"needs_more_search={assessment.needs_more_search}"
)
return assessment
except Exception as e:
logger.error(f"质量评估失败: {e}")
return self._default_assessment(answer)
def _default_assessment(self, answer: Answer) -> QualityAssessment:
"""默认评估结果"""
# 根据答案置信度估计完整性
confidence_score = {
"high": 0.9,
"medium": 0.7,
"low": 0.4
}.get(answer.confidence, 0.5)
return QualityAssessment(
completeness=confidence_score,
missing_aspects=[],
needs_more_search=confidence_score < self.COMPLETENESS_THRESHOLD,
suggested_queries=[]
)
def should_continue(
self,
assessment: QualityAssessment,
current_iteration: int
) -> bool:
"""
判断是否应该继续迭代
Args:
assessment: 质量评估结果
current_iteration: 当前迭代次数
Returns:
是否继续迭代
"""
# 达到最大迭代次数
if current_iteration >= self.config.max_iterations:
logger.info(f"达到最大迭代次数 ({self.config.max_iterations}),停止迭代")
return False
# 完整性达标
if assessment.completeness >= self.COMPLETENESS_THRESHOLD:
logger.info(f"完整性达标 ({assessment.completeness:.2f}),停止迭代")
return False
# 没有建议的补充搜索
if not assessment.suggested_queries:
logger.info("没有建议的补充搜索,停止迭代")
return False
return assessment.needs_more_search
@@ -0,0 +1,107 @@
"""
结果处理模块
负责结果去重、相关性排序、筛选
"""
from typing import List
from loguru import logger
from search_agent.config import Config
from models.schemas import Document, RankedDocument
from tools.jina_reranker import JinaRerankerClient
from utils.helpers import deduplicate_by_url
class ResultProcessor:
"""结果处理模块"""
def __init__(self, config: Config):
"""
初始化结果处理器
Args:
config: 配置对象
"""
self.config = config
self.reranker = JinaRerankerClient(
api_key=config.jina_api_key,
timeout=config.timeout
)
async def process(
self,
query: str,
documents: List[Document],
top_k: int = 5
) -> List[RankedDocument]:
"""
处理文档:去重 + 重排序 + 筛选
Args:
query: 原始查询
documents: 文档列表
top_k: 返回前k个结果
Returns:
排序后的RankedDocument列表
"""
if not documents:
logger.warning("没有文档需要处理")
return []
logger.info(f"开始处理 {len(documents)} 个文档")
# 1. 去重
unique_docs = self._deduplicate(documents)
logger.debug(f"去重后: {len(unique_docs)} 个文档")
# 2. 过滤空内容
valid_docs = [d for d in unique_docs if d.content and len(d.content.strip()) > 50]
logger.debug(f"有效文档: {len(valid_docs)} 个")
if not valid_docs:
logger.warning("没有有效文档")
return []
# 3. 重排序
ranked_docs = await self.reranker.rerank(
query=query,
documents=valid_docs,
top_k=top_k,
content_max_length=self.config.content_max_length // 5 # 使用较短内容进行排序
)
logger.info(f"处理完成,返回 {len(ranked_docs)} 个排序结果")
return ranked_docs
def _deduplicate(self, documents: List[Document]) -> List[Document]:
"""去重文档"""
return deduplicate_by_url(documents, "url")
async def process_without_rerank(
self,
documents: List[Document],
top_k: int = 5
) -> List[RankedDocument]:
"""
处理文档(不进行重排序)
Args:
documents: 文档列表
top_k: 返回前k个结果
Returns:
RankedDocument列表(按原始顺序)
"""
unique_docs = self._deduplicate(documents)
valid_docs = [d for d in unique_docs if d.content and len(d.content.strip()) > 50]
return [
RankedDocument(
document=doc,
relevance_score=1.0 - (i * 0.1),
rank=i + 1
)
for i, doc in enumerate(valid_docs[:top_k])
]
@@ -0,0 +1,89 @@
"""
搜索执行模块
执行搜索计划,调用Serper API
"""
import asyncio
from typing import List
from loguru import logger
from search_agent.config import Config
from models.schemas import SearchPlan, SearchTask, SearchResult
from tools.serper import SerperClient
class SearchExecutor:
"""搜索执行模块"""
def __init__(self, config: Config):
"""
初始化搜索执行器
Args:
config: 配置对象
"""
self.config = config
self.serper = SerperClient(
api_key=config.serper_api_key,
timeout=config.timeout
)
async def execute(self, plan: SearchPlan) -> List[SearchResult]:
"""
执行搜索计划
Args:
plan: 搜索计划
Returns:
搜索结果列表
"""
logger.info(f"开始执行搜索计划: {len(plan.tasks)} 个任务")
if plan.strategy == "parallel":
results = await self._execute_parallel(plan.tasks)
else:
results = await self._execute_sequential(plan.tasks)
logger.info(f"搜索完成,共获取 {len(results)} 条结果")
return results
async def _execute_parallel(self, tasks: List[SearchTask]) -> List[SearchResult]:
"""并行执行搜索任务"""
coroutines = [self._execute_task(task) for task in tasks]
results_list = await asyncio.gather(*coroutines, return_exceptions=True)
# 合并结果
all_results = []
for results in results_list:
if isinstance(results, list):
all_results.extend(results)
elif isinstance(results, Exception):
logger.warning(f"搜索任务失败: {results}")
return all_results
async def _execute_sequential(self, tasks: List[SearchTask]) -> List[SearchResult]:
"""串行执行搜索任务"""
all_results = []
for task in tasks:
try:
results = await self._execute_task(task)
all_results.extend(results)
except Exception as e:
logger.warning(f"搜索任务失败: {e}")
return all_results
async def _execute_task(self, task: SearchTask) -> List[SearchResult]:
"""执行单个搜索任务"""
logger.debug(f"执行搜索: {task.query} [{task.source.value}]")
return await self.serper.search(
query=task.query,
source=task.source,
num_results=task.num_results,
time_filter=task.time_filter
)
@@ -0,0 +1,137 @@
"""
搜索规划模块
根据查询分析结果制定搜索计划
"""
from typing import List
from loguru import logger
from search_agent.config import Config
from models.schemas import (
QueryAnalysis,
SearchPlan,
SearchTask,
SearchSource,
Intent
)
class SearchPlanner:
"""搜索规划模块"""
def __init__(self, config: Config):
"""
初始化搜索规划器
Args:
config: 配置对象
"""
self.config = config
self.max_results = config.max_results_per_query
async def plan(self, analysis: QueryAnalysis) -> SearchPlan:
"""
根据查询分析制定搜索计划
Args:
analysis: 查询分析结果
Returns:
SearchPlan对象
"""
logger.info(f"开始制定搜索计划: intent={analysis.intent.value}")
tasks = []
# 根据意图确定搜索策略
strategy = self._determine_strategy(analysis)
# 构建搜索任务
tasks.extend(self._create_web_tasks(analysis))
if analysis.need_news:
tasks.extend(self._create_news_tasks(analysis))
plan = SearchPlan(
tasks=tasks,
strategy=strategy
)
logger.info(f"搜索计划: {len(tasks)} 个任务, 策略={strategy}")
return plan
def _determine_strategy(self, analysis: QueryAnalysis) -> str:
"""确定执行策略"""
# 大多数情况使用并行策略
if analysis.intent == Intent.COMPARISON:
# 对比类查询可能需要串行以获取更相关的结果
return "parallel"
return "parallel"
def _create_web_tasks(self, analysis: QueryAnalysis) -> List[SearchTask]:
"""创建Web搜索任务"""
tasks = []
# 原始查询
tasks.append(SearchTask(
query=analysis.original_query,
source=SearchSource.WEB,
time_filter=analysis.time_filter,
num_results=self.max_results
))
# 扩展查询(限制数量避免过多请求)
for query in analysis.expanded_queries[:2]:
if query != analysis.original_query:
tasks.append(SearchTask(
query=query,
source=SearchSource.WEB,
time_filter=analysis.time_filter,
num_results=self.max_results
))
return tasks
def _create_news_tasks(self, analysis: QueryAnalysis) -> List[SearchTask]:
"""创建新闻搜索任务"""
tasks = []
# 新闻搜索使用原始查询
tasks.append(SearchTask(
query=analysis.original_query,
source=SearchSource.NEWS,
time_filter=analysis.time_filter or "qdr:m", # 默认过去一个月
num_results=self.max_results
))
return tasks
def plan_supplementary(
self,
original_query: str,
suggested_queries: List[str]
) -> SearchPlan:
"""
创建补充搜索计划
Args:
original_query: 原始查询
suggested_queries: 建议的补充查询
Returns:
SearchPlan对象
"""
tasks = []
for query in suggested_queries[:3]: # 限制补充搜索数量
tasks.append(SearchTask(
query=query,
source=SearchSource.WEB,
num_results=self.max_results
))
return SearchPlan(
tasks=tasks,
strategy="parallel"
)
@@ -0,0 +1,14 @@
"""
外部API工具封装模块
"""
from .serper import SerperClient
from .jina_reader import JinaReaderClient
from .jina_reranker import JinaRerankerClient
__all__ = [
"SerperClient",
"JinaReaderClient",
"JinaRerankerClient",
]
@@ -0,0 +1,180 @@
"""
Jina Reader API封装
提供网页内容提取功能
"""
import asyncio
from typing import List, Optional
import aiohttp
from loguru import logger
from models.schemas import Document, SearchSource
class JinaReaderClient:
"""Jina Reader API客户端"""
BASE_URL = "https://r.jina.ai"
def __init__(
self,
api_key: str,
timeout: int = 30,
max_concurrent: int = 5,
max_content_length: int = 5000
):
"""
初始化Jina Reader客户端
Args:
api_key: Jina API密钥
timeout: 请求超时时间(秒)
max_concurrent: 最大并发请求数
max_content_length: 最大内容长度
"""
self.api_key = api_key
self.timeout = timeout
self.max_concurrent = max_concurrent
self.max_content_length = max_content_length
self._semaphore = asyncio.Semaphore(max_concurrent)
async def extract_content(
self,
url: str,
source: SearchSource = SearchSource.WEB
) -> Optional[Document]:
"""
提取单个URL的内容
Args:
url: 要提取的网页URL
source: 来源类型
Returns:
Document对象,如果提取失败则返回None
"""
reader_url = f"{self.BASE_URL}/{url}"
headers = {
"Authorization": f"Bearer {self.api_key}",
"Accept": "application/json"
}
try:
async with self._semaphore:
async with aiohttp.ClientSession() as session:
async with session.get(
reader_url,
headers=headers,
timeout=aiohttp.ClientTimeout(total=self.timeout)
) as response:
if response.status != 200:
logger.warning(f"Jina Reader提取失败 [{response.status}]: {url}")
return None
# Jina Reader可能返回JSON或纯文本
content_type = response.headers.get("Content-Type", "")
if "application/json" in content_type:
result = await response.json()
# 处理嵌套的data字段
if "data" in result:
result = result["data"]
content = result.get("content", "")
title = result.get("title", "")
else:
# 纯文本响应(Markdown格式)
content = await response.text()
# 从内容中提取标题(第一行通常是标题)
lines = content.strip().split("\n")
title = lines[0].lstrip("#").strip() if lines else ""
# 限制内容长度
if len(content) > self.max_content_length:
content = content[:self.max_content_length]
logger.debug(f"提取成功: {url[:50]}... 内容长度: {len(content)}")
return Document(
url=url,
title=title,
content=content,
source=source
)
except aiohttp.ClientError as e:
logger.warning(f"Jina Reader网络错误 [{url}]: {e}")
return None
except asyncio.TimeoutError:
logger.warning(f"Jina Reader超时: {url}")
return None
except Exception as e:
logger.warning(f"Jina Reader异常 [{url}]: {e}")
return None
async def extract_batch(
self,
urls: List[str],
source: SearchSource = SearchSource.WEB
) -> List[Document]:
"""
批量提取多个URL的内容
Args:
urls: URL列表
source: 来源类型
Returns:
成功提取的Document列表
"""
logger.info(f"批量提取 {len(urls)} 个URL的内容")
tasks = [
self.extract_content(url, source)
for url in urls
]
results = await asyncio.gather(*tasks, return_exceptions=True)
# 过滤掉失败的结果
documents = []
for result in results:
if isinstance(result, Document):
documents.append(result)
elif isinstance(result, Exception):
logger.warning(f"提取异常: {result}")
logger.info(f"成功提取 {len(documents)}/{len(urls)} 个文档")
return documents
async def extract_with_retry(
self,
url: str,
source: SearchSource = SearchSource.WEB,
max_retries: int = 2,
retry_delay: float = 1.0
) -> Optional[Document]:
"""
带重试的内容提取
Args:
url: 要提取的网页URL
source: 来源类型
max_retries: 最大重试次数
retry_delay: 重试延迟(秒)
Returns:
Document对象,如果最终失败则返回None
"""
for attempt in range(max_retries + 1):
result = await self.extract_content(url, source)
if result is not None:
return result
if attempt < max_retries:
logger.debug(f"重试提取 [{attempt + 1}/{max_retries}]: {url}")
await asyncio.sleep(retry_delay)
return None
@@ -0,0 +1,191 @@
"""
Jina Reranker API封装
提供搜索结果重排序功能
"""
from typing import List, Tuple
import aiohttp
from loguru import logger
from models.schemas import Document, RankedDocument
class JinaRerankerClient:
"""Jina Reranker API客户端"""
BASE_URL = "https://api.jina.ai/v1/rerank"
MODEL = "jina-reranker-v2-base-multilingual"
def __init__(self, api_key: str, timeout: int = 30):
"""
初始化Jina Reranker客户端
Args:
api_key: Jina API密钥
timeout: 请求超时时间(秒)
"""
self.api_key = api_key
self.timeout = timeout
async def rerank(
self,
query: str,
documents: List[Document],
top_k: int = 5,
content_max_length: int = 1000
) -> List[RankedDocument]:
"""
对文档进行相关性重排序
Args:
query: 查询字符串
documents: 文档列表
top_k: 返回前k个结果
content_max_length: 用于排序的内容最大长度
Returns:
排序后的RankedDocument列表
"""
if not documents:
return []
# 准备文档内容(截断到合适长度)
doc_contents = [
doc.content[:content_max_length] if doc.content else doc.title
for doc in documents
]
headers = {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json"
}
payload = {
"model": self.MODEL,
"query": query,
"documents": doc_contents,
"top_n": min(top_k, len(documents))
}
try:
async with aiohttp.ClientSession() as session:
async with session.post(
self.BASE_URL,
headers=headers,
json=payload,
timeout=aiohttp.ClientTimeout(total=self.timeout)
) as response:
if response.status != 200:
error_text = await response.text()
logger.error(f"Jina Reranker API错误: {response.status} - {error_text}")
# 如果重排序失败,返回原始顺序
return self._fallback_ranking(documents, top_k)
result = await response.json()
return self._parse_rerank_results(documents, result, top_k)
except aiohttp.ClientError as e:
logger.error(f"Jina Reranker网络错误: {e}")
return self._fallback_ranking(documents, top_k)
except Exception as e:
logger.error(f"Jina Reranker异常: {e}")
return self._fallback_ranking(documents, top_k)
def _parse_rerank_results(
self,
documents: List[Document],
response: dict,
top_k: int
) -> List[RankedDocument]:
"""解析重排序结果"""
results = []
reranked = response.get("results", [])
for rank, item in enumerate(reranked[:top_k], 1):
index = item.get("index", 0)
score = item.get("relevance_score", 0.0)
if 0 <= index < len(documents):
ranked_doc = RankedDocument(
document=documents[index],
relevance_score=score,
rank=rank
)
results.append(ranked_doc)
logger.debug(f"重排序返回 {len(results)} 个结果")
return results
def _fallback_ranking(
self,
documents: List[Document],
top_k: int
) -> List[RankedDocument]:
"""后备排序:保持原始顺序"""
logger.warning("使用后备排序(原始顺序)")
return [
RankedDocument(
document=doc,
relevance_score=1.0 - (i * 0.1), # 模拟递减分数
rank=i + 1
)
for i, doc in enumerate(documents[:top_k])
]
async def rerank_texts(
self,
query: str,
texts: List[str],
top_k: int = 5
) -> List[Tuple[int, float]]:
"""
对纯文本列表进行重排序
Args:
query: 查询字符串
texts: 文本列表
top_k: 返回前k个结果
Returns:
(原始索引, 相关性分数) 的列表
"""
if not texts:
return []
headers = {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json"
}
payload = {
"model": self.MODEL,
"query": query,
"documents": texts,
"top_n": min(top_k, len(texts))
}
try:
async with aiohttp.ClientSession() as session:
async with session.post(
self.BASE_URL,
headers=headers,
json=payload,
timeout=aiohttp.ClientTimeout(total=self.timeout)
) as response:
if response.status != 200:
logger.error(f"Reranker API错误: {response.status}")
return [(i, 1.0 - i * 0.1) for i in range(min(top_k, len(texts)))]
result = await response.json()
return [
(item["index"], item["relevance_score"])
for item in result.get("results", [])[:top_k]
]
except Exception as e:
logger.error(f"Reranker异常: {e}")
return [(i, 1.0 - i * 0.1) for i in range(min(top_k, len(texts)))]
@@ -0,0 +1,212 @@
"""
Serper API封装
提供Google搜索和新闻搜索功能
"""
from typing import List, Optional, Dict, Any
import aiohttp
from loguru import logger
from models.schemas import SearchResult, SearchSource
class SerperClient:
"""Serper API客户端"""
BASE_URL = "https://google.serper.dev"
ENDPOINTS = {
"web": "/search",
"news": "/news"
}
def __init__(self, api_key: str, timeout: int = 30):
"""
初始化Serper客户端
Args:
api_key: Serper API密钥
timeout: 请求超时时间(秒)
"""
self.api_key = api_key
self.timeout = timeout
async def _request(
self,
endpoint: str,
payload: Dict[str, Any]
) -> Dict[str, Any]:
"""
发送请求到Serper API
Args:
endpoint: API端点
payload: 请求体
Returns:
API响应
"""
url = f"{self.BASE_URL}{endpoint}"
headers = {
"X-API-KEY": self.api_key,
"Content-Type": "application/json"
}
try:
async with aiohttp.ClientSession() as session:
async with session.post(
url,
headers=headers,
json=payload,
timeout=aiohttp.ClientTimeout(total=self.timeout)
) as response:
if response.status != 200:
error_text = await response.text()
logger.error(f"Serper API错误: {response.status} - {error_text}")
raise Exception(f"Serper API请求失败: {response.status}")
return await response.json()
except aiohttp.ClientError as e:
logger.error(f"Serper请求网络错误: {e}")
raise
async def search_web(
self,
query: str,
num_results: int = 10,
gl: str = "cn",
hl: str = "zh-cn",
time_filter: Optional[str] = None
) -> List[SearchResult]:
"""
执行Web搜索
Args:
query: 搜索查询
num_results: 返回结果数量
gl: 地区代码
hl: 语言代码
time_filter: 时间过滤器 (qdr:d/qdr:w/qdr:m/qdr:y)
Returns:
搜索结果列表
"""
payload = {
"q": query,
"num": num_results,
"gl": gl,
"hl": hl
}
if time_filter:
payload["tbs"] = time_filter
logger.info(f"执行Web搜索: {query}")
result = await self._request(self.ENDPOINTS["web"], payload)
return self._parse_web_results(result)
async def search_news(
self,
query: str,
num_results: int = 10,
gl: str = "cn",
hl: str = "zh-cn",
time_filter: Optional[str] = None
) -> List[SearchResult]:
"""
执行新闻搜索
Args:
query: 搜索查询
num_results: 返回结果数量
gl: 地区代码
hl: 语言代码
time_filter: 时间过滤器
Returns:
搜索结果列表
"""
payload = {
"q": query,
"num": num_results,
"gl": gl,
"hl": hl
}
if time_filter:
payload["tbs"] = time_filter
logger.info(f"执行新闻搜索: {query}")
result = await self._request(self.ENDPOINTS["news"], payload)
return self._parse_news_results(result)
def _parse_web_results(self, response: Dict[str, Any]) -> List[SearchResult]:
"""解析Web搜索结果"""
results = []
organic = response.get("organic", [])
for item in organic:
result = SearchResult(
title=item.get("title", ""),
url=item.get("link", ""),
snippet=item.get("snippet", ""),
source=SearchSource.WEB,
position=item.get("position", 0),
date=None
)
results.append(result)
logger.debug(f"Web搜索返回 {len(results)} 条结果")
return results
def _parse_news_results(self, response: Dict[str, Any]) -> List[SearchResult]:
"""解析新闻搜索结果"""
results = []
news = response.get("news", [])
for i, item in enumerate(news, 1):
result = SearchResult(
title=item.get("title", ""),
url=item.get("link", ""),
snippet=item.get("snippet", ""),
source=SearchSource.NEWS,
position=i,
date=item.get("date")
)
results.append(result)
logger.debug(f"新闻搜索返回 {len(results)} 条结果")
return results
async def search(
self,
query: str,
source: SearchSource,
num_results: int = 10,
time_filter: Optional[str] = None
) -> List[SearchResult]:
"""
统一搜索接口
Args:
query: 搜索查询
source: 搜索来源类型
num_results: 返回结果数量
time_filter: 时间过滤器
Returns:
搜索结果列表
"""
if source == SearchSource.NEWS:
return await self.search_news(query, num_results, time_filter=time_filter)
else:
return await self.search_web(query, num_results, time_filter=time_filter)
@@ -0,0 +1,22 @@
"""
工具函数模块
"""
from .llm_client import LLMClient
from .helpers import (
flatten,
deduplicate_by_url,
truncate_text,
extract_json_from_text,
format_documents_for_prompt,
)
__all__ = [
"LLMClient",
"flatten",
"deduplicate_by_url",
"truncate_text",
"extract_json_from_text",
"format_documents_for_prompt",
]
@@ -0,0 +1,197 @@
"""
通用工具函数
"""
import re
import json
from typing import List, TypeVar, Optional, Dict, Any
T = TypeVar('T')
def flatten(nested_list: List[List[T]]) -> List[T]:
"""
将嵌套列表展平为一维列表
Args:
nested_list: 嵌套列表
Returns:
展平后的一维列表
"""
return [item for sublist in nested_list for item in sublist]
def deduplicate_by_url(items: List[Any], url_attr: str = "url") -> List[Any]:
"""
根据URL去重
Args:
items: 包含URL属性的对象列表
url_attr: URL属性名
Returns:
去重后的列表
"""
seen_urls = set()
unique_items = []
for item in items:
url = getattr(item, url_attr, None) or item.get(url_attr)
if url and url not in seen_urls:
seen_urls.add(url)
unique_items.append(item)
return unique_items
def truncate_text(text: str, max_length: int, suffix: str = "...") -> str:
"""
截断文本到指定长度
Args:
text: 原始文本
max_length: 最大长度
suffix: 截断后缀
Returns:
截断后的文本
"""
if len(text) <= max_length:
return text
return text[:max_length - len(suffix)] + suffix
def extract_json_from_text(text: str) -> Optional[Dict[str, Any]]:
"""
从文本中提取JSON对象
Args:
text: 可能包含JSON的文本
Returns:
提取的JSON字典,如果提取失败则返回None
"""
# 尝试直接解析
try:
return json.loads(text)
except json.JSONDecodeError:
pass
# 尝试提取```json ... ```块
json_block_pattern = r'```(?:json)?\s*([\s\S]*?)```'
matches = re.findall(json_block_pattern, text)
for match in matches:
try:
return json.loads(match.strip())
except json.JSONDecodeError:
continue
# 尝试提取{ ... }块
brace_pattern = r'\{[\s\S]*\}'
matches = re.findall(brace_pattern, text)
for match in matches:
try:
return json.loads(match)
except json.JSONDecodeError:
continue
return None
def format_documents_for_prompt(documents: List[Any], max_length: int = 2000) -> str:
"""
格式化文档列表为Prompt中使用的文本
Args:
documents: 文档列表(RankedDocument或Document对象)
max_length: 每个文档的最大内容长度
Returns:
格式化后的文本
"""
formatted_parts = []
for i, doc in enumerate(documents, 1):
# 支持RankedDocument和Document两种类型
if hasattr(doc, 'document'):
# RankedDocument
actual_doc = doc.document
score = f" (相关性: {doc.relevance_score:.2f})"
else:
# Document
actual_doc = doc
score = ""
content = truncate_text(actual_doc.content, max_length)
part = f"""### 来源 [{i}]{score}
**标题**: {actual_doc.title}
**URL**: {actual_doc.url}
**内容**:
{content}
"""
formatted_parts.append(part)
return "\n---\n".join(formatted_parts)
def clean_url(url: str) -> str:
"""
清理和标准化URL
Args:
url: 原始URL
Returns:
清理后的URL
"""
# 移除末尾的斜杠
url = url.rstrip("/")
# 移除锚点
if "#" in url:
url = url.split("#")[0]
return url
def is_valid_url(url: str) -> bool:
"""
验证URL是否有效
Args:
url: URL字符串
Returns:
是否有效
"""
url_pattern = re.compile(
r'^https?://' # http:// or https://
r'(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+[A-Z]{2,6}\.?|' # domain
r'localhost|' # localhost
r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})' # IP
r'(?::\d+)?' # optional port
r'(?:/?|[/?]\S+)$', re.IGNORECASE)
return bool(url_pattern.match(url))
def merge_dicts(base: Dict, override: Dict) -> Dict:
"""
合并两个字典,override中的值会覆盖base中的值
Args:
base: 基础字典
override: 覆盖字典
Returns:
合并后的字典
"""
result = base.copy()
result.update(override)
return result
@@ -0,0 +1,179 @@
"""
LLM客户端模块 - 使用LiteLLM SDK
封装与LLM的交互(通过LiteLLM)
"""
import json
from typing import Optional, List, Dict, Any
import httpx
from loguru import logger
class LLMClient:
"""LLM客户端,使用LiteLLM SDK方式调用"""
def __init__(
self,
base_url: str,
api_key: str,
model: str = "gpt-4",
timeout: int = 60
):
"""
初始化LLM客户端
Args:
base_url: LiteLLM服务的基础URL
api_key: API密钥
model: 模型名称
timeout: 超时时间(秒)
"""
self.base_url = base_url.rstrip("/")
self.api_key = api_key
self.model = model
self.timeout = timeout
# LiteLLM的chat completions端点
self.chat_endpoint = f"{self.base_url}/chat/completions"
# HTTP客户端
self._client: Optional[httpx.AsyncClient] = None
async def _get_client(self) -> httpx.AsyncClient:
"""获取或创建HTTP客户端"""
if self._client is None or self._client.is_closed:
self._client = httpx.AsyncClient(
timeout=httpx.Timeout(self.timeout),
headers={
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json"
}
)
return self._client
async def close(self):
"""关闭HTTP客户端"""
if self._client and not self._client.is_closed:
await self._client.aclose()
async def chat(
self,
messages: List[Dict[str, str]],
temperature: float = 0.7,
max_tokens: int = 4096,
response_format: Optional[Dict[str, str]] = None
) -> str:
"""
发送聊天请求到LLM
Args:
messages: 消息列表,格式 [{"role": "user", "content": "..."}]
temperature: 温度参数
max_tokens: 最大token数
response_format: 响应格式(如 {"type": "json_object"})
Returns:
LLM的响应文本
"""
client = await self._get_client()
payload = {
"model": self.model,
"messages": messages,
"temperature": temperature,
"max_tokens": max_tokens
}
if response_format:
payload["response_format"] = response_format
try:
response = await client.post(
self.chat_endpoint,
json=payload
)
if response.status_code != 200:
error_text = response.text
logger.error(f"LLM API错误: {response.status_code} - {error_text}")
raise Exception(f"LLM API请求失败: {response.status_code}")
result = response.json()
return result["choices"][0]["message"]["content"]
except httpx.HTTPError as e:
logger.error(f"LLM请求网络错误: {e}")
raise
except Exception as e:
logger.error(f"LLM请求异常: {e}")
raise
async def chat_with_system(
self,
system_prompt: str,
user_message: str,
temperature: float = 0.7,
max_tokens: int = 4096,
response_format: Optional[Dict[str, str]] = None
) -> str:
"""
使用系统提示和用户消息进行对话
Args:
system_prompt: 系统提示
user_message: 用户消息
temperature: 温度参数
max_tokens: 最大token数
response_format: 响应格式
Returns:
LLM的响应文本
"""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_message}
]
return await self.chat(
messages=messages,
temperature=temperature,
max_tokens=max_tokens,
response_format=response_format
)
async def chat_json(
self,
system_prompt: str,
user_message: str,
temperature: float = 0.3
) -> Dict[str, Any]:
"""
请求JSON格式的响应
Args:
system_prompt: 系统提示
user_message: 用户消息
temperature: 温度参数(JSON响应建议使用较低温度)
Returns:
解析后的JSON字典
"""
from .helpers import extract_json_from_text
response = await self.chat_with_system(
system_prompt=system_prompt,
user_message=user_message,
temperature=temperature,
response_format={"type": "json_object"}
)
try:
return json.loads(response)
except json.JSONDecodeError:
# 尝试从文本中提取JSON
extracted = extract_json_from_text(response)
if extracted:
return extracted
logger.error(f"无法解析LLM响应为JSON: {response[:200]}")
raise ValueError("LLM响应不是有效的JSON格式")
@@ -0,0 +1,60 @@
FROM python:3.11-slim
WORKDIR /app
# 安装系统依赖
RUN apt-get update && apt-get install -y \
curl \
&& rm -rf /var/lib/apt/lists/*
# 复制requirements文件
COPY agents/search_agent/search_agent_A2A/requirements.txt /app/requirements.txt
COPY agents/search_agent/search_agent/search_agent/requirements.txt /app/search_agent_requirements.txt
# 安装Python依赖
# 先安装基础依赖(a2a-sdk的依赖)
RUN pip install --no-cache-dir \
fastapi>=0.109.0 \
uvicorn[standard]>=0.27.0 \
pydantic>=2.5.3 \
httpx>=0.27.0 \
python-dotenv>=1.0.0 \
orjson>=3.9.0 \
typing-extensions>=4.9.0 \
loguru>=0.7.0 \
asyncio-throttle>=1.0.2 \
aiohttp>=3.9.0 \
requests>=2.31.0
# 安装A2A SDK(包含http-server支持)
RUN pip install --no-cache-dir "a2a-sdk[http-server]>=0.3.0"
# 安装search_agent核心依赖(如果存在)
RUN if [ -f /app/search_agent_requirements.txt ]; then \
pip install --no-cache-dir -r /app/search_agent_requirements.txt; \
fi
# 复制search_agent_A2A目录
COPY agents/search_agent/search_agent_A2A/ /app/
# 复制search_agent核心代码
COPY agents/search_agent/search_agent/search_agent/ /app/search_agent/
# 设置环境变量
ENV PYTHONUNBUFFERED=1
ENV SERVICE_HOST=0.0.0.0
ENV SERVICE_PORT=8080
ENV PYTHONPATH=/app
# 默认 API 密钥(硬编码)
ENV JINA_API_KEY=jina_e26dc30420a44a1e859216528065b203TkMRmsoz-FgMDQC5FZX9jr5oF2CI
ENV SERPER_API_KEY=8253b4f240b520194065312f90e85f9be0fa205f
# 健康检查 - A2A SDK默认提供根路径和/.well-known/agent.json
# 先尝试根路径,如果不可用则尝试health端点
HEALTHCHECK --interval=30s --timeout=10s --start-period=40s --retries=3 \
CMD python3 -c "import urllib.request; import json; response = urllib.request.urlopen('http://localhost:8080/'); response.read()" || exit 1
# 运行agent (直接使用Python,避免shell)
CMD ["python3", "-u", "main.py"]
@@ -0,0 +1,167 @@
# MCP Search Agent
基于 Model Context Protocol (MCP) 协议的智能搜索 Agent,使用 LiteLLM SDK 进行模型调用。
## 特性
- 基于 MCP 协议标准
- 使用 LiteLLM SDK 进行模型调用
- 支持从请求传入或环境变量获取 API key 和模型名称
- 支持流式响应(SSE)
- 完整的搜索功能,包括查询理解、搜索规划、多源搜索等
## 目录结构
```
search_agent_MCP/
├── __init__.py
├── agent.py # SearchAgent包装器
├── config.py # 配置管理
├── mcp_server.py # MCP服务器实现
├── main.py # 主入口
├── requirements.txt # Python依赖
└── search_agent_MCP.Dockerfile # Docker构建文件
```
## 环境变量
- `LITELLM_API_KEY`: LiteLLM API密钥(可选,可通过请求传入)
- `MODEL_NAME` 或 `LITELLM_MODEL`: 模型名称(可选,可通过请求传入)
- `LITELLM_BASE_URL`: LiteLLM服务基础URL
- `SERPER_API_KEY`: Serper搜索API密钥
- `JINA_API_KEY`: Jina Reader API密钥
- `SERVICE_HOST`: 服务主机地址(默认:0.0.0.0)
- `SERVICE_PORT`: 服务端口(默认:8080)
## API端点
### 健康检查
```
GET /health
```
### MCP搜索
```
POST /mcp/v1/search
Content-Type: application/json
{
"jsonrpc": "2.0",
"id": "1",
"method": "search",
"params": {
"query": "搜索查询",
"api_key": "可选,LiteLLM API密钥",
"model": "可选,模型名称"
}
}
```
### MCP流式搜索
```
POST /mcp/v1/search/stream
Content-Type: application/json
{
"jsonrpc": "2.0",
"id": "1",
"method": "search/stream",
"params": {
"query": "搜索查询",
"api_key": "可选,LiteLLM API密钥",
"model": "可选,模型名称"
}
}
```
### 通用MCP调用
```
POST /mcp/v1/call
Content-Type: application/json
{
"jsonrpc": "2.0",
"id": "1",
"method": "search",
"params": {
"query": "搜索查询"
}
}
```
## 响应格式
### 成功响应
```json
{
"jsonrpc": "2.0",
"id": "1",
"result": {
"query": "搜索查询",
"answer": "答案内容(Markdown格式)",
"sources": [
{
"index": 1,
"title": "来源标题",
"url": "来源URL"
}
],
"confidence": "high",
"iterations": 2,
"total_sources": 5,
"search_queries": ["查询1", "查询2"],
"timestamp": "2024-01-01T00:00:00.000000Z"
}
}
```
### 错误响应
```json
{
"jsonrpc": "2.0",
"id": "1",
"error": {
"code": -32602,
"message": "Invalid params: 'query' is required"
}
}
```
## 运行
### 本地运行
```bash
python main.py
```
### 使用Docker
```bash
docker build -f search_agent_MCP.Dockerfile -t search-agent-mcp .
docker run -p 8080:8080 \
-e LITELLM_API_KEY=your-key \
-e MODEL_NAME=your-model \
search-agent-mcp
```
### 使用uvicorn
```bash
uvicorn mcp_server:app --host 0.0.0.0 --port 8080
```
## 使用LiteLLM SDK
本实现使用 LiteLLM SDK 进行模型调用。虽然核心 SearchAgent 仍然使用 LLMClient(基于 HTTP),但可以通过配置 LiteLLM 的基础 URL 和 API 密钥来使用 LiteLLM 服务。
如需完全使用 LiteLLM SDK 调用,可以修改 `search_agent/utils/llm_client.py` 以使用 `litellm.completion()` 而不是直接的 HTTP 请求。
## 与A2A版本的区别
- 协议:使用 MCP (Model Context Protocol) 而不是 A2A (Agent2Agent)
- API端点:使用 `/mcp/v1/` 前缀而不是 A2A 的端点
- 请求格式:使用 MCP 标准格式
- SDK:明确支持 LiteLLM SDK(在 requirements.txt 中包含 litellm)
## 许可证
与主项目相同
@@ -0,0 +1,420 @@
# MCP Search Agent 用户调用指南
## 概述
MCP Search Agent 是基于 Model Context Protocol (MCP) 协议的智能搜索服务,能够理解用户查询意图、自动规划搜索策略、从多个来源获取信息,并生成高质量、有来源引用的答案。
**重要提示**:本服务部署在 AKS 后,模型、基础URL等配置已通过环境变量在部署时配置,用户调用时**不需要**传递这些参数。
---
## API 接口说明
### 基础信息
- **协议**: MCP (Model Context Protocol)
- **通信格式**: JSON-RPC 2.0
- **Content-Type**: `application/json`
- **基础URL**: 部署后提供的服务地址
---
## 核心接口
### 1. 搜索接口
#### POST /mcp/v1/search
执行智能搜索,根据查询返回答案和相关来源。
**请求格式**:
```json
{
"jsonrpc": "2.0",
"id": "request-001",
"method": "search",
"params": {
"query": "什么是人工智能?",
"api_key": "your-llm-api-key"
}
}
```
**或者使用 `llm_api_key`(与API格式保持一致)**:
```json
{
"jsonrpc": "2.0",
"id": "request-001",
"method": "search",
"params": {
"query": "什么是人工智能?",
"llm_api_key": "your-llm-api-key"
}
}
```
**请求参数说明**:
| 字段 | 类型 | 必填 | 说明 |
|------|------|------|------|
| jsonrpc | string | 是 | JSON-RPC版本,固定为 "2.0" |
| id | string | 是 | 请求ID,用于关联请求和响应 |
| method | string | 是 | 方法名,固定为 "search" |
| params | object | 是 | 请求参数对象 |
| params.query | string | 是 | 搜索查询内容 |
| params.api_key | string | 是 | LLM API密钥(**必填**,等同于API格式版本的`llm_api_key`) |
| params.llm_api_key | string | 是 | LLM API密钥(**必填**,`api_key`的别名,与API格式保持一致) |
**响应示例**:
```json
{
"jsonrpc": "2.0",
"id": "request-001",
"result": {
"query": "什么是人工智能?",
"answer": "人工智能(AI)是计算机科学的一个分支,致力于创建能够执行通常需要人类智能的任务的系统...",
"sources": [
{
"index": 1,
"title": "人工智能 - 维基百科",
"url": "https://zh.wikipedia.org/wiki/人工智能"
},
{
"index": 2,
"title": "什么是AI?",
"url": "https://example.com/ai-introduction"
}
],
"confidence": "high",
"iterations": 2,
"total_sources": 5,
"search_queries": [
"什么是人工智能",
"AI定义"
],
"timestamp": "2024-01-01T00:00:00.000000Z"
}
}
```
**响应字段说明**:
| 字段 | 类型 | 说明 |
|------|------|------|
| jsonrpc | string | JSON-RPC版本 |
| id | string | 请求ID |
| result | object | 搜索结果对象 |
| result.query | string | 原始查询 |
| result.answer | string | 生成的答案内容(Markdown格式) |
| result.sources | array | 来源列表 |
| result.sources[].index | integer | 来源索引 |
| result.sources[].title | string | 来源标题 |
| result.sources[].url | string | 来源URL |
| result.confidence | string | 置信度:"high" / "medium" / "low" |
| result.iterations | integer | 迭代次数 |
| result.total_sources | integer | 参考来源总数 |
| result.search_queries | array[string] | 使用的搜索查询列表 |
| result.timestamp | string | 时间戳(ISO格式) |
---
### 2. 流式搜索接口
#### POST /mcp/v1/search/stream
执行智能搜索,通过 Server-Sent Events (SSE) 流式返回结果。
**请求格式**:
```json
{
"jsonrpc": "2.0",
"id": "request-002",
"method": "search/stream",
"params": {
"query": "Python编程语言的特点",
"api_key": "your-llm-api-key"
}
}
```
**响应格式** (SSE流):
```
data: {"jsonrpc":"2.0","id":"request-002","method":"task/start","params":{"task_id":"..."}}
data: {"jsonrpc":"2.0","id":"request-002","method":"result/delta","params":{"task_id":"...","delta":"答案内容片段"}}
data: {"jsonrpc":"2.0","id":"request-002","method":"task/complete","params":{"task_id":"...","result":{...}}}
```
---
### 3. 通用调用接口
#### POST /mcp/v1/call
通用MCP调用接口,支持所有MCP方法。
**请求格式**:
```json
{
"jsonrpc": "2.0",
"id": "request-003",
"method": "search",
"params": {
"query": "搜索查询",
"api_key": "your-llm-api-key"
}
}
```
**响应格式**: 与对应的具体方法响应相同。
---
### 4. 健康检查
#### GET /health
检查服务健康状态。
**请求示例**:
```
GET /health
```
**响应示例**:
```json
{
"status": "healthy",
"pod_name": "search-agent-mcp",
"template_type": "search_agent_MCP",
"configured": true,
"llm_base_url": "https://litellm.example.com",
"llm_model": "gpt-4",
"timestamp": "2024-01-01T00:00:00.000000"
}
```
---
## 错误响应
### 请求格式错误
```json
{
"jsonrpc": "2.0",
"id": "request-001",
"error": {
"code": -32600,
"message": "Invalid Request: 错误详情"
}
}
```
### 方法不存在
```json
{
"jsonrpc": "2.0",
"id": "request-001",
"error": {
"code": -32601,
"message": "Method not found: 方法名"
}
}
```
### 参数错误
```json
{
"jsonrpc": "2.0",
"id": "request-001",
"error": {
"code": -32602,
"message": "Invalid params: 'query' is required"
}
}
```
### 服务器错误
```json
{
"jsonrpc": "2.0",
"id": "request-001",
"error": {
"code": -32000,
"message": "Server error: 错误详情"
}
}
```
---
## 调用示例
### cURL 示例
```bash
# 基础搜索
curl -X POST https://your-service-url/mcp/v1/search \
-H "Content-Type: application/json" \
-d '{
"jsonrpc": "2.0",
"id": "test-1",
"method": "search",
"params": {
"query": "什么是人工智能?",
"api_key": "your-llm-api-key"
}
}'
# 流式搜索
curl -X POST https://your-service-url/mcp/v1/search/stream \
-H "Content-Type: application/json" \
-d '{
"jsonrpc": "2.0",
"id": "test-2",
"method": "search/stream",
"params": {
"query": "Python编程语言的特点",
"api_key": "your-llm-api-key"
}
}'
```
### Python 示例
```python
import requests
import json
# 搜索接口
url = "https://your-service-url/mcp/v1/search"
payload = {
"jsonrpc": "2.0",
"id": "python-request-1",
"method": "search",
"params": {
"query": "什么是人工智能?",
"api_key": "your-llm-api-key"
}
}
response = requests.post(url, json=payload)
result = response.json()
if "result" in result:
print(f"查询: {result['result']['query']}")
print(f"答案: {result['result']['answer']}")
print(f"来源数: {result['result']['total_sources']}")
else:
print(f"错误: {result.get('error', {}).get('message', '未知错误')}")
```
### JavaScript 示例
```javascript
// 搜索接口
const url = 'https://your-service-url/mcp/v1/search';
const payload = {
jsonrpc: '2.0',
id: 'js-request-1',
method: 'search',
params: {
query: '什么是人工智能?',
api_key: 'your-llm-api-key'
}
};
fetch(url, {
method: 'POST',
headers: {
'Content-Type': 'application/json'
},
body: JSON.stringify(payload)
})
.then(response => response.json())
.then(data => {
if (data.result) {
console.log('查询:', data.result.query);
console.log('答案:', data.result.answer);
console.log('来源数:', data.result.total_sources);
} else {
console.error('错误:', data.error?.message || '未知错误');
}
});
```
---
## 重要说明
### 1. 环境变量配置(部署时已配置)
以下环境变量在 AKS 部署时已配置,用户**不需要**在请求中传递:
- `LLM_BASE_URL` - LLM服务基础URL
- `MODEL_NAME` 或 `LLM_MODEL` - 模型名称
- `SERPER_API_KEY` - Serper搜索API密钥
- `JINA_API_KEY` - Jina Reader API密钥
### 2. 用户请求参数
用户调用时**必须**传递:
- `query` - 搜索查询内容(**必填**)
- `api_key` 或 `llm_api_key` - LLM API密钥(**必填**,等同于API格式版本的`llm_api_key`,用于计费和身份验证)
**注意**:
- **`api_key`/`llm_api_key`是必需的**,必须在每次请求中传递(这是用户的LLM API密钥,用于计费和身份验证)
- **`model`不需要传递**,模型名称已在AKS部署时通过环境变量(`MODEL_NAME`或`LLM_MODEL`)配置
- 与API格式版本保持一致:用户的API密钥必须传递,其他配置由部署时通过环境变量配置
### 3. 请求格式
- 所有请求必须使用 JSON-RPC 2.0 格式
- `jsonrpc` 字段必须为 `"2.0"`
- `id` 字段用于关联请求和响应,可以是任意字符串
- `method` 字段指定要调用的方法
- `params` 字段包含方法参数
---
## 错误码说明
| 错误码 | 说明 |
|--------|------|
| -32600 | Invalid Request - 请求格式错误 |
| -32601 | Method not found - 方法不存在 |
| -32602 | Invalid params - 参数错误 |
| -32000 | Server error - 服务器内部错误 |
---
## 常见问题
### Q1: 请求中是否需要传递 api_key?
**A**: **必须传递**。`api_key`(或`llm_api_key`)是用户的LLM API密钥,必须在每次请求中传递(等同于API格式版本的`llm_api_key`)。这是用于计费和身份验证的,不能从环境变量获取。
### Q2: 请求中是否需要传递 model?
**A**: **不需要传递**。模型名称已在AKS部署时通过环境变量(`MODEL_NAME`或`LLM_MODEL`)配置,用户请求中不需要传递。
### Q3: 如何处理超时?
**A**: 默认超时时间为 30 秒,如果搜索查询较复杂可能需要更长时间。建议在客户端设置合理的超时时间(建议 120 秒)。
### Q4: 如何获取流式响应?
**A**: 使用 `/mcp/v1/search/stream` 接口,客户端需要支持 SSE (Server-Sent Events) 格式的流式响应处理。
---
## 技术支持
如有问题或需要帮助,请联系技术支持团队。
@@ -0,0 +1,4 @@
"""
MCP Search Agent
"""
@@ -0,0 +1,127 @@
"""
Search Agent 核心模块 - MCP版本
基于LiteLLM SDK和MCP协议的搜索Agent实现
"""
import sys
import os
# 先处理路径,再导入模块
# 添加search_agent目录到Python路径的最前面,确保优先导入
search_agent_dir = os.path.join(os.path.dirname(__file__), 'search_agent')
if search_agent_dir not in sys.path:
# 将search_agent目录放在路径最前面
sys.path.insert(0, search_agent_dir)
# 导入当前目录的config(使用绝对路径避免冲突)
# 临时移除当前目录,避免导入search_agent时冲突
current_dir = os.path.dirname(os.path.abspath(__file__))
current_dir_in_path = current_dir in sys.path
if current_dir_in_path:
sys.path.remove(current_dir)
# 先导入search_agent模块(此时当前目录不在路径中)
from search_agent.config import Config as SearchAgentConfig
from search_agent.agent.search_agent import SearchAgent as CoreSearchAgent
# 现在可以安全地导入当前目录的config
if not current_dir_in_path:
sys.path.append(current_dir)
else:
sys.path.append(current_dir)
# 导入当前目录的MCP config模块(使用重命名后的模块名)
from typing import Optional
from loguru import logger
from mcp_config import LiteLLMConfig, AgentConfig, get_config
# 为了避免与当前目录的config混淆,重命名
Config = SearchAgentConfig
class SearchAgentWrapper:
"""
Search Agent包装器
用于适配MCP框架,将SearchAgent包装为可配置的Agent实例
使用LiteLLM SDK进行模型调用
"""
def __init__(
self,
litellm_config: Optional[LiteLLMConfig] = None,
agent_config: Optional[AgentConfig] = None,
api_key: Optional[str] = None,
model: Optional[str] = None
):
"""
初始化Search Agent
Args:
litellm_config: LiteLLM配置对象
agent_config: Agent配置对象
api_key: LiteLLM API密钥(可选,优先使用)
model: 模型名称(可选,优先使用)
"""
# 获取配置
if not litellm_config:
llm_config, _, _ = get_config(api_key=api_key, model=model)
else:
llm_config = litellm_config
if not agent_config:
_, agent_config, _ = get_config(api_key=api_key, model=model)
self.litellm_config = llm_config
self.agent_config = agent_config
# 验证配置
self.litellm_config.validate()
# 创建SearchAgent配置(使用litellm的base_url和api_key)
# 需要从环境变量获取其他配置
serper_api_key = os.getenv("SERPER_API_KEY", "")
jina_api_key = os.getenv("JINA_API_KEY", "")
self.search_config = SearchAgentConfig(
llm_base_url=llm_config.base_url,
llm_api_key=llm_config.api_key,
llm_model=llm_config.model,
serper_api_key=serper_api_key,
jina_api_key=jina_api_key,
max_iterations=int(os.getenv("MAX_ITERATIONS", "3")),
max_results_per_query=int(os.getenv("MAX_RESULTS_PER_QUERY", "10")),
content_max_length=int(os.getenv("CONTENT_MAX_LENGTH", "5000")),
log_level=os.getenv("LOG_LEVEL", "INFO"),
timeout=int(os.getenv("TIMEOUT", "30"))
)
# 注意:SearchAgent内部使用LLMClient,它直接调用HTTP API
# 如果需要完全使用litellm SDK,需要修改SearchAgent的LLMClient实现
# 但为了保持兼容性,这里仍然使用原有的Config和SearchAgent
# 创建SearchAgent实例
self.agent = CoreSearchAgent(self.search_config)
logger.info(
"SearchAgent初始化完成 (MCP版本)",
agent_name=self.agent_config.name,
model=self.litellm_config.model,
base_url=self.litellm_config.base_url
)
async def search(self, query: str):
"""
执行搜索
Args:
query: 搜索查询
Returns:
AgentResponse对象
"""
return await self.agent.search(query)
async def close(self):
"""关闭资源(SearchAgent不需要特殊清理)"""
pass
@@ -0,0 +1,47 @@
"""
MCP Search Agent 主入口
支持从环境变量或请求传入 API key
"""
import os
import uvicorn
from mcp_server import create_app
# 环境变量配置
SERVICE_HOST = os.getenv("SERVICE_HOST", "0.0.0.0")
SERVICE_PORT = int(os.getenv("SERVICE_PORT", "8080"))
POD_NAME = os.getenv("POD_NAME", "search-agent-mcp")
TEMPLATE_TYPE = os.getenv("TEMPLATE_TYPE", "search_agent_MCP")
# 从环境变量获取默认配置(可选)
# 支持多种环境变量名称(向后兼容,与API格式和AKS部署保持一致)
# API密钥:优先使用 LITELLM_API_KEY(LiteLLM约定),也支持 LLM_API_KEY(AKS部署)
default_api_key = os.getenv("LITELLM_API_KEY") or os.getenv("LLM_API_KEY")
# 模型名称:优先使用 MODEL_NAME(API格式),也支持 LLM_MODEL(AKS部署)、LITELLM_MODEL
default_model = os.getenv("MODEL_NAME") or os.getenv("LLM_MODEL") or os.getenv("LITELLM_MODEL")
# 创建应用
app = create_app(api_key=default_api_key, model=default_model)
def main():
"""主函数"""
print(f"🚀 启动 MCP Search Agent")
print(f" - Pod名称: {POD_NAME}")
print(f" - 模板类型: {TEMPLATE_TYPE}")
print(f" - 服务地址: http://{SERVICE_HOST}:{SERVICE_PORT}")
if default_api_key:
print(f" - 已配置默认 API key(可通过请求覆盖)")
else:
print(f" - 未配置默认 API key,需在请求中传入")
uvicorn.run(
app,
host=SERVICE_HOST,
port=SERVICE_PORT,
log_level="info"
)
if __name__ == "__main__":
main()
@@ -0,0 +1,116 @@
"""
MCP Search Agent 配置模块
支持用户传入密钥和模型名称,同时支持从环境变量获取
"""
import os
from dataclasses import dataclass, field
from typing import Optional
from dotenv import load_dotenv
# 加载环境变量
load_dotenv()
@dataclass
class LiteLLMConfig:
"""LiteLLM 配置"""
# 基础URL - 用户提供的LiteLLM服务地址
# 优先使用 LLM_BASE_URL(与API格式和AKS部署保持一致)
# 也支持 LITELLM_BASE_URL(向后兼容)
base_url: str = field(default_factory=lambda: os.getenv(
"LLM_BASE_URL"
) or os.getenv(
"LITELLM_BASE_URL",
"https://litellm.graystone-fb459c5d.southeastasia.azurecontainerapps.io"
))
# API密钥 - 优先使用传入的,否则从环境变量获取
api_key: Optional[str] = None
# 模型名称 - 优先使用传入的,否则从环境变量获取
model: Optional[str] = None
# 请求超时时间(秒)
timeout: int = 120
# 温度参数
temperature: float = 0.7
# 最大token数
max_tokens: int = 4096
def __post_init__(self):
# 从环境变量读取(如果未直接提供)
# API密钥:优先使用 LITELLM_API_KEY(LiteLLM约定),也支持 LLM_API_KEY(AKS部署)
if self.api_key is None:
self.api_key = os.getenv("LITELLM_API_KEY") or os.getenv("LLM_API_KEY")
# 模型名称:优先使用 MODEL_NAME(API格式),也支持 LLM_MODEL(AKS部署)、LITELLM_MODEL
if self.model is None:
self.model = os.getenv("MODEL_NAME") or os.getenv("LLM_MODEL") or os.getenv("LITELLM_MODEL", "gpt-4")
def validate(self) -> bool:
"""验证配置是否完整"""
if not self.api_key:
raise ValueError("API密钥未设置! 请设置 LITELLM_API_KEY 环境变量或直接传入 api_key")
if not self.model:
raise ValueError("模型名称未设置! 请设置 MODEL_NAME 环境变量或直接传入 model")
return True
@dataclass
class AgentConfig:
"""Agent 配置"""
# Agent名称
name: str = "search-agent"
# Agent描述
description: str = "智能AI搜索Agent,基于LiteLLM和MCP协议,能够理解用户查询意图、自动规划搜索策略、从多个来源获取信息,并生成高质量、有来源引用的答案"
# Agent版本
version: str = "1.0.0"
# 服务端口
port: int = 8080
# 服务主机
host: str = "0.0.0.0"
# 是否启用流式响应
enable_streaming: bool = True
# 系统提示词
system_prompt: str = "你是一个智能搜索助手。"
@dataclass
class MCPConfig:
"""MCP协议配置"""
# MCP协议版本
protocol_version: str = "2024-11-05"
# 服务器信息
server_name: str = "search-agent-mcp"
server_version: str = "1.0.0"
def get_config(
api_key: Optional[str] = None,
model: Optional[str] = None
) -> tuple[LiteLLMConfig, AgentConfig, MCPConfig]:
"""
获取完整配置
Args:
api_key: LiteLLM API密钥(可选,优先使用,否则从环境变量获取)
model: 模型名称(可选,优先使用,否则从环境变量获取)
Returns:
(LiteLLMConfig, AgentConfig, MCPConfig) 配置元组
"""
litellm_config = LiteLLMConfig(api_key=api_key, model=model)
agent_config = AgentConfig()
mcp_config = MCPConfig()
return litellm_config, agent_config, mcp_config
@@ -0,0 +1,540 @@
"""
MCP协议兼容的Search Agent服务
实现Model Context Protocol协议规范
支持从请求传入 API key,也支持从环境变量获取
使用LiteLLM SDK进行模型调用
"""
import asyncio
import json
import uuid
import os
from typing import Optional, Dict, Any, AsyncGenerator
from datetime import datetime
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException, Request, Response
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from loguru import logger
from agent import SearchAgentWrapper
from mcp_config import get_config, AgentConfig, MCPConfig
# 环境变量配置
SERVICE_HOST = os.getenv("SERVICE_HOST", "0.0.0.0")
SERVICE_PORT = int(os.getenv("SERVICE_PORT", "8080"))
POD_NAME = os.getenv("POD_NAME", "search-agent-mcp")
TEMPLATE_TYPE = os.getenv("TEMPLATE_TYPE", "search_agent_MCP")
# ============== MCP 协议数据模型 ==============
class MCPMessage(BaseModel):
"""MCP消息"""
role: str
content: str
toolCalls: Optional[list[Dict[str, Any]]] = None
class MCPRequest(BaseModel):
"""MCP JSON-RPC请求"""
jsonrpc: str = "2.0"
id: Optional[str] = None
method: str
params: Optional[Dict[str, Any]] = None
class MCPResponse(BaseModel):
"""MCP JSON-RPC响应"""
jsonrpc: str = "2.0"
id: Optional[str] = None
result: Optional[Dict[str, Any]] = None
error: Optional[Dict[str, Any]] = None
class MCPSearchParams(BaseModel):
"""MCP搜索参数"""
query: str = Field(..., description="搜索查询")
api_key: Optional[str] = Field(None, description="LiteLLM API密钥(可选,优先使用,否则从环境变量获取)")
model: Optional[str] = Field(None, description="模型名称(可选,优先使用,否则从环境变量获取)")
class MCPSearchResult(BaseModel):
"""MCP搜索结果"""
query: str
answer: str
sources: list[Dict[str, Any]]
confidence: str
iterations: int
total_sources: int
search_queries: list[str]
timestamp: str = Field(default_factory=lambda: datetime.utcnow().isoformat() + "Z")
# ============== MCP Server ==============
class MCPSearchAgentServer:
"""MCP协议Search Agent服务器"""
def __init__(
self,
api_key: Optional[str] = None,
model: Optional[str] = None
):
"""
初始化MCP Search Agent服务器
Args:
api_key: LiteLLM API密钥(可选,优先使用,否则从环境变量获取)
model: 模型名称(可选,优先使用,否则从环境变量获取)
"""
# 获取配置
self.llm_config, self.agent_config, self.mcp_config = get_config(api_key, model)
# 创建默认Agent(使用默认配置)
self.default_agent = SearchAgentWrapper(
litellm_config=self.llm_config,
agent_config=self.agent_config
)
# 任务存储
self.tasks: Dict[str, Dict[str, Any]] = {}
# 创建FastAPI应用
self.app = self._create_app()
def _create_app(self) -> FastAPI:
"""创建FastAPI应用"""
@asynccontextmanager
async def lifespan(app: FastAPI):
logger.info("MCP Search Agent服务启动", agent_name=self.agent_config.name)
yield
await self.default_agent.close()
logger.info("MCP Search Agent服务关闭")
app = FastAPI(
title=f"{self.agent_config.name} - MCP Agent",
description=self.agent_config.description,
version=self.agent_config.version,
lifespan=lifespan
)
# CORS中间件
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# 注册路由
self._register_routes(app)
return app
def _get_agent(self, api_key: Optional[str] = None, model: Optional[str] = None) -> SearchAgentWrapper:
"""
获取Agent实例
如果提供了api_key或model,创建新的Agent实例
否则使用默认Agent
"""
if api_key or model:
# 创建新的配置和Agent
llm_config, agent_config, _ = get_config(api_key, model)
return SearchAgentWrapper(
litellm_config=llm_config,
agent_config=agent_config,
api_key=api_key,
model=model
)
return self.default_agent
def _register_routes(self, app: FastAPI):
"""注册MCP协议路由"""
@app.get("/")
async def root():
"""服务根路径"""
return {
"name": self.agent_config.name,
"version": self.agent_config.version,
"protocol": "MCP",
"status": "running",
"pod_name": POD_NAME,
"template_type": TEMPLATE_TYPE
}
@app.get("/health")
async def health_check():
"""健康检查"""
return {
"status": "healthy",
"pod_name": POD_NAME,
"template_type": TEMPLATE_TYPE,
"configured": self.llm_config.api_key is not None,
"llm_base_url": self.llm_config.base_url,
"llm_model": self.llm_config.model,
"timestamp": datetime.utcnow().isoformat()
}
@app.post("/mcp/v1/search")
async def mcp_search(request: Request):
"""MCP搜索端点"""
body = await request.json()
# 解析JSON-RPC请求
try:
rpc_request = MCPRequest(**body)
except Exception as e:
return JSONResponse({
"jsonrpc": "2.0",
"id": body.get("id", "unknown"),
"error": {
"code": -32600,
"message": f"Invalid Request: {str(e)}"
}
})
# 处理 search 方法
if rpc_request.method == "search":
return await self._handle_search(rpc_request)
else:
return JSONResponse({
"jsonrpc": "2.0",
"id": rpc_request.id,
"error": {
"code": -32601,
"message": f"Method not found: {rpc_request.method}"
}
})
@app.post("/mcp/v1/search/stream")
async def mcp_search_stream(request: Request):
"""MCP流式搜索端点 (SSE)"""
body = await request.json()
try:
rpc_request = MCPRequest(**body)
except Exception as e:
return JSONResponse({
"jsonrpc": "2.0",
"id": body.get("id", "unknown"),
"error": {
"code": -32600,
"message": f"Invalid Request: {str(e)}"
}
})
return await self._handle_search_stream(rpc_request)
@app.post("/mcp/v1/call")
async def mcp_call(request: Request):
"""MCP通用调用端点(JSON-RPC兼容)"""
body = await request.json()
try:
rpc_request = MCPRequest(**body)
except Exception as e:
return JSONResponse({
"jsonrpc": "2.0",
"id": body.get("id", "unknown"),
"error": {
"code": -32600,
"message": f"Invalid Request: {str(e)}"
}
})
# 根据方法名路由
if rpc_request.method == "search":
return await self._handle_search(rpc_request)
elif rpc_request.method == "search/stream":
return await self._handle_search_stream(rpc_request)
else:
return JSONResponse({
"jsonrpc": "2.0",
"id": rpc_request.id,
"error": {
"code": -32601,
"message": f"Method not found: {rpc_request.method}"
}
})
async def _handle_search(self, request: MCPRequest) -> JSONResponse:
"""处理 search 请求"""
params = request.params or {}
# 提取搜索查询
query = params.get("query", "")
if not query:
return JSONResponse({
"jsonrpc": "2.0",
"id": request.id,
"error": {
"code": -32602,
"message": "Invalid params: 'query' is required"
}
})
# 提取API key(必须从请求参数中获取,等同于API格式版本的llm_api_key)
api_key = params.get("api_key") or params.get("llm_api_key")
if not api_key:
return JSONResponse({
"jsonrpc": "2.0",
"id": request.id,
"error": {
"code": -32602,
"message": "Invalid params: 'api_key' or 'llm_api_key' is required"
}
})
# 提取model(从环境变量获取,不支持在请求中传递,与API格式版本保持一致)
# 支持多种环境变量名称:MODEL_NAME(优先)、LLM_MODEL(AKS部署)、LITELLM_MODEL
model = os.getenv("MODEL_NAME") or os.getenv("LLM_MODEL") or os.getenv("LITELLM_MODEL")
if not model:
return JSONResponse({
"jsonrpc": "2.0",
"id": request.id,
"error": {
"code": -32000,
"message": "Model not configured: MODEL_NAME or LLM_MODEL environment variable is required"
}
})
# 创建任务ID
task_id = params.get("task_id", uuid.uuid4().hex)
try:
# 获取Agent实例
agent = self._get_agent(api_key, model)
# 调用Agent获取响应
logger.info("处理搜索请求", task_id=task_id, query_preview=query[:50])
response = await agent.search(query=query)
# 如果创建了新Agent,关闭它
if api_key or model:
await agent.close()
# 构建响应数据
sources = []
if response.answer.sources:
for i, source in enumerate(response.answer.sources, 1):
sources.append({
"index": i,
"title": source.title,
"url": source.url
})
result = MCPSearchResult(
query=query,
answer=response.answer.content,
sources=sources,
confidence=response.answer.confidence,
iterations=response.iterations,
total_sources=response.total_sources_consulted,
search_queries=response.search_queries_used
)
return JSONResponse({
"jsonrpc": "2.0",
"id": request.id,
"result": result.model_dump()
})
except Exception as e:
logger.error("处理搜索请求失败", error=str(e))
return JSONResponse({
"jsonrpc": "2.0",
"id": request.id,
"error": {
"code": -32000,
"message": f"Server error: {str(e)}"
}
})
async def _handle_search_stream(self, request: MCPRequest) -> StreamingResponse:
"""处理 search/stream 请求 (SSE)"""
params = request.params or {}
# 提取搜索查询
query = params.get("query", "")
if not query:
return JSONResponse({
"jsonrpc": "2.0",
"id": request.id,
"error": {
"code": -32602,
"message": "Invalid params: 'query' is required"
}
})
# 提取API key(必须从请求参数中获取,等同于API格式版本的llm_api_key)
api_key = params.get("api_key") or params.get("llm_api_key")
if not api_key:
return JSONResponse({
"jsonrpc": "2.0",
"id": request.id,
"error": {
"code": -32602,
"message": "Invalid params: 'api_key' or 'llm_api_key' is required"
}
})
# 提取model(从环境变量获取,不支持在请求中传递,与API格式版本保持一致)
# 支持多种环境变量名称:MODEL_NAME(优先)、LLM_MODEL(AKS部署)、LITELLM_MODEL
model = os.getenv("MODEL_NAME") or os.getenv("LLM_MODEL") or os.getenv("LITELLM_MODEL")
if not model:
return JSONResponse({
"jsonrpc": "2.0",
"id": request.id,
"error": {
"code": -32000,
"message": "Model not configured: MODEL_NAME or LLM_MODEL environment variable is required"
}
})
task_id = params.get("task_id", uuid.uuid4().hex)
async def event_generator() -> AsyncGenerator[str, None]:
"""生成SSE事件流"""
agent = None
try:
# 获取Agent实例
agent = self._get_agent(api_key, model)
# 发送任务开始事件
start_event = {
"jsonrpc": "2.0",
"id": request.id,
"method": "task/start",
"params": {
"task_id": task_id
}
}
yield f"data: {json.dumps(start_event)}\n\n"
# 执行搜索
response = await agent.search(query=query)
# 构建答案文本
answer_parts = [response.answer.content]
if response.answer.sources:
answer_parts.append("\n\n## 来源")
for i, source in enumerate(response.answer.sources, 1):
answer_parts.append(f"{i}. [{source.title}]({source.url})")
answer_text = "\n".join(answer_parts)
# 发送完整答案(作为增量发送,以便显示进度)
chunk_size = 100
for i in range(0, len(answer_text), chunk_size):
chunk = answer_text[i:i + chunk_size]
delta_event = {
"jsonrpc": "2.0",
"id": request.id,
"method": "result/delta",
"params": {
"task_id": task_id,
"delta": chunk
}
}
yield f"data: {json.dumps(delta_event)}\n\n"
await asyncio.sleep(0.01)
# 发送完成事件
sources = []
if response.answer.sources:
for i, source in enumerate(response.answer.sources, 1):
sources.append({
"index": i,
"title": source.title,
"url": source.url
})
complete_event = {
"jsonrpc": "2.0",
"id": request.id,
"method": "task/complete",
"params": {
"task_id": task_id,
"result": {
"query": query,
"answer": response.answer.content,
"sources": sources,
"confidence": response.answer.confidence,
"iterations": response.iterations,
"total_sources": response.total_sources_consulted,
"search_queries": response.search_queries_used
}
}
}
yield f"data: {json.dumps(complete_event)}\n\n"
except Exception as e:
# 发送错误事件
error_event = {
"jsonrpc": "2.0",
"id": request.id,
"method": "task/error",
"params": {
"task_id": task_id,
"error": str(e)
}
}
yield f"data: {json.dumps(error_event)}\n\n"
finally:
# 如果创建了新Agent,关闭它
if agent and (api_key or model):
await agent.close()
return StreamingResponse(
event_generator(),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no"
}
)
def run(self, host: Optional[str] = None, port: Optional[int] = None):
"""运行服务器"""
import uvicorn
host = host or self.agent_config.host
port = port or self.agent_config.port
logger.info(f"启动MCP Search Agent服务", host=host, port=port)
uvicorn.run(self.app, host=host, port=port)
def create_app(api_key: Optional[str] = None, model: Optional[str] = None) -> FastAPI:
"""
创建FastAPI应用(用于uvicorn启动)
使用方式:
uvicorn mcp_server:app --host 0.0.0.0 --port 8080
或设置环境变量后:
export LITELLM_API_KEY="your-key"
export MODEL_NAME="your-model"
uvicorn mcp_server:app --host 0.0.0.0 --port 8080
"""
server = MCPSearchAgentServer(api_key=api_key, model=model)
return server.app
# uvicorn 启动入口
# 环境变量: LITELLM_API_KEY, MODEL_NAME (或 LITELLM_MODEL)
# 注意: app 只在 main.py 中创建,避免导入时立即执行验证
# app = create_app()
@@ -0,0 +1,30 @@
# FastAPI 和 Web 服务器
fastapi>=0.109.0
uvicorn[standard]>=0.27.0
pydantic>=2.5.3
# LiteLLM SDK - 用于模型调用
litellm>=1.40.0
# HTTP客户端 - 用于LiteLLM SDK调用和其他API调用
httpx>=0.27.0
# 环境变量
python-dotenv>=1.0.0
# JSON处理
orjson>=3.9.0
# 类型提示
typing-extensions>=4.9.0
# 日志
loguru>=0.7.0
# 异步工具
asyncio-throttle>=1.0.2
# HTTP客户端 - 用于搜索和其他API调用
aiohttp>=3.9.0
requests>=2.31.0
@@ -0,0 +1 @@
../../search_agent/search_agent
@@ -0,0 +1,44 @@
FROM python:3.11-slim
WORKDIR /app
# 安装系统依赖
RUN apt-get update && apt-get install -y \
curl \
&& rm -rf /var/lib/apt/lists/*
# 复制requirements文件
COPY agents/search_agent/search_agent_MCP/requirements.txt /app/requirements.txt
COPY agents/search_agent/search_agent/search_agent/requirements.txt /app/search_agent_requirements.txt
# 安装Python依赖
RUN pip install --no-cache-dir \
fastapi==0.109.0 \
uvicorn[standard]==0.27.0 \
pydantic==2.5.3 \
&& pip install --no-cache-dir -r /app/requirements.txt \
&& pip install --no-cache-dir -r /app/search_agent_requirements.txt
# 复制search_agent_MCP目录
COPY agents/search_agent/search_agent_MCP/ /app/
# 复制search_agent核心代码
COPY agents/search_agent/search_agent/search_agent/ /app/search_agent/
# 设置环境变量
ENV PYTHONUNBUFFERED=1
ENV SERVICE_HOST=0.0.0.0
ENV SERVICE_PORT=8080
ENV PYTHONPATH=/app
# 默认 API 密钥(硬编码)
ENV JINA_API_KEY=jina_e26dc30420a44a1e859216528065b203TkMRmsoz-FgMDQC5FZX9jr5oF2CI
ENV SERPER_API_KEY=8253b4f240b520194065312f90e85f9be0fa205f
# 健康检查 - 使用Python避免僵尸进程
HEALTHCHECK --interval=30s --timeout=10s --start-period=40s --retries=3 \
CMD python3 -c "import urllib.request; urllib.request.urlopen('http://localhost:8080/health').read()" || exit 1
# 运行agent (直接使用Python,避免shell)
CMD ["python3", "-u", "main.py"]
@@ -0,0 +1,305 @@
"""
MCP Search Agent 测试脚本
测试MCP格式的API调用
"""
import asyncio
import json
import sys
import os
from pathlib import Path
# 添加当前目录到Python路径
current_dir = Path(__file__).parent
sys.path.insert(0, str(current_dir))
import httpx
from loguru import logger
# 配置日志
logger.remove()
logger.add(sys.stderr, level="INFO", format="<green>{time:HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{message}</cyan>")
# 服务地址
BASE_URL = os.getenv("TEST_BASE_URL", "http://localhost:8080")
def test_health():
"""测试健康检查"""
print("\n" + "=" * 70)
print("测试健康检查")
print("=" * 70)
try:
response = httpx.get(f"{BASE_URL}/health", timeout=10.0)
print(f"状态码: {response.status_code}")
print(f"响应: {json.dumps(response.json(), indent=2, ensure_ascii=False)}")
assert response.status_code == 200
print("✓ 健康检查通过")
return True
except Exception as e:
print(f"✗ 健康检查失败: {e}")
return False
def test_mcp_search():
"""测试MCP搜索接口"""
print("\n" + "=" * 70)
print("测试MCP搜索接口")
print("=" * 70)
request_data = {
"jsonrpc": "2.0",
"id": "test-1",
"method": "search",
"params": {
"query": "什么是人工智能?",
# 可选:可以从环境变量或请求中传入
# "api_key": os.getenv("LITELLM_API_KEY"),
# "model": os.getenv("MODEL_NAME")
}
}
print(f"请求: {json.dumps(request_data, indent=2, ensure_ascii=False)}")
try:
response = httpx.post(
f"{BASE_URL}/mcp/v1/search",
json=request_data,
timeout=120.0 # 搜索可能需要较长时间
)
print(f"\n状态码: {response.status_code}")
if response.status_code == 200:
result = response.json()
print(f"响应: {json.dumps(result, indent=2, ensure_ascii=False)}")
# 验证响应格式
assert "jsonrpc" in result
assert result["jsonrpc"] == "2.0"
assert "id" in result
assert result["id"] == "test-1"
if "result" in result:
assert "query" in result["result"]
assert "answer" in result["result"]
assert "sources" in result["result"]
print("\n✓ MCP搜索测试通过")
return True
elif "error" in result:
print(f"\n✗ MCP搜索返回错误: {result['error']}")
return False
else:
print(f"错误响应: {response.text}")
return False
except Exception as e:
print(f"✗ MCP搜索测试失败: {e}")
import traceback
traceback.print_exc()
return False
def test_mcp_call():
"""测试MCP通用调用接口"""
print("\n" + "=" * 70)
print("测试MCP通用调用接口")
print("=" * 70)
request_data = {
"jsonrpc": "2.0",
"id": "test-2",
"method": "search",
"params": {
"query": "Python编程语言的特点"
}
}
print(f"请求: {json.dumps(request_data, indent=2, ensure_ascii=False)}")
try:
response = httpx.post(
f"{BASE_URL}/mcp/v1/call",
json=request_data,
timeout=120.0
)
print(f"\n状态码: {response.status_code}")
if response.status_code == 200:
result = response.json()
print(f"响应: {json.dumps(result, indent=2, ensure_ascii=False)}")
# 验证响应格式
assert "jsonrpc" in result
assert result["jsonrpc"] == "2.0"
if "result" in result:
print("\n✓ MCP通用调用测试通过")
return True
elif "error" in result:
print(f"\n✗ MCP通用调用返回错误: {result['error']}")
return False
else:
print(f"错误响应: {response.text}")
return False
except Exception as e:
print(f"✗ MCP通用调用测试失败: {e}")
import traceback
traceback.print_exc()
return False
def test_mcp_search_stream():
"""测试MCP流式搜索接口"""
print("\n" + "=" * 70)
print("测试MCP流式搜索接口")
print("=" * 70)
request_data = {
"jsonrpc": "2.0",
"id": "test-3",
"method": "search/stream",
"params": {
"query": "机器学习的基本概念"
}
}
print(f"请求: {json.dumps(request_data, indent=2, ensure_ascii=False)}")
print("\n流式响应:")
try:
with httpx.stream(
"POST",
f"{BASE_URL}/mcp/v1/search/stream",
json=request_data,
timeout=120.0
) as response:
print(f"状态码: {response.status_code}")
if response.status_code == 200:
chunks = []
for line in response.iter_lines():
if line.startswith("data: "):
data = line[6:] # 去掉 "data: " 前缀
try:
event = json.loads(data)
chunks.append(event)
print(f"事件: {json.dumps(event, indent=2, ensure_ascii=False)}")
except json.JSONDecodeError:
print(f"非JSON数据: {data}")
if chunks:
print(f"\n✓ 收到 {len(chunks)} 个事件")
return True
else:
print("\n✗ 未收到任何事件")
return False
else:
text = response.read().decode()
print(f"错误响应: {text}")
return False
except Exception as e:
print(f"✗ MCP流式搜索测试失败: {e}")
import traceback
traceback.print_exc()
return False
def test_mcp_error_handling():
"""测试MCP错误处理"""
print("\n" + "=" * 70)
print("测试MCP错误处理")
print("=" * 70)
# 测试缺少query参数
request_data = {
"jsonrpc": "2.0",
"id": "test-error-1",
"method": "search",
"params": {}
}
print(f"请求(缺少query): {json.dumps(request_data, indent=2, ensure_ascii=False)}")
try:
response = httpx.post(
f"{BASE_URL}/mcp/v1/search",
json=request_data,
timeout=10.0
)
print(f"\n状态码: {response.status_code}")
result = response.json()
print(f"响应: {json.dumps(result, indent=2, ensure_ascii=False)}")
# 应该返回错误
if "error" in result:
assert result["error"]["code"] != 0
print("\n✓ 错误处理测试通过(正确返回错误)")
return True
else:
print("\n✗ 错误处理测试失败(应该返回错误但没有)")
return False
except Exception as e:
print(f"✗ 错误处理测试失败: {e}")
return False
def main():
"""运行所有测试"""
print("\n" + "=" * 70)
print("MCP Search Agent 测试")
print("=" * 70)
print(f"服务地址: {BASE_URL}")
print("=" * 70)
# 检查服务是否可访问
try:
httpx.get(f"{BASE_URL}/health", timeout=5.0)
except Exception as e:
print(f"\n✗ 无法连接到服务: {e}")
print("请确保服务已启动:python main.py")
return
results = []
# 运行测试
results.append(("健康检查", test_health()))
results.append(("MCP搜索", test_mcp_search()))
results.append(("MCP通用调用", test_mcp_call()))
results.append(("MCP流式搜索", test_mcp_search_stream()))
results.append(("错误处理", test_mcp_error_handling()))
# 汇总结果
print("\n" + "=" * 70)
print("测试结果汇总")
print("=" * 70)
passed = 0
total = len(results)
for name, result in results:
status = "✓ 通过" if result else "✗ 失败"
print(f"{name:20s}: {status}")
if result:
passed += 1
print("=" * 70)
print(f"总计: {passed}/{total} 通过")
print("=" * 70)
if passed == total:
print("\n✓ 所有测试通过!")
return 0
else:
print(f"\n✗ {total - passed} 个测试失败")
return 1
if __name__ == "__main__":
sys.exit(main())
@@ -0,0 +1,129 @@
"""
测试 MCP 格式调用结构
"""
import os
import sys
import json
import asyncio
from datetime import datetime
# 添加当前目录到路径
sys.path.insert(0, os.path.dirname(__file__))
async def test_mcp_structure():
"""测试 MCP 格式调用结构"""
# 设置环境变量(如果需要)
test_api_key = os.getenv("LITELLM_API_KEY", "test-api-key")
test_model = os.getenv("MODEL_NAME", "gpt-4")
print("=" * 60)
print("测试 MCP Search Agent")
print("=" * 60)
print(f"API Key: {test_api_key[:10]}..." if test_api_key else "未设置")
print(f"Model: {test_model}")
print()
try:
# 创建服务器实例(不验证 API key,仅测试结构)
print("1. 创建 MCP 服务器实例...")
# 先检查是否能导入
from mcp_config import get_config
from agent import SearchAgentWrapper
print(" ✓ 模块导入成功")
# 测试配置
print("2. 测试配置加载...")
llm_config, agent_config, mcp_config = get_config(
api_key=test_api_key,
model=test_model
)
print(f" ✓ LLM Base URL: {llm_config.base_url}")
print(f" ✓ LLM Model: {llm_config.model}")
print(f" ✓ Agent Name: {agent_config.name}")
# 测试 MCP 请求结构
print("\n3. 测试 MCP 请求结构...")
# 模拟 MCP 请求
test_request = {
"jsonrpc": "2.0",
"id": "test-request-001",
"method": "search",
"params": {
"query": "什么是人工智能?",
"api_key": test_api_key,
"model": test_model
}
}
print(" ✓ MCP 请求格式正确")
print(f" 请求示例: {json.dumps(test_request, indent=2, ensure_ascii=False)}")
# 检查 MCP 数据模型
print("\n4. 检查 MCP 数据模型...")
# 导入 MCP 数据模型(不触发服务器创建)
import mcp_server
from mcp_server import (
MCPMessage,
MCPRequest,
MCPResponse,
MCPSearchParams,
MCPSearchResult
)
# 创建测试消息
message = MCPMessage(role="user", content="测试消息")
request = MCPRequest(
jsonrpc="2.0",
id="test-1",
method="search",
params={"query": "测试查询"}
)
print(f" ✓ MCPMessage 创建成功: {message.role}")
print(f" ✓ MCPRequest 创建成功: {request.method}")
# 测试 SearchResult
print("\n5. 测试 MCPSearchResult...")
result = MCPSearchResult(
query="测试查询",
answer="测试答案",
sources=[],
confidence="high",
iterations=1,
total_sources=0,
search_queries=["测试查询"]
)
print(f" ✓ MCPSearchResult 创建成功")
print(f" 结果示例: {json.dumps(result.model_dump(), indent=2, ensure_ascii=False)}")
print("\n" + "=" * 60)
print("✅ 所有测试通过!MCP 格式调用结构正常")
print("=" * 60)
print("\n注意: 要进行实际的搜索调用,需要:")
print(" 1. 设置有效的 LITELLM_API_KEY 环境变量")
print(" 2. 设置 SERPER_API_KEY 和 JINA_API_KEY 环境变量")
print(" 3. 启动服务器: python main.py")
print(" 4. 发送 MCP 格式的 HTTP 请求到 /mcp/v1/search")
print("\n示例请求:")
print(json.dumps(test_request, indent=2, ensure_ascii=False))
return True
except Exception as e:
print(f"\n❌ 测试失败: {str(e)}")
import traceback
traceback.print_exc()
return False
if __name__ == "__main__":
success = asyncio.run(test_mcp_structure())
sys.exit(0 if success else 1)
@@ -1,10 +1,13 @@
"""
智能搜索 AI Agent - FastAPI版本
通过HTTP API接收搜索请求,提供智能搜索功能
使用场景:多个用户使用不同的 API key,但共享固定的 model 和 endpoint
"""
import os
import sys
import logging
import copy
from typing import Optional, Dict, Any, List
from datetime import datetime
from fastapi import FastAPI, HTTPException
@@ -20,7 +23,12 @@ if search_agent_dir not in sys.path:
# 直接导入,避免与文件名冲突
from config import Config
from agent.search_agent import SearchAgent
from agent_callback_utils import AgentCallbackHandler, CallbackContextManager
# 尝试从不同路径导入回调工具
try:
from common.agent_callback_utils import AgentCallbackHandler, CallbackContextManager
except ImportError:
from agent_callback_utils import AgentCallbackHandler, CallbackContextManager
# 配置日志
logging.basicConfig(
@@ -35,11 +43,13 @@ SERVICE_PORT = int(os.getenv("SERVICE_PORT", "8080"))
POD_NAME = os.getenv("POD_NAME", "search-agent")
TEMPLATE_TYPE = os.getenv("TEMPLATE_TYPE", "search_agent")
# 全局搜索Agent和回调处理器
search_agent: Optional[SearchAgent] = None
config: Optional[Config] = None
# 基础配置(启动时加载,不包含 llm_api_key,是不可变的)
base_config: Optional[Config] = None
callback_handler: Optional[AgentCallbackHandler] = None
# 环境变量
USER_ID = os.getenv("USER_ID", "")
# FastAPI应用
app = FastAPI(
title="Intelligent Search AI Agent",
@@ -64,11 +74,10 @@ class ConfigRequest(BaseModel):
class SearchRequest(BaseModel):
"""搜索请求"""
"""搜索请求 - 简化版:只需传入query和llm_api_key,其他从环境变量获取"""
query: str = Field(..., description="搜索查询")
llm_api_key: str = Field(..., description="LLM API密钥")
user_id: Optional[str] = Field(None, description="用户ID(用于计费回调)")
auto_configure: bool = Field(default=False, description="是否自动从环境变量配置")
class Source(BaseModel):
@@ -107,31 +116,93 @@ class ErrorResponse(BaseModel):
# ==================== Agent操作函数 ====================
def initialize_agent_from_env():
"""从环境变量初始化Agent"""
global search_agent, config
def load_base_config():
"""
从环境变量加载基础配置(不包含 llm_api_key)
必须的环境变量:
- LLM_BASE_URL: LLM 服务地址(固定)
- MODEL_NAME: 模型名称(固定)
- SERPER_API_KEY: Serper 搜索 API 密钥(固定)
- JINA_API_KEY: Jina Reader API 密钥(固定)
llm_api_key 在每次请求时由用户传入
"""
global base_config
try:
config = Config.from_env()
config.validate()
search_agent = SearchAgent(config)
logger.info("Search Agent从环境变量初始化成功")
# 检查必须的环境变量
required_vars = ["LLM_BASE_URL", "SERPER_API_KEY", "JINA_API_KEY"]
missing = [v for v in required_vars if not os.getenv(v)]
if missing:
raise ValueError(f"缺少必要的环境变量: {', '.join(missing)}")
# 创建基础配置(llm_api_key 使用占位符,每次请求时会被替换)
base_config = Config(
llm_base_url=os.getenv("LLM_BASE_URL", ""),
llm_api_key="__PLACEHOLDER__", # 占位符,每次请求时替换
llm_model=os.getenv("MODEL_NAME", "gpt-4o-mini"),
serper_api_key=os.getenv("SERPER_API_KEY", ""),
jina_api_key=os.getenv("JINA_API_KEY", ""),
max_iterations=int(os.getenv("MAX_ITERATIONS", "3")),
max_results_per_query=int(os.getenv("MAX_RESULTS_PER_QUERY", "10")),
content_max_length=int(os.getenv("CONTENT_MAX_LENGTH", "5000")),
log_level=os.getenv("LOG_LEVEL", "INFO"),
timeout=int(os.getenv("TIMEOUT", "30"))
)
logger.info("基础配置加载成功")
logger.info(f" LLM_BASE_URL: {base_config.llm_base_url}")
logger.info(f" MODEL_NAME: {base_config.llm_model}")
logger.info(f" SERPER_API_KEY: {'已设置' if base_config.serper_api_key else '未设置'}")
logger.info(f" JINA_API_KEY: {'已设置' if base_config.jina_api_key else '未设置'}")
return True
except Exception as e:
logger.error(f"从环境变量初始化Agent失败: {str(e)}")
logger.error(f"加载基础配置失败: {str(e)}")
return False
def create_agent_for_request(llm_api_key: str) -> SearchAgent:
"""
为单次请求创建 SearchAgent(使用用户的 API key)
Args:
llm_api_key: 用户的 LLM API 密钥
Returns:
配置了用户 API key 的 SearchAgent 实例
"""
global base_config
if not base_config:
raise ValueError("基础配置未加载")
# 创建配置副本,设置用户的 API key(线程安全)
request_config = Config(
llm_base_url=base_config.llm_base_url,
llm_api_key=llm_api_key,
llm_model=base_config.llm_model,
serper_api_key=base_config.serper_api_key,
jina_api_key=base_config.jina_api_key,
max_iterations=base_config.max_iterations,
max_results_per_query=base_config.max_results_per_query,
content_max_length=base_config.content_max_length,
log_level=base_config.log_level,
timeout=base_config.timeout
)
return SearchAgent(request_config)
def initialize_agent_from_config(config_data: Dict[str, Any]):
"""从配置数据初始化Agent"""
global search_agent, config
"""从配置数据初始化基础配置(用于 /configure 端点)"""
global base_config
try:
# 创建配置对象
config = Config(
base_config = Config(
llm_base_url=config_data.get("llm_base_url", ""),
llm_api_key=config_data.get("llm_api_key", ""),
llm_model=config_data.get("llm_model", "xchat52"),
llm_api_key="__PLACEHOLDER__", # 占位符
llm_model=config_data.get("llm_model", "gpt-4o-mini"),
serper_api_key=config_data.get("serper_api_key", ""),
jina_api_key=config_data.get("jina_api_key", ""),
max_iterations=config_data.get("max_iterations", 3),
@@ -141,12 +212,10 @@ def initialize_agent_from_config(config_data: Dict[str, Any]):
timeout=config_data.get("timeout", 30)
)
config.validate()
search_agent = SearchAgent(config)
logger.info("Search Agent从配置初始化成功")
logger.info("基础配置更新成功")
return True
except Exception as e:
logger.error(f"从配置初始化Agent失败: {str(e)}")
logger.error(f"配置更新失败: {str(e)}")
raise
@@ -159,7 +228,9 @@ async def health_check():
"status": "healthy",
"pod_name": POD_NAME,
"template_type": TEMPLATE_TYPE,
"configured": search_agent is not None,
"configured": base_config is not None,
"llm_base_url": base_config.llm_base_url if base_config else None,
"llm_model": base_config.llm_model if base_config else None,
"timestamp": datetime.utcnow().isoformat()
}
@@ -168,98 +239,130 @@ async def health_check():
async def get_status():
"""获取状态"""
return StatusResponse(
status="running" if search_agent else "not_configured",
status="running" if base_config else "not_configured",
pod_name=POD_NAME,
template_type=TEMPLATE_TYPE,
configured=search_agent is not None,
configured=base_config is not None,
timestamp=datetime.utcnow().isoformat()
)
@app.post("/configure")
async def configure_agent(config_req: ConfigRequest):
"""配置Agent"""
"""配置基础参数(model、endpoint、serper_key、jina_key)"""
try:
initialize_agent_from_config(config_req.dict())
return {
"status": "success",
"message": "Agent配置成功",
"message": "基础配置更新成功",
"llm_base_url": base_config.llm_base_url,
"llm_model": base_config.llm_model,
"timestamp": datetime.utcnow().isoformat()
}
except Exception as e:
logger.error(f"配置Agent失败: {str(e)}")
logger.error(f"配置更新失败: {str(e)}")
raise HTTPException(status_code=400, detail=f"配置失败: {str(e)}")
@app.post("/search", response_model=SearchResponse)
async def search(request: SearchRequest):
"""执行搜索"""
global search_agent, callback_handler, config
"""
执行搜索
# 如果未配置且需要自动配置
if not search_agent and request.auto_configure:
if not initialize_agent_from_env():
请求参数:
- query: 搜索查询(必须)
- llm_api_key: 用户的 LLM API 密钥(必须)
- user_id: 用户ID(可选,用于计费回调)
固定配置(从环境变量获取,所有用户共享):
- LLM_BASE_URL: LLM 服务地址
- MODEL_NAME: 模型名称
- SERPER_API_KEY: Serper 搜索 API 密钥
- JINA_API_KEY: Jina Reader API 密钥
"""
global base_config, callback_handler
# 自动加载基础配置(如果尚未加载)
if not base_config:
if not load_base_config():
raise HTTPException(
status_code=400,
detail="Agent未配置且自动配置失败,请先调用/configure接口"
status_code=500,
detail="基础配置加载失败,请检查环境变量: LLM_BASE_URL, SERPER_API_KEY, JINA_API_KEY"
)
if not search_agent:
raise HTTPException(
status_code=400,
detail="Agent未配置,请先调用/configure接口"
# 初始化回调处理器(全局单例)
if not callback_handler:
callback_handler = AgentCallbackHandler(
agent_name=POD_NAME,
user_id=USER_ID
)
# 初始化回调处理器(如果尚未初始化)
if not callback_handler:
callback_handler = AgentCallbackHandler()
# 验证必须的参数
if not request.llm_api_key:
raise HTTPException(
status_code=400,
detail="llm_api_key 是必须的参数"
)
# 使用上下文管理器自动处理回调
# 为本次请求创建独立的 SearchAgent(使用用户的 API key,线程安全)
try:
with CallbackContextManager(
handler=callback_handler,
user_id=request.user_id,
request_id=f"search-{int(datetime.utcnow().timestamp())}"
) as ctx:
# 临时更新API key
original_api_key = config.llm_api_key if config else None
if config:
config.llm_api_key = request.llm_api_key
search_agent.config.llm_api_key = request.llm_api_key
search_agent = create_agent_for_request(request.llm_api_key)
except Exception as e:
logger.error(f"创建 SearchAgent 失败: {str(e)}")
raise HTTPException(status_code=500, detail=f"Agent 创建失败: {str(e)}")
try:
# 执行搜索
# 确定用于回调的 user_id
effective_user_id = request.user_id or USER_ID
try:
# 执行搜索(如果有 user_id,使用回调上下文管理器)
if effective_user_id:
with CallbackContextManager(
handler=callback_handler,
user_id=effective_user_id,
request_id=f"search-{int(datetime.utcnow().timestamp())}"
) as ctx:
ctx.add_tool("web_search")
ctx.add_tool("content_reader")
ctx.add_tool("llm_completion")
result = await search_agent.search(request.query)
else:
result = await search_agent.search(request.query)
# 转换响应
sources = [
Source(
index=s.index,
title=s.title,
url=s.url
)
for s in result.answer.sources
]
# 转换响应
sources = [
Source(
index=s.index,
title=s.title,
url=s.url
)
for s in result.answer.sources
]
return SearchResponse(
query=request.query,
answer=result.answer.content,
sources=sources,
confidence=result.answer.confidence,
iterations=result.iterations,
total_sources=result.total_sources_consulted,
search_queries=result.search_queries_used,
timestamp=datetime.utcnow().isoformat()
)
finally:
# 恢复原始API key
if config and original_api_key:
config.llm_api_key = original_api_key
search_agent.config.llm_api_key = original_api_key
return SearchResponse(
query=request.query,
answer=result.answer.content,
sources=sources,
confidence=result.answer.confidence,
iterations=result.iterations,
total_sources=result.total_sources_consulted,
search_queries=result.search_queries_used,
timestamp=datetime.utcnow().isoformat()
)
except Exception as e:
logger.error(f"搜索失败: {str(e)}")
# 输出 API key 用于调试(仅显示前8位和后4位)
api_key = request.llm_api_key
if api_key and len(api_key) > 12:
masked_key = f"{api_key[:8]}...{api_key[-4:]}"
else:
masked_key = api_key if api_key else "未提供"
logger.error(f"使用的 LLM API Key: {masked_key}")
logger.error(f"LLM Base URL: {os.getenv('LLM_BASE_URL', '未配置')}")
logger.error(f"MODEL_NAME: {os.getenv('MODEL_NAME', '未配置')}")
import traceback
logger.error(traceback.format_exc())
raise HTTPException(status_code=500, detail=f"搜索失败: {str(e)}")
@@ -291,13 +394,14 @@ def main():
"""主函数"""
logger.info(f"启动 Search Agent - {POD_NAME}")
logger.info(f"Template Type: {TEMPLATE_TYPE}")
logger.info(f"使用模式: 多用户共享(固定 model/endpoint,每次请求传入 API key)")
# 尝试从环境变量初始化
if os.getenv("LLM_API_KEY"):
logger.info("检测到环境变量配置,尝试自动初始化...")
initialize_agent_from_env()
# 预加载基础配置
if os.getenv("LLM_BASE_URL") and os.getenv("SERPER_API_KEY"):
logger.info("预加载基础配置...")
load_base_config()
else:
logger.info("未检测到环境变量配置,等待通过API配置...")
logger.warning("环境变量未完全配置,将在首次请求时加载...")
# 启动服务
uvicorn.run(
Submodule agent_templates/aks_agent deleted from b45aa748ee
@@ -32,7 +32,7 @@ class AgentCallbackHandler:
self.user_id = user_id or os.getenv("USER_ID", "")
self.callback_url = callback_url or os.getenv(
"AGENT_CALLBACK_URL",
"http://mcp-server:8002/api/v1/billing/agent-callback"
"http://mcp-server.taiji-ai.svc.cluster.local:8000/api/v1/billing/agent-callback"
)
self.start_time: Optional[datetime] = None
@@ -0,0 +1,6 @@
# Requirements for Azure Blob Agent - A2A Version
fastapi==0.109.0
uvicorn[standard]==0.27.0
pydantic==2.5.3
azure-storage-blob==12.19.0
httpx==0.26.0
@@ -0,0 +1,5 @@
# Requirements for Azure Blob Agent - MCP Version
fastapi==0.109.0
uvicorn[standard]==0.27.0
pydantic==2.5.3
azure-storage-blob==12.19.0
+6 -6
View File
@@ -40,13 +40,13 @@ fi
docker buildx use multiarch-builder
# Agent 列表
# Agent 列表(更新路径)
declare -A AGENTS=(
["search-agent"]="search_agent.Dockerfile"
["jina-search-agent"]="jina_search_agent.Dockerfile"
["mysql-agent"]="mysql_agent.Dockerfile"
["postgresql-agent"]="postgresql_agent.Dockerfile"
["azure-blob-agent"]="azure_blob_agent.Dockerfile"
["search-agent"]="agents/search_agent/search_agent.Dockerfile"
["azure-blob-agent"]="agents/azure_blob_agent/azure_blob_agent.Dockerfile"
["azure-blob-agent-a2a"]="agents/azure_blob_agent_a2a/azure_blob_agent_a2a.Dockerfile"
["azure-blob-agent-mcp"]="agents/azure_blob_agent_mcp/azure_blob_agent_mcp.Dockerfile"
["a2a-litellm-agent"]="agents/a2a_litellm_agent/a2a_litellm_agent.Dockerfile"
)
# 构建函数
@@ -11,8 +11,8 @@ IMAGE_NAME="ai-agents/search-agent"
TAG="${1:-latest}"
FULL_IMAGE="${ACR_NAME}/${IMAGE_NAME}:${TAG}"
# 支持的平台
PLATFORMS="linux/amd64,linux/arm64"
# 支持的平台(K8s 使用 ARM64)
PLATFORMS="linux/arm64"
echo "=========================================="
echo "构建 Intelligent Search Agent"
@@ -47,10 +47,10 @@ if [[ $REPLY =~ ^[Yy]$ ]]; then
echo "登录到 ACR..."
az acr login --name $(echo ${ACR_NAME} | cut -d'.' -f1)
# 构建并推送镜像(多架构)
# 构建并推送镜像(ARM64)
docker buildx build \
--platform "${PLATFORMS}" \
-f search_agent.Dockerfile \
-f agents/search_agent/search_agent.Dockerfile \
-t "${FULL_IMAGE}" \
--push \
.
@@ -64,7 +64,7 @@ else
echo "⏭️ 只构建本地镜像 (linux/arm64)..."
docker buildx build \
--platform "linux/arm64" \
-f search_agent.Dockerfile \
-f agents/search_agent/search_agent.Dockerfile \
-t "${FULL_IMAGE}" \
--load \
.
+85
View File
@@ -0,0 +1,85 @@
apiVersion: apps/v1
kind: Deployment
metadata:
name: test-search-agent
namespace: agent-test-search
labels:
app: test-search-agent
managed-by: manual-test
spec:
replicas: 1
selector:
matchLabels:
app: test-search-agent
template:
metadata:
labels:
app: test-search-agent
managed-by: manual-test
spec:
# ARM 架构节点选择器
nodeSelector:
kubernetes.io/arch: arm64
imagePullSecrets:
- name: acr-secret
containers:
- name: search-agent
image: agnettaiji.azurecr.io/ai-agents/search-agent:latest
imagePullPolicy: Always
ports:
- containerPort: 8080
name: http
env:
- name: POD_NAME
valueFrom:
fieldRef:
fieldPath: metadata.name
- name: TEMPLATE_TYPE
value: "search_agent"
- name: SERVICE_HOST
value: "0.0.0.0"
- name: SERVICE_PORT
value: "8080"
resources:
requests:
cpu: "500m"
memory: "512Mi"
limits:
cpu: "1000m"
memory: "1Gi"
livenessProbe:
httpGet:
path: /health
port: 8080
initialDelaySeconds: 60
periodSeconds: 30
timeoutSeconds: 10
failureThreshold: 3
readinessProbe:
httpGet:
path: /health
port: 8080
initialDelaySeconds: 30
periodSeconds: 10
timeoutSeconds: 5
failureThreshold: 3
---
apiVersion: v1
kind: Service
metadata:
name: test-search-agent
namespace: agent-test-search
labels:
app: test-search-agent
managed-by: manual-test
spec:
type: LoadBalancer
selector:
app: test-search-agent
ports:
- port: 80
targetPort: 8080
protocol: TCP
name: http
+4 -4
View File
@@ -300,7 +300,7 @@ async def create_agent(request: CreateAgentRequest, db: Session = Depends(get_db
logger.info(f"收到创建Agent请求: {request.name}, 模板: {request.template}")
# 验证模板类型
valid_templates = ["echo_agent", "chat_agent", "code_agent", "search_agent", "mysql_agent", "postgresql_agent", "jina_search_agent", "azure_blob_agent", "azure_blob_agent_mcp", "azure_blob_agent_a2a"]
valid_templates = ["echo_agent", "search_agent", "mysql_agent", "postgresql_agent", "jina_search_agent", "azure_blob_agent", "azure_blob_agent_mcp", "azure_blob_agent_a2a", "a2a_litellm_agent"]
if request.template not in valid_templates:
raise HTTPException(
status_code=400,
@@ -689,7 +689,7 @@ async def list_templates():
Returns:
模板列表及其配置信息
"""
valid_templates = ["echo_agent", "chat_agent", "code_agent", "search_agent", "mysql_agent", "postgresql_agent", "jina_search_agent", "azure_blob_agent", "azure_blob_agent_mcp", "azure_blob_agent_a2a"]
valid_templates = ["echo_agent", "search_agent", "mysql_agent", "postgresql_agent", "jina_search_agent", "azure_blob_agent", "azure_blob_agent_mcp", "azure_blob_agent_a2a", "a2a_litellm_agent"]
templates_info = []
for template in valid_templates:
@@ -711,7 +711,7 @@ async def list_platform_templates():
平台提供的Agent模板列表
"""
# 平台 Agent 是预定义的标准模板
platform_templates = ["echo_agent", "chat_agent", "code_agent", "search_agent", "jina_search_agent", "azure_blob_agent", "azure_blob_agent_mcp", "azure_blob_agent_a2a"]
platform_templates = ["echo_agent", "search_agent", "jina_search_agent", "azure_blob_agent", "azure_blob_agent_mcp", "azure_blob_agent_a2a", "a2a_litellm_agent"]
templates_info = []
for template in platform_templates:
@@ -761,7 +761,7 @@ async def get_template_info(template_name: str):
Returns:
模板详细信息(端口、所需环境变量等)
"""
valid_templates = ["echo_agent", "chat_agent", "code_agent", "search_agent", "mysql_agent", "postgresql_agent", "jina_search_agent", "azure_blob_agent", "azure_blob_agent_mcp", "azure_blob_agent_a2a"]
valid_templates = ["echo_agent", "search_agent", "mysql_agent", "postgresql_agent", "jina_search_agent", "azure_blob_agent", "azure_blob_agent_mcp", "azure_blob_agent_a2a", "a2a_litellm_agent"]
if template_name not in valid_templates:
raise HTTPException(
-1033
View File
File diff suppressed because it is too large Load Diff
+35 -7
View File
@@ -313,15 +313,14 @@ class K8sManager:
# 模板端口映射
TEMPLATE_PORTS = {
"echo_agent": 8000,
"chat_agent": 8000,
"code_agent": 8000,
"search_agent": 8000,
"search_agent": 8080,
"mysql_agent": 8000,
"postgresql_agent": 8000,
"jina_search_agent": 8080,
"azure_blob_agent": 8080,
"azure_blob_agent_mcp": 8080,
"azure_blob_agent_a2a": 8080,
"a2a_litellm_agent": 8080,
}
# 模板所需环境变量说明
@@ -409,6 +408,36 @@ class K8sManager:
"TENANT_ID": "租户标识",
"NAMESPACE": "Kubernetes 命名空间"
}
},
"search_agent": {
"required": {
"LLM_BASE_URL": "LLM 服务地址,如 https://api.openai.com/v1",
"SERPER_API_KEY": "Serper 搜索 API 密钥,从 https://serper.dev 获取",
"JINA_API_KEY": "Jina Reader API 密钥,从 https://jina.ai 获取"
},
"optional": {
"LLM_API_KEY": "LLM API 密钥(可在搜索请求中传入)",
"LLM_MODEL": "LLM 模型名称,默认 gpt-4o-mini",
"MAX_ITERATIONS": "最大搜索迭代次数,默认 3",
"MAX_RESULTS_PER_QUERY": "每次搜索最大结果数,默认 10",
"CONTENT_MAX_LENGTH": "内容最大长度,默认 5000",
"TIMEOUT": "超时时间(秒),默认 30",
"SERVICE_PORT": "HTTP服务端口,默认 8080",
"SERVICE_HOST": "HTTP服务监听地址,默认 0.0.0.0"
}
},
"a2a_litellm_agent": {
"required": {
"LITELLM_API_BASE": "LiteLLM 服务地址",
"LITELLM_MODEL": "LiteLLM 模型名称"
},
"optional": {
"LITELLM_API_KEY": "LiteLLM API 密钥(可在请求中传入)",
"AGENT_NAME": "Agent 名称",
"AGENT_DESCRIPTION": "Agent 描述",
"SERVICE_PORT": "HTTP服务端口,默认 8080",
"SERVICE_HOST": "HTTP服务监听地址,默认 0.0.0.0"
}
}
}
@@ -492,9 +521,7 @@ class K8sManager:
# 根据模板类型选择镜像
image_map = {
"echo_agent": "agnettaiji.azurecr.io/ai-agents/echo-agent:latest",
"chat_agent": "agnettaiji.azurecr.io/ai-agents/chat-agent:latest",
"code_agent": "agnettaiji.azurecr.io/ai-agents/code-agent:latest",
"echo_agent": "agnettaiji.azurecr.io/echo-agent:latest",
"search_agent": "agnettaiji.azurecr.io/ai-agents/search-agent:latest",
"mysql_agent": "agnettaiji.azurecr.io/ai-agents/mysql-agent:latest",
"postgresql_agent": "agnettaiji.azurecr.io/ai-agents/postgresql-agent:latest",
@@ -502,8 +529,9 @@ class K8sManager:
"azure_blob_agent": "agnettaiji.azurecr.io/ai-agents/azure-blob-agent:latest",
"azure_blob_agent_mcp": "agnettaiji.azurecr.io/ai-agents/azure-blob-agent-mcp:latest",
"azure_blob_agent_a2a": "agnettaiji.azurecr.io/ai-agents/azure-blob-agent-a2a:latest",
"a2a_litellm_agent": "agnettaiji.azurecr.io/ai-agents/a2a-litellm-agent:latest",
}
image = image_map.get(template, image_map["echo_agent"])
image = image_map.get(template, image_map["search_agent"])
# 构建环境变量列表
env_vars = [
-626
View File
@@ -1,626 +0,0 @@
"""
Enhanced Kubernetes Manager - 支持Deployment、Service、HPA和Secrets
"""
from kubernetes import client, config
from kubernetes.client.rest import ApiException
from datetime import datetime
import logging
import os
import base64
logger = logging.getLogger(__name__)
class K8sManager:
"""Kubernetes资源管理器 - 增强版"""
def __init__(self, namespace="ai-agents", kubeconfig_path=None):
"""
初始化K8s管理器
Args:
namespace: 命名空间
kubeconfig_path: kubeconfig文件路径(可选,用于本地开发)
"""
self.namespace = namespace
try:
if kubeconfig_path and os.path.exists(kubeconfig_path):
config.load_kube_config(kubeconfig_path)
logger.info(f"使用kubeconfig: {kubeconfig_path}")
else:
config.load_incluster_config()
logger.info("使用集群内ServiceAccount")
except Exception as e:
logger.error(f"K8s配置加载失败: {str(e)}")
raise
self.core_v1 = client.CoreV1Api()
self.apps_v1 = client.AppsV1Api()
self.autoscaling_v2 = client.AutoscalingV2Api()
self._ensure_namespace()
def _ensure_namespace(self):
"""确保命名空间存在"""
try:
self.core_v1.read_namespace(self.namespace)
logger.info(f"命名空间 {self.namespace} 已存在")
except ApiException as e:
if e.status == 404:
namespace = client.V1Namespace(
metadata=client.V1ObjectMeta(name=self.namespace)
)
self.core_v1.create_namespace(namespace)
logger.info(f"创建命名空间: {self.namespace}")
else:
raise
def create_secret(self, name: str, data: dict) -> dict:
"""
创建Kubernetes Secret存储敏感数据
Args:
name: Secret名称
data: 敏感数据字典
Returns:
Secret信息
"""
try:
# 编码数据为base64
encoded_data = {}
for key, value in data.items():
if isinstance(value, str):
encoded_data[key] = base64.b64encode(value.encode()).decode()
else:
encoded_data[key] = base64.b64encode(str(value).encode()).decode()
secret = client.V1Secret(
metadata=client.V1ObjectMeta(
name=name,
namespace=self.namespace,
labels={
"managed-by": "agent-manager",
"type": "agent-secret"
}
),
type="Opaque",
data=encoded_data
)
result = self.core_v1.create_namespaced_secret(self.namespace, secret)
logger.info(f"Created secret: {name}")
return {"name": name, "namespace": self.namespace}
except ApiException as e:
if e.status == 409:
# Secret已存在,更新它
logger.info(f"Secret {name} exists, updating...")
result = self.core_v1.replace_namespaced_secret(name, self.namespace, secret)
return {"name": name, "namespace": self.namespace}
else:
logger.error(f"Failed to create secret: {e}")
raise
def delete_secret(self, name: str):
"""删除Secret"""
try:
self.core_v1.delete_namespaced_secret(name, self.namespace)
logger.info(f"Deleted secret: {name}")
except ApiException as e:
if e.status != 404:
logger.error(f"Failed to delete secret: {e}")
def create_deployment_and_service(self, name: str, template, agent, env_vars: dict) -> dict:
"""
创建Deployment和Service
Args:
name: Agent名称
template: Template数据库对象
agent: Agent数据库对象
env_vars: 环境变量字典
Returns:
部署结果信息
"""
try:
deployment_name = f"{name}-deployment"
service_name = f"{name}-service"
# 1. 如果有敏感环境变量,创建Secret
secret_name = None
if env_vars:
secret_name = f"{name}-secret"
self.create_secret(secret_name, env_vars)
# 2. 创建Deployment
deployment = self._build_deployment(
name=deployment_name,
image=template.image,
port=template.port,
secret_name=secret_name,
agent=agent,
labels={
"app": name,
"managed-by": "agent-manager",
"template": template.name,
"agent-type": agent.agent_type.value,
"owner": agent.owner_id
}
)
self.apps_v1.create_namespaced_deployment(self.namespace, deployment)
logger.info(f"Created deployment: {deployment_name}")
# 3. 创建Service(如果模板定义了端口)
service_url = None
if template.port:
service = self._build_service(
name=service_name,
port=template.port,
selector={"app": name}
)
self.core_v1.create_namespaced_service(self.namespace, service)
logger.info(f"Created service: {service_name}")
# 生成服务URL(集群内访问)
service_url = f"http://{service_name}.{self.namespace}.svc.cluster.local:{template.port}"
# 4. 创建HPA(如果配置了弹性伸缩)
if agent.max_replicas > agent.min_replicas:
self.create_hpa(
name=f"{name}-hpa",
deployment_name=deployment_name,
min_replicas=agent.min_replicas,
max_replicas=agent.max_replicas,
target_cpu_utilization=agent.target_cpu_utilization
)
return {
"deployment_name": deployment_name,
"service_name": service_name,
"service_url": service_url,
"secret_name": secret_name
}
except Exception as e:
logger.error(f"Failed to create deployment and service: {str(e)}")
# 清理已创建的资源
self._cleanup_resources(deployment_name, service_name, secret_name)
raise
def _build_deployment(self, name: str, image: str, port: int, secret_name: str,
agent, labels: dict) -> client.V1Deployment:
"""构建Deployment对象"""
# 环境变量配置
env_vars = []
if secret_name:
# 从Secret引用环境变量
for key in agent.environment_vars.keys():
env_vars.append(client.V1EnvVar(
name=key,
value_from=client.V1EnvVarSource(
secret_key_ref=client.V1SecretKeySelector(
name=secret_name,
key=key
)
)
))
# 容器配置
container = client.V1Container(
name="agent",
image=image,
image_pull_policy="Always",
env=env_vars if env_vars else None,
resources=client.V1ResourceRequirements(
requests={
"cpu": agent.cpu_request or "100m",
"memory": agent.memory_request or "128Mi"
},
limits={
"cpu": agent.cpu_limit or "500m",
"memory": agent.memory_limit or "512Mi"
}
)
)
# 如果有端口,添加端口配置
if port:
container.ports = [client.V1ContainerPort(container_port=port)]
# Pod模板
template = client.V1PodTemplateSpec(
metadata=client.V1ObjectMeta(
labels=labels
),
spec=client.V1PodSpec(
containers=[container],
image_pull_secrets=[client.V1LocalObjectReference(name="acr-secret")]
)
)
# Deployment规格
spec = client.V1DeploymentSpec(
replicas=agent.min_replicas,
selector=client.V1LabelSelector(
match_labels={"app": labels["app"]}
),
template=template
)
# Deployment对象
deployment = client.V1Deployment(
api_version="apps/v1",
kind="Deployment",
metadata=client.V1ObjectMeta(
name=name,
namespace=self.namespace,
labels=labels
),
spec=spec
)
return deployment
def _build_service(self, name: str, port: int, selector: dict) -> client.V1Service:
"""构建Service对象"""
service = client.V1Service(
api_version="v1",
kind="Service",
metadata=client.V1ObjectMeta(
name=name,
namespace=self.namespace,
labels={
"managed-by": "agent-manager"
}
),
spec=client.V1ServiceSpec(
selector=selector,
ports=[client.V1ServicePort(
port=port,
target_port=port,
protocol="TCP"
)],
type="ClusterIP"
)
)
return service
def create_hpa(self, name: str, deployment_name: str, min_replicas: int,
max_replicas: int, target_cpu_utilization: int) -> dict:
"""
创建HorizontalPodAutoscaler
Args:
name: HPA名称
deployment_name: 目标Deployment名称
min_replicas: 最小副本数
max_replicas: 最大副本数
target_cpu_utilization: 目标CPU利用率(百分比)
Returns:
HPA信息
"""
try:
hpa = client.V2HorizontalPodAutoscaler(
api_version="autoscaling/v2",
kind="HorizontalPodAutoscaler",
metadata=client.V1ObjectMeta(
name=name,
namespace=self.namespace
),
spec=client.V2HorizontalPodAutoscalerSpec(
scale_target_ref=client.V2CrossVersionObjectReference(
api_version="apps/v1",
kind="Deployment",
name=deployment_name
),
min_replicas=min_replicas,
max_replicas=max_replicas,
metrics=[
client.V2MetricSpec(
type="Resource",
resource=client.V2ResourceMetricSource(
name="cpu",
target=client.V2MetricTarget(
type="Utilization",
average_utilization=target_cpu_utilization
)
)
)
]
)
)
result = self.autoscaling_v2.create_namespaced_horizontal_pod_autoscaler(
self.namespace, hpa
)
logger.info(f"Created HPA: {name}")
return {"name": name, "namespace": self.namespace}
except ApiException as e:
logger.error(f"Failed to create HPA: {e}")
raise
def delete_hpa(self, name: str):
"""删除HPA"""
try:
self.autoscaling_v2.delete_namespaced_horizontal_pod_autoscaler(
name, self.namespace
)
logger.info(f"Deleted HPA: {name}")
except ApiException as e:
if e.status != 404:
logger.error(f"Failed to delete HPA: {e}")
def update_deployment_env(self, deployment_name: str, env_vars: dict):
"""
更新Deployment的环境变量(通过更新Secret)
Args:
deployment_name: Deployment名称
env_vars: 新的环境变量字典
"""
try:
# 获取Deployment
deployment = self.apps_v1.read_namespaced_deployment(
deployment_name, self.namespace
)
# 查找Secret名称
secret_name = None
for env in deployment.spec.template.spec.containers[0].env or []:
if env.value_from and env.value_from.secret_key_ref:
secret_name = env.value_from.secret_key_ref.name
break
if secret_name:
# 更新Secret
self.create_secret(secret_name, env_vars)
# 触发Pod重启(通过添加annotation)
if not deployment.spec.template.metadata.annotations:
deployment.spec.template.metadata.annotations = {}
deployment.spec.template.metadata.annotations["kubectl.kubernetes.io/restartedAt"] = \
datetime.utcnow().isoformat()
self.apps_v1.replace_namespaced_deployment(
deployment_name, self.namespace, deployment
)
logger.info(f"Updated deployment env: {deployment_name}")
else:
raise ValueError("No secret found in deployment")
except ApiException as e:
logger.error(f"Failed to update deployment env: {e}")
raise
def delete_deployment_and_service(self, deployment_name: str, service_name: str):
"""
删除Deployment、Service和相关资源
Args:
deployment_name: Deployment名称
service_name: Service名称
"""
try:
# 删除Deployment
try:
self.apps_v1.delete_namespaced_deployment(
deployment_name, self.namespace,
propagation_policy='Foreground'
)
logger.info(f"Deleted deployment: {deployment_name}")
except ApiException as e:
if e.status != 404:
logger.error(f"Failed to delete deployment: {e}")
# 删除Service
try:
self.core_v1.delete_namespaced_service(service_name, self.namespace)
logger.info(f"Deleted service: {service_name}")
except ApiException as e:
if e.status != 404:
logger.error(f"Failed to delete service: {e}")
# 删除HPA
hpa_name = deployment_name.replace("-deployment", "-hpa")
self.delete_hpa(hpa_name)
# 删除Secret
secret_name = deployment_name.replace("-deployment", "-secret")
self.delete_secret(secret_name)
except Exception as e:
logger.error(f"Failed to delete resources: {str(e)}")
raise
def _cleanup_resources(self, deployment_name: str, service_name: str, secret_name: str):
"""清理资源(用于错误恢复)"""
if deployment_name:
try:
self.apps_v1.delete_namespaced_deployment(deployment_name, self.namespace)
except:
pass
if service_name:
try:
self.core_v1.delete_namespaced_service(service_name, self.namespace)
except:
pass
if secret_name:
try:
self.delete_secret(secret_name)
except:
pass
def get_deployment_status(self, deployment_name: str) -> dict:
"""获取Deployment状态"""
try:
deployment = self.apps_v1.read_namespaced_deployment(
deployment_name, self.namespace
)
# 获取 Deployment 对应的 Pods 实际状态
label_selector = f"app={deployment_name.replace('-deployment', '')}"
pods = self.core_v1.list_namespaced_pod(
self.namespace,
label_selector=label_selector
)
# 检查 Pod 的健康状态
health_status = "healthy"
pod_details = []
for pod in pods.items:
pod_health = "healthy"
container_statuses = pod.status.container_statuses or []
for container_status in container_statuses:
container_info = {
"name": container_status.name,
"ready": container_status.ready,
"restart_count": container_status.restart_count
}
# 检查容器状态
if container_status.state.waiting:
container_info["state"] = "waiting"
container_info["reason"] = container_status.state.waiting.reason
pod_health = "unhealthy"
elif container_status.state.terminated:
container_info["state"] = "terminated"
container_info["reason"] = container_status.state.terminated.reason
container_info["exit_code"] = container_status.state.terminated.exit_code
pod_health = "unhealthy"
elif container_status.state.running:
container_info["state"] = "running"
# 检查是否就绪
if not container_status.ready:
pod_health = "unhealthy"
# 检查重启次数
if container_status.restart_count > 5:
pod_health = "degraded"
pod_details.append({
"name": pod.metadata.name,
"phase": pod.status.phase,
"health": pod_health,
"containers": [container_info]
})
# 更新整体健康状态
if pod_health == "unhealthy":
health_status = "unhealthy"
elif pod_health == "degraded" and health_status != "unhealthy":
health_status = "degraded"
return {
"name": deployment_name,
"namespace": self.namespace,
"status": "Running" if deployment.status.available_replicas else "Pending",
"health_status": health_status, # 新增:真实健康状态
"replicas": deployment.status.replicas or 0,
"ready_replicas": deployment.status.ready_replicas or 0,
"available_replicas": deployment.status.available_replicas or 0,
"pods": pod_details, # 新增:Pod详细信息
"conditions": [
{
"type": c.type,
"status": c.status,
"reason": c.reason,
"message": c.message
}
for c in (deployment.status.conditions or [])
]
}
except ApiException as e:
if e.status == 404:
return {"status": "not_found", "message": f"Deployment {deployment_name} not found"}
raise
def get_pod_logs(self, deployment_name: str, lines: int = 100) -> str:
"""获取Pod日志"""
try:
# 查找Deployment对应的Pods
label_selector = f"app={deployment_name.replace('-deployment', '')}"
pods = self.core_v1.list_namespaced_pod(
self.namespace,
label_selector=label_selector
)
if not pods.items:
return "No pods found"
# 获取第一个Pod的日志
pod_name = pods.items[0].metadata.name
logs = self.core_v1.read_namespaced_pod_log(
pod_name, self.namespace,
tail_lines=lines
)
return logs
except ApiException as e:
logger.error(f"Failed to get pod logs: {e}")
raise
# ==================== 向后兼容的方法 ====================
def create_pod(self, pod_name: str, template: str, config_data: dict) -> dict:
"""创建Pod(旧方法,保留向后兼容)"""
# 这个方法现在已被create_deployment_and_service替代
# 但为了兼容性保留
raise NotImplementedError("Use create_deployment_and_service instead")
def delete_pod(self, pod_name: str) -> dict:
"""删除Pod(旧方法)"""
raise NotImplementedError("Use delete_deployment_and_service instead")
def get_pod_status(self, pod_name: str) -> dict:
"""获取Pod状态(旧方法)"""
# 尝试查找对应的Deployment
deployment_name = f"{pod_name}-deployment"
return self.get_deployment_status(deployment_name)
def list_pods(self, label_selector: str = None) -> list:
"""列出Pods"""
try:
if label_selector:
deployments = self.apps_v1.list_namespaced_deployment(
self.namespace,
label_selector=label_selector
)
else:
deployments = self.apps_v1.list_namespaced_deployment(self.namespace)
result = []
for deployment in deployments.items:
result.append({
"name": deployment.metadata.name,
"namespace": self.namespace,
"replicas": deployment.status.replicas or 0,
"ready_replicas": deployment.status.ready_replicas or 0,
"labels": deployment.metadata.labels
})
return result
except ApiException as e:
logger.error(f"Failed to list deployments: {e}")
raise
+248
View File
@@ -0,0 +1,248 @@
# API Key 未正确使用 - 代码问题分析
## 问题现象
请求参数:
```json
{
"query": "什么是杜鹃花",
"llm_api_key": "sk-rxegkFOciNmQLhOHr3qP3A"
}
```
但 LiteLLM 收到的是 `placeholder`,而不是真实的 API Key。
## 代码逻辑分析
### 1. 请求处理流程(`search_agent_main.py:233-282`)
```python
@app.post("/search", response_model=SearchResponse)
async def search(request: SearchRequest):
global search_agent, callback_handler, config
# 第 251 行:应该记录请求信息
logger.info(f"收到搜索请求: query={request.query}, user_id={request.user_id}")
# 第 252 行:应该记录 API Key 状态
logger.info(f"LLM API Key: {'已提供' if request.llm_api_key else '未提供'}")
# 第 255-261 行:如果 Agent 未初始化,从环境变量初始化
if not search_agent:
logger.info("Agent未初始化,从环境变量自动配置...")
if not initialize_agent_from_env():
raise HTTPException(...)
# 第 274-282 行:关键!更新 API Key 的逻辑
original_api_key = config.llm_api_key if config else None
if config and request.llm_api_key: # ← 问题可能在这里
logger.info(f"使用请求中的 LLM API Key: {request.llm_api_key[:10]}...")
config.llm_api_key = request.llm_api_key
search_agent = SearchAgent(config) # 重新初始化
logger.info("SearchAgent 已使用新的 API Key 重新初始化")
```
### 2. 初始化逻辑(`search_agent_main.py:117-161`)
```python
def initialize_agent_from_env():
global search_agent, config
# 第 140-151 行:创建配置
config = Config(
llm_base_url=os.getenv("LLM_BASE_URL", ""),
llm_api_key=os.getenv("LLM_API_KEY", "placeholder"), # ← 默认是 placeholder
llm_model=os.getenv("LLM_MODEL", "gpt-4o-mini"),
# ...
)
search_agent = SearchAgent(config) # 使用 placeholder 初始化
return True
```
## 问题定位
### 关键发现
1. **日志中没有 "收到搜索请求"**:
- 说明第 251 行的日志没有被执行
- 或者日志被过滤/丢失了
2. **日志中没有 "LLM API Key: 已提供"**:
- 说明第 252 行的日志没有被执行
- 或者 `request.llm_api_key` 是空的
3. **日志中没有 "使用请求中的 LLM API Key"**:
- 说明第 277 行的条件判断 `if config and request.llm_api_key:` 为 False
- 可能的原因:
- `config` 是 None
- `request.llm_api_key` 是空字符串、None 或 False
4. **但搜索确实执行了**:
- 日志显示 "开始搜索: 什么是杜鹃花"
- 说明请求确实到达了 SearchAgent
## 可能的问题点
### 问题 1:日志系统不一致
**发现**:
- `search_agent_main.py` 使用标准 `logging` 模块
- `search_agent` 模块使用 `loguru`
- 日志格式不一致,可能导致日志丢失
**证据**:
- 标准 logging 格式:`2026-01-16 12:46:08,456 - __main__ - INFO - ...`
- loguru 格式:`2026-01-16 12:46:08.456 | INFO | agent.search_agent:search:61 - ...`
**影响**:
- 如果 uvicorn 的日志级别设置不当,可能过滤掉标准 logging 的日志
### 问题 2:条件判断失败
**代码**:
```python
if config and request.llm_api_key:
```
**可能的原因**:
1. `config` 是 None(但不太可能,因为搜索执行了)
2. `request.llm_api_key` 是:
- 空字符串 `""`(Python 中空字符串是 False)
- None
- 其他 falsy 值
### 问题 3:请求验证失败(422 错误)
**日志显示**:
```
INFO: 10.224.0.5:46411 - "POST /search HTTP/1.1" 422 Unprocessable Entity
INFO: 10.224.0.7:60577 - "POST /search HTTP/1.1" 422 Unprocessable Entity
```
**422 错误**:表示请求格式正确,但语义验证失败(FastAPI 的 Pydantic 验证)
**可能的原因**:
- `llm_api_key` 字段验证失败
- 字段类型不匹配
- 必填字段缺失
### 问题 4:SearchRequest 模型定义
**代码**(`search_agent_main.py:74-78`):
```python
class SearchRequest(BaseModel):
"""搜索请求 - 简化版:只需传入query和llm_api_key,其他从环境变量获取"""
query: str = Field(..., description="搜索查询")
llm_api_key: str = Field(..., description="LLM API密钥")
user_id: Optional[str] = Field(None, description="用户ID(用于计费回调)")
```
**分析**:
- `llm_api_key: str` 是必填字段(没有 Optional)
- 如果请求中没有这个字段,FastAPI 会返回 422 错误
- 如果字段值是空字符串,Pydantic 可能也会验证失败(取决于配置)
## 诊断步骤
### 1. 检查请求是否到达函数
在代码中添加更详细的日志:
```python
@app.post("/search", response_model=SearchResponse)
async def search(request: SearchRequest):
# 添加这行,确保日志被记录
print(f"DEBUG: 收到搜索请求: query={request.query}, llm_api_key={request.llm_api_key[:10] if request.llm_api_key else 'None'}...")
logger.info(f"收到搜索请求: query={request.query}, user_id={request.user_id}")
logger.info(f"LLM API Key: {'已提供' if request.llm_api_key else '未提供'}")
logger.info(f"LLM API Key 值: {request.llm_api_key[:10] if request.llm_api_key else 'None'}...")
# 检查 config
logger.info(f"Config 状态: {config is not None}, Config.llm_api_key: {config.llm_api_key[:10] if config and config.llm_api_key else 'None'}...")
# 检查条件判断
condition_result = bool(config and request.llm_api_key)
logger.info(f"条件判断结果: config={config is not None}, request.llm_api_key={bool(request.llm_api_key)}, 结果={condition_result}")
```
### 2. 检查 FastAPI 请求验证
查看 FastAPI 的自动生成的文档:
```bash
curl http://<pod-ip>:8080/docs
```
或者直接测试请求:
```bash
curl -X POST http://<pod-ip>:8080/search \
-H "Content-Type: application/json" \
-d '{
"query": "测试",
"llm_api_key": "sk-rxegkFOciNmQLhOHr3qP3A"
}'
```
### 3. 检查日志配置
确认 uvicorn 的日志级别和格式:
```python
uvicorn.run(
app,
host=SERVICE_HOST,
port=SERVICE_PORT,
log_level="info" # 这里可能过滤了某些日志
)
```
## 最可能的问题
基于代码分析,**最可能的问题是**:
1. **日志被过滤**:标准 logging 的日志被 uvicorn 过滤掉了
2. **条件判断失败**:`request.llm_api_key` 可能是空字符串,导致 `if config and request.llm_api_key:` 为 False
3. **请求验证问题**:虽然有些请求返回 200,但可能请求体解析有问题
## 建议的修复方法(不改代码的情况下)
### 方法 1:检查实际请求
从 Pod 内部检查请求日志:
```bash
kubectl exec -n agent-search-agent-bd115e1f-fc89c0 search-agent-bd115e1f-fc89c0 -- \
curl -X POST http://localhost:8080/search \
-H "Content-Type: application/json" \
-d '{"query":"测试","llm_api_key":"sk-rxegkFOciNmQLhOHr3qP3A"}' \
-v
```
### 方法 2:检查环境变量
确认环境变量中的 LLM_API_KEY:
```bash
kubectl exec -n agent-search-agent-bd115e1f-fc89c0 search-agent-bd115e1f-fc89c0 -- env | grep LLM_API_KEY
```
### 方法 3:直接测试 API Key 更新逻辑
在代码中添加调试输出(虽然不能改代码,但可以检查现有日志):
- 检查是否有 "Search Agent从环境变量初始化成功" 的日志
- 检查是否有 "LLM_API_KEY: 需在请求中传入" 的日志
## 结论
**问题很可能出在第 276 行的条件判断**:
```python
if config and request.llm_api_key:
```
这个条件可能因为:
1. `request.llm_api_key` 是空字符串(falsy)
2. 或者请求根本没有正确解析 `llm_api_key` 字段
**需要进一步检查**:
1. 实际的 HTTP 请求体
2. FastAPI 的请求验证日志
3. Pydantic 模型的验证结果
@@ -0,0 +1,347 @@
# Search Agent LiteLLM 问题分析报告
## 一、Search Agent 模型和密钥部署方式分析
### 1.1 配置来源
Search Agent 的模型和密钥配置通过以下方式部署:
#### 环境变量配置(主要方式)
从 `agent_templates/agents/search_agent/search_agent/config.py` 和 `search_agent_main.py` 可以看到:
```python
# 必须的环境变量
LLM_BASE_URL: str # LLM API基础URL(如:https://litellm.graystone-fb459c5d.southeastasia.azurecontainerapps.io/)
LLM_API_KEY: str # LLM API密钥(可选,可在请求中传入)
LLM_MODEL: str # 模型名称(默认:gpt-4o-mini,实际使用:taiji/gpt-4o-mini)
SERPER_API_KEY: str # Serper 搜索 API 密钥
JINA_API_KEY: str # Jina Reader API 密钥
```
#### 部署配置示例
从 `k8s-test-deployment.yaml` 可以看到实际部署配置:
```yaml
env:
- name: LLM_BASE_URL
value: "https://litellm.graystone-fb459c5d.southeastasia.azurecontainerapps.io/"
- name: LLM_MODEL
value: "taiji/gpt-4o-mini"
- name: LLM_API_KEY
valueFrom:
secretKeyRef:
name: search-agent-secrets
key: LLM_API_KEY
```
### 1.2 密钥使用流程
1. **初始化阶段**:
- Agent 启动时从环境变量读取 `LLM_BASE_URL` 和 `LLM_MODEL`
- `LLM_API_KEY` 可以为空(使用占位符),等待请求时传入
2. **请求处理阶段**(`search_agent_main.py:233-282`):
```python
# 临时更新API key - 重要:在搜索前设置
if config and request.llm_api_key:
logger.info(f"使用请求中的 LLM API Key: {request.llm_api_key[:10]}...")
config.llm_api_key = request.llm_api_key
# 重新初始化整个 SearchAgent 以使用新的 API key
search_agent = SearchAgent(config)
```
3. **LLM 调用**(`llm_client.py:49-89`):
```python
# Azure OpenAI 风格的URL
url = f"{self.base_url}/chat/completions?api-version={self.API_VERSION}"
# Azure OpenAI 使用 api-key 头
headers = {
"api-key": self.api_key, # 使用传入的 API Key
"Content-Type": "application/json"
}
```
### 1.3 关键发现
1. **模型名称**:使用的是 `taiji/gpt-4o-mini`,这是一个**模型别名**(model alias),不是原始模型名
2. **API Key 传递方式**:
- 通过 HTTP Header `api-key` 传递(Azure OpenAI 风格)
- 支持在请求中动态传入,覆盖环境变量
3. **LiteLLM Proxy 地址**:`https://litellm.graystone-fb459c5d.southeastasia.azurecontainerapps.io/`
---
## 二、LiteLLM 回调报错原因分析
### 2.1 错误现象
根据提供的日志信息:
```
"call_type": "/chat/completions",
"model": "taiji/gpt-4o-mini",
"status": "failure",
"response_time": 0.00039 # 极短的响应时间,说明在鉴权阶段就被拒绝
异常位置:
ProxyException
File ".../auth_checks.py", line 1756
can_team_access_model(model=_model, team_model_aliases=...)
```
### 2.2 根本原因
**LiteLLM Proxy 在鉴权阶段无法将当前请求使用的 API Key 映射到任何租户(tenant/team),因此直接拒绝了模型调用。**
#### 问题链路:
1. **请求到达 LiteLLM Proxy**
- Search Agent 发送请求到 `https://litellm.graystone-fb459c5d.southeastasia.azurecontainerapps.io/chat/completions`
- Header: `api-key: <用户提供的API Key>`
- Model: `taiji/gpt-4o-mini`
2. **LiteLLM 鉴权流程**
```
LiteLLM Proxy 收到请求
↓
查找 API Key 对应的 team/tenant
↓
检查该 team/tenant 是否有权限访问 "taiji/gpt-4o-mini"
↓
❌ 失败:无法找到 API Key 对应的租户,或租户没有该模型的访问权限
↓
抛出 ProxyException,拒绝请求
```
3. **为什么 Agent 还有回复?**
- 可能的原因:
- Agent 使用了**备用 API Key**(环境变量中的默认值)
- 或者 LiteLLM Proxy 配置了**降级策略**(fallback)
- 或者请求被**重试**,使用了不同的 API Key
### 2.3 具体问题点
#### 问题 1:API Key 未正确映射到租户
LiteLLM Proxy 需要知道:
- 哪个 API Key 属于哪个 team/tenant
- 该 team/tenant 可以访问哪些模型
**可能的原因**:
- API Key 未在 LiteLLM 的数据库中注册
- API Key 没有关联 `team_id` 或 `tenant_id`
- API Key 的 `auth_metadata` 中缺少租户信息
#### 问题 2:模型别名权限配置缺失
模型 `taiji/gpt-4o-mini` 是一个别名,需要:
- 在 LiteLLM 中配置该别名映射到实际模型
- 配置哪些 team/tenant 可以访问该别名
**可能的原因**:
- 模型别名 `taiji/gpt-4o-mini` 未在 LiteLLM 中配置
- 或者配置了,但当前租户没有访问权限
#### 问题 3:回调时缺少租户信息
从回调文档可以看到,LiteLLM 回调需要包含租户信息:
```json
{
"metadata": {
"user_api_key_auth_metadata": {
"tenant_id": "tenant-123",
"channel_id": "channel-456"
}
}
}
```
**如果回调中缺少这些信息,计费系统无法识别租户,导致计费失败。**
---
## 三、修复方案
### 3.1 立即修复(LiteLLM Proxy 配置)
#### 方案 A:在 LiteLLM 中正确配置 API Key 和租户映射
1. **检查 LiteLLM 数据库中的 API Key 配置**
确保每个 API Key 都有:
```yaml
# LiteLLM config.yaml 或数据库记录
api_keys:
- sk-xxx:
team_id: "team-123"
metadata:
tenant_id: "tenant-123"
channel_id: "channel-456"
```
2. **配置模型别名和团队访问权限**
```yaml
# LiteLLM config.yaml
model_list:
- model_name: taiji/gpt-4o-mini
litellm_params:
model: gpt-4o-mini
api_key: os.environ/OPENAI_API_KEY
# 团队模型访问配置
team_settings:
team-123:
team_model_aliases:
- taiji/gpt-4o-mini
```
#### 方案 B:在请求中添加租户信息(如果 LiteLLM 支持)
如果 LiteLLM Proxy 支持通过 Header 传递租户信息:
```python
# 在 llm_client.py 中修改
headers = {
"api-key": self.api_key,
"Content-Type": "application/json",
"x-tenant-id": tenant_id, # 如果 LiteLLM 支持
"x-team-id": team_id # 如果 LiteLLM 支持
}
```
### 3.2 代码层面修复(Search Agent)
#### 修复 1:在 LLM 调用时传递租户信息
修改 `search_agent/utils/llm_client.py`,在请求中包含租户信息:
```python
class LLMClient:
def __init__(
self,
base_url: str,
api_key: str,
model: str = "xchat52",
timeout: int = 60,
tenant_id: Optional[str] = None, # 新增
team_id: Optional[str] = None # 新增
):
self.base_url = base_url.rstrip("/")
self.api_key = api_key
self.model = model
self.timeout = timeout
self.tenant_id = tenant_id
self.team_id = team_id
async def chat(self, ...):
headers = {
"api-key": self.api_key,
"Content-Type": "application/json"
}
# 如果 LiteLLM 支持通过 Header 传递租户信息
if self.tenant_id:
headers["x-tenant-id"] = self.tenant_id
if self.team_id:
headers["x-team-id"] = self.team_id
# ... 其余代码
```
#### 修复 2:从请求中提取并传递租户信息
修改 `search_agent_main.py`,在搜索请求中包含租户信息:
```python
class SearchRequest(BaseModel):
query: str
llm_api_key: str
user_id: Optional[str] = None
tenant_id: Optional[str] = None # 新增
team_id: Optional[str] = None # 新增
# 在初始化 SearchAgent 时传递租户信息
config = Config(
...
tenant_id=request.tenant_id, # 传递租户信息
team_id=request.team_id
)
```
### 3.3 LiteLLM Proxy 配置检查清单
1. ✅ **API Key 配置**
- [ ] 所有使用的 API Key 都在 LiteLLM 数据库中
- [ ] 每个 API Key 都关联了 `team_id` 或 `tenant_id`
- [ ] API Key 的 `auth_metadata` 包含租户信息
2. ✅ **模型别名配置**
- [ ] `taiji/gpt-4o-mini` 在 `model_list` 中定义
- [ ] 模型别名正确映射到实际模型
3. ✅ **团队/租户访问权限**
- [ ] 每个 team/tenant 的 `team_model_aliases` 包含 `taiji/gpt-4o-mini`
- [ ] 或者使用通配符允许所有模型
4. ✅ **回调配置**
- [ ] LiteLLM 的 `webhook_url` 指向正确的回调地址
- [ ] 回调中包含 `metadata.user_api_key_auth_metadata` 信息
### 3.4 验证步骤
1. **测试 API Key 映射**
```bash
# 使用 curl 测试
curl -X POST https://litellm.graystone-fb459c5d.southeastasia.azurecontainerapps.io/chat/completions \
-H "api-key: <your-api-key>" \
-H "Content-Type: application/json" \
-d '{
"model": "taiji/gpt-4o-mini",
"messages": [{"role": "user", "content": "test"}]
}'
```
2. **检查 LiteLLM 日志**
- 查看 LiteLLM Proxy 的日志,确认 API Key 是否被正确识别
- 检查是否有 `can_team_access_model` 相关的错误
3. **验证回调数据**
- 检查回调请求中的 `metadata` 字段
- 确认包含 `user_api_key_auth_metadata.tenant_id`
---
## 四、总结
### 核心问题
**LiteLLM Proxy 的多租户鉴权机制无法将 API Key 映射到租户,导致模型调用被拒绝。**
### 修复优先级
1. **高优先级**:修复 LiteLLM Proxy 配置
- 确保 API Key 正确映射到租户
- 配置模型别名和访问权限
2. **中优先级**:代码层面改进
- 在请求中传递租户信息(如果 LiteLLM 支持)
- 改进错误处理和日志记录
3. **低优先级**:长期优化
- 统一租户信息管理
- 添加更详细的监控和告警
### 为什么 Agent 还有回复?
可能的原因:
1. Agent 使用了环境变量中的备用 API Key(有权限的)
2. LiteLLM Proxy 配置了降级策略
3. 请求被重试,使用了不同的 API Key
**建议**:检查 LiteLLM Proxy 的日志,确认实际使用的 API Key 和租户信息。
+363
View File
@@ -0,0 +1,363 @@
# Search Agent 问题诊断报告
## Pod 信息
- **Pod 名称**: `search-agent-bd115e1f-fc89c0`
- **命名空间**: `agent-search-agent-bd115e1f-fc89c0`
- **状态**: Running
---
## 一、发现的问题
### 问题 0:为什么 Agent 还能返回内容?(即使 LiteLLM 返回 401)
#### 现象
虽然 LiteLLM 返回 401 错误,但 Agent 仍然返回了内容,例如:
```json
{
"query": "什么是人工智能",
"answer": "关于「什么是人工智能」,以下是搜索到的相关信息:\n### 来源 [1]: 人工智能 (AI)\n..."
}
```
#### 原因:降级策略(Fallback Mechanism)
Agent 有一个**降级策略**,当 LLM 调用失败时,会直接从搜索结果中提取内容:
1. **正常流程**(LLM 成功时):
```
搜索 → 提取内容 → 重排序 → LLM 生成答案 → 返回
```
2. **降级流程**(LLM 失败时):
```
搜索 → 提取内容 → 重排序 → LLM 失败 → 使用 _fallback_answer() → 返回
```
3. **降级答案生成逻辑**(`answer_generator.py:125-150`):
```python
def _fallback_answer(self, query: str, documents: List[RankedDocument]) -> Answer:
"""后备答案生成(LLM失败时)"""
# 简单汇总文档内容
content_parts = [f"关于「{query}」,以下是搜索到的相关信息:\n"]
for i, doc in enumerate(documents[:5], 1):
content_parts.append(f"### 来源 [{i}]: {doc.title}\n")
content_parts.append(f"{doc.content[:500]}...\n\n")
return Answer(
content="".join(content_parts),
sources=sources,
confidence="low" # 注意:置信度是 low
)
```
4. **日志证据**:
```
2026-01-16 12:37:00.009 | ERROR | modules.answer_generator:generate:114 - 答案生成失败: LLM API请求失败: 401
2026-01-16 12:37:00.009 | INFO | agent.search_agent:search:131 - 答案生成完成: confidence=low
```
注意:虽然 LLM 失败了,但答案生成"完成"了,只是 `confidence=low`。
#### 这意味着什么?
- ✅ **搜索功能正常**:Serper 搜索和 Jina 内容提取都成功了
- ✅ **内容提取正常**:从网页提取了内容并进行了重排序
- ❌ **LLM 生成失败**:无法使用 LLM 生成高质量答案
- ⚠️ **返回降级答案**:直接返回搜索结果的简单汇总,质量较低
**所以 Agent 返回的内容是降级答案,不是 LLM 生成的,质量会明显下降。**
---
### 问题 1:API Key 未正确使用(导致 LiteLLM 回调失败)
#### 现象
日志显示 LiteLLM 收到的 API Key 是 `placeholder`,而不是用户传入的真实 API Key:
```
LLM API错误: 401 - {"error":{"message":"Authentication Error, LiteLLM Virtual Key expected. Received=placeholder, expected to start with 'sk-'.","type":"auth_error"}}
```
#### 原因分析
1. **代码逻辑**:`search_agent_main.py` 第 276-282 行有更新 API Key 的逻辑:
```python
if config and request.llm_api_key:
logger.info(f"使用请求中的 LLM API Key: {request.llm_api_key[:10]}...")
config.llm_api_key = request.llm_api_key
search_agent = SearchAgent(config) # 重新初始化
```
2. **问题**:日志中**没有看到** "使用请求中的 LLM API Key" 这条日志,说明:
- 要么 `request.llm_api_key` 为空/None
- 要么请求中的字段名不是 `llm_api_key`
3. **SearchAgent 初始化时机**:
- 各个模块(QueryAnalyzer, AnswerGenerator 等)在 `__init__` 时创建了 `LLMClient`
- 即使重新初始化 `SearchAgent`,如果模块内部已经缓存了旧的 `LLMClient`,仍会使用 placeholder
#### 验证方法
检查最近的请求日志,看是否有 "使用请求中的 LLM API Key" 这条日志:
```bash
kubectl logs -n agent-search-agent-bd115e1f-fc89c0 search-agent-bd115e1f-fc89c0 --tail=1000 | grep "使用请求中的"
```
**如果没有这条日志,说明 API Key 没有被正确传入或处理。**
---
### 问题 2:Agent 回调 URL 无法解析(Agent Manager 回调失败)
#### 现象
日志显示回调请求失败,无法解析 `mcp-server` 主机名:
```
Failed to send callback: HTTPConnectionPool(host='mcp-server', port=8002): Max retries exceeded with url: /api/v1/billing/agent-callback
(Caused by NameResolutionError("HTTPConnection(host='mcp-server', port=8002): Failed to resolve 'mcp-server' ([Errno -2] Name or service not known)"))
```
#### 回调 URL 信息
- **当前配置的 URL**: `http://mcp-server:8002/api/v1/billing/agent-callback`
- **来源**: `agent_callback_utils.py` 第 33-35 行
```python
self.callback_url = callback_url or os.getenv(
"AGENT_CALLBACK_URL",
"http://mcp-server:8002/api/v1/billing/agent-callback"
)
```
#### 原因分析
1. **服务发现问题**:
- `mcp-server` 服务在 `taiji-ai` 命名空间中
- 当前 Pod 在 `agent-search-agent-bd115e1f-fc89c0` 命名空间中
- 跨命名空间访问需要使用完整的服务名:`mcp-server.taiji-ai.svc.cluster.local`
2. **端口问题**:
- mcp-server 服务实际端口是 **8000**(不是 8002)
- 代码中默认使用的是 8002 端口
3. **正确的回调 URL 应该是**:
```
http://mcp-server.taiji-ai.svc.cluster.local:8000/api/v1/billing/agent-callback
```
或者(如果 mcp-server 在同一个集群中):
```
http://mcp-server.taiji-ai:8000/api/v1/billing/agent-callback
```
**注意**:端口是 8000,不是 8002!
#### 回调请求详情
从日志中可以看到回调请求的 payload:
```json
{
"agentName": "search-agent-bd115e1f-fc89c0",
"userId": "bd115e1f-e2de-4bd2-a641-1ed74ac1a34a",
"podRunningTimeSeconds": 22,
"toolsUsed": ["web_search", "content_reader"],
"startTime": "2026-01-16T12:02:39.355978+00:00",
"endTime": "2026-01-16T12:03:01.701388+00:00",
"requestId": "search-1768564959"
}
```
---
### 问题 3:LiteLLM 回调失败(但直接测试密钥正常)
#### 现象
- 用户使用 `sk-rxegkFOciNmQLhOHr3qP3A` 直接测试模型回调正常
- 但通过 search_agent 调用时,LiteLLM 回调报错
#### 原因分析
1. **API Key 传递问题**:
- Agent 使用的是 `placeholder`,不是真实的 `sk-rxegkFOciNmQLhOHr3qP3A`
- LiteLLM 无法识别 `placeholder`,导致鉴权失败
- 鉴权失败后,LiteLLM 的回调也会失败(因为无法识别租户)
2. **为什么直接测试正常**:
- 直接测试时使用的是真实的 API Key `sk-rxegkFOciNmQLhOHr3qP3A`
- 该 API Key 在 LiteLLM 中正确配置了租户信息
- 所以回调正常
3. **为什么 Agent 调用失败**:
- Agent 实际使用的是 `placeholder`
- LiteLLM 无法识别 `placeholder`,返回 401 错误
- 回调时也无法识别租户,导致回调失败
---
## 二、问题根源总结
### 核心问题链
```
用户传入 llm_api_key: "sk-rxegkFOciNmQLhOHr3qP3A"
↓
代码应该更新 config.llm_api_key 并重新初始化 SearchAgent
↓
❌ 但日志显示没有执行更新逻辑(没有 "使用请求中的 LLM API Key" 日志)
↓
SearchAgent 继续使用 placeholder
↓
LLMClient 使用 placeholder 调用 LiteLLM
↓
LiteLLM 返回 401: "Received=placeholder, expected to start with 'sk-'"
↓
LiteLLM 回调失败(无法识别租户)
```
---
## 三、验证步骤
### 1. 验证 API Key 是否被传入
检查最近的请求日志:
```bash
kubectl logs -n agent-search-agent-bd115e1f-fc89c0 search-agent-bd115e1f-fc89c0 --tail=1000 | grep -E "使用请求中的|llm_api_key|POST /search"
```
### 2. 验证回调 URL 配置
检查 Pod 的环境变量:
```bash
kubectl exec -n agent-search-agent-bd115e1f-fc89c0 search-agent-bd115e1f-fc89c0 -- env | grep -i "CALLBACK\|MCP"
```
### 3. 验证 mcp-server 服务
确认 mcp-server 服务的完整地址:
```bash
kubectl get svc -n taiji-ai mcp-server
kubectl get endpoints -n taiji-ai mcp-server
```
### 4. 测试回调 URL 连通性
从 Pod 内部测试回调 URL:
```bash
# 测试健康检查接口(端口 8000)
kubectl exec -n agent-search-agent-bd115e1f-fc89c0 search-agent-bd115e1f-fc89c0 -- curl -v http://mcp-server.taiji-ai.svc.cluster.local:8000/api/v1/billing/agent-callback/health
# 测试 DNS 解析
kubectl exec -n agent-search-agent-bd115e1f-fc89c0 search-agent-bd115e1f-fc89c0 -- nslookup mcp-server.taiji-ai.svc.cluster.local
```
---
## 四、修复建议(不改代码)
### 修复 1:配置正确的回调 URL
通过环境变量或 ConfigMap 设置正确的回调 URL:
```yaml
env:
- name: AGENT_CALLBACK_URL
value: "http://mcp-server.taiji-ai.svc.cluster.local:8000/api/v1/billing/agent-callback"
```
或者如果 mcp-server 在同一个集群中:
```yaml
env:
- name: AGENT_CALLBACK_URL
value: "http://mcp-server.taiji-ai:8000/api/v1/billing/agent-callback"
```
**重要**:
- 服务名:`mcp-server.taiji-ai.svc.cluster.local`(跨命名空间访问)
- 端口:**8000**(不是 8002)
- 路径:`/api/v1/billing/agent-callback`
### 修复 2:确认 API Key 传递方式
检查请求格式是否正确:
```json
{
"query": "什么是人工智能",
"llm_api_key": "sk-rxegkFOciNmQLhOHr3qP3A",
"user_id": "bd115e1f-e2de-4bd2-a641-1ed74ac1a34a"
}
```
**注意**:字段名必须是 `llm_api_key`(不是 `LLM_API_KEY` 或其他)。
### 修复 3:检查 LiteLLM 配置
确认 LiteLLM Proxy 中 `sk-rxegkFOciNmQLhOHr3qP3A` 的配置:
- API Key 是否正确注册
- 是否关联了 `team_id` 或 `tenant_id`
- 是否有权限访问 `taiji/gpt-4o-mini` 模型
---
## 五、关键日志位置
### Agent 回调日志
```
Sending callback: {'agentName': '...', 'userId': '...', ...}
Failed to send callback: HTTPConnectionPool(host='mcp-server', port=8002): ...
```
### LiteLLM 调用日志
```
LLM API错误: 401 - {"error":{"message":"Authentication Error, LiteLLM Virtual Key expected. Received=placeholder, ..."}}
```
### API Key 更新日志(应该出现但没出现)
```
使用请求中的 LLM API Key: sk-rxegkFO...
SearchAgent 已使用新的 API Key 重新初始化
```
---
## 六、Agent 执行流程说明
### 完整执行流程
```
1. 接收请求(POST /search)
↓
2. 查询分析(QueryAnalyzer)→ ❌ LLM 失败,使用默认 intent
↓
3. 搜索规划(SearchPlanner)→ ✅ 成功
↓
4. 执行搜索(SearchExecutor + Serper)→ ✅ 成功,获取 10 条结果
↓
5. 提取内容(ContentExtractor + Jina Reader)→ ✅ 成功,提取 10 个文档
↓
6. 结果处理(ResultProcessor + Jina Reranker)→ ✅ 成功,排序后返回 5 个
↓
7. 生成答案(AnswerGenerator + LLM)→ ❌ LLM 失败(401)
↓
8. 降级处理(_fallback_answer)→ ✅ 从文档中提取内容,生成简单答案
↓
9. 质量评估(Reflector + LLM)→ ❌ LLM 失败(401),跳过评估
↓
10. 返回结果 → ✅ 返回降级答案(confidence=low)
```
### 为什么还能返回内容?
- **搜索和内容提取不依赖 LLM**:使用 Serper API 和 Jina Reader API,这些服务都有独立的 API Key
- **降级策略**:当 LLM 失败时,直接从搜索结果中提取内容并格式化返回
- **质量下降**:降级答案只是简单汇总,没有 LLM 的智能整合和结构化
### 如何判断返回的是降级答案?
1. **检查 confidence**:降级答案的 `confidence` 是 `"low"`
2. **检查答案格式**:降级答案通常以 "关于「xxx」,以下是搜索到的相关信息:" 开头
3. **检查日志**:日志中会有 "答案生成失败" 和 "confidence=low" 的记录
---
## 七、下一步行动
1. **立即检查**:确认请求中是否真的传入了 `llm_api_key` 字段
2. **修复回调 URL**:通过环境变量设置正确的 `AGENT_CALLBACK_URL`
3. **验证修复**:重新发送请求,检查日志中是否出现 "使用请求中的 LLM API Key"
4. **监控回调**:确认 Agent 回调和 LiteLLM 回调都成功
5. **验证答案质量**:修复后,答案的 `confidence` 应该是 `"high"` 或 `"medium"`,而不是 `"low"`