Coverage for src/commands/models.py: 82%

17 statements  

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

1""" 

2Command handler for basic model listing. 

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 .base import CommandResult 

15 

16 

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

18 """List available models.""" 

19 try: 

20 models_list = list_models_api(client) 

21 if models_list: 

22 return CommandResult(True, data=models_list) 

23 else: 

24 no_models_help = """No models found. To set up a model in Databricks: 

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

262. Click 'Create Model'. 

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

284. Configure the model settings and deploy the model. 

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

30 return CommandResult(True, data=[], message=no_models_help) 

31 except Exception as e: 

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

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

34 

35 

36DEFINITION = CommandDefinition( 

37 name="list-models", 

38 description="List available language models in the Databricks workspace", 

39 handler=handle_command, 

40 parameters={}, 

41 required_params=[], 

42 tui_aliases=["/models"], 

43 visible_to_user=True, 

44 visible_to_agent=True, 

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

46)