diff --git a/dsLightRag/Test/TestGoogleGemini.py b/dsLightRag/Test/TestGoogleGemini.py index 6034300c..853ea2fa 100644 --- a/dsLightRag/Test/TestGoogleGemini.py +++ b/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)