'commit'

dsLightRag/Test/TestGoogleGemini.py

@@ -4,14 +4,16 @@ import requests
 
 from Config.Config import GPTNB_API_KEY
 
+
 class ModelInteractor:
-    def __init__(self):
-        self.api_url = "https://goapi.gptnb.ai/v1/chat/completions",
+    def __init__(self, api_key=GPTNB_API_KEY, api_url="https://goapi.gptnb.ai/v1/chat/completions"):
+        self.api_key = api_key
+        self.api_url = api_url
         self.headers = {
             "Content-Type": "application/json",
-            "Authorization": f"Bearer {GPTNB_API_KEY}"
+            "Authorization": f"Bearer {self.api_key}"
         }
-    
+
     def stream_request(self, model, prompt, temperature=0.7):
         """
         发送流式请求到模型API
@@ -33,7 +35,7 @@ class ModelInteractor:
             "temperature": temperature,
             "stream": True
         }
-        
+
         try:
             response = requests.post(
                 self.api_url,
@@ -43,24 +45,24 @@ class ModelInteractor:
                 timeout=30
             )
             response.raise_for_status()
-            
+
             print(f"使用模型 {model} 的流式响应内容: ")
             for chunk in response.iter_content(chunk_size=None):
                 if chunk:
                     chunk_data = chunk.decode('utf-8', errors='replace')
-                    
+
                     for line in chunk_data.splitlines():
                         line = line.strip()
                         if not line:
                             continue
-                            
+
                         if line == 'data: [DONE]':
                             print("\n流式响应结束")
                             return
-                            
+
                         if line.startswith('data: '):
                             line = line[6:]
-                            
+
                         try:
                             data = json.loads(line)
                             if 'choices' in data and len(data['choices']) > 0:
@@ -70,18 +72,19 @@ class ModelInteractor:
                                     print(content, end='', flush=True)
                         except json.JSONDecodeError as e:
                             print(f"[调试] JSON解析错误: {e}, 内容: {line}")
-                            
+
         except requests.exceptions.RequestException as e:
             print(f"请求发生错误: {e}")
 
+
 # 示例使用
 if __name__ == "__main__":
     # 创建模型交互器实例
     interactor = ModelInteractor()
-    
+
     # 使用不同的模型和提示词
     model_name = "gemini-2.5-pro"
     prompt_text = "请详细介绍一下你自己，分成几个段落来说明"
-    
+
     # 发送流式请求
     interactor.stream_request(model_name, prompt_text)