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