修改后端支持流式响应时的缓存
This commit is contained in:
parent
5d15567460
commit
2c0ea3467b
@ -65,29 +65,47 @@ def build_chat_cache_key(chat_id: str, question: str) -> str:
|
||||
@ragflowController.post('/converse_with_chat_assistant')
|
||||
async def converse_with_chat_assistant(
|
||||
request: Request,
|
||||
converse_params: ConverseWithChatAssistantModel, # 使用正确的Pydantic模型
|
||||
converse_params: ConverseWithChatAssistantModel,
|
||||
):
|
||||
"""
|
||||
与聊天助手进行对话 - 集成语义缓存版本
|
||||
与聊天助手进行对话 - 集成语义缓存版本（支持流式和非流式）
|
||||
"""
|
||||
start_time = time.perf_counter()
|
||||
|
||||
# 获取redis实例
|
||||
redis = getattr(request.app.state, 'redis', None)
|
||||
cache_key = None
|
||||
cached_answer = None
|
||||
cache_similarity = 0.0
|
||||
|
||||
# 语义缓存查找（非流式对话）
|
||||
if not converse_params.stream and redis:
|
||||
# ========== 语义缓存查找（流式和非流式都支持）==========
|
||||
if redis:
|
||||
cache_result = await lookup_question(
|
||||
converse_params.chat_id,
|
||||
converse_params.question,
|
||||
redis
|
||||
)
|
||||
if cache_result:
|
||||
cached_answer, similarity = cache_result
|
||||
logger.info(f'[SemanticCache] 命中缓存 (相似度={similarity:.2f}): chat_id={converse_params.chat_id}')
|
||||
# 返回缓存的答案，包装成标准响应格式
|
||||
return ResponseUtil.success(data={'answer': cached_answer, 'from_cache': True, 'similarity': similarity})
|
||||
cached_answer, cache_similarity = cache_result
|
||||
logger.info(f'[SemanticCache] 命中缓存 (相似度={cache_similarity:.2f}): chat_id={converse_params.chat_id}')
|
||||
|
||||
# 非流式：直接返回缓存答案
|
||||
if not converse_params.stream:
|
||||
return ResponseUtil.success(data={'answer': cached_answer, 'from_cache': True, 'similarity': cache_similarity})
|
||||
|
||||
# 流式：使用缓存答案进行流式响应
|
||||
if converse_params.stream:
|
||||
logger.info(f'[SemanticCache] 流式响应使用缓存答案，chat_id={converse_params.chat_id}')
|
||||
return StreamingResponse(
|
||||
stream_cached_response(cached_answer, converse_params.chat_id, start_time),
|
||||
media_type='text/event-stream',
|
||||
headers={
|
||||
'Cache-Control': 'no-cache',
|
||||
'Connection': 'keep-alive',
|
||||
'X-Accel-Buffering': 'no',
|
||||
'Transfer-Encoding': 'chunked'
|
||||
}
|
||||
)
|
||||
|
||||
# 检查是否应该使用搜索服务
|
||||
try:
|
||||
@ -99,7 +117,7 @@ async def converse_with_chat_assistant(
|
||||
except Exception as e:
|
||||
logger.warning(f"意图分类失败，使用RAG服务: {e}")
|
||||
|
||||
# 缓存检查（原有的精确缓存）
|
||||
# 缓存检查（原有的精确缓存）- 非流式
|
||||
if not converse_params.stream and redis:
|
||||
cache_key = build_chat_cache_key(converse_params.chat_id, converse_params.question)
|
||||
try:
|
||||
@ -310,4 +328,64 @@ async def delete_datasets(
|
||||
):
|
||||
"""删除数据集"""
|
||||
result = RAGFlowService.delete_datasets_services(delete_params)
|
||||
return parse_result(result)
|
||||
return parse_result(result)
|
||||
|
||||
|
||||
def stream_cached_response(cached_answer: str, chat_id: str, start_time: float) -> Generator[str, None, None]:
|
||||
"""
|
||||
流式返回缓存的答案
|
||||
|
||||
Args:
|
||||
cached_answer: 缓存的答案文本
|
||||
chat_id: 会话ID
|
||||
start_time: 请求开始时间
|
||||
"""
|
||||
import time
|
||||
|
||||
server_stream_start = time.time()
|
||||
logger.info(f"[CACHE_STREAM {server_stream_start:.3f}] 🚀 开始流式返回缓存答案，chat_id: {chat_id}")
|
||||
|
||||
# 立即发送连接建立消息
|
||||
connection_time = time.time()
|
||||
ping_message = "event: ping\ndata: {\"status\": \"connected\"}\n\n"
|
||||
yield ping_message
|
||||
logger.info(f"[CACHE_STREAM {connection_time:.3f}] 📡 连接建立ping发送，耗时: {connection_time - server_stream_start:.3f}s")
|
||||
|
||||
# 模拟流式输出效果（可选：如果是完整答案，可以选择立即返回或模拟流式）
|
||||
# 这里我们选择模拟流式输出，每次发送一小部分，模拟打字效果
|
||||
if cached_answer:
|
||||
# 将答案分块发送，模拟打字效果
|
||||
chunk_size = 10 # 每个块10个字符
|
||||
answer_chunks = [cached_answer[i:i+chunk_size] for i in range(0, len(cached_answer), chunk_size)]
|
||||
|
||||
first_token_time = time.time()
|
||||
for i, chunk in enumerate(answer_chunks):
|
||||
chunk_time = time.time()
|
||||
|
||||
# 首token延迟
|
||||
if i == 0:
|
||||
latency = time.perf_counter() - start_time
|
||||
logger.info(f"[CACHE_STREAM {chunk_time:.3f}] 🎯 首token到达（缓存），耗时: {latency:.3f}s")
|
||||
|
||||
body = {'answer': chunk, 'chunk_index': i, 'total_chunks': len(answer_chunks)}
|
||||
yield format_sse(body)
|
||||
|
||||
# 小延迟模拟打字效果（可调整或移除）
|
||||
# time.sleep(0.01)
|
||||
|
||||
logger.info(f"[CACHE_STREAM {chunk_time:.3f}] 📝 缓存答案流式发送完成，共 {len(answer_chunks)} 个chunk")
|
||||
|
||||
# 流结束
|
||||
stream_end_time = time.time()
|
||||
end_message = format_sse({
|
||||
'status': 'completed',
|
||||
'from_cache': True,
|
||||
'total_time': stream_end_time - server_stream_start
|
||||
}, event='end')
|
||||
yield end_message
|
||||
|
||||
logger.info(f"[CACHE_STREAM {stream_end_time:.3f}] 🏁 缓存流式响应完成")
|
||||
logger.info(f"[CACHE_STREAM {stream_end_time:.3f}] 📊 缓存答案长度: {len(cached_answer)} 字符")
|
||||
|
||||
total_time = time.perf_counter() - start_time
|
||||
logger.info(f'[CACHE_STREAM {time.time():.3f}] ⏱️ Total Cache Stream Duration ({chat_id}): {total_time:.3f}s')
|
||||
@ -2,8 +2,8 @@ import argparse # 解析命令行参数
|
||||
import asyncio # 异步事件循环
|
||||
import json # 处理 JSON 文本
|
||||
import os # 读取环境变量
|
||||
import threading # 后台线程
|
||||
import time # 简单的等待/演示主线程未阻塞
|
||||
import sys # 系统参数
|
||||
import time # 时间统计
|
||||
import uuid # 生成随机 chat/session id
|
||||
|
||||
import aiohttp # 异步 HTTP 客户端
|
||||
@ -22,7 +22,7 @@ DEFAULT_CHAT_ID = os.environ.get(
|
||||
"RAGFLOW_CHAT_ID", "a455abcadbcb11f0884b0242ac130006"
|
||||
)
|
||||
DEFAULT_SESSION_ID = os.environ.get(
|
||||
"RAGFLOW_SESSION_ID", "38d765e48a3811f0be310242ac130006"
|
||||
"RAGFLOW_SESSION_ID", "3d3ec1e2b9344591a1486fc95c1e01ed"
|
||||
)
|
||||
DEFAULT_QUESTION = os.environ.get(
|
||||
"RAGFLOW_QUESTION", "你好，请用简洁的语言介绍你自己"
|
||||
@ -32,36 +32,114 @@ LOGIN_URL = os.environ.get(
|
||||
"http://10.0.0.202:9099/login",
|
||||
)
|
||||
|
||||
# 登录凭据
|
||||
LOGIN_USERNAME = os.environ.get("RAGFLOW_USERNAME", "admin")
|
||||
LOGIN_PASSWORD = os.environ.get("RAGFLOW_PASSWORD", "admin123")
|
||||
|
||||
|
||||
class PerformanceTimer:
|
||||
"""性能计时器类"""
|
||||
|
||||
def __init__(self):
|
||||
self.start_time = None
|
||||
self.end_time = None
|
||||
self.elapsed_ms = 0
|
||||
|
||||
def start(self):
|
||||
self.start_time = time.perf_counter()
|
||||
return self
|
||||
|
||||
def stop(self):
|
||||
if self.start_time is None:
|
||||
raise ValueError("计时器未启动")
|
||||
self.end_time = time.perf_counter()
|
||||
self.elapsed_ms = (self.end_time - self.start_time) * 1000
|
||||
return self
|
||||
|
||||
def get_elapsed_ms(self) -> float:
|
||||
return self.elapsed_ms
|
||||
|
||||
def get_elapsed_seconds(self) -> float:
|
||||
return self.elapsed_ms / 1000
|
||||
|
||||
def format(self) -> str:
|
||||
"""格式化输出时间"""
|
||||
if self.elapsed_ms < 1000:
|
||||
return f"{self.elapsed_ms:.2f}ms"
|
||||
else:
|
||||
return f"{self.get_elapsed_seconds():.2f}s ({self.elapsed_ms:.0f}ms)"
|
||||
|
||||
|
||||
async def login_get_token() -> str:
|
||||
"""
|
||||
登录获取token
|
||||
"""
|
||||
print("🔐 正在登录获取token...")
|
||||
|
||||
payload = {
|
||||
"username": "admin",
|
||||
"password": "admin123"
|
||||
}
|
||||
headers = {
|
||||
"Authorization": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VyX2lkIjoiMSIsInVzZXJfbmFtZSI6ImFkbWluIiwiZGVwdF9uYW1lIjoiXHU3ODE0XHU1M2QxXHU5MGU4XHU5NWU4Iiwic2Vzc2lvbl9pZCI6IjYwNjdlMzkxLTRiNGMtNGUwYy1hZjM5LTVmNzczZjIzMmE3OSIsImxvZ2luX2luZm8iOnsiaXBhZGRyIjpudWxsLCJsb2dpbkxvY2F0aW9uIjoiXHU2NzJhXHU3N2U1IiwiYnJvd3NlciI6Ik90aGVyIiwib3MiOiJPdGhlciIsImxvZ2luVGltZSI6IjIwMjUtMTItMDIgMTA6NTk6NDAifSwiZXhwIjoxNzY3MjM2MzgwfQ.NM9j0emNz8trxU1DX87PVZZaWWHlFFwF7o0Pxop6B5c"
|
||||
"username": LOGIN_USERNAME,
|
||||
"password": LOGIN_PASSWORD
|
||||
}
|
||||
|
||||
timer = PerformanceTimer().start()
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
# 使用data参数发送form-data而不是json
|
||||
async with session.post(
|
||||
LOGIN_URL, data=payload, headers=headers, timeout=60
|
||||
LOGIN_URL, data=payload, timeout=60
|
||||
) as resp:
|
||||
response_data = await resp.json() # 将响应转换为JSON
|
||||
token = response_data.get("token") # 获取token字段
|
||||
timer.stop()
|
||||
|
||||
if resp.status != 200:
|
||||
print(f"❌ 登录失败，HTTP状态码: {resp.status}")
|
||||
error_text = await resp.text()
|
||||
print(f"错误响应: {error_text}")
|
||||
raise Exception(f"登录失败: {resp.status}")
|
||||
|
||||
response_data = await resp.json()
|
||||
|
||||
# 检查响应格式
|
||||
if response_data.get("code") != 200:
|
||||
print(f"❌ 登录失败，响应码: {response_data.get('code')}")
|
||||
print(f"响应信息: {response_data.get('msg', '未知错误')}")
|
||||
raise Exception(f"登录失败: {response_data.get('msg')}")
|
||||
|
||||
token = response_data.get("token")
|
||||
|
||||
if not token:
|
||||
print("❌ 登录成功但未获取到token")
|
||||
print(f"完整响应: {response_data}")
|
||||
raise Exception("登录成功但未获取到token")
|
||||
|
||||
print(f"✅ 登录成功，耗时: {timer.format()}")
|
||||
print(f"📝 Token预览: {token[:50]}...")
|
||||
|
||||
return token
|
||||
|
||||
|
||||
async def stream_chat_async(question: str, chat_id: str, session_id: str) -> str:
|
||||
async def stream_chat_async(
|
||||
question: str,
|
||||
chat_id: str,
|
||||
session_id: str,
|
||||
request_number: int = 1
|
||||
) -> tuple[str, PerformanceTimer]:
|
||||
"""
|
||||
异步版（aiohttp），适合已有事件循环的应用，不阻塞主线程。
|
||||
返回完整回答和耗时
|
||||
"""
|
||||
print(f"\n{'='*60}")
|
||||
print(f"📤 请求 #{request_number}")
|
||||
print(f"❓ 问题: {question}")
|
||||
print(f"🆔 Chat ID: {chat_id}")
|
||||
print(f"🔗 Session ID: {session_id}")
|
||||
print(f"{'='*60}")
|
||||
|
||||
token = await login_get_token()
|
||||
|
||||
payload = {
|
||||
"chatId": chat_id,
|
||||
"chat_id": chat_id,
|
||||
"question": question,
|
||||
"stream": True,
|
||||
"sessionId": session_id,
|
||||
"session_id": session_id,
|
||||
}
|
||||
headers = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
@ -70,10 +148,14 @@ async def stream_chat_async(question: str, chat_id: str, session_id: str) -> str
|
||||
|
||||
answer_parts = []
|
||||
event_type = "message"
|
||||
|
||||
|
||||
# 性能计时
|
||||
timer = PerformanceTimer().start()
|
||||
first_token_time = None
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(
|
||||
API_URL, json=payload, headers=headers, timeout=60
|
||||
API_URL, json=payload, headers=headers, timeout=120
|
||||
) as resp:
|
||||
print(
|
||||
f"# status={resp.status} content-type={resp.headers.get('content-type')}")
|
||||
@ -102,6 +184,10 @@ async def stream_chat_async(question: str, chat_id: str, session_id: str) -> str
|
||||
event_type = "message"
|
||||
continue
|
||||
|
||||
# 记录首Token时间
|
||||
if first_token_time is None and "answer" in payload_obj:
|
||||
first_token_time = timer.get_elapsed_ms()
|
||||
|
||||
if event_type == "end" or payload_obj.get("status") == "completed":
|
||||
print("\n[stream] completed")
|
||||
break
|
||||
@ -111,21 +197,137 @@ async def stream_chat_async(question: str, chat_id: str, session_id: str) -> str
|
||||
if "answer" in payload_obj:
|
||||
piece = payload_obj.get("answer", "")
|
||||
answer_parts.append(piece)
|
||||
# 流式输出：每个片段都单独打印，并且添加延时让效果更明显
|
||||
# 流式输出：每个片段都单独打印
|
||||
print(piece, end="", flush=True)
|
||||
time.sleep(0.1) # 每个片段间隔100ms，让流式效果更明显
|
||||
time.sleep(0.05) # 每个片段间隔50ms，让流式效果更明显
|
||||
else:
|
||||
print(f"[{event_type}] {payload_obj}")
|
||||
event_type = "message"
|
||||
|
||||
print("\n\n[stream completed] - Answer received in real-time via streaming!")
|
||||
timer.stop()
|
||||
print(f"\n\n{'='*60}")
|
||||
print(f"✅ 请求 #{request_number} 完成")
|
||||
print(f"⏱️ 总耗时: {timer.format()}")
|
||||
if first_token_time:
|
||||
print(f"🚀 首Token延迟: {first_token_time:.0f}ms")
|
||||
print(f"📝 回答长度: {len(''.join(answer_parts))} 字符")
|
||||
print(f"{'='*60}")
|
||||
|
||||
full_answer = "".join(answer_parts)
|
||||
return full_answer
|
||||
return full_answer, timer
|
||||
|
||||
|
||||
async def run_performance_test(
|
||||
question: str,
|
||||
chat_id: str,
|
||||
session_id: str,
|
||||
num_requests: int = 3,
|
||||
delay_between_requests: float = 2.0
|
||||
):
|
||||
"""
|
||||
运行性能测试，多次请求以测试缓存效果
|
||||
"""
|
||||
print("\n" + "🚀"*30)
|
||||
print("🎯 性能测试开始")
|
||||
print(f"📊 测试配置:")
|
||||
print(f" - 请求次数: {num_requests}")
|
||||
print(f" - 请求间隔: {delay_between_requests}s")
|
||||
print(f" - 问题: {question}")
|
||||
print("🚀"*30)
|
||||
|
||||
results = []
|
||||
|
||||
for i in range(1, num_requests + 1):
|
||||
try:
|
||||
answer, timer = await stream_chat_async(
|
||||
question=question,
|
||||
chat_id=chat_id,
|
||||
session_id=session_id,
|
||||
request_number=i
|
||||
)
|
||||
results.append({
|
||||
"request_number": i,
|
||||
"elapsed_ms": timer.get_elapsed_ms(),
|
||||
"elapsed_seconds": timer.get_elapsed_seconds(),
|
||||
"answer_length": len(answer)
|
||||
})
|
||||
|
||||
# 如果不是最后一次请求，等待一段时间
|
||||
if i < num_requests:
|
||||
print(f"\n⏳ 等待 {delay_between_requests}s 后进行下一次请求...")
|
||||
await asyncio.sleep(delay_between_requests)
|
||||
|
||||
except Exception as e:
|
||||
print(f"\n❌ 请求 #{i} 失败: {str(e)}")
|
||||
results.append({
|
||||
"request_number": i,
|
||||
"elapsed_ms": -1,
|
||||
"elapsed_seconds": -1,
|
||||
"answer_length": 0,
|
||||
"error": str(e)
|
||||
})
|
||||
|
||||
# 打印性能测试汇总
|
||||
print("\n" + "📊"*30)
|
||||
print("📈 性能测试结果汇总")
|
||||
print("📊"*30)
|
||||
|
||||
valid_results = [r for r in results if r["elapsed_ms"] > 0]
|
||||
|
||||
if not valid_results:
|
||||
print("❌ 所有请求都失败了")
|
||||
return
|
||||
|
||||
elapsed_times = [r["elapsed_ms"] for r in valid_results]
|
||||
avg_time = sum(elapsed_times) / len(elapsed_times)
|
||||
min_time = min(elapsed_times)
|
||||
max_time = max(elapsed_times)
|
||||
|
||||
print(f"\n{'请求#':<8} {'耗时':<15} {'速度评估'}")
|
||||
print("-" * 50)
|
||||
|
||||
for r in results:
|
||||
if r["elapsed_ms"] < 0:
|
||||
print(f"#{r['request_number']:<6} {'失败':<15} ❌")
|
||||
else:
|
||||
# 速度评估
|
||||
elapsed = r["elapsed_ms"]
|
||||
if elapsed < 500:
|
||||
speed = "🚀 极快 (缓存?)"
|
||||
elif elapsed < 1500:
|
||||
speed = "⚡ 较快"
|
||||
elif elapsed < 3000:
|
||||
speed = "⏱️ 正常"
|
||||
else:
|
||||
speed = "🐢 较慢"
|
||||
|
||||
print(f"#{r['request_number']:<6} {r['elapsed_ms']:.0f}ms".ljust(15) + speed)
|
||||
|
||||
print("-" * 50)
|
||||
print(f"\n📊 统计信息:")
|
||||
print(f" 平均耗时: {avg_time:.0f}ms ({avg_time/1000:.2f}s)")
|
||||
print(f" 最快耗时: {min_time:.0f}ms")
|
||||
print(f" 最慢耗时: {max_time:.0f}ms")
|
||||
|
||||
# 缓存效果分析
|
||||
if len(valid_results) >= 2:
|
||||
first_time = valid_results[0]["elapsed_ms"]
|
||||
last_time = valid_results[-1]["elapsed_ms"]
|
||||
|
||||
print(f"\n💡 缓存效果分析:")
|
||||
if last_time < first_time * 0.7:
|
||||
print(f" ✅ 缓存可能生效! 后端请求({first_time:.0f}ms)比首次请求快 {((first_time - last_time) / first_time * 100):.1f}%")
|
||||
elif last_time > first_time * 1.3:
|
||||
print(f" ⚠️ 缓存可能未生效，后续请求比首次请求慢")
|
||||
else:
|
||||
print(f" 🤔 响应时间相近，缓存效果不确定")
|
||||
|
||||
print("📊"*30 + "\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Call ragflow chat endpoint with SSE streaming."
|
||||
description="Call ragflow chat endpoint with SSE streaming and performance testing."
|
||||
)
|
||||
parser.add_argument(
|
||||
"question",
|
||||
@ -153,6 +355,25 @@ if __name__ == "__main__":
|
||||
default=API_URL,
|
||||
help="接口地址，默认取 RAGFLOW_URL。",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-n",
|
||||
"--num-requests",
|
||||
type=int,
|
||||
default=3,
|
||||
help="性能测试的请求次数，默认3次。",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-d",
|
||||
"--delay",
|
||||
type=float,
|
||||
default=2.0,
|
||||
help="请求之间的间隔时间(秒)，默认2秒。",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--single",
|
||||
action="store_true",
|
||||
help="单次请求模式(不进行性能测试)。",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Allow either positional or -q/--question to override env default.
|
||||
@ -161,10 +382,23 @@ if __name__ == "__main__":
|
||||
# Override globals if provided via CLI.
|
||||
API_URL = args.url # type: ignore
|
||||
|
||||
asyncio.run(
|
||||
stream_chat_async(
|
||||
question=question,
|
||||
chat_id=args.chat_id,
|
||||
session_id=args.session_id,
|
||||
if args.single:
|
||||
# 单次请求模式
|
||||
asyncio.run(
|
||||
stream_chat_async(
|
||||
question=question,
|
||||
chat_id=args.chat_id,
|
||||
session_id=args.session_id,
|
||||
)
|
||||
)
|
||||
else:
|
||||
# 性能测试模式
|
||||
asyncio.run(
|
||||
run_performance_test(
|
||||
question=question,
|
||||
chat_id=args.chat_id,
|
||||
session_id=args.session_id,
|
||||
num_requests=args.num_requests,
|
||||
delay_between_requests=args.delay
|
||||
)
|
||||
)
|
||||
)
|
||||
Loading…
Reference in New Issue
Block a user