feat(monitoring): add llm oriented metrics to grafana dashboard

api/app.py

@@ -4,7 +4,6 @@
 from fastapi import FastAPI, Request
 import sentry_sdk
 from starlette.middleware.sessions import SessionMiddleware
-from starlette.responses import JSONResponse
 
 from api.endpoints.monitoring import setup_prometheus
 from api.schemas.core.context import RequestContext
@@ -89,7 +88,3 @@ def _setup_monitoring(app: FastAPI, configuration: Configuration) -> None:
 
     if configuration.settings.monitoring_prometheus_enabled:
         setup_prometheus(app, include_in_schema=include_in_schema)
-
-    @app.get(path="/health", tags=[RouterName.MONITORING.title()], include_in_schema=include_in_schema)
-    def health() -> JSONResponse:
-        return JSONResponse(content={"status": "ok"}, status_code=200)

api/clients/model/_basemodelprovider.py

@@ -118,7 +118,6 @@ def _get_usage(self, request_content: RequestContent, response_data: dict | list
         tokenizer = getattr(global_context, "tokenizer", None)
         if tokenizer and request_content.endpoint in tokenizer.USAGE_ENDPOINTS:
             try:
-                completion_tokens = 0
                 prompt_tokens = tokenizer.get_prompt_tokens(endpoint=request_content.endpoint, body=request_content.json)
                 completion_tokens = tokenizer.get_completion_tokens(endpoint=request_content.endpoint, response_data=response_data)
                 total_tokens = prompt_tokens + completion_tokens
@@ -154,7 +153,7 @@ def _format_request(self, request_content: RequestContent) -> RequestContent:
         Format a request to a provider model. This method can be overridden by a subclass to add additional headers or parameters. This method format the requested endpoint thanks the ENDPOINT_TABLE attribute.
 
         Args:
-            content(RequestContent): The request content to format.
+            request_content(RequestContent): The request content to format.
 
         Returns:
             content(RequestContent): The formatted request content.
@@ -226,7 +225,8 @@ def _format_response(self, request_content: RequestContent, response: httpx.Resp
 
         return response
 
-    async def _ensure_timeseries_exists(self, redis_client: AsyncRedis, key: str) -> None:
+    @staticmethod
+    async def _ensure_timeseries_exists(redis_client: AsyncRedis, key: str) -> None:
         """
         Ensure a time series exists with proper retention configuration.
 

api/endpoints/health.py

@@ -0,0 +1,11 @@
+from fastapi import APIRouter
+from starlette.responses import JSONResponse
+
+from api.utils.variables import RouterName
+
+router = APIRouter(tags=[RouterName.HEALTH.title()])
+
+
+@router.get(path="/health")
+def health() -> JSONResponse:
+    return JSONResponse(content={"status": "ok"}, status_code=200)

api/endpoints/monitoring.py

@@ -7,12 +7,25 @@
 from starlette.responses import Response
 
 from api.helpers._accesscontroller import AccessController
+from api.helpers._metricsmiddleware import (
+    inference_requests_duration_seconds,
+    inference_requests_total,
+    inference_tokens_total,
+    inference_ttft_milliseconds,
+)
 from api.schemas.admin.roles import PermissionType
 from api.utils.variables import RouterName
 
 
 def setup_prometheus(app: FastAPI, include_in_schema: bool = True) -> None:
-    app.instrumentator = Instrumentator().instrument(app=app)
+    app.instrumentator = (
+        Instrumentator()
+        .instrument(app=app)
+        .add(inference_requests_total())
+        .add(inference_requests_duration_seconds())
+        .add(inference_ttft_milliseconds())
+        .add(inference_tokens_total())
+    )
 
     @app.get(
         path="/metrics",

api/helpers/_metricsmiddleware.py

@@ -0,0 +1,104 @@
+from collections.abc import Callable
+
+from prometheus_client import Counter, Histogram
+from prometheus_fastapi_instrumentator.metrics import Info
+
+from api.utils.context import request_context
+
+
+def inference_requests_total() -> Callable[[Info], None]:
+    metric = Counter(
+        "inference_requests_total",
+        "Total number of LLM requests.",
+        labelnames=("endpoint", "model", "status_code"),
+    )
+
+    def instrumentation(info: Info) -> None:
+        try:
+            context = request_context.get()
+            model = context.router_name
+            endpoint = context.endpoint
+            if model and endpoint:
+                metric.labels(
+                    endpoint=endpoint,
+                    model=model,
+                    status_code=info.modified_status,
+                ).inc()
+        except Exception:
+            pass
+
+    return instrumentation
+
+
+def inference_requests_duration_seconds() -> Callable[[Info], None]:
+    metric = Histogram(
+        "inference_requests_duration_seconds",
+        "Duration of LLM requests in seconds.",
+        labelnames=("endpoint", "model", "status_code"),
+    )
+
+    def instrumentation(info: Info) -> None:
+        try:
+            context = request_context.get()
+            model = context.router_name
+            endpoint = context.endpoint
+            latency = context.latency
+            if model and endpoint and latency is not None:
+                metric.labels(
+                    endpoint=endpoint,
+                    model=model,
+                    status_code=info.modified_status,
+                ).observe(latency / 1000)
+        except Exception:
+            pass
+
+    return instrumentation
+
+
+def inference_ttft_milliseconds() -> Callable[[Info], None]:
+    metric = Histogram(
+        "inference_ttft_milliseconds",
+        "Time to first token for streaming LLM responses in milliseconds.",
+        labelnames=("endpoint", "model", "status_code"),
+    )
+
+    def instrumentation(info: Info) -> None:
+        try:
+            context = request_context.get()
+            model = context.router_name
+            endpoint = context.endpoint
+            ttft = context.ttft
+            if model and endpoint and ttft is not None:
+                metric.labels(
+                    endpoint=endpoint,
+                    model=model,
+                    status_code=info.modified_status,
+                ).observe(ttft)
+        except Exception:
+            pass
+
+    return instrumentation
+
+
+def inference_tokens_total() -> Callable[[Info], None]:
+    metric = Counter(
+        "inference_tokens_total",
+        "Total number of tokens consumed (prompt and completion).",
+        labelnames=("endpoint", "model", "type"),
+    )
+
+    def instrumentation(info: Info) -> None:
+        try:
+            context = request_context.get()
+            model = context.router_name
+            endpoint = context.endpoint
+            usage = context.usage
+            if model and endpoint and usage is not None:
+                if usage.prompt_tokens:
+                    metric.labels(endpoint=endpoint, model=model, type="prompt").inc(usage.prompt_tokens)
+                if usage.completion_tokens:
+                    metric.labels(endpoint=endpoint, model=model, type="completion").inc(usage.completion_tokens)
+        except Exception:
+            pass
+
+    return instrumentation

api/helpers/load_balancing/_leastbusyloadbalancingstrategy.py

@@ -20,7 +20,6 @@ def __init__(self, redis_client: AsyncRedis | Redis, load_balancing_metric: Metr
         Get a provider to handle the request based on the specified routing strategy.
 
         Args:
-            candidates (list[int]): The list of provider candidates (provider IDs) to choose from
             redis_client (AsyncRedis): Redis client instance, required for least busy strategy
             load_balancing_metric (Metric): The type of metric to use for performance evaluation
 

api/utils/hooks_decorator.py

@@ -189,7 +189,7 @@ async def update_budget(usage: Usage):
                 # Update the budget
                 update_stmt = update(User).where(User.id == user_id).values(budget=new_budget, updated=func.now()).returning(User.budget)
 
-                result = await postgres_session.execute(update_stmt)
+                await postgres_session.execute(update_stmt)
 
         except Exception as e:
             logger.exception(f"Failed to update budget for user {user_id}: {e}")

api/utils/variables.py

@@ -19,6 +19,7 @@ class RouterName(StrEnum):
     COLLECTIONS = ("collections", "api.endpoints.collections")
     DOCUMENTS = ("documents", "api.endpoints.documents")
     EMBEDDINGS = ("embeddings", "api.endpoints.embeddings")
+    HEALTH = ("health", "api.endpoints.health")
     ME = ("me", "api.endpoints.me")
     MODELS = ("models", "api.infrastructure.fastapi.endpoints.models")
     MONITORING = ("monitoring", None)