321 lines
10 KiB
Python
321 lines
10 KiB
Python
"""
|
|
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()
|