diff --git a/lightrag/api/lightrag_ollama.py b/lightrag/api/lightrag_ollama.py index 004c2739..3b92902f 100644 --- a/lightrag/api/lightrag_ollama.py +++ b/lightrag/api/lightrag_ollama.py @@ -231,6 +231,7 @@ class SearchMode(str, Enum): class OllamaMessage(BaseModel): role: str content: str + images: Optional[List[str]] = None class OllamaChatRequest(BaseModel): model: str = LIGHTRAG_MODEL @@ -712,7 +713,8 @@ def create_app(args): "created_at": LIGHTRAG_CREATED_AT, "message": { "role": "assistant", - "content": response + "content": response, + "images": None }, "done": True } @@ -726,21 +728,24 @@ def create_app(args): "created_at": LIGHTRAG_CREATED_AT, "message": { "role": "assistant", - "content": chunk + "content": chunk, + "images": None }, "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 + "done": True, + "total_duration": 0, # 由于我们没有实际统计这些指标,暂时使用默认值 + "load_duration": 0, + "prompt_eval_count": 0, + "prompt_eval_duration": 0, + "eval_count": 0, + "eval_duration": 0 } yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n" return # 确保生成器在发送完成标记后立即结束 @@ -777,7 +782,8 @@ def create_app(args): created_at=LIGHTRAG_CREATED_AT, message=OllamaMessage( role="assistant", - content=str(response_text) # 确保转换为字符串 + content=str(response_text), # 确保转换为字符串 + images=None ), done=True ) diff --git a/test_lightrag_ollama_chat.py b/test_lightrag_ollama_chat.py index 067b8877..b941ee27 100644 --- a/test_lightrag_ollama_chat.py +++ b/test_lightrag_ollama_chat.py @@ -35,7 +35,7 @@ def test_stream_chat(): "messages": [ { "role": "user", - "content": "/naive 孙悟空有什么法力,性格特征是什么" + "content": "孙悟空有什么法力,性格特征是什么" } ], "stream": True @@ -51,12 +51,16 @@ def test_stream_chat(): for event in client.events(): try: data = json.loads(event.data) - message = data.get("message", {}) - content = message.get("content", "") - if content: # 只收集非空内容 - output_buffer.append(content) - if data.get("done", False): # 如果收到完成标记,退出循环 - break + if data.get("done", False): # 如果是完成标记 + if "total_duration" in data: # 最终的性能统计消息 + print("\n=== 性能统计 ===") + print(json.dumps(data, ensure_ascii=False, indent=2)) + break + else: # 正常的内容消息 + message = data.get("message", {}) + content = message.get("content", "") + if content: # 只收集非空内容 + output_buffer.append(content) except json.JSONDecodeError: print("Error decoding JSON from SSE event") finally: