313 lines
9.6 KiB
Python
313 lines
9.6 KiB
Python
"""
|
|
智能搜索 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()
|