dsProject/dsLightRag/Test/TestGoogleGemini.py

import json

import requests

from Config.Config import GPTNB_API_KEY

class ModelInteractor:
    def __init__(self):
        self.api_url = "https://goapi.gptnb.ai/v1/chat/completions",
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {GPTNB_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}")

# 示例使用
if __name__ == "__main__":
    # 创建模型交互器实例
    interactor = ModelInteractor()
    
    # 使用不同的模型和提示词
    model_name = "gemini-2.5-pro"
    prompt_text = "请详细介绍一下你自己，分成几个段落来说明"
    
    # 发送流式请求
    interactor.stream_request(model_name, prompt_text)