Merge pull request #293 from umago/ISSUE-291-metrics-as-library

[ISSUE-291] Fix setup_model_metrics() as a library

Author: manstis

src/app/main.py

@@ -74,13 +74,13 @@ async def rest_api_metrics(
     if isinstance(route, (Mount, Route, WebSocketRoute))
 ]
 
-setup_model_metrics()
-
 
 @app.on_event("startup")
 async def startup_event() -> None:
     """Perform logger setup on service startup."""
     logger.info("Registering MCP servers")
     await register_mcp_servers_async(logger, configuration.configuration)
     get_logger("app.endpoints.handlers")
+    logger.info("Setting up model metrics")
+    await setup_model_metrics()
     logger.info("App startup complete")

src/metrics/utils.py

@@ -1,19 +1,24 @@
 """Utility functions for metrics handling."""
 
 from configuration import configuration
-from client import LlamaStackClientHolder
+from client import LlamaStackClientHolder, AsyncLlamaStackClientHolder
 from log import get_logger
 import metrics
 
 logger = get_logger(__name__)
 
 
-def setup_model_metrics() -> None:
+async def setup_model_metrics() -> None:
     """Perform setup of all metrics related to LLM model and provider."""
-    client = LlamaStackClientHolder().get_client()
+    model_list = []
+    if configuration.llama_stack_configuration.use_as_library_client:
+        model_list = await AsyncLlamaStackClientHolder().get_client().models.list()
+    else:
+        model_list = LlamaStackClientHolder().get_client().models.list()
+
     models = [
         model
-        for model in client.models.list()
+        for model in model_list
         if model.model_type == "llm"  # pyright: ignore[reportAttributeAccessIssue]
     ]
 

tests/unit/metrics/test_utis.py

@@ -3,7 +3,7 @@
 from metrics.utils import setup_model_metrics
 
 
-def test_setup_model_metrics(mocker):
+async def test_setup_model_metrics(mocker):
     """Test the setup_model_metrics function."""
 
     # Mock the LlamaStackAsLibraryClient
@@ -51,7 +51,7 @@ def test_setup_model_metrics(mocker):
         model_1,
     ]
 
-    setup_model_metrics()
+    await setup_model_metrics()
 
     # Check that the provider_model_configuration metric was set correctly
     # The default model should have a value of 1, others should be 0