yangdx
commited on
Commit
·
fda27b8
1
Parent(s):
ead5d24
优化流式响应处理并添加测试用例
Browse files- 修复流式响应中的完成标记逻辑
- 添加非流式调用测试
- 添加流式调用测试
- 优化JSON序列化,支持非ASCII字符
- 确保生成器在完成标记后立即结束
- lightrag/api/lightrag_ollama.py +26 -32
- test_lightrag_ollama_chat.py +73 -0
lightrag/api/lightrag_ollama.py
CHANGED
@@ -706,50 +706,44 @@ def create_app(args):
|
|
706 |
try:
|
707 |
# 确保 response 是异步生成器
|
708 |
if isinstance(response, str):
|
|
|
709 |
data = {
|
710 |
-
|
711 |
-
|
712 |
-
|
713 |
-
|
714 |
-
|
715 |
},
|
716 |
-
|
717 |
}
|
718 |
-
yield f"data: {json.dumps(data)}\n\n"
|
719 |
else:
|
|
|
720 |
async for chunk in response:
|
721 |
-
|
722 |
-
|
723 |
-
|
724 |
-
|
725 |
-
"
|
726 |
-
|
727 |
-
|
728 |
-
|
729 |
-
|
730 |
-
|
|
|
|
|
|
|
731 |
data = {
|
732 |
"model": LIGHTRAG_MODEL,
|
733 |
"created_at": LIGHTRAG_CREATED_AT,
|
734 |
"message": {
|
735 |
"role": "assistant",
|
736 |
-
"content":
|
737 |
},
|
738 |
-
"done":
|
739 |
}
|
740 |
-
yield f"data: {json.dumps(data)}\n\n"
|
741 |
-
|
742 |
-
# 发送完成标记
|
743 |
-
data = {
|
744 |
-
"model": LIGHTRAG_MODEL,
|
745 |
-
"created_at": LIGHTRAG_CREATED_AT,
|
746 |
-
"message": {
|
747 |
-
"role": "assistant",
|
748 |
-
"content": ""
|
749 |
-
},
|
750 |
-
"done": True
|
751 |
-
}
|
752 |
-
yield f"data: {json.dumps(data)}\n\n"
|
753 |
except Exception as e:
|
754 |
logging.error(f"Error in stream_generator: {str(e)}")
|
755 |
raise
|
|
|
706 |
try:
|
707 |
# 确保 response 是异步生成器
|
708 |
if isinstance(response, str):
|
709 |
+
# 如果是字符串,作为单个完整响应发送
|
710 |
data = {
|
711 |
+
"model": LIGHTRAG_MODEL,
|
712 |
+
"created_at": LIGHTRAG_CREATED_AT,
|
713 |
+
"message": {
|
714 |
+
"role": "assistant",
|
715 |
+
"content": response
|
716 |
},
|
717 |
+
"done": True
|
718 |
}
|
719 |
+
yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
|
720 |
else:
|
721 |
+
# 流式响应
|
722 |
async for chunk in response:
|
723 |
+
if chunk: # 只发送非空内容
|
724 |
+
data = {
|
725 |
+
"model": LIGHTRAG_MODEL,
|
726 |
+
"created_at": LIGHTRAG_CREATED_AT,
|
727 |
+
"message": {
|
728 |
+
"role": "assistant",
|
729 |
+
"content": chunk
|
730 |
+
},
|
731 |
+
"done": False
|
732 |
+
}
|
733 |
+
yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
|
734 |
+
|
735 |
+
# 发送完成标记
|
736 |
data = {
|
737 |
"model": LIGHTRAG_MODEL,
|
738 |
"created_at": LIGHTRAG_CREATED_AT,
|
739 |
"message": {
|
740 |
"role": "assistant",
|
741 |
+
"content": ""
|
742 |
},
|
743 |
+
"done": True
|
744 |
}
|
745 |
+
yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
|
746 |
+
return # 确保生成器在发送完成标记后立即结束
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
747 |
except Exception as e:
|
748 |
logging.error(f"Error in stream_generator: {str(e)}")
|
749 |
raise
|
test_lightrag_ollama_chat.py
ADDED
@@ -0,0 +1,73 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import requests
|
2 |
+
import json
|
3 |
+
import sseclient
|
4 |
+
|
5 |
+
def test_non_stream_chat():
|
6 |
+
"""测试非流式调用 /api/chat 接口"""
|
7 |
+
url = "http://localhost:9621/api/chat"
|
8 |
+
|
9 |
+
# 构造请求数据
|
10 |
+
data = {
|
11 |
+
"model": "lightrag:latest",
|
12 |
+
"messages": [
|
13 |
+
{
|
14 |
+
"role": "user",
|
15 |
+
"content": "孙悟空"
|
16 |
+
}
|
17 |
+
],
|
18 |
+
"stream": False
|
19 |
+
}
|
20 |
+
|
21 |
+
# 发送请求
|
22 |
+
response = requests.post(url, json=data)
|
23 |
+
|
24 |
+
# 打印响应
|
25 |
+
print("\n=== 非流式调用响应 ===")
|
26 |
+
print(json.dumps(response.json(), ensure_ascii=False, indent=2))
|
27 |
+
|
28 |
+
def test_stream_chat():
|
29 |
+
"""测试流式调用 /api/chat 接口"""
|
30 |
+
url = "http://localhost:9621/api/chat"
|
31 |
+
|
32 |
+
# 构造请求数据
|
33 |
+
data = {
|
34 |
+
"model": "lightrag:latest",
|
35 |
+
"messages": [
|
36 |
+
{
|
37 |
+
"role": "user",
|
38 |
+
"content": "/naive 孙悟空有什么法力,性格特征是什么"
|
39 |
+
}
|
40 |
+
],
|
41 |
+
"stream": True
|
42 |
+
}
|
43 |
+
|
44 |
+
# 发送请求并获取 SSE 流
|
45 |
+
response = requests.post(url, json=data, stream=True)
|
46 |
+
client = sseclient.SSEClient(response)
|
47 |
+
|
48 |
+
print("\n=== 流式调用响应 ===")
|
49 |
+
output_buffer = []
|
50 |
+
try:
|
51 |
+
for event in client.events():
|
52 |
+
try:
|
53 |
+
data = json.loads(event.data)
|
54 |
+
message = data.get("message", {})
|
55 |
+
content = message.get("content", "")
|
56 |
+
if content: # 只收集非空内容
|
57 |
+
output_buffer.append(content)
|
58 |
+
if data.get("done", False): # 如果收到完成标记,退出循环
|
59 |
+
break
|
60 |
+
except json.JSONDecodeError:
|
61 |
+
print("Error decoding JSON from SSE event")
|
62 |
+
finally:
|
63 |
+
response.close() # 确保关闭响应连接
|
64 |
+
|
65 |
+
# 一次性打印所有收集到的内容
|
66 |
+
print("".join(output_buffer))
|
67 |
+
|
68 |
+
if __name__ == "__main__":
|
69 |
+
# 先测试非流式调用
|
70 |
+
test_non_stream_chat()
|
71 |
+
|
72 |
+
# 再测试流式调用
|
73 |
+
test_stream_chat()
|