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:
@@ -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}"""
|
||||
|
||||
+209
@@ -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",
|
||||
]
|
||||
|
||||
+151
@@ -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"
|
||||
)
|
||||
|
||||
+102
@@ -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)
|
||||
|
||||
+117
@@ -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
|
||||
)
|
||||
|
||||
+166
@@ -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
|
||||
|
||||
+107
@@ -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])
|
||||
]
|
||||
|
||||
+89
@@ -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
|
||||
)
|
||||
|
||||
+137
@@ -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",
|
||||
]
|
||||
|
||||
+180
@@ -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
|
||||
|
||||
+191
@@ -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
|
||||
|
||||
+179
@@ -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
|
||||
@@ -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 \
|
||||
.
|
||||
|
||||
@@ -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
|
||||
@@ -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
File diff suppressed because it is too large
Load Diff
+35
-7
@@ -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 = [
|
||||
|
||||
@@ -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
|
||||
@@ -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 和租户信息。
|
||||
|
||||
@@ -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"`
|
||||
|
||||
Reference in New Issue
Block a user