From 9632a8f0dc539f0b733d93a5c939e7a4a484d175 Mon Sep 17 00:00:00 2001
From: yangdx <yangdx@znipower.com>
Date: Wed, 15 Jan 2025 23:09:50 +0800
Subject: [PATCH] =?UTF-8?q?=E8=A7=A3=E5=86=B3=E6=9F=A5=E8=AF=A2=E5=91=BD?=
 =?UTF-8?q?=E4=B8=AD=E7=BC=93=E5=AD=98=E6=97=B6=E6=B5=81=E5=BC=8F=E5=93=8D?=
 =?UTF-8?q?=E5=BA=94=E6=9C=AA=E9=81=B5=E5=BE=AAOllma=E8=A7=84=E8=8C=83?=
 =?UTF-8?q?=E7=9A=84=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- rag返回结果未字符串时，响应分两次发送
- 第一次发送查询内容
- 第二次发送统计信息
---
 lightrag/api/lightrag_ollama.py | 21 ++++++++++++++++++---
 test_lightrag_ollama_chat.py    |  2 +-
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/lightrag/api/lightrag_ollama.py b/lightrag/api/lightrag_ollama.py
index 714731eb..5768fc42 100644
--- a/lightrag/api/lightrag_ollama.py
+++ b/lightrag/api/lightrag_ollama.py
@@ -709,16 +709,31 @@ def create_app(args):
                     try:
                         # 确保 response 是异步生成器
                         if isinstance(response, str):
-                            # 如果是字符串，作为单个完整响应发送
+                            # 如果是字符串,分两次发送
+                            # 第一次发送查询内容
                             data = {
                                 "model": LIGHTRAG_MODEL,
                                 "created_at": LIGHTRAG_CREATED_AT,
                                 "message": {
-                                    "role": "assistant",
+                                    "role": "assistant", 
                                     "content": response,
                                     "images": None
                                 },
-                                "done": True
+                                "done": False
+                            }
+                            yield f"{json.dumps(data, ensure_ascii=False)}\n"
+                            
+                            # 第二次发送统计信息
+                            data = {
+                                "model": LIGHTRAG_MODEL,
+                                "created_at": LIGHTRAG_CREATED_AT,
+                                "done": True,
+                                "total_duration": 1,
+                                "load_duration": 1,
+                                "prompt_eval_count": 999,
+                                "prompt_eval_duration": 1,
+                                "eval_count": 999,
+                                "eval_duration": 1
                             }
                             yield f"{json.dumps(data, ensure_ascii=False)}\n"
                         else:
diff --git a/test_lightrag_ollama_chat.py b/test_lightrag_ollama_chat.py
index b9902e77..02b51b22 100644
--- a/test_lightrag_ollama_chat.py
+++ b/test_lightrag_ollama_chat.py
@@ -119,7 +119,7 @@ DEFAULT_CONFIG = {
     "test_cases": {
         "basic": {
             "query": "孙悟空",
-            "stream_query": "孙悟空有什么法力，性格特征是什么"
+            "stream_query": "孙悟空"
         }
     }
 }