Coverage for src/models.py: 100%

19 statements  

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

1""" 

2Module for interacting with Databricks model serving endpoints. 

3""" 

4 

5import logging 

6 

7 

8def list_models(client): 

9 """ 

10 Fetch a list of models from the Databricks Serving API. 

11 

12 Args: 

13 client: DatabricksAPIClient instance 

14 

15 Returns: 

16 List of available model endpoints 

17 """ 

18 try: 

19 return client.list_models() 

20 except ValueError as e: 

21 logging.error(f"Failed to list models: {e}") 

22 raise ValueError(f"Model serving API error: {e}") 

23 except ConnectionError as e: 

24 logging.error(f"Connection error when listing models: {e}") 

25 raise ConnectionError(f"Failed to connect to serving endpoint: {e}") 

26 

27 

28def get_model(client, model_name): 

29 """ 

30 Get details of a specific model from Databricks Serving API. 

31 

32 Args: 

33 client: DatabricksAPIClient instance 

34 model_name: Name of the model to retrieve 

35 

36 Returns: 

37 Model details if found, None otherwise 

38 """ 

39 try: 

40 return client.get_model(model_name) 

41 except ValueError as e: 

42 logging.error(f"Failed to get model: {e}") 

43 raise ValueError(f"Model serving API error: {e}") 

44 except ConnectionError as e: 

45 logging.error(f"Connection error when getting model: {e}") 

46 raise ConnectionError(f"Failed to connect to serving endpoint: {e}")