580 lines
18 KiB
Python
580 lines
18 KiB
Python
"""
|
|
FastAPI服务 - Facebook搜索智能Agent
|
|
提供搜索API接口和MCP协议支持
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import uuid
|
|
from typing import Dict, Any, Optional, AsyncGenerator
|
|
from fastapi import FastAPI, HTTPException, Request, Header, Depends
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from fastapi.responses import StreamingResponse, JSONResponse
|
|
from pydantic import BaseModel, Field
|
|
from datetime import datetime
|
|
import sys
|
|
from typing import Optional
|
|
from loguru import logger
|
|
|
|
# 支持相对导入和绝对导入
|
|
try:
|
|
from .config import Config
|
|
from .agent import FacebookAgent
|
|
from .models.schemas import SearchRequest, SearchResponse
|
|
from .mcp_server import search_facebook, initialize_agent
|
|
except ImportError:
|
|
# 如果相对导入失败，尝试绝对导入
|
|
from config import Config
|
|
from agent import FacebookAgent
|
|
from models.schemas import SearchRequest, SearchResponse
|
|
from mcp_server import search_facebook, initialize_agent
|
|
|
|
try:
|
|
from common.agent_callback_utils import AgentCallbackHandler, CallbackContextManager
|
|
CALLBACK_ENABLED = True
|
|
except ImportError:
|
|
CALLBACK_ENABLED = False
|
|
AgentCallbackHandler = None
|
|
CallbackContextManager = None
|
|
|
|
|
|
# ==================== FastAPI应用 ====================
|
|
|
|
# 创建FastAPI应用
|
|
app = FastAPI(
|
|
title="Facebook搜索智能Agent API",
|
|
description="提供Facebook内容搜索服务和MCP协议支持，基于Pydantic AI框架和LiteLLM Gateway",
|
|
version="1.0.0"
|
|
)
|
|
|
|
# 添加CORS中间件
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
# 全局变量
|
|
config: Optional[Config] = None
|
|
agent: Optional[FacebookAgent] = None
|
|
callback_handler: Optional[AgentCallbackHandler] = None
|
|
POD_NAME = os.getenv("POD_NAME", "facebook-agent")
|
|
USER_ID = os.getenv("USER_ID", "")
|
|
|
|
# MCP 工具映射
|
|
TOOL_MAP = {
|
|
'search_facebook': search_facebook,
|
|
}
|
|
|
|
# Session 管理
|
|
sessions: Dict[str, Dict[str, Any]] = {}
|
|
|
|
|
|
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[Any] = None
|
|
error: Optional[Dict[str, Any]] = None
|
|
|
|
|
|
async def handle_mcp_request(request_data: Dict[str, Any], session_id: Optional[str] = None, api_key: Optional[str] = None) -> Dict[str, Any]:
|
|
"""处理 MCP 请求"""
|
|
method = request_data.get("method")
|
|
params = request_data.get("params", {})
|
|
request_id = request_data.get("id")
|
|
|
|
# 对于 tools/call 方法，需要验证 API key
|
|
if method == "tools/call":
|
|
if not api_key or api_key.strip() == "" or api_key.strip() == "sk":
|
|
return {
|
|
"jsonrpc": "2.0",
|
|
"id": request_id,
|
|
"error": {
|
|
"code": -32001,
|
|
"message": "缺少 API key。请在请求头中提供 'api-key' 或 'Authorization: Bearer <token>'。"
|
|
}
|
|
}
|
|
|
|
try:
|
|
if method == "initialize":
|
|
# 初始化会话
|
|
if not session_id:
|
|
session_id = str(uuid.uuid4())
|
|
sessions[session_id] = {
|
|
"initialized": True,
|
|
"capabilities": {}
|
|
}
|
|
|
|
return {
|
|
"jsonrpc": "2.0",
|
|
"id": request_id,
|
|
"result": {
|
|
"protocolVersion": "2024-11-05",
|
|
"capabilities": {
|
|
"tools": {},
|
|
"resources": {}
|
|
},
|
|
"serverInfo": {
|
|
"name": "Facebook搜索Agent",
|
|
"version": "1.0.0"
|
|
}
|
|
}
|
|
}
|
|
|
|
elif method == "tools/list":
|
|
# 列出所有工具
|
|
tools = [
|
|
{
|
|
"name": "search_facebook",
|
|
"description": "搜索Facebook内容，返回相关帖子和AI生成的总结。支持搜索关键词，返回帖子标题、链接、摘要、作者、点赞数等信息，并提供AI生成的总结。",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {
|
|
"type": "string",
|
|
"description": "搜索关键词，例如：technology news、travel tips、food recipes等"
|
|
},
|
|
"limit": {
|
|
"type": "integer",
|
|
"description": "返回结果数量，默认5，最大20",
|
|
"default": 5,
|
|
"minimum": 1,
|
|
"maximum": 20
|
|
}
|
|
},
|
|
"required": ["query"]
|
|
}
|
|
}
|
|
]
|
|
|
|
return {
|
|
"jsonrpc": "2.0",
|
|
"id": request_id,
|
|
"result": {
|
|
"tools": tools
|
|
}
|
|
}
|
|
|
|
elif method == "tools/call":
|
|
# 调用工具
|
|
tool_name = params.get("name")
|
|
arguments = params.get("arguments", {})
|
|
|
|
if tool_name not in TOOL_MAP:
|
|
raise ValueError(f"Tool '{tool_name}' not found")
|
|
|
|
# 如果提供了 API key，临时更新配置
|
|
if api_key and agent is not None:
|
|
old_api_key = agent.config.litellm_api_key
|
|
agent.config.litellm_api_key = api_key
|
|
# 重新创建 LLM 客户端以使用新的 API key
|
|
try:
|
|
from .clients.litellm_client import LiteLLMClient
|
|
except ImportError:
|
|
from clients.litellm_client import LiteLLMClient
|
|
agent.deps.llm_client = LiteLLMClient(agent.config)
|
|
|
|
try:
|
|
# 获取工具函数
|
|
tool_func = TOOL_MAP[tool_name]
|
|
|
|
# 调用工具（异步）
|
|
if CALLBACK_ENABLED and callback_handler:
|
|
with CallbackContextManager(
|
|
handler=callback_handler,
|
|
user_id=USER_ID,
|
|
request_id=f"facebook-mcp-{tool_name}-{request_id or uuid.uuid4().hex}"
|
|
) as ctx:
|
|
ctx.add_tool(tool_name)
|
|
result = await tool_func(**arguments)
|
|
else:
|
|
result = await tool_func(**arguments)
|
|
finally:
|
|
# 恢复原来的 API key
|
|
if api_key and agent is not None and 'old_api_key' in locals():
|
|
agent.config.litellm_api_key = old_api_key
|
|
try:
|
|
from .clients.litellm_client import LiteLLMClient
|
|
except ImportError:
|
|
from clients.litellm_client import LiteLLMClient
|
|
agent.deps.llm_client = LiteLLMClient(agent.config)
|
|
|
|
return {
|
|
"jsonrpc": "2.0",
|
|
"id": request_id,
|
|
"result": {
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": str(result)
|
|
}
|
|
]
|
|
}
|
|
}
|
|
|
|
elif method == "ping":
|
|
return {
|
|
"jsonrpc": "2.0",
|
|
"id": request_id,
|
|
"result": {}
|
|
}
|
|
|
|
else:
|
|
raise ValueError(f"Unknown method: {method}")
|
|
|
|
except Exception as e:
|
|
return {
|
|
"jsonrpc": "2.0",
|
|
"id": request_id,
|
|
"error": {
|
|
"code": -32603,
|
|
"message": str(e)
|
|
}
|
|
}
|
|
|
|
|
|
def setup_logger():
|
|
"""配置日志"""
|
|
logger.remove()
|
|
logger.add(
|
|
sys.stderr,
|
|
level=config.log_level if config else "INFO",
|
|
format="<green>{time:HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{message}</cyan>"
|
|
)
|
|
|
|
|
|
@app.on_event("startup")
|
|
async def startup_event():
|
|
"""应用启动时初始化"""
|
|
global config, agent, callback_handler
|
|
|
|
try:
|
|
# 加载配置
|
|
config = Config.from_env()
|
|
config.validate()
|
|
|
|
# 配置日志
|
|
setup_logger()
|
|
|
|
# 创建Agent
|
|
agent = FacebookAgent(config)
|
|
if CALLBACK_ENABLED and AgentCallbackHandler:
|
|
callback_handler = AgentCallbackHandler(agent_name=POD_NAME, user_id=USER_ID)
|
|
|
|
logger.info("=" * 60)
|
|
logger.info("Facebook搜索智能Agent API 启动成功")
|
|
logger.info("=" * 60)
|
|
logger.info(f"LiteLLM Gateway: {config.litellm_gateway_url}")
|
|
logger.info(f"模型: {config.litellm_model}")
|
|
logger.info(f"Facebook API Host: {config.facebook_api_host}")
|
|
logger.info("MCP服务器已就绪，支持HTTP/SSE传输")
|
|
|
|
# 初始化MCP Agent（如果还没有初始化）
|
|
try:
|
|
initialize_agent()
|
|
except Exception:
|
|
pass # 如果已经初始化，忽略错误
|
|
|
|
except Exception as e:
|
|
logger.error(f"启动失败: {e}")
|
|
raise
|
|
|
|
|
|
@app.get("/", tags=["健康检查"])
|
|
async def root():
|
|
"""根路径 - 服务信息"""
|
|
return {
|
|
"service": "Facebook搜索智能Agent API",
|
|
"status": "running",
|
|
"version": "1.0.0",
|
|
"timestamp": datetime.now().isoformat(),
|
|
"endpoints": {
|
|
"api": {
|
|
"search": "/search",
|
|
"config": "/config",
|
|
"health": "/health"
|
|
},
|
|
"mcp": {
|
|
"http": "/mcp",
|
|
"sse": "/mcp/sse"
|
|
}
|
|
},
|
|
"tools": list(TOOL_MAP.keys())
|
|
}
|
|
|
|
|
|
@app.get("/health", tags=["健康检查"])
|
|
async def health_check():
|
|
"""健康检查接口"""
|
|
return {
|
|
"status": "healthy",
|
|
"service": "Facebook搜索智能Agent API + MCP Server",
|
|
"agent_initialized": agent is not None,
|
|
"tools_count": len(TOOL_MAP),
|
|
"timestamp": datetime.now().isoformat()
|
|
}
|
|
|
|
|
|
# ==================== API Key 验证 ====================
|
|
|
|
async def verify_api_key(
|
|
api_key: Optional[str] = Header(None, alias="api-key"),
|
|
authorization: Optional[str] = Header(None)
|
|
) -> str:
|
|
"""
|
|
验证 API Key
|
|
|
|
支持从以下位置获取 API key:
|
|
1. api-key 请求头
|
|
2. Authorization: Bearer <token> 请求头
|
|
|
|
如果没有提供 API key，返回 401 错误
|
|
"""
|
|
# 从 api-key 请求头获取
|
|
if api_key:
|
|
if not api_key.strip() or api_key.strip() == "sk":
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="无效的 API key。请提供有效的 API key。"
|
|
)
|
|
return api_key.strip()
|
|
|
|
# 从 Authorization 请求头获取
|
|
if authorization:
|
|
if authorization.startswith("Bearer "):
|
|
api_key = authorization[7:].strip()
|
|
else:
|
|
api_key = authorization.strip()
|
|
|
|
if not api_key or api_key == "sk":
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="无效的 API key。请提供有效的 API key。"
|
|
)
|
|
return api_key
|
|
|
|
# 如果没有提供 API key，返回错误
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="缺少 API key。请在请求头中提供 'api-key' 或 'Authorization: Bearer <token>'。"
|
|
)
|
|
|
|
|
|
@app.post("/search", response_model=SearchResponse, tags=["搜索"])
|
|
async def search(request: SearchRequest, api_key: str = Depends(verify_api_key)):
|
|
"""
|
|
执行Facebook搜索
|
|
|
|
- **query**: 搜索关键词（必填）
|
|
- **limit**: 返回结果数量（可选，默认5，最大20）
|
|
|
|
返回：
|
|
- 搜索结果列表
|
|
- AI生成的总结
|
|
- 统计信息
|
|
"""
|
|
if agent is None:
|
|
raise HTTPException(status_code=503, detail="服务未初始化")
|
|
|
|
try:
|
|
# 如果提供了 API key，临时更新配置
|
|
if api_key:
|
|
old_api_key = agent.config.litellm_api_key
|
|
agent.config.litellm_api_key = api_key
|
|
# 重新创建 LLM 客户端以使用新的 API key
|
|
try:
|
|
from .clients.litellm_client import LiteLLMClient
|
|
except ImportError:
|
|
from clients.litellm_client import LiteLLMClient
|
|
agent.deps.llm_client = LiteLLMClient(agent.config)
|
|
|
|
try:
|
|
if CALLBACK_ENABLED and callback_handler:
|
|
with CallbackContextManager(
|
|
handler=callback_handler,
|
|
user_id=USER_ID,
|
|
request_id=f"facebook-search-{uuid.uuid4().hex}"
|
|
) as ctx:
|
|
ctx.add_tool("search_facebook")
|
|
response = await agent.search(request)
|
|
else:
|
|
response = await agent.search(request)
|
|
finally:
|
|
# 恢复原来的 API key
|
|
if api_key and 'old_api_key' in locals():
|
|
agent.config.litellm_api_key = old_api_key
|
|
try:
|
|
from .clients.litellm_client import LiteLLMClient
|
|
except ImportError:
|
|
from clients.litellm_client import LiteLLMClient
|
|
agent.deps.llm_client = LiteLLMClient(agent.config)
|
|
|
|
if not response.success:
|
|
raise HTTPException(status_code=500, detail=response.message or "搜索失败")
|
|
|
|
return response
|
|
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"搜索失败: {e}")
|
|
raise HTTPException(status_code=500, detail=f"搜索失败: {str(e)}")
|
|
|
|
|
|
@app.get("/config", tags=["配置"])
|
|
async def get_config():
|
|
"""获取当前配置信息（隐藏敏感信息）"""
|
|
if config is None:
|
|
raise HTTPException(status_code=503, detail="服务未初始化")
|
|
|
|
return {
|
|
"litellm_model": config.litellm_model,
|
|
"facebook_api_host": config.facebook_api_host,
|
|
"max_results": config.max_results,
|
|
"log_level": config.log_level,
|
|
"timeout": config.timeout
|
|
}
|
|
|
|
|
|
# ==================== MCP 协议端点 ====================
|
|
|
|
@app.post("/mcp", tags=["MCP"])
|
|
async def mcp_http_endpoint(request: Request):
|
|
"""MCP HTTP 端点 - Streamable HTTP"""
|
|
try:
|
|
body = await request.json()
|
|
session_id = request.headers.get("x-mcp-session-id")
|
|
|
|
# 从请求头获取 API key
|
|
api_key = request.headers.get("api-key") or request.headers.get("api_key")
|
|
if not api_key:
|
|
auth_header = request.headers.get("Authorization")
|
|
if auth_header:
|
|
if auth_header.startswith("Bearer "):
|
|
api_key = auth_header[7:]
|
|
else:
|
|
api_key = auth_header
|
|
|
|
response = await handle_mcp_request(body, session_id, api_key=api_key)
|
|
|
|
# 如果创建了新会话，返回 session ID
|
|
if "result" in response and isinstance(response["result"], dict):
|
|
if "sessionId" not in response["result"] and session_id:
|
|
response["result"]["sessionId"] = session_id
|
|
|
|
return JSONResponse(
|
|
content=response,
|
|
headers={"x-mcp-session-id": session_id or ""}
|
|
)
|
|
|
|
except Exception as e:
|
|
return JSONResponse(
|
|
status_code=400,
|
|
content={
|
|
"jsonrpc": "2.0",
|
|
"error": {
|
|
"code": -32700,
|
|
"message": f"Parse error: {str(e)}"
|
|
}
|
|
}
|
|
)
|
|
|
|
|
|
@app.get("/mcp/sse", tags=["MCP"])
|
|
async def mcp_sse_endpoint(request: Request):
|
|
"""MCP SSE 端点 - Server-Sent Events"""
|
|
session_id = request.headers.get("x-mcp-session-id") or str(uuid.uuid4())
|
|
|
|
async def event_stream() -> AsyncGenerator[str, None]:
|
|
# 发送初始连接消息
|
|
yield f"data: {json.dumps({'type': 'connection', 'sessionId': session_id})}\n\n"
|
|
|
|
# 发送工具列表
|
|
tools = list(TOOL_MAP.keys())
|
|
yield f"data: {json.dumps({'type': 'tools', 'tools': tools})}\n\n"
|
|
|
|
# 保持连接
|
|
import asyncio
|
|
while True:
|
|
await asyncio.sleep(30) # 心跳
|
|
yield f"data: {json.dumps({'type': 'ping'})}\n\n"
|
|
|
|
return StreamingResponse(
|
|
event_stream(),
|
|
media_type="text/event-stream",
|
|
headers={
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"x-mcp-session-id": session_id
|
|
}
|
|
)
|
|
|
|
|
|
@app.post("/mcp/sse", tags=["MCP"])
|
|
async def mcp_sse_post(request: Request):
|
|
"""MCP SSE POST 端点 - 处理 SSE 请求"""
|
|
try:
|
|
body = await request.json()
|
|
session_id = request.headers.get("x-mcp-session-id") or str(uuid.uuid4())
|
|
|
|
# 从请求头获取 API key
|
|
api_key = request.headers.get("api-key") or request.headers.get("api_key")
|
|
if not api_key:
|
|
auth_header = request.headers.get("Authorization")
|
|
if auth_header:
|
|
if auth_header.startswith("Bearer "):
|
|
api_key = auth_header[7:]
|
|
else:
|
|
api_key = auth_header
|
|
|
|
async def response_stream() -> AsyncGenerator[str, None]:
|
|
response = await handle_mcp_request(body, session_id, api_key=api_key)
|
|
yield f"data: {json.dumps(response)}\n\n"
|
|
|
|
return StreamingResponse(
|
|
response_stream(),
|
|
media_type="text/event-stream",
|
|
headers={
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"x-mcp-session-id": session_id
|
|
}
|
|
)
|
|
|
|
except Exception as e:
|
|
return JSONResponse(
|
|
status_code=400,
|
|
content={
|
|
"jsonrpc": "2.0",
|
|
"error": {
|
|
"code": -32700,
|
|
"message": f"Parse error: {str(e)}"
|
|
}
|
|
}
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
|
|
# 运行服务
|
|
uvicorn.run(
|
|
"facebook_agent.api:app",
|
|
host="0.0.0.0",
|
|
port=8000,
|
|
reload=True,
|
|
log_level="info"
|
|
)
|
|
|