'commit'

dsLightRag/Test/TestGoogleGemini.py

@@ -1,5 +1,6 @@
 import requests
 import json
+import sys
 
 # API配置
 API_URL = "https://goapi.gptnb.ai/v1/chat/completions"
@@ -11,30 +12,67 @@ headers = {
     "Authorization": f"Bearer {API_KEY}"
 }
 
-# 请求体
+# 请求体 - 添加stream: true参数启用流式响应
 payload = {
     "model": "gemini-2.5-pro",
-    "messages": [{"role": "user", "content": "你是谁?"}],
-    "temperature": 0.7
+    "messages": [{
+        "role": "user",
+        "content": "请详细介绍一下你自己，分成几个段落来说明"
+    }],
+    "temperature": 0.7,
+    "stream": True  # 启用流式响应
 }
 
 try:
-    # 发送POST请求
-    response = requests.post(API_URL, headers=headers, data=json.dumps(payload))
+    # 发送POST请求，设置stream=True
+    response = requests.post(
+        API_URL,
+        headers=headers,
+        data=json.dumps(payload),
+        stream=True,
+        timeout=30  # 设置超时时间
+    )
     response.raise_for_status()  # 检查请求是否成功
 
-    # 解析响应
-    result = response.json()
+    print("流式响应内容:\n")
+    # 逐块处理响应
+    for chunk in response.iter_content(chunk_size=None):
+        if chunk:
+            # 解码chunk
+            chunk_data = chunk.decode('utf-8', errors='replace')
+            print(f"[调试] 原始chunk: {chunk_data}", file=sys.stderr)
 
-    # 提取并打印返回的文本内容
-    try:
-        # 从响应中获取文本内容
-        text_content = result['choices'][0]['message']['content']
-        print("\n提取的文本内容:")
-        print(text_content)
-    except (KeyError, IndexError) as e:
-        print(f"\n提取文本内容时出错: {e}")
-        print("响应结构可能与预期不符。")
+            # 处理可能的多部分响应
+            for line in chunk_data.splitlines():
+                line = line.strip()
+                if not line:
+                    continue
+
+                # 检查是否结束
+                if line == '[DONE]':
+                    print("\n[调试] 接收到结束信号", file=sys.stderr)
+                    break
+
+                # 去除可能的前缀
+                if line.startswith('data: '):
+                    line = line[6:]
+
+                try:
+                    # 解析JSON
+                    data = json.loads(line)
+                    # 提取文本内容
+                    if 'choices' in data and len(data['choices']) > 0:
+                        delta = data['choices'][0].get('delta', {})
+                        content = delta.get('content', '')
+                        if content:
+                            # 实时输出内容，不换行
+                            print(content, end='', flush=True)
+                except json.JSONDecodeError as e:
+                    print(f"[调试] JSON解析错误: {e}, 内容: {line}", file=sys.stderr)
+                except Exception as e:
+                    print(f"[调试] 处理错误: {e}", file=sys.stderr)
+
+    print("\n\n流式响应结束")
 
 except requests.exceptions.RequestException as e:
     print(f"请求发生错误: {e}")