feat(): Support multiple large model APIs

Author: Robbings

config/config.py

@@ -2,16 +2,18 @@
 # model list
 model_quester_anster = "text-davinci-003"
 model_gpt_35_turbo = "gpt-3.5-turbo"
+model_gpt_4o = "gpt-4o"
 model_programming_translate = "code-davinci-002"
 
-# gpt key
-openai_api_key = "your openai key"
+# api 接口封装类
+llm_api_impl = "llm_api_default.LLMApiDefault"
 
-# openai api
-openai_baseurl = "https://api.openai.com/v1"
-
-# gpt model
-openai_model_name = model_gpt_35_turbo
+# api 配置，参考LiteLLM文档：https://docs.litellm.ai/docs
+api_config = {
+    "OPENAI_API_KEY": "Your OpenAI API Key",
+    "OPENAI_API_BASE": "https://api.openai.com/v1",
+    "MODEL_NAME": model_gpt_4o,
+}
 
 # 2. 提示词
 gpt_message = """

llm_api/llm_api_default.py

@@ -0,0 +1,57 @@
+import os
+from math import trunc
+
+from litellm import completion
+from config.config import api_config as out_config
+from llm_api.llm_api_interface import LLMApiInterface
+from llm_api.load_api import create_llm_api_instance
+
+
+class LLMApiDefault(LLMApiInterface):
+
+    def __init__(self):
+        self.model_name = None
+        self.response = None
+
+    def set_config(self, api_config: dict) -> bool:
+        if api_config is None:
+            raise ValueError("api_config is None")
+        for key in api_config:
+            if key == "MODEL_NAME":
+                self.model_name = api_config[key]
+                continue
+            os.environ[key] = api_config[key]
+        return True
+
+    def generate_text(self, messages: list) -> bool:
+
+        self.response = completion(
+            model=self.model_name,
+            messages=messages,
+        )
+        return True
+
+    def get_respond_content(self) -> str:
+        return self.response['choices'][0]['message']['content']
+
+    def get_respond_tokens(self) -> int:
+        return trunc(int(self.response['usage']['total_tokens']))
+
+
+
+
+# 示例使用
+if __name__ == "__main__":
+    api = create_llm_api_instance()
+    api.set_config(out_config)
+    api.generate_text([
+        {"role": "system",
+         "content": "你是一位作家"
+         },
+        {"role": "user",
+         "content": "请写一首抒情的诗",
+         }
+    ])
+    print(api.get_respond_content())
+    print(api.get_respond_tokens())
+

llm_api/llm_api_interface.py

@@ -0,0 +1,23 @@
+from abc import ABC, abstractmethod
+
+class LLMApiInterface(ABC):
+
+    @abstractmethod
+    def set_config(self, api_config: dict) -> bool:
+        """设置模型配置"""
+        pass
+
+    @abstractmethod
+    def generate_text(self, messages: str) -> bool:
+        """根据提示生成文本"""
+        pass
+
+    @abstractmethod
+    def get_respond_content(self) -> str:
+        """获取模型返回内容"""
+        pass
+
+    @abstractmethod
+    def get_respond_tokens(self) -> int:
+        """获取模型返回token数"""
+        pass

llm_api/load_api.py

@@ -0,0 +1,14 @@
+import importlib
+
+from config.config import llm_api_impl
+
+def get_llm_api_class():
+    module_name, class_name = llm_api_impl.rsplit('.', 1)
+    module = importlib.import_module(module_name)
+    cls = getattr(module, class_name)
+    return cls
+
+# 使用工厂函数获取类实例
+def create_llm_api_instance():
+    cls = get_llm_api_class()
+    return cls()

requirements.txt

@@ -20,13 +20,14 @@ importlib-metadata==6.8.0
 itsdangerous==2.1.2
 javalang==0.13.0
 Jinja2==3.1.2
+litellm==1.58.2
 markdown-it-py==3.0.0
 MarkupSafe==2.1.3
 mdurl==0.1.2
 multidict==6.0.4
 netifaces==0.10.6
 numpy==1.26.0
-openai==0.27.0
+openai==1.59.7
 openpyxl==3.1.2
 outcome==1.2.0
 pandas==2.1.1

service/chat_review.py

@@ -3,7 +3,8 @@
 import openai
 from openai import OpenAIError
 from app.gitlab_utils import *
-from config.config import gitlab_server_url, gitlab_private_token, openai_api_key, openai_baseurl, openai_model_name
+from config.config import gitlab_server_url, gitlab_private_token, api_config
+from llm_api.load_api import create_llm_api_instance
 from service.content_handle import filter_diff_content
 from utils.logger import log
 from utils.dingding import send_dingtalk_message_by_sign
@@ -45,8 +46,6 @@ def wait_and_retry(exception):
 def generate_review_note(change):
     try:
         content = filter_diff_content(change['diff'])
-        openai.api_key = openai_api_key
-        openai.api_base = openai_baseurl
         messages = [
             {"role": "system",
              "content": gpt_message
@@ -56,14 +55,13 @@ def generate_review_note(change):
              },
         ]
         log.info(f"发送给gpt 内容如下：{messages}")
-        response = openai.ChatCompletion.create(
-            model=openai_model_name,
-            messages=messages,
-        )
+        api = create_llm_api_instance()
+        api.set_config(api_config)
+        api.generate_text(messages)
         new_path = change['new_path']
         log.info(f'对 {new_path} review中...')
-        response_content = response['choices'][0]['message']['content'].replace('\n\n', '\n')
-        total_tokens = response['usage']['total_tokens']
+        response_content = api.get_respond_content().replace('\n\n', '\n')
+        total_tokens = api.get_respond_tokens()
         review_note = f'# 📚`{new_path}`' + '\n\n'
         review_note += f'({total_tokens} tokens) {"AI review 意见如下:"}' + '\n\n'
         review_note += response_content + """