大修LLMReq
This commit is contained in:
92
src/llm_models/model_manager_bak.py
Normal file
92
src/llm_models/model_manager_bak.py
Normal file
@@ -0,0 +1,92 @@
|
||||
import importlib
|
||||
from typing import Dict
|
||||
|
||||
from src.config.config import model_config
|
||||
from src.config.api_ada_configs import ModuleConfig, ModelUsageArgConfig
|
||||
from src.common.logger import get_logger
|
||||
|
||||
from .model_client import ModelRequestHandler, BaseClient
|
||||
|
||||
logger = get_logger("模型管理器")
|
||||
|
||||
class ModelManager:
|
||||
# TODO: 添加读写锁,防止异步刷新配置时发生数据竞争
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: ModuleConfig,
|
||||
):
|
||||
self.config: ModuleConfig = config
|
||||
"""配置信息"""
|
||||
|
||||
self.api_client_map: Dict[str, BaseClient] = {}
|
||||
"""API客户端映射表"""
|
||||
|
||||
self._request_handler_cache: Dict[str, ModelRequestHandler] = {}
|
||||
"""ModelRequestHandler缓存,避免重复创建"""
|
||||
|
||||
for provider_name, api_provider in self.config.api_providers.items():
|
||||
# 初始化API客户端
|
||||
try:
|
||||
# 根据配置动态加载实现
|
||||
client_module = importlib.import_module(
|
||||
f".model_client.{api_provider.client_type}_client", __package__
|
||||
)
|
||||
client_class = getattr(
|
||||
client_module, f"{api_provider.client_type.capitalize()}Client"
|
||||
)
|
||||
if not issubclass(client_class, BaseClient):
|
||||
raise TypeError(
|
||||
f"'{client_class.__name__}' is not a subclass of 'BaseClient'"
|
||||
)
|
||||
self.api_client_map[api_provider.name] = client_class(
|
||||
api_provider
|
||||
) # 实例化,放入api_client_map
|
||||
except ImportError as e:
|
||||
logger.error(f"Failed to import client module: {e}")
|
||||
raise ImportError(
|
||||
f"Failed to import client module for '{provider_name}': {e}"
|
||||
) from e
|
||||
|
||||
def __getitem__(self, task_name: str) -> ModelRequestHandler:
|
||||
"""
|
||||
获取任务所需的模型客户端(封装)
|
||||
使用缓存机制避免重复创建ModelRequestHandler
|
||||
:param task_name: 任务名称
|
||||
:return: 模型客户端
|
||||
"""
|
||||
if task_name not in self.config.task_model_arg_map:
|
||||
raise KeyError(f"'{task_name}' not registered in ModelManager")
|
||||
|
||||
# 检查缓存中是否已存在
|
||||
if task_name in self._request_handler_cache:
|
||||
logger.debug(f"🚀 [性能优化] 从缓存获取ModelRequestHandler: {task_name}")
|
||||
return self._request_handler_cache[task_name]
|
||||
|
||||
# 创建新的ModelRequestHandler并缓存
|
||||
logger.debug(f"🔧 [性能优化] 创建并缓存ModelRequestHandler: {task_name}")
|
||||
handler = ModelRequestHandler(
|
||||
task_name=task_name,
|
||||
config=self.config,
|
||||
api_client_map=self.api_client_map,
|
||||
)
|
||||
self._request_handler_cache[task_name] = handler
|
||||
return handler
|
||||
|
||||
def __setitem__(self, task_name: str, value: ModelUsageArgConfig):
|
||||
"""
|
||||
注册任务的模型使用配置
|
||||
:param task_name: 任务名称
|
||||
:param value: 模型使用配置
|
||||
"""
|
||||
self.config.task_model_arg_map[task_name] = value
|
||||
|
||||
def __contains__(self, task_name: str):
|
||||
"""
|
||||
判断任务是否已注册
|
||||
:param task_name: 任务名称
|
||||
:return: 是否在模型列表中
|
||||
"""
|
||||
return task_name in self.config.task_model_arg_map
|
||||
|
||||
|
||||
Reference in New Issue
Block a user