gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

优化流式响应处理并添加测试用例

Browse Source 
- 修复流式响应中的完成标记逻辑
- 添加非流式调用测试
- 添加流式调用测试
- 优化JSON序列化，支持非ASCII字符
- 确保生成器在完成标记后立即结束
This commit is contained in:
yangdx
2025-01-15 20:18:17 +08:00
parent f15f97a51d
commit 23f838ec94
2 changed files with 103 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

66
lightrag/api/lightrag_ollama.py
View File

@@ -706,50 +706,44 @@ def create_app(args):
try:
# 确保 response 是异步生成器
if isinstance(response, str):
data = {
'model': LIGHTRAG_MODEL,
'created_at': LIGHTRAG_CREATED_AT,
'message': {
'role': 'assistant',
'content': response
},
'done': True
}
yield f"data: {json.dumps(data)}\n\n"
else:
async for chunk in response:
data = {
"model": LIGHTRAG_MODEL,
"created_at": LIGHTRAG_CREATED_AT,
"message": {
"role": "assistant",
"content": chunk
},
"done": False
}
yield f"data: {json.dumps(data)}\n\n"
# 如果是字符串，作为单个完整响应发送
data = {
"model": LIGHTRAG_MODEL,
"created_at": LIGHTRAG_CREATED_AT,
"message": {
"role": "assistant",
"content": chunk
"content": response
},
"done": False
"done": True
}
yield f"data: {json.dumps(data)}\n\n"
yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
else:
# 流式响应
async for chunk in response:
if chunk: # 只发送非空内容
data = {
"model": LIGHTRAG_MODEL,
"created_at": LIGHTRAG_CREATED_AT,
"message": {
"role": "assistant",
"content": chunk
},
"done": False
}
yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
# 发送完成标记
data = {
"model": LIGHTRAG_MODEL,
"created_at": LIGHTRAG_CREATED_AT,
"message": {
"role": "assistant",
"content": ""
},
"done": True
}
yield f"data: {json.dumps(data)}\n\n"
# 发送完成标记
data = {
"model": LIGHTRAG_MODEL,
"created_at": LIGHTRAG_CREATED_AT,
"message": {
"role": "assistant",
"content": ""
},
"done": True
}
yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
return # 确保生成器在发送完成标记后立即结束
except Exception as e:
logging.error(f"Error in stream_generator: {str(e)}")
raise