420 lines
13 KiB
Python
420 lines
13 KiB
Python
"""
|
||
RAG API 生产环境服务 - Redis缓存版
|
||
基于原始 rag_api_server_production.py 添加Redis缓存
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import json
|
||
import uvicorn
|
||
import hashlib
|
||
import redis
|
||
from pathlib import Path
|
||
from datetime import datetime
|
||
from typing import Dict, Any, List, Optional
|
||
from contextlib import asynccontextmanager
|
||
from fastapi import FastAPI, HTTPException
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
from fastapi.responses import FileResponse
|
||
from pydantic import BaseModel, Field
|
||
from dotenv import load_dotenv
|
||
|
||
# 设置环境变量
|
||
os.environ['RAG_CONFIG_PATH'] = 'rag_config_production.yaml'
|
||
|
||
# 项目根目录
|
||
project_root = Path(__file__).parent
|
||
sys.path.append(str(project_root))
|
||
|
||
# 加载环境变量
|
||
env_path = project_root / '.env'
|
||
if env_path.exists():
|
||
load_dotenv(env_path)
|
||
|
||
from retriver.langsmith.langsmith_retriever import create_langsmith_retriever, check_langsmith_connection
|
||
|
||
# ============= Redis配置 =============
|
||
REDIS_HOST = os.getenv('REDIS_HOST', '172.18.0.3') # 使用Docker内部IP
|
||
REDIS_PORT = int(os.getenv('REDIS_PORT', 6379))
|
||
REDIS_PASSWORD = os.getenv('REDIS_PASSWORD', None) # Redis密码
|
||
CACHE_TTL = int(os.getenv('CACHE_TTL', 3600)) # 缓存1小时
|
||
CACHE_ENABLED = True
|
||
|
||
# ============= 生产环境配置 =============
|
||
OUTPUT_DIR = project_root / "api_outputs"
|
||
|
||
# ============= 请求/响应模型 =============
|
||
class RetrieveRequest(BaseModel):
|
||
"""检索请求模型"""
|
||
query: str = Field(..., description="查询问题")
|
||
mode: str = Field(default="0", description="调试模式: 0=自动")
|
||
save_output: bool = Field(default=True, description="是否保存输出到文件")
|
||
use_cache: bool = Field(default=True, description="是否使用缓存")
|
||
|
||
|
||
class RetrieveResponse(BaseModel):
|
||
"""检索响应模型"""
|
||
success: bool = Field(..., description="是否成功")
|
||
query: str = Field(..., description="原始查询")
|
||
answer: str = Field(..., description="生成的答案")
|
||
supporting_facts: List[List[str]] = Field(default_factory=list, description="支撑段落")
|
||
supporting_events: List[List[str]] = Field(default_factory=list, description="支撑事件")
|
||
metadata: Dict[str, Any] = Field(default_factory=dict, description="元数据")
|
||
output_file: Optional[str] = Field(None, description="输出文件路径")
|
||
error: Optional[str] = Field(None, description="错误信息")
|
||
|
||
|
||
class HealthResponse(BaseModel):
|
||
"""健康检查响应"""
|
||
status: str = "healthy"
|
||
environment: str = "production"
|
||
langsmith_connected: bool = False
|
||
redis_connected: bool = False
|
||
timestamp: str = ""
|
||
|
||
|
||
# ============= Redis缓存管理 =============
|
||
class CacheManager:
|
||
"""Redis缓存管理器"""
|
||
|
||
def __init__(self):
|
||
"""初始化Redis连接"""
|
||
self.enabled = False
|
||
self.client = None
|
||
self.connect()
|
||
|
||
def connect(self):
|
||
"""连接Redis"""
|
||
if not CACHE_ENABLED:
|
||
print("[INFO] 缓存已禁用")
|
||
return
|
||
|
||
try:
|
||
self.client = redis.Redis(
|
||
host=REDIS_HOST,
|
||
port=REDIS_PORT,
|
||
password=REDIS_PASSWORD,
|
||
db=0,
|
||
decode_responses=True,
|
||
socket_connect_timeout=5,
|
||
socket_timeout=5
|
||
)
|
||
# 测试连接
|
||
self.client.ping()
|
||
self.enabled = True
|
||
print(f"[OK] Redis缓存已连接: {REDIS_HOST}:{REDIS_PORT}")
|
||
except Exception as e:
|
||
print(f"[WARNING] Redis连接失败: {e}")
|
||
print("[INFO] 系统将在无缓存模式下运行")
|
||
self.enabled = False
|
||
|
||
def get_key(self, query: str, mode: str) -> str:
|
||
"""生成缓存键"""
|
||
content = f"{query}:{mode}"
|
||
return f"rag:{hashlib.md5(content.encode()).hexdigest()}"
|
||
|
||
def get(self, query: str, mode: str) -> Optional[Dict]:
|
||
"""获取缓存"""
|
||
if not self.enabled:
|
||
return None
|
||
|
||
try:
|
||
key = self.get_key(query, mode)
|
||
cached_data = self.client.get(key)
|
||
if cached_data:
|
||
print(f"[CACHE HIT] 查询: {query[:50]}...")
|
||
return json.loads(cached_data)
|
||
except Exception as e:
|
||
print(f"[CACHE ERROR] 读取失败: {e}")
|
||
|
||
return None
|
||
|
||
def set(self, query: str, mode: str, data: Dict):
|
||
"""设置缓存"""
|
||
if not self.enabled:
|
||
return
|
||
|
||
try:
|
||
key = self.get_key(query, mode)
|
||
self.client.setex(
|
||
key,
|
||
CACHE_TTL,
|
||
json.dumps(data, ensure_ascii=False)
|
||
)
|
||
print(f"[CACHE SET] 已缓存: {query[:50]}... (TTL: {CACHE_TTL}秒)")
|
||
except Exception as e:
|
||
print(f"[CACHE ERROR] 写入失败: {e}")
|
||
|
||
def get_stats(self) -> Dict:
|
||
"""获取缓存统计"""
|
||
if not self.enabled:
|
||
return {"enabled": False}
|
||
|
||
try:
|
||
info = self.client.info()
|
||
dbsize = self.client.dbsize()
|
||
return {
|
||
"enabled": True,
|
||
"keys": dbsize,
|
||
"used_memory": info.get("used_memory_human"),
|
||
"hits": info.get("keyspace_hits", 0),
|
||
"misses": info.get("keyspace_misses", 0),
|
||
"hit_rate": round(
|
||
info.get("keyspace_hits", 0) /
|
||
max(info.get("keyspace_hits", 0) + info.get("keyspace_misses", 0), 1) * 100,
|
||
2
|
||
)
|
||
}
|
||
except:
|
||
return {"enabled": False, "error": "无法获取统计"}
|
||
|
||
def clear(self):
|
||
"""清空缓存"""
|
||
if not self.enabled:
|
||
return
|
||
|
||
try:
|
||
self.client.flushdb()
|
||
print("[CACHE] 缓存已清空")
|
||
except Exception as e:
|
||
print(f"[CACHE ERROR] 清空失败: {e}")
|
||
|
||
|
||
# ============= 全局状态管理 =============
|
||
class AppState:
|
||
"""应用状态管理"""
|
||
retriever = None
|
||
langsmith_connected = False
|
||
cache_manager = None
|
||
request_count = 0
|
||
cache_hits = 0
|
||
last_request_time = None
|
||
|
||
|
||
# ============= 生命周期管理 =============
|
||
def initialize_retriever():
|
||
"""初始化检索器"""
|
||
try:
|
||
AppState.langsmith_connected = check_langsmith_connection()
|
||
AppState.retriever = create_langsmith_retriever(keyword="test")
|
||
print("[OK] 检索器初始化成功")
|
||
except Exception as e:
|
||
print(f"[ERROR] 检索器初始化失败: {e}")
|
||
AppState.retriever = None
|
||
|
||
|
||
def initialize_cache():
|
||
"""初始化缓存"""
|
||
AppState.cache_manager = CacheManager()
|
||
|
||
|
||
@asynccontextmanager
|
||
async def lifespan(app: FastAPI):
|
||
"""应用生命周期管理"""
|
||
print("="*60)
|
||
print("[STARTING] 启动RAG API服务（Redis缓存版）")
|
||
print("="*60)
|
||
|
||
# 创建输出目录
|
||
OUTPUT_DIR.mkdir(exist_ok=True)
|
||
|
||
# 初始化组件
|
||
initialize_retriever()
|
||
initialize_cache()
|
||
|
||
print(f"[INFO] Redis配置: {REDIS_HOST}:{REDIS_PORT}")
|
||
print(f"[INFO] 缓存TTL: {CACHE_TTL}秒")
|
||
print("="*60)
|
||
|
||
yield
|
||
|
||
print("\n[INFO] 服务关闭")
|
||
|
||
|
||
# ============= FastAPI应用 =============
|
||
app = FastAPI(
|
||
title="RAG API Service (Redis Cache)",
|
||
description="RAG检索服务（Redis缓存版）",
|
||
version="2.0.0",
|
||
lifespan=lifespan
|
||
)
|
||
|
||
# 配置CORS
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=["*"],
|
||
allow_credentials=True,
|
||
allow_methods=["*"],
|
||
allow_headers=["*"],
|
||
)
|
||
|
||
|
||
# ============= API端点 =============
|
||
@app.get("/health", response_model=HealthResponse)
|
||
async def health_check():
|
||
"""健康检查"""
|
||
return HealthResponse(
|
||
status="healthy" if AppState.retriever else "degraded",
|
||
environment="production",
|
||
langsmith_connected=AppState.langsmith_connected,
|
||
redis_connected=AppState.cache_manager.enabled if AppState.cache_manager else False,
|
||
timestamp=datetime.now().isoformat()
|
||
)
|
||
|
||
|
||
@app.post("/retrieve", response_model=RetrieveResponse)
|
||
async def retrieve(request: RetrieveRequest):
|
||
"""RAG查询接口（带缓存）"""
|
||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
AppState.request_count += 1
|
||
AppState.last_request_time = datetime.now()
|
||
|
||
print(f"\n{'='*60}")
|
||
print(f"[IN] 请求 #{AppState.request_count}")
|
||
print(f" 查询: {request.query}")
|
||
print(f" 使用缓存: {request.use_cache}")
|
||
|
||
if not AppState.retriever:
|
||
raise HTTPException(status_code=503, detail="检索器未初始化")
|
||
|
||
try:
|
||
# 1. 检查缓存
|
||
if request.use_cache and AppState.cache_manager and AppState.cache_manager.enabled:
|
||
cached_result = AppState.cache_manager.get(request.query, request.mode)
|
||
if cached_result:
|
||
AppState.cache_hits += 1
|
||
print(f"[OK] 缓存命中 (命中率: {AppState.cache_hits}/{AppState.request_count})")
|
||
|
||
# 直接返回缓存结果
|
||
return RetrieveResponse(
|
||
success=True,
|
||
query=request.query,
|
||
answer=cached_result.get("answer", ""),
|
||
supporting_facts=cached_result.get("supporting_facts", []),
|
||
supporting_events=cached_result.get("supporting_events", []),
|
||
metadata={
|
||
"environment": "production",
|
||
"timestamp": timestamp,
|
||
"cached": True,
|
||
"cache_hit_rate": f"{AppState.cache_hits}/{AppState.request_count}"
|
||
}
|
||
)
|
||
|
||
# 2. 执行检索
|
||
print("[PROCESSING] 执行RAG检索...")
|
||
result = AppState.retriever.retrieve(
|
||
request.query,
|
||
request.mode
|
||
)
|
||
|
||
# 3. 存入缓存
|
||
if request.use_cache and AppState.cache_manager and AppState.cache_manager.enabled:
|
||
cache_data = {
|
||
"answer": result.get("answer", ""),
|
||
"supporting_facts": result.get("supporting_facts", []),
|
||
"supporting_events": result.get("supporting_events", [])
|
||
}
|
||
AppState.cache_manager.set(request.query, request.mode, cache_data)
|
||
|
||
# 4. 保存输出（可选）
|
||
output_file = None
|
||
if request.save_output:
|
||
output_file = OUTPUT_DIR / f"output_{timestamp}.json"
|
||
output_data = {
|
||
"timestamp": timestamp,
|
||
"query": request.query,
|
||
"answer": result.get("answer", ""),
|
||
"supporting_facts": result.get("supporting_facts", []),
|
||
"supporting_events": result.get("supporting_events", []),
|
||
"metadata": {
|
||
"cached": False,
|
||
"mode": request.mode
|
||
}
|
||
}
|
||
with open(output_file, 'w', encoding='utf-8') as f:
|
||
json.dump(output_data, f, ensure_ascii=False, indent=2)
|
||
|
||
response = RetrieveResponse(
|
||
success=True,
|
||
query=request.query,
|
||
answer=result.get("answer", ""),
|
||
supporting_facts=result.get("supporting_facts", []),
|
||
supporting_events=result.get("supporting_events", []),
|
||
metadata={
|
||
"environment": "production",
|
||
"timestamp": timestamp,
|
||
"cached": False
|
||
},
|
||
output_file=str(output_file) if output_file else None
|
||
)
|
||
|
||
print(f"[OK] 查询完成 (答案长度: {len(response.answer)} 字符)")
|
||
return response
|
||
|
||
except Exception as e:
|
||
print(f"[ERROR] 查询失败: {e}")
|
||
return RetrieveResponse(
|
||
success=False,
|
||
query=request.query,
|
||
answer="",
|
||
error=str(e),
|
||
metadata={"environment": "production", "timestamp": timestamp}
|
||
)
|
||
|
||
|
||
@app.get("/cache/stats")
|
||
async def cache_stats():
|
||
"""获取缓存统计"""
|
||
if not AppState.cache_manager:
|
||
return {"error": "缓存未初始化"}
|
||
|
||
stats = AppState.cache_manager.get_stats()
|
||
stats.update({
|
||
"total_requests": AppState.request_count,
|
||
"cache_hits": AppState.cache_hits,
|
||
"hit_rate_session": f"{round(AppState.cache_hits/max(AppState.request_count,1)*100, 2)}%"
|
||
})
|
||
return stats
|
||
|
||
|
||
@app.delete("/cache/clear")
|
||
async def clear_cache():
|
||
"""清空缓存"""
|
||
if not AppState.cache_manager:
|
||
return {"error": "缓存未初始化"}
|
||
|
||
AppState.cache_manager.clear()
|
||
return {"success": True, "message": "缓存已清空"}
|
||
|
||
|
||
# ============= 兼容旧接口 =============
|
||
@app.post("/query", response_model=RetrieveResponse)
|
||
async def query(request: RetrieveRequest):
|
||
"""查询接口（兼容旧版）"""
|
||
return await retrieve(request)
|
||
|
||
|
||
# ============= 主函数 =============
|
||
def main():
|
||
"""启动服务"""
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser(description="RAG API服务（Redis缓存版）")
|
||
parser.add_argument("--host", default="0.0.0.0", help="服务地址")
|
||
parser.add_argument("--port", type=int, default=8000, help="服务端口")
|
||
parser.add_argument("--workers", type=int, default=4, help="Worker进程数")
|
||
|
||
args = parser.parse_args()
|
||
|
||
uvicorn.run(
|
||
"rag_api_server_with_cache:app",
|
||
host=args.host,
|
||
port=args.port,
|
||
reload=False,
|
||
workers=args.workers,
|
||
log_level="info"
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |