feat: add Prometheus support to Tabby.

Makefile

@@ -24,8 +24,9 @@ bump-release-version:
 	cargo ws version --allow-branch "r*" --no-individual-tags --force "*"
 
 update-openapi-doc:
-	curl http://localhost:8080/api-docs/openapi.json | jq '                                                       \
-	  delpaths([                                                                                                  \
+	curl http://localhost:8080/api-docs/openapi.json | jq '                                                             \
+	  delpaths([                                                                                                        \
+		  ["paths", "/v1/metrics"],                                                                                 \
 		  ["paths", "/v1beta/chat/completions"],                                                                    \
 		  ["paths", "/v1beta/search"],                                                                              \
 		  ["components", "schemas", "CompletionRequest", "properties", "prompt"],                                   \
@@ -37,4 +38,4 @@ update-openapi-doc:
 			    > website/static/openapi.json
 
 update-graphql-schema:
-	cargo run --package tabby-webserver --example update-schema
+	cargo run --package tabby-webserver --example update-schema

crates/tabby/Cargo.toml

@@ -49,6 +49,7 @@ async-trait.workspace = true
 tabby-webserver = { path = "../../ee/tabby-webserver", optional = true }
 thiserror.workspace = true
 chrono = "0.4.31"
+axum-prometheus = "0.4.0"
 
 [dependencies.uuid]
 version = "1.3.3"

crates/tabby/src/routes/metrics.rs

@@ -0,0 +1,15 @@
+use std::sync::Arc;
+use axum::extract::State;
+use axum_prometheus::metrics_exporter_prometheus::PrometheusHandle;
+
+#[utoipa::path(
+    get,
+    path = "/v1/metrics",
+    tag = "v1",
+    responses(
+        (status = 200, description = "Success", body = String, content_type = "text/plain"),
+    )
+)]
+pub async fn metrics(State(state): State<Arc<PrometheusHandle>>) -> String {
+    state.render()
+}

crates/tabby/src/routes/mod.rs

@@ -3,9 +3,12 @@ mod completions;
 mod events;
 mod health;
 mod search;
+mod metrics;
+
 
 pub use chat::*;
 pub use completions::*;
 pub use events::*;
 pub use health::*;
 pub use search::*;
+pub use metrics::*;

crates/tabby/src/serve.rs

@@ -5,6 +5,7 @@ use std::{
 };
 
 use axum::{routing, Router, Server};
+use axum_prometheus::{PrometheusMetricLayer, metrics_exporter_prometheus::PrometheusHandle};
 use axum_tracing_opentelemetry::opentelemetry_tracing_layer;
 use clap::Args;
 use tabby_common::{
@@ -49,7 +50,7 @@ Install following IDE / Editor extensions to get started with [Tabby](https://gi
     servers(
         (url = "/", description = "Server"),
     ),
-    paths(routes::log_event, routes::completions, routes::completions, routes::health, routes::search),
+    paths(routes::log_event, routes::completions, routes::completions, routes::health, routes::search, routes::metrics),
     components(schemas(
         api::event::LogEventRequest,
         completion::CompletionRequest,
@@ -108,9 +109,11 @@ pub async fn main(config: &Config, args: &ServeArgs) {
 
     let logger = Arc::new(create_logger());
     let code = Arc::new(create_code_search());
+    let (prometheus_layer, prometheus_handle) = PrometheusMetricLayer::pair();
+    let metrics_handle = Arc::new(prometheus_handle);
 
     let app = Router::new()
-        .merge(api_router(args, config, logger.clone(), code.clone()).await)
+        .merge(api_router(args, config, logger.clone(), code.clone(), metrics_handle).await)
         .merge(SwaggerUi::new("/swagger-ui").url("/api-docs/openapi.json", ApiDoc::openapi()));
 
     #[cfg(feature = "ee")]
@@ -121,7 +124,8 @@ pub async fn main(config: &Config, args: &ServeArgs) {
 
     let app = app
         .layer(CorsLayer::permissive())
-        .layer(opentelemetry_tracing_layer());
+        .layer(opentelemetry_tracing_layer())
+        .layer(prometheus_layer);
 
     let address = SocketAddr::from((Ipv4Addr::UNSPECIFIED, args.port));
     info!("Listening at {}", address);
@@ -148,6 +152,7 @@ async fn api_router(
     config: &Config,
     logger: Arc<dyn EventLogger>,
     code: Arc<dyn CodeSearch>,
+    metrics_handle: Arc<PrometheusHandle>
 ) -> Router {
     let completion_state = if let Some(model) = &args.model {
         Some(Arc::new(
@@ -179,6 +184,7 @@ async fn api_router(
         args.chat_model.as_deref(),
         &args.device,
     ));
+
     routers.push({
         Router::new()
             .route(
@@ -193,6 +199,10 @@ async fn api_router(
                 "/v1/health",
                 routing::get(routes::health).with_state(health_state),
             )
+            .route(
+                "/v1/metrics",
+                routing::get(routes::metrics).with_state(metrics_handle),
+            )
     });
 
     if let Some(completion_state) = completion_state {

crates/tabby/src/worker.rs

@@ -6,6 +6,7 @@ use std::{
 
 use anyhow::Result;
 use axum::{routing, Router};
+use axum_prometheus::{PrometheusMetricLayer, metrics_exporter_prometheus::PrometheusHandle};
 use axum_tracing_opentelemetry::opentelemetry_tracing_layer;
 use clap::Args;
 use hyper::Server;
@@ -63,19 +64,25 @@ async fn make_chat_route(context: WorkerContext, args: &WorkerArgs) -> Router {
     )
 }
 
-async fn make_completion_route(context: WorkerContext, args: &WorkerArgs) -> Router {
+async fn make_completion_route(context: WorkerContext, args: &WorkerArgs, metrics_handle: PrometheusHandle) -> Router {
     context.register(WorkerKind::Completion, args).await;
 
     let code = Arc::new(context.client.clone());
     let logger = Arc::new(context.client);
     let completion_state = Arc::new(
         create_completion_service(code, logger, &args.model, &args.device, args.parallelism).await,
     );
-
-    Router::new().route(
-        "/v1/completions",
-        routing::post(routes::completions).with_state(completion_state),
-    )
+    let metrics_handle = Arc::new(metrics_handle);
+
+    Router::new()
+        .route(
+            "/v1/completions",
+            routing::post(routes::completions).with_state(completion_state),
+        )
+        .route(
+            "/v1/metrics",
+            routing::get(routes::metrics).with_state(metrics_handle),
+        )
 }
 
 pub async fn main(kind: WorkerKind, args: &WorkerArgs) {
@@ -84,14 +91,18 @@ pub async fn main(kind: WorkerKind, args: &WorkerArgs) {
     info!("Starting worker, this might takes a few minutes...");
 
     let context = WorkerContext::new(&args.url).await;
+
+    let (prometheus_layer, prometheus_handle) = PrometheusMetricLayer::pair();
+
     let app = match kind {
-        WorkerKind::Completion => make_completion_route(context, args).await,
+        WorkerKind::Completion => make_completion_route(context, args, prometheus_handle).await,
         WorkerKind::Chat => make_chat_route(context, args).await,
     };
 
     let app = app
         .layer(CorsLayer::permissive())
-        .layer(opentelemetry_tracing_layer());
+        .layer(opentelemetry_tracing_layer())
+        .layer(prometheus_layer);
 
     let address = SocketAddr::from((Ipv4Addr::UNSPECIFIED, args.port));
     info!("Listening at {}", address);