From 118fbe3f7d66acf64b28193c3a6b4535f2c5c8dd Mon Sep 17 00:00:00 2001
From: Rock Chin <1010553892@qq.com>
Date: Thu, 2 Mar 2023 19:50:31 +0800
Subject: [PATCH] =?UTF-8?q?perf(modelmgr.py):=20=E7=B1=BB=E5=90=8D?=
 =?UTF-8?q?=E7=A7=B0=E5=BC=BA=E8=B0=83=E5=85=B6=E4=B8=BA=E4=B8=80=E4=B8=AA?=
 =?UTF-8?q?=E8=AF=B7=E6=B1=82=E5=AF=B9=E8=B1=A1?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pkg/openai/manager.py  | 4 ++--
 pkg/openai/modelmgr.py | 9 +++++----
 pkg/openai/session.py  | 2 --
 3 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/pkg/openai/manager.py b/pkg/openai/manager.py
index 7320d221..563e14ad 100644
--- a/pkg/openai/manager.py
+++ b/pkg/openai/manager.py
@@ -5,7 +5,7 @@ import openai
 import pkg.openai.keymgr
 import pkg.utils.context
 import pkg.audit.gatherer
-from pkg.openai.modelmgr import Model, ChatCompletionModel, OpenaiModel
+from pkg.openai.modelmgr import ModelRequest, create_openai_model_request
 
 # 为其他模块提供与OpenAI交互的接口
 class OpenAIInteract:
@@ -35,7 +35,7 @@ class OpenAIInteract:
     def request_completion(self, messages):
         config = pkg.utils.context.get_config()
 
-        ai:Model = OpenaiModel(config.completion_api_params['model'], 'user')
+        ai:ModelRequest = create_openai_model_request(config.completion_api_params['model'], 'user')
         ai.request(
             messages,
             **config.completion_api_params
diff --git a/pkg/openai/modelmgr.py b/pkg/openai/modelmgr.py
index dbd44054..8942b4c6 100644
--- a/pkg/openai/modelmgr.py
+++ b/pkg/openai/modelmgr.py
@@ -32,7 +32,7 @@ IMAGE_MODELS = {
 }
 
 
-class Model():
+class ModelRequest():
 
     can_chat = False
 
@@ -62,7 +62,7 @@ class Model():
         return self.ret
     
 
-class ChatCompletionModel(Model):
+class ChatCompletionModel(ModelRequest):
     """ChatCompletion接口实现"""
     Chat_role = ['system', 'user', 'assistant']
     def __init__(self, model_name, user_name):
@@ -87,7 +87,7 @@ class ChatCompletionModel(Model):
         return self.message
     
 
-class CompletionModel(Model):
+class CompletionModel(ModelRequest):
     """Completion接口实现"""
     def __init__(self, model_name, user_name):
         request_fun = openai.Completion.create
@@ -111,7 +111,8 @@ class CompletionModel(Model):
         return self.message
     
 
-def OpenaiModel(model_name:str, user_name='user'):
+def create_openai_model_request(model_name: str, user_name: str = 'user') -> ModelRequest:
+    """使用给定的模型名称创建模型请求对象"""
     if model_name in CHAT_COMPLETION_MODELS:
         model = ChatCompletionModel(model_name, user_name)
     elif model_name in COMPLETION_MODELS:
diff --git a/pkg/openai/session.py b/pkg/openai/session.py
index 37e8a241..dcaac9ff 100644
--- a/pkg/openai/session.py
+++ b/pkg/openai/session.py
@@ -94,8 +94,6 @@ class Session:
 
     prompt = {}
 
-    import config
-
     create_timestamp = 0
 
     last_interact_timestamp = 0