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
« prev ^ index » next coverage.py v7.8.0, created at 2025-06-05 22:56 -0700
1"""
2Command handler for model selection.
4This module contains the handler for selecting an active model
5for use in a Databricks workspace.
6"""
8import logging
9from typing import Optional
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
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))
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)