From 350e080ec139e881c4359b53247097bbaf4a7edf Mon Sep 17 00:00:00 2001
From: yangdx <yangdx@znipower.com>
Date: Thu, 16 Jan 2025 01:11:59 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E6=B5=8B=E8=AF=95=E8=84=9A?=
 =?UTF-8?q?=E6=9C=AC=E9=85=8D=E7=BD=AE=E5=92=8C=E5=91=BD=E4=BB=A4=E8=A1=8C?=
 =?UTF-8?q?=E5=8F=82=E6=95=B0=E5=A4=84=E7=90=86?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 移除冗余的stream_query配置
- 统一使用query作为测试查询内容
- 新增--ask参数覆盖查询内容
- 调整命令行参数顺序
---
 test_lightrag_ollama_chat.py | 32 ++++++++++++++++++++------------
 1 file changed, 20 insertions(+), 12 deletions(-)

diff --git a/test_lightrag_ollama_chat.py b/test_lightrag_ollama_chat.py
index 02b51b22..e0ed5342 100644
--- a/test_lightrag_ollama_chat.py
+++ b/test_lightrag_ollama_chat.py
@@ -118,8 +118,7 @@ DEFAULT_CONFIG = {
     },
     "test_cases": {
         "basic": {
-            "query": "孙悟空",
-            "stream_query": "孙悟空"
+            "query": "唐僧有几个徒弟"
         }
     }
 }
@@ -292,7 +291,7 @@ def test_stream_chat():
     """
     url = get_base_url()
     data = create_request_data(
-        CONFIG["test_cases"]["basic"]["stream_query"],
+        CONFIG["test_cases"]["basic"]["query"],
         stream=True
     )
     
@@ -344,7 +343,7 @@ def test_query_modes():
         if OutputControl.is_verbose():
             print(f"\n=== 测试 /{mode} 模式 ===")
         data = create_request_data(
-            f"/{mode} 孙悟空的特点",
+            f"/{mode} {CONFIG['test_cases']['basic']['query']}",
             stream=False
         )
         
@@ -534,11 +533,14 @@ def parse_args() -> argparse.Namespace:
 """
     )
     parser.add_argument(
-        "--tests",
-        nargs="+",
-        choices=list(get_test_cases().keys()) + ["all"],
-        default=["all"],
-        help="要运行的测试用例，可选: %(choices)s。使用 all 运行所有测试"
+        "-q", "--quiet",
+        action="store_true",
+        help="静默模式，只显示测试结果摘要"
+    )
+    parser.add_argument(
+        "-a", "--ask",
+        type=str,
+        help="指定查询内容，会覆盖配置文件中的查询设置"
     )
     parser.add_argument(
         "--init-config",
@@ -552,9 +554,11 @@ def parse_args() -> argparse.Namespace:
         help="测试结果输出文件路径"
     )
     parser.add_argument(
-        "-q", "--quiet",
-        action="store_true",
-        help="静默模式，只显示测试结果摘要"
+        "--tests",
+        nargs="+",
+        choices=list(get_test_cases().keys()) + ["all"],
+        default=["all"],
+        help="要运行的测试用例，可选: %(choices)s。使用 all 运行所有测试"
     )
     return parser.parse_args()
 
@@ -564,6 +568,10 @@ if __name__ == "__main__":
     # 设置输出模式
     OutputControl.set_verbose(not args.quiet)
     
+    # 如果指定了查询内容，更新配置
+    if args.ask:
+        CONFIG["test_cases"]["basic"]["query"] = args.ask
+    
     # 如果指定了创建配置文件
     if args.init_config:
         create_default_config()