102 lines
3.5 KiB
Python
102 lines
3.5 KiB
Python
"""
|
||
搜索服务测试 - 直接调用 SearchService
|
||
"""
|
||
import asyncio
|
||
import sys
|
||
import os
|
||
import json
|
||
import time
|
||
|
||
# 添加项目路径
|
||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||
|
||
from module_admin.service.search_service import SearchService
|
||
from module_admin.entity.vo.ragflow_vo import ConverseWithChatAssistantModel
|
||
from fastapi.responses import StreamingResponse
|
||
|
||
async def test_search_service():
|
||
"""测试 SearchService.handle_search_chat"""
|
||
|
||
print("=" * 60)
|
||
print("搜索服务测试 (调用 SearchService)")
|
||
print("=" * 60)
|
||
|
||
query = "今天有什么新闻"
|
||
print(f"\n测试查询: {query}")
|
||
print("-" * 60)
|
||
|
||
# 构造请求参数
|
||
converse_params = ConverseWithChatAssistantModel(
|
||
chat_id="test-chat",
|
||
question=query,
|
||
stream=True
|
||
)
|
||
|
||
print("调用 handle_search_chat (stream=True)...")
|
||
try:
|
||
start_time = time.time()
|
||
last_time = start_time
|
||
|
||
response = await SearchService.handle_search_chat(converse_params, redis=None)
|
||
|
||
if isinstance(response, StreamingResponse):
|
||
print("\n✓ 获取到 StreamingResponse! 开始接收流式数据...")
|
||
print("\n【流式输出】")
|
||
|
||
async for chunk in response.body_iterator:
|
||
current_time = time.time()
|
||
delay = current_time - last_time
|
||
last_time = current_time
|
||
|
||
# chunk 是 bytes，格式为 "data: {...}\n\n"
|
||
chunk_str = chunk.decode('utf-8') if isinstance(chunk, bytes) else chunk
|
||
|
||
if not chunk_str.strip():
|
||
continue
|
||
|
||
print(f"[{current_time - start_time:.3f}s] (+{delay:.3f}s) RAW: {chunk_str.strip()}")
|
||
|
||
# 尝试解析并显示内容
|
||
for line in chunk_str.split('\n'):
|
||
if line.startswith('data:'):
|
||
data_str = line[5:].strip()
|
||
try:
|
||
# 忽略控制消息
|
||
if data_str in ('[DONE]'):
|
||
continue
|
||
|
||
data_json = json.loads(data_str)
|
||
# 我们的 format_sse 返回的是 {'data': content}
|
||
content = data_json.get('data', '')
|
||
|
||
# 如果是 status 消息
|
||
if 'status' in data_json:
|
||
print(f"\n[Status: {data_json['status']}]")
|
||
continue
|
||
|
||
if content:
|
||
# print(f"Content: {content}")
|
||
pass
|
||
|
||
except:
|
||
pass
|
||
|
||
print("\n✓ 流式接收完成!")
|
||
else:
|
||
print(f"\n✗ 预期返回 StreamingResponse，实际返回: {type(response)}")
|
||
# 打印非流式结果
|
||
if hasattr(response, 'body'):
|
||
print(response.body)
|
||
|
||
except Exception as e:
|
||
print(f"\n✗ 发生错误: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
print("\n" + "=" * 60)
|
||
print("测试完成")
|
||
print("=" * 60)
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(test_search_service())
|