Coverage for src/chuck_data/commands/model_selection.py: 0%

22 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-06-05 22:56 -0700

1""" 

2Command handler for model selection. 

3 

4This module contains the handler for selecting an active model 

5for use in a Databricks workspace. 

6""" 

7 

8import logging 

9from typing import Optional 

10 

11from ..clients.databricks import DatabricksAPIClient 

12from ..models import list_models as list_models_api 

13from ..config import set_active_model 

14from ..command_registry import CommandDefinition 

15from .base import CommandResult 

16 

17 

18def handle_command(client: Optional[DatabricksAPIClient], **kwargs) -> CommandResult: 

19 """Set the active model. 

20 Args: 

21 client: API client instance 

22 **kwargs: model_name (str) 

23 """ 

24 model_name: str = kwargs.get("model_name") 

25 if not model_name: 

26 return CommandResult(False, message="model_name parameter is required.") 

27 try: 

28 models_list = list_models_api(client) 

29 model_names = [m["name"] for m in models_list] 

30 if model_name not in model_names: 

31 return CommandResult(False, message=f"Model '{model_name}' not found.") 

32 set_active_model(model_name) 

33 return CommandResult( 

34 True, message=f"Active model is now set to '{model_name}'." 

35 ) 

36 except Exception as e: 

37 logging.error(f"Failed to set model '{model_name}': {e}", exc_info=True) 

38 return CommandResult(False, error=e, message=str(e)) 

39 

40 

41DEFINITION = CommandDefinition( 

42 name="select-model", 

43 description="Set the active model for agent operations", 

44 handler=handle_command, 

45 parameters={ 

46 "model_name": { 

47 "type": "string", 

48 "description": "Name of the model to set as active", 

49 } 

50 }, 

51 required_params=["model_name"], 

52 tui_aliases=["/select-model"], 

53 visible_to_user=True, 

54 visible_to_agent=True, 

55)