CodeBridge/llm.py at main · SMAT-Lab/CodeBridge · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
from openai import OpenAI
import os
from dotenv import load_dotenv, find_dotenv
import time
import logging

def get_openai_answer(ques, model_name="gpt-4o-mini", system_prompt=None, max_retries=5, retry_delay=1):
    _ = load_dotenv(find_dotenv())  # read local .env file
    api_key = os.environ['OPENAI_API_KEY']
    api_base = os.environ['OPENAI_API_BASE']

    client = OpenAI(api_key=api_key, base_url=api_base)

    messages = [{"role": "user", "content": ques}]

    if system_prompt:
        messages.insert(0, {"role": "system", "content": system_prompt})

    for attempt in range(max_retries):
        try:
            response = client.chat.completions.create(
                model=model_name,
                messages=messages
            )
            return response.choices[0].message.content
        except Exception as e:
            if attempt == max_retries - 1:
                logging.error(f"最大重试次数已达到")
                raise
            logging.warning(f"第{attempt + 1}次调用失败: {str(e)}, {retry_delay}秒后重试...")
            time.sleep(retry_delay)

def get_ollama_answer(ques, model_name, system_prompt=None, max_retries=3, retry_delay=1, **kwargs):
    _ = load_dotenv(find_dotenv())
    api_key = os.environ['OLLAMA_API_KEY']
    api_base = os.environ["OLLAMA_API_BASE"]

    client = OpenAI(api_key=api_key, base_url=api_base)

    messages = [{"role": "user", "content": ques}]

    if system_prompt:
        messages.insert(0, {"role": "system", "content": system_prompt})

    # 设置默认参数
    default_params = {
        "model": model_name,
        "messages": messages
    }

    # 使用传入的kwargs更新默认参数
    default_params.update(kwargs)

    for attempt in range(max_retries):
        try:
            response = client.chat.completions.create(**default_params)
            return response.choices[0].message.content
        except Exception as e:
            if attempt == max_retries - 1:
                logging.error(f"最大重试次数已达到,最终错误: {str(e)}")
                raise
            logging.warning(f"第{attempt + 1}次调用失败: {str(e)}, {retry_delay}秒后重试...")
            time.sleep(retry_delay)

def get_gptgod_answer(ques, model_name, system_prompt=None, max_retries=3, retry_delay=1):
    _ = load_dotenv(find_dotenv())
    api_key = os.environ['GPTGOD_API_KEY']
    api_base = os.environ["GPTGOD_API_BASE"]

    client = OpenAI(api_key=api_key, base_url=api_base)

    messages = [{"role": "user", "content": ques}]

    if system_prompt:
        messages.insert(0, {"role": "system", "content": system_prompt})

    for attempt in range(max_retries):
        try:
            response = client.chat.completions.create(
                model=model_name,
                messages=messages
            )
            return response.choices[0].message.content
        except Exception as e:
            if attempt == max_retries - 1:
                logging.error(f"最大重试次数已达到,最终错误: {str(e)}")
                raise
            logging.warning(f"第{attempt + 1}次调用失败: {str(e)}, {retry_delay}秒后重试...")
            time.sleep(retry_delay)