yangdx
commited on
Commit
·
d5d330d
1
Parent(s):
fda27b8
为Ollama API返回结果添加图像字段和性能统计信息
Browse files- 在OllamaMessage中添加images字段
- 响应消息中增加images字段
- 完成标记中添加性能统计信息
- 更新测试用例以处理性能统计
- 移除测试用例中的/naive前缀
- lightrag/api/lightrag_ollama.py +15 -9
- test_lightrag_ollama_chat.py +11 -7
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 |
-
"
|
740 |
-
|
741 |
-
|
742 |
-
|
743 |
-
"
|
|
|
|
|
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": "
|
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 |
-
|
55 |
-
|
56 |
-
|
57 |
-
|
58 |
-
|
59 |
-
|
|
|
|
|
|
|
|
|
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:
|