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

为Ollama API返回结果添加图像字段和性能统计信息

Browse files

- 在OllamaMessage中添加images字段
- 响应消息中增加images字段
- 完成标记中添加性能统计信息
- 更新测试用例以处理性能统计
- 移除测试用例中的/naive前缀

lightrag/api/lightrag_ollama.py CHANGED
@@ -231,6 +231,7 @@ class SearchMode(str, Enum):
231
  class OllamaMessage(BaseModel):
232
  role: str
233
  content: str
 
234
 
235
  class OllamaChatRequest(BaseModel):
236
  model: str = LIGHTRAG_MODEL
@@ -712,7 +713,8 @@ def create_app(args):
712
  "created_at": LIGHTRAG_CREATED_AT,
713
  "message": {
714
  "role": "assistant",
715
- "content": response
 
716
  },
717
  "done": True
718
  }
@@ -726,21 +728,24 @@ def create_app(args):
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 # 确保生成器在发送完成标记后立即结束
@@ -777,7 +782,8 @@ def create_app(args):
777
  created_at=LIGHTRAG_CREATED_AT,
778
  message=OllamaMessage(
779
  role="assistant",
780
- content=str(response_text) # 确保转换为字符串
 
781
  ),
782
  done=True
783
  )
 
231
  class OllamaMessage(BaseModel):
232
  role: str
233
  content: str
234
+ images: Optional[List[str]] = None
235
 
236
  class OllamaChatRequest(BaseModel):
237
  model: str = LIGHTRAG_MODEL
 
713
  "created_at": LIGHTRAG_CREATED_AT,
714
  "message": {
715
  "role": "assistant",
716
+ "content": response,
717
+ "images": None
718
  },
719
  "done": True
720
  }
 
728
  "created_at": LIGHTRAG_CREATED_AT,
729
  "message": {
730
  "role": "assistant",
731
+ "content": chunk,
732
+ "images": None
733
  },
734
  "done": False
735
  }
736
  yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
737
 
738
+ # 发送完成标记,包含性能统计信息
739
  data = {
740
  "model": LIGHTRAG_MODEL,
741
  "created_at": LIGHTRAG_CREATED_AT,
742
+ "done": True,
743
+ "total_duration": 0, # 由于我们没有实际统计这些指标,暂时使用默认值
744
+ "load_duration": 0,
745
+ "prompt_eval_count": 0,
746
+ "prompt_eval_duration": 0,
747
+ "eval_count": 0,
748
+ "eval_duration": 0
749
  }
750
  yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
751
  return # 确保生成器在发送完成标记后立即结束
 
782
  created_at=LIGHTRAG_CREATED_AT,
783
  message=OllamaMessage(
784
  role="assistant",
785
+ content=str(response_text), # 确保转换为字符串
786
+ images=None
787
  ),
788
  done=True
789
  )
test_lightrag_ollama_chat.py CHANGED
@@ -35,7 +35,7 @@ def test_stream_chat():
35
  "messages": [
36
  {
37
  "role": "user",
38
- "content": "/naive 孙悟空有什么法力,性格特征是什么"
39
  }
40
  ],
41
  "stream": True
@@ -51,12 +51,16 @@ def test_stream_chat():
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:
 
35
  "messages": [
36
  {
37
  "role": "user",
38
+ "content": "孙悟空有什么法力,性格特征是什么"
39
  }
40
  ],
41
  "stream": True
 
51
  for event in client.events():
52
  try:
53
  data = json.loads(event.data)
54
+ if data.get("done", False): # 如果是完成标记
55
+ if "total_duration" in data: # 最终的性能统计消息
56
+ print("\n=== 性能统计 ===")
57
+ print(json.dumps(data, ensure_ascii=False, indent=2))
58
+ break
59
+ else: # 正常的内容消息
60
+ message = data.get("message", {})
61
+ content = message.get("content", "")
62
+ if content: # 只收集非空内容
63
+ output_buffer.append(content)
64
  except json.JSONDecodeError:
65
  print("Error decoding JSON from SSE event")
66
  finally: