ck_pro/tests/test_agent_model_inheritance.py

#!/usr/bin/env python3
"""
Test agent model inheritance - verify WebAgent and FileAgent properly inherit model configs
"""
import os
import sys
import types
import pytest

# Ensure package root is on path
sys.path.insert(0, os.path.abspath('.'))

# Stub heavy dependencies to avoid import overhead
stub_model_mod = types.ModuleType('ck_pro.agents.model')
class _StubLLM:
    def __init__(self, _default_init=False, **kwargs):
        self.call_target = kwargs.get('call_target', 'https://api.openai.com/v1/chat/completions')
        self.api_key = kwargs.get('api_key', 'default-key')
        self.model = kwargs.get('model', 'gpt-4o-mini')
        self.extract_body = kwargs.get('extract_body', {})
        self._default_init = _default_init
        
    def __call__(self, messages):
        return "test response"
        
stub_model_mod.LLM = _StubLLM
sys.modules['ck_pro.agents.model'] = stub_model_mod

# Stub other heavy modules
stub_utils_mod = types.ModuleType('ck_pro.agents.utils')
stub_utils_mod.zwarn = lambda x: None
stub_utils_mod.zlog = lambda x: None
stub_utils_mod.have_images_in_messages = lambda x: False
stub_utils_mod.rprint = lambda x, **kwargs: None
stub_utils_mod.TemplatedString = lambda x: type('T', (), {'format': lambda **k: x})()
stub_utils_mod.parse_response = lambda x: {'code': 'print("ok")'}
stub_utils_mod.CodeExecutor = lambda: type('CE', (), {'run': lambda *a, **k: None, 'get_print_results': lambda: 'ok'})()
stub_utils_mod.KwargsInitializable = object
stub_utils_mod.ActionResult = lambda x: x
sys.modules['ck_pro.agents.utils'] = stub_utils_mod

# Stub agent base
stub_agent_mod = types.ModuleType('ck_pro.agents.agent')
class _StubMultiStepAgent:
    def __init__(self, **kwargs):
        # Simulate MultiStepAgent behavior: use model from kwargs or default
        if 'model' in kwargs:
            self.model = _StubLLM(**kwargs['model'])
        else:
            self.model = _StubLLM(_default_init=True)
        self.ACTIVE_FUNCTIONS = {}
stub_agent_mod.MultiStepAgent = _StubMultiStepAgent
stub_agent_mod.register_template = lambda x: None
stub_agent_mod.ActionResult = lambda x: x
sys.modules['ck_pro.agents.agent'] = stub_agent_mod

stub_session_mod = types.ModuleType('ck_pro.agents.session')
stub_session_mod.AgentSession = object
sys.modules['ck_pro.agents.session'] = stub_session_mod

stub_tool_mod = types.ModuleType('ck_pro.agents.tool')
stub_tool_mod.Tool = object
stub_tool_mod.SimpleSearchTool = lambda **kwargs: type('SST', (), {})()
sys.modules['ck_pro.agents.tool'] = stub_tool_mod

# Stub file utils
stub_file_utils_mod = types.ModuleType('ck_pro.ck_file.utils')
stub_file_utils_mod.FileEnv = lambda **kwargs: type('FE', (), {})()
sys.modules['ck_pro.ck_file.utils'] = stub_file_utils_mod

# Stub file prompts
stub_file_prompts_mod = types.ModuleType('ck_pro.ck_file.prompts')
stub_file_prompts_mod.PROMPTS = {}
sys.modules['ck_pro.ck_file.prompts'] = stub_file_prompts_mod

# Stub web prompts
stub_web_prompts_mod = types.ModuleType('ck_pro.ck_web.prompts')
stub_web_prompts_mod.PROMPTS = {}
sys.modules['ck_pro.ck_web.prompts'] = stub_web_prompts_mod

# Import after stubbing
from ck_pro.config.settings import Settings, LLMConfig
from ck_pro.ck_file.agent import FileAgent
from ck_pro.ck_web.agent import WebAgent


class TestAgentModelInheritance:
    """Test that WebAgent and FileAgent properly inherit model configurations"""
    
    def test_file_agent_inherits_main_model_from_kwargs(self):
        """Test FileAgent inherits main model config through kwargs -> super().__init__"""
        # Create model config that should be inherited
        model_config = {
            'call_target': 'https://test.modelscope.cn/v1/chat/completions',
            'api_key': 'test-key-123',
            'model': 'test-model-456',
            'extract_body': {'temperature': 0.3}
        }
        
        # Create FileAgent with model config
        agent = FileAgent(settings=None, model=model_config)
        
        # Verify main model inherited the config
        assert agent.model.call_target == 'https://test.modelscope.cn/v1/chat/completions'
        assert agent.model.api_key == 'test-key-123'
        assert agent.model.model == 'test-model-456'
        assert agent.model.extract_body == {'temperature': 0.3}
    
    def test_file_agent_inherits_multimodal_model_from_kwargs(self):
        """Test FileAgent inherits multimodal model config from model_multimodal kwargs"""
        # Create multimodal model config
        mm_config = {
            'call_target': 'https://test-mm.modelscope.cn/v1/chat/completions',
            'api_key': 'test-mm-key',
            'model': 'test-mm-model',
            'extract_body': {'temperature': 0.0}
        }
        
        # Create FileAgent with multimodal config
        agent = FileAgent(settings=None, model_multimodal=mm_config)
        
        # Verify multimodal model inherited the config
        assert agent.model_multimodal.call_target == 'https://test-mm.modelscope.cn/v1/chat/completions'
        assert agent.model_multimodal.api_key == 'test-mm-key'
        assert agent.model_multimodal.model == 'test-mm-model'
        assert agent.model_multimodal.extract_body == {'temperature': 0.0}
    
    def test_web_agent_inherits_main_model_from_kwargs(self):
        """Test WebAgent inherits main model config through kwargs -> super().__init__"""
        # Create model config that should be inherited
        model_config = {
            'call_target': 'https://test.modelscope.cn/v1/chat/completions',
            'api_key': 'test-key-789',
            'model': 'test-model-web',
            'extract_body': {'temperature': 0.0}
        }
        
        # Create WebAgent with model config
        agent = WebAgent(settings=None, model=model_config)
        
        # Verify main model inherited the config
        assert agent.model.call_target == 'https://test.modelscope.cn/v1/chat/completions'
        assert agent.model.api_key == 'test-key-789'
        assert agent.model.model == 'test-model-web'
        assert agent.model.extract_body == {'temperature': 0.0}
    
    def test_web_agent_inherits_multimodal_model_from_kwargs(self):
        """Test WebAgent inherits multimodal model config from model kwargs (reused)"""
        # WebAgent reuses main model config for multimodal
        model_config = {
            'call_target': 'https://test-web-mm.modelscope.cn/v1/chat/completions',
            'api_key': 'test-web-mm-key',
            'model': 'test-web-mm-model',
            'extract_body': {'temperature': 0.1}
        }
        
        # Create WebAgent with model config
        agent = WebAgent(settings=None, model=model_config)
        
        # Verify multimodal model inherited the same config
        assert agent.model_multimodal.call_target == 'https://test-web-mm.modelscope.cn/v1/chat/completions'
        assert agent.model_multimodal.api_key == 'test-web-mm-key'
        assert agent.model_multimodal.model == 'test-web-mm-model'
        assert agent.model_multimodal.extract_body == {'temperature': 0.1}
    
    def test_file_agent_defaults_when_no_model_config(self):
        """Test FileAgent falls back to defaults when no model config provided"""
        # Create FileAgent without model config
        agent = FileAgent(settings=None)
        
        # Should use default LLM(_default_init=True) behavior
        assert agent.model._default_init == True
        assert agent.model_multimodal._default_init == True
    
    def test_web_agent_defaults_when_no_model_config(self):
        """Test WebAgent falls back to defaults when no model config provided"""
        # Create WebAgent without model config
        agent = WebAgent(settings=None)
        
        # Should use default LLM(_default_init=True) behavior
        assert agent.model._default_init == True
        assert agent.model_multimodal._default_init == True

    def test_full_config_chain_settings_to_agents(self):
        """Test complete config chain: Settings -> CKAgent kwargs -> sub-agents"""
        # Create settings with ModelScope endpoints
        settings = Settings()
        settings.ck.model = LLMConfig(
            call_target='https://api-inference.modelscope.cn/v1/chat/completions',
            api_key='parent-key',
            model='Qwen3-235B-A22B-Instruct-2507'
        )
        settings.file.model = LLMConfig(
            call_target='https://file.modelscope.cn/v1/chat/completions',
            api_key='file-key',
            model='file-model'
        )
        settings.web.model = LLMConfig(
            call_target='https://web.modelscope.cn/v1/chat/completions',
            api_key='web-key',
            model='web-model'
        )

        # Convert to CKAgent kwargs
        kwargs = settings.to_ckagent_kwargs()

        # Extract sub-agent configs
        web_kwargs = kwargs.get('web_agent', {})
        file_kwargs = kwargs.get('file_agent', {})

        # Create agents with extracted configs
        web_agent = WebAgent(settings=settings, **web_kwargs)
        file_agent = FileAgent(settings=settings, **file_kwargs)

        # Verify web agent got correct config
        assert web_agent.model.call_target == 'https://web.modelscope.cn/v1/chat/completions'
        assert web_agent.model.api_key == 'web-key'
        assert web_agent.model.model == 'web-model'

        # Verify file agent got correct config
        assert file_agent.model.call_target == 'https://file.modelscope.cn/v1/chat/completions'
        assert file_agent.model.api_key == 'file-key'
        assert file_agent.model.model == 'file-model'


if __name__ == "__main__":
    pytest.main([__file__, "-v"])