yangdx commited on
Commit
fda27b8
·
1 Parent(s): ead5d24

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

Browse files

- 修复流式响应中的完成标记逻辑
- 添加非流式调用测试
- 添加流式调用测试
- 优化JSON序列化,支持非ASCII字符
- 确保生成器在完成标记后立即结束

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
- 'model': LIGHTRAG_MODEL,
711
- 'created_at': LIGHTRAG_CREATED_AT,
712
- 'message': {
713
- 'role': 'assistant',
714
- 'content': response
715
  },
716
- 'done': True
717
  }
718
- yield f"data: {json.dumps(data)}\n\n"
719
  else:
 
720
  async for chunk in response:
721
- data = {
722
- "model": LIGHTRAG_MODEL,
723
- "created_at": LIGHTRAG_CREATED_AT,
724
- "message": {
725
- "role": "assistant",
726
- "content": chunk
727
- },
728
- "done": False
729
- }
730
- yield f"data: {json.dumps(data)}\n\n"
 
 
 
731
  data = {
732
  "model": LIGHTRAG_MODEL,
733
  "created_at": LIGHTRAG_CREATED_AT,
734
  "message": {
735
  "role": "assistant",
736
- "content": chunk
737
  },
738
- "done": False
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()