refactor(llm): rewrite large model related code

Author: Robbings

config/config.py

@@ -6,7 +6,7 @@
 model_programming_translate = "code-davinci-002"
 
 # api 接口封装类
-llm_api_impl = "llm_api.llm_api_default.LLMApiDefault"
+llm_api_impl = "large_model.api.default_api.DefaultApi"
 
 # api 配置方式参考 docs/llm_api.md
 # 默认使用认UnionLLM，参考：https://github.com/EvalsOne/UnionLLM/tree/main/docs
@@ -56,4 +56,4 @@
 # ------------- Message notification --------------------
 # dingding notification （un necessary）
 dingding_bot_webhook = "https://oapi.dingtalk.com/robot/send?access_token=*****************************************"
-dingding_secret = "S********************************950f"
+dingding_secret = "S********************************950f"

large_model/__init__.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 
-class LLMApiInterface(ABC):
+class AbstractApi(ABC):
 
     @abstractmethod
     def set_config(self, api_config: dict) -> bool:

llm_api/llm_api_interface.py renamed to large_model/abstract_api.py

@@ -1,14 +1,12 @@
 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
 from unionllm import unionchat
 
+from large_model.abstract_api import AbstractApi
 
-class LLMApiDefault(LLMApiInterface):
+
+class DefaultApi(AbstractApi):
 
     def __init__(self):
         self.params = {}
@@ -36,19 +34,3 @@ def get_respond_content(self) -> str:
 
     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())

large_model/api/__init__.py

@@ -0,0 +1,30 @@
+import importlib
+import warnings
+
+from config.config import llm_api_impl, api_config
+
+
+class LLMGenerator:
+
+    @classmethod
+    def new_model(cls, config = api_config):
+        api = cls.create_model_instance()
+        api.set_config(config)
+        return api
+
+    @classmethod
+    def get_llm_api_class(cls):
+        module_name, class_name = llm_api_impl.rsplit('.', 1)
+        module = importlib.import_module(module_name)
+        llm_class = getattr(module, class_name)
+        return llm_class
+
+    @classmethod
+    def create_model_instance(cls):
+        """
+        使用工厂函数获取类实例
+        """
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", category=UserWarning)
+            llm_class = cls.get_llm_api_class()
+            return llm_class()

llm_api/llm_api_default.py renamed to large_model/api/default_api.py

@@ -7,5 +7,5 @@ class ReviewHandle(object):
     def __init__(self):
         pass
 
-    def merge_handle(self, changes, merge_info, hook_info, reply):
+    def merge_handle(self, changes, merge_info, hook_info, reply, model):
         pass

large_model/llm_generator.py

@@ -4,20 +4,19 @@
 from retrying import retry
 
 from config.config import api_config, gpt_message
-from llm_api.load_api import create_llm_api_instance
 from review_engine.handler.abstract_handler import ReviewHandle
 from utils.gitlab_parser import filter_diff_content
 from utils.logger import log
 
 
-def chat_review(changes):
+def chat_review(changes, model):
     log.info('开始code review')
     with concurrent.futures.ThreadPoolExecutor() as executor:
         review_results = []
         result_lock = threading.Lock()
 
         def process_change(change):
-            result = generate_review_note(change)
+            result = generate_review_note(change, model)
             with result_lock:
                 review_results.append(result)
 
@@ -34,7 +33,7 @@ def process_change(change):
     return "\n\n".join(review_results) if review_results else ""
 
 @retry(stop_max_attempt_number=3, wait_fixed=60000)
-def generate_review_note(change):
+def generate_review_note(change, model):
     try:
         content = filter_diff_content(change['diff'])
         messages = [
@@ -46,13 +45,11 @@ def generate_review_note(change):
              },
         ]
         log.info(f"发送给gpt 内容如下：{messages}")
-        api = create_llm_api_instance()
-        api.set_config(api_config)
-        api.generate_text(messages)
+        model.generate_text(messages)
         new_path = change['new_path']
         log.info(f'对 {new_path} review中...')
-        response_content = api.get_respond_content().replace('\n\n', '\n')
-        total_tokens = api.get_respond_tokens()
+        response_content = model.get_respond_content().replace('\n\n', '\n')
+        total_tokens = model.get_respond_tokens()
         review_note = f'# 📚`{new_path}`' + '\n\n'
         review_note += f'({total_tokens} tokens) {"AI review 意见如下:"}' + '\n\n'
         review_note += response_content + """
@@ -71,14 +68,14 @@ def generate_review_note(change):
 
 
 class MainReviewHandle(ReviewHandle):
-    def merge_handle(self, changes, merge_info, hook_info, reply):
-        self.default_handle(changes, merge_info, hook_info, reply)
+    def merge_handle(self, changes, merge_info, hook_info, reply, model):
+        self.default_handle(changes, merge_info, hook_info, reply, model)
 
-    def default_handle(self, changes, merge_info, hook_info, reply):
+    def default_handle(self, changes, merge_info, hook_info, reply, model):
         maximum_files = 50
         if changes and len(changes) <= maximum_files:
             # Code Review 信息
-            review_info = chat_review(changes)
+            review_info = chat_review(changes, model)
             if review_info:
                 reply.add_reply({
                     'title': '__MAIN_REVIEW__',
@@ -150,6 +147,8 @@ def default_handle(self, changes, merge_info, hook_info, reply):
     reply = Reply({'type': 'merge_request',
                      'project_id': 9885,
                      'merge_request_iid': 18})
+    from large_model.llm_generator import LLMGenerator
+    model = LLMGenerator.new_model()
     hook_info = {
         "object_kind": "merge_request",
         "event_type": "merge_request",
@@ -198,4 +197,4 @@ def default_handle(self, changes, merge_info, hook_info, reply):
             ]
         }
     }
-    main_handle.merge_handle(changes, info, hook_info, reply)
+    main_handle.merge_handle(changes, info, hook_info, reply, model)

review_engine/handler/abstract_handler.py

@@ -1,5 +1,6 @@
 import threading
 
+from large_model.llm_generator import LLMGenerator
 from utils.tools import import_submodules
 
 import_submodules('review_engine.handler')
@@ -15,7 +16,8 @@ def __init__(self, reply):
 
     def handle_merge(self, changes, merge_info, webhook_info):
         # 多线程处理
-        threads = [threading.Thread(target=handle.merge_handle, args=(changes, merge_info, webhook_info, self.reply))
+        threads = [threading.Thread(target=handle.merge_handle, args=(changes, merge_info, webhook_info,
+                                                                      self.reply, LLMGenerator.new_model()))
                    for handle in self.handles]
         for thread in threads:
             thread.start()

review_engine/handler/default_handler.py

@@ -54,8 +54,9 @@ def check_dingding_config(config):
     """
     result = {'passed': True, 'errors': []}
     try:
-        from utils.dingding import send_dingtalk_message_by_sign
-        response = send_dingtalk_message_by_sign("连通性测试：测试消息，请勿回复。")
+        from reply_module.reply_target.dingtalk_reply import DingtalkReply
+        dingtalk_reply = DingtalkReply({'type': 'merge_request', 'project_id': 1, 'merge_request_iid': 1})
+        response = dingtalk_reply.send("连通性测试：测试消息，请勿回复。")
         if not response:
             error_msg = "Dingding configuration is invalid"
             result['errors'].append(error_msg)
@@ -100,9 +101,8 @@ def check_api_config(config):
     """
     result = {'passed': True, 'errors': []}
     try:
-        from llm_api.load_api import create_llm_api_instance
-        api = create_llm_api_instance()
-        api.set_config(config.api_config)
+        from large_model.llm_generator import LLMGenerator
+        api = LLMGenerator.new_model()
         api.generate_text([
             {"role": "system",
              "content": "你是一个有用的助手"