'commit'

dsLightRag/Test/TestGoogleGemini.py

@@ -4,12 +4,14 @@ 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):
@@ -74,6 +76,7 @@ class ModelInteractor:
         except requests.exceptions.RequestException as e:
             print(f"请求发生错误: {e}")
 
+
 # 示例使用
 if __name__ == "__main__":
     # 创建模型交互器实例