193 lines
6.0 KiB
Python
193 lines
6.0 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
简化后的RAGFlow实现测试脚本
|
||
验证同步Generator在异步环境中的工作
|
||
"""
|
||
|
||
import asyncio
|
||
import time
|
||
from typing import Generator, Any
|
||
|
||
# 模拟同步RAGFlow客户端
|
||
class MockSyncRAGFlowClient:
|
||
"""模拟同步RAGFlow客户端，返回Generator"""
|
||
|
||
def __init__(self, base_url: str, api_key: str):
|
||
self.base_url = base_url
|
||
self.api_key = api_key
|
||
|
||
def converse_with_chat_assistant(self, **kwargs) -> Generator[dict, None, None]:
|
||
"""模拟流式对话，返回同步Generator"""
|
||
print(f"Mock client: 开始生成流式数据...")
|
||
|
||
# 模拟流式数据生成
|
||
responses = [
|
||
{'data': {'answer': 'Hello'}},
|
||
{'data': {'answer': 'Hello, I am'}},
|
||
{'data': {'answer': 'Hello, I am an AI'}},
|
||
{'data': {'answer': 'Hello, I am an AI assistant'}},
|
||
{'data': {'answer': 'Hello, I am an AI assistant.'}},
|
||
]
|
||
|
||
for i, response in enumerate(responses):
|
||
print(f"Mock client: 生成第{i+1}个数据块")
|
||
time.sleep(0.5) # 模拟网络延迟
|
||
yield response
|
||
|
||
print(f"Mock client: 流式数据生成完成")
|
||
|
||
|
||
# 模拟同步RAGFlowService
|
||
class MockRAGFlowService:
|
||
"""模拟同步RAGFlowService"""
|
||
|
||
@staticmethod
|
||
def converse_with_chat_assistant_services(converse_params) -> Generator[dict, None, None]:
|
||
"""返回同步Generator"""
|
||
print("MockService: 调用converse_with_chat_assistant_services")
|
||
client = MockSyncRAGFlowClient("http://localhost:9099", "test_key")
|
||
return client.converse_with_chat_assistant(
|
||
chat_id=converse_params.get('chat_id'),
|
||
question=converse_params.get('question'),
|
||
stream=True,
|
||
session_id=converse_params.get('session_id')
|
||
)
|
||
|
||
|
||
# 模拟异步控制器
|
||
async def async_controller_test():
|
||
"""测试异步控制器中消费同步Generator"""
|
||
|
||
# 模拟参数
|
||
params = {
|
||
'chat_id': 'test_chat_123',
|
||
'question': '你好，请介绍一下自己',
|
||
'stream': True,
|
||
'session_id': 'session_456'
|
||
}
|
||
|
||
print("=" * 60)
|
||
print("测试：异步控制器消费同步Generator")
|
||
print("=" * 60)
|
||
|
||
start_time = time.perf_counter()
|
||
first_token_received = False
|
||
|
||
try:
|
||
# 调用服务层（同步方法）
|
||
print("1. 调用RAGFlowService.converse_with_chat_assistant_services...")
|
||
result = MockRAGFlowService.converse_with_chat_assistant_services(params)
|
||
print(f" 返回类型: {type(result)}")
|
||
|
||
if not isinstance(result, Generator):
|
||
raise TypeError(f"期望Generator类型，但得到 {type(result)}")
|
||
|
||
# 在异步上下文中消费同步Generator
|
||
print("2. 开始消费同步Generator...")
|
||
chunk_count = 0
|
||
|
||
try:
|
||
for chunk in result:
|
||
chunk_count += 1
|
||
print(f" 接收到第{chunk_count}个数据块: {chunk}")
|
||
|
||
# 检查第一个token延迟
|
||
if not first_token_received:
|
||
first_token_received = True
|
||
latency = time.perf_counter() - start_time
|
||
print(f" 首Token延迟: {latency:.3f}s")
|
||
|
||
# 模拟处理每个chunk
|
||
await asyncio.sleep(0.01) # 让出控制权
|
||
|
||
except Exception as e:
|
||
print(f" 消费数据时出错: {e}")
|
||
raise
|
||
|
||
total_time = time.perf_counter() - start_time
|
||
print(f"3. 流式处理完成，总耗时: {total_time:.3f}s")
|
||
print(f" 总共接收数据块: {chunk_count}")
|
||
|
||
# 验证结果
|
||
if chunk_count == 5:
|
||
print("✅ 测试通过：成功接收到所有5个数据块")
|
||
else:
|
||
print(f"❌ 测试失败：期望5个数据块，实际收到{chunk_count}个")
|
||
|
||
except Exception as e:
|
||
print(f"❌ 测试失败：{e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
|
||
# 测试同步消费 vs 异步消费
|
||
def sync_vs_async_test():
|
||
"""测试同步消费和异步消费的差异"""
|
||
|
||
print("\n" + "=" * 60)
|
||
print("测试：同步消费 vs 异步消费")
|
||
print("=" * 60)
|
||
|
||
# 创建同步Generator
|
||
def sync_generator():
|
||
for i in range(5):
|
||
time.sleep(0.1)
|
||
yield f"数据块 {i+1}"
|
||
|
||
generator = sync_generator()
|
||
|
||
# 1. 同步消费
|
||
print("1. 同步消费测试:")
|
||
start_time = time.perf_counter()
|
||
for item in generator:
|
||
print(f" {item}")
|
||
sync_time = time.perf_counter() - start_time
|
||
print(f" 同步消费耗时: {sync_time:.3f}s")
|
||
|
||
# 2. 异步消费
|
||
print("\n2. 异步消费测试:")
|
||
generator2 = sync_generator()
|
||
start_time = time.perf_counter()
|
||
|
||
async def async_consumer(gen):
|
||
count = 0
|
||
for item in gen:
|
||
count += 1
|
||
print(f" {item}")
|
||
await asyncio.sleep(0.01) # 让出控制权
|
||
return count
|
||
|
||
async def run_async_test():
|
||
return await async_consumer(generator2)
|
||
|
||
try:
|
||
count = asyncio.run(run_async_test())
|
||
async_time = time.perf_counter() - start_time
|
||
print(f" 异步消费耗时: {async_time:.3f}s")
|
||
print(f" 处理了{count}个项目")
|
||
except Exception as e:
|
||
print(f" 异步消费失败: {e}")
|
||
|
||
|
||
async def main():
|
||
"""主测试函数"""
|
||
print("RAGFlow简化实现测试")
|
||
print("=" * 60)
|
||
|
||
# 运行主要测试
|
||
await async_controller_test()
|
||
|
||
# 运行对比测试
|
||
sync_vs_async_test()
|
||
|
||
print("\n" + "=" * 60)
|
||
print("测试总结:")
|
||
print("1. 同步Generator可以在异步环境中正常工作")
|
||
print("2. 使用for循环可以自动处理同步Generator")
|
||
print("3. 异步消费需要适当让出控制权(await)")
|
||
print("4. 简化后的架构避免了复杂的async/await链")
|
||
print("=" * 60)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main()) |