diff --git a/dsLightRag/Test/TestGoogleGemini.py b/dsLightRag/Test/TestGoogleGemini.py index 853ea2fa..a69f587f 100644 --- a/dsLightRag/Test/TestGoogleGemini.py +++ b/dsLightRag/Test/TestGoogleGemini.py @@ -1,81 +1,4 @@ -import json - -import requests - -from Config.Config import GPTNB_API_KEY - - -class ModelInteractor: - 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 {self.api_key}" - } - - def stream_request(self, model, prompt, temperature=0.7): - """ - 发送流式请求到模型API - - 参数: - - model: 模型名称 - - prompt: 用户提示词 - - temperature: 温度参数,控制输出的随机性 - - 返回: - - 无返回值,直接打印流式响应 - """ - payload = { - "model": model, - "messages": [{ - "role": "user", - "content": prompt - }], - "temperature": temperature, - "stream": True - } - - try: - response = requests.post( - self.api_url, - headers=self.headers, - data=json.dumps(payload), - stream=True, - 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: - delta = data['choices'][0].get('delta', {}) - content = delta.get('content', '') - if content and content != '\n': - print(content, end='', flush=True) - except json.JSONDecodeError as e: - print(f"[调试] JSON解析错误: {e}, 内容: {line}") - - except requests.exceptions.RequestException as e: - print(f"请求发生错误: {e}") - +from Util.GoApiUtil import ModelInteractor # 示例使用 if __name__ == "__main__": diff --git a/dsLightRag/Util/GoApiUtil.py b/dsLightRag/Util/GoApiUtil.py new file mode 100644 index 00000000..eb99ba37 --- /dev/null +++ b/dsLightRag/Util/GoApiUtil.py @@ -0,0 +1,78 @@ +import json + +import requests + +from Config.Config import GPTNB_API_KEY + + +class ModelInteractor: + 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 {self.api_key}" + } + + def stream_request(self, model, prompt, temperature=0.7): + """ + 发送流式请求到模型API + + 参数: + - model: 模型名称 + - prompt: 用户提示词 + - temperature: 温度参数,控制输出的随机性 + + 返回: + - 无返回值,直接打印流式响应 + """ + payload = { + "model": model, + "messages": [{ + "role": "user", + "content": prompt + }], + "temperature": temperature, + "stream": True + } + + try: + response = requests.post( + self.api_url, + headers=self.headers, + data=json.dumps(payload), + stream=True, + 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: + delta = data['choices'][0].get('delta', {}) + content = delta.get('content', '') + if content and content != '\n': + print(content, end='', flush=True) + except json.JSONDecodeError as e: + print(f"[调试] JSON解析错误: {e}, 内容: {line}") + + except requests.exceptions.RequestException as e: + print(f"请求发生错误: {e}") + diff --git a/dsLightRag/Util/__pycache__/GoApiUtil.cpython-310.pyc b/dsLightRag/Util/__pycache__/GoApiUtil.cpython-310.pyc new file mode 100644 index 00000000..54f36ab9 Binary files /dev/null and b/dsLightRag/Util/__pycache__/GoApiUtil.cpython-310.pyc differ