609 lines
21 KiB
Plaintext
609 lines
21 KiB
Plaintext
"""
|
||
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()
|
||
|