AgentCoreService/pqai_agent/utils/agent_abtest_utils.py

from typing import Optional

from pqai_agent.abtest.utils import get_abtest_info
from pqai_agent.data_models.agent_configuration import AgentConfiguration
from pqai_agent.service_module_manager import ServiceModuleManager
from pqai_agent.agent_config_manager import AgentConfigManager

def get_agent_abtest_config(module_name: str, uid: str,
                            service_module_manager: ServiceModuleManager,
                            agent_config_manager: AgentConfigManager) -> Optional[AgentConfiguration]:
    abtest_info = get_abtest_info(uid)
    module_config = service_module_manager.get_module_config(f'{module_name}_module')
    agent_id = module_config['default_agent_id']
    param_key = f'module_{module_name}_agent_id'
    if param_key in abtest_info.params:
        agent_id = int(abtest_info.params[param_key])
    agent_config = agent_config_manager.get_config(agent_id)
    return agent_config