Coverage for src/commands/list_models.py: 100%

31 statements  

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

1""" 

2Command handler for listing models. 

3 

4This module contains the handler for listing available models 

5in a Databricks workspace. 

6""" 

7 

8import logging 

9from typing import Optional 

10 

11from src.clients.databricks import DatabricksAPIClient 

12from src.models import list_models as list_models_api 

13from src.command_registry import CommandDefinition 

14from src.config import get_active_model 

15from .base import CommandResult 

16 

17 

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

19 """ 

20 List available models with optional filtering and detailed information. 

21 Args: 

22 client: API client instance 

23 **kwargs: 

24 detailed (bool): Whether to show detailed information. Defaults to False. 

25 filter (str, optional): Filter string for model names. 

26 """ 

27 detailed: bool = kwargs.get("detailed", False) 

28 filter_str: Optional[str] = kwargs.get("filter") 

29 

30 try: 

31 models_list = list_models_api(client) 

32 if filter_str: 

33 normalized_filter = filter_str.lower() 

34 models_list = [ 

35 m for m in models_list if normalized_filter in m.get("name", "").lower() 

36 ] 

37 

38 if detailed and models_list: 

39 for model_item in models_list: 

40 from src.models import get_model 

41 

42 model_details = get_model(client, model_item["name"]) 

43 if model_details: 

44 model_item["details"] = model_details 

45 

46 active_model_name = get_active_model() 

47 result_data = { 

48 "models": models_list, 

49 "active_model": active_model_name, 

50 "detailed": detailed, 

51 "filter": filter_str, 

52 } 

53 

54 message = None 

55 if not models_list: 

56 message = """No models found. To set up a model in Databricks: 

571. Go to the Databricks Model Serving page in your workspace. 

582. Click 'Create Model'. 

593. Choose a model (e.g., Claude, OpenAI, or another supported LLM). 

604. Configure the model settings and deploy the model. 

61After deployment, run the models command again to verify availability.""" 

62 return CommandResult(True, data=result_data, message=message) 

63 except Exception as e: 

64 logging.error(f"Failed to list models: {e}", exc_info=True) 

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

66 

67 

68DEFINITION = CommandDefinition( 

69 name="detailed-models", 

70 description="List available models with filtering and detailed information", 

71 handler=handle_command, 

72 parameters={ 

73 "detailed": { 

74 "type": "boolean", 

75 "description": "Show detailed information about each model", 

76 }, 

77 "filter": { 

78 "type": "string", 

79 "description": "Filter string to match against model names", 

80 }, 

81 }, 

82 required_params=[], 

83 tui_aliases=["/list-models"], 

84 visible_to_user=True, 

85 visible_to_agent=True, 

86 agent_display="full", # Show full model list in tables 

87)