Add --root-path (and normalize_root_path)

Signed-off-by: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>

Author: alvarobartt

router/src/http/server.rs

@@ -1632,6 +1632,7 @@ pub async fn run(
     payload_limit: usize,
     api_key: Option<String>,
     cors_allow_origin: Option<Vec<String>>,
+    root_path: Option<String>,
 ) -> Result<(), anyhow::Error> {
     // OpenAPI documentation
     #[derive(OpenApi)]
@@ -1872,6 +1873,14 @@ pub async fn run(
         .layer(DefaultBodyLimit::max(payload_limit))
         .layer(cors_layer);
 
+    let app = match normalize_root_path(root_path.as_deref())? {
+        Some(root_path) => {
+            tracing::info!("Serving HTTP routes under root path: {root_path}");
+            Router::new().nest(&root_path, app)
+        }
+        None => app,
+    };
+
     // Run server
     let listener = tokio::net::TcpListener::bind(&addr)
         .await
@@ -1888,6 +1897,24 @@ pub async fn run(
     Ok(())
 }
 
+fn normalize_root_path(root_path: Option<&str>) -> Result<Option<String>, anyhow::Error> {
+    let Some(root_path) = root_path else {
+        return Ok(None);
+    };
+
+    let root_path = root_path.trim();
+    if root_path.is_empty() || root_path == "/" {
+        return Ok(None);
+    }
+
+    if root_path.contains('?') || root_path.contains('#') {
+        anyhow::bail!("`--root-path` must be a path and cannot contain `?` or `#`");
+    }
+
+    let root_path = format!("/{}", root_path.trim_matches('/'));
+    Ok(Some(root_path))
+}
+
 impl From<&ErrorType> for StatusCode {
     fn from(value: &ErrorType) -> Self {
         match value {
@@ -1932,3 +1959,34 @@ impl From<serde_json::Error> for ErrorResponse {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::normalize_root_path;
+
+    #[test]
+    fn normalize_root_path_ignores_empty_and_root() {
+        assert_eq!(normalize_root_path(None).unwrap(), None);
+        assert_eq!(normalize_root_path(Some("")).unwrap(), None);
+        assert_eq!(normalize_root_path(Some("   ")).unwrap(), None);
+        assert_eq!(normalize_root_path(Some("/")).unwrap(), None);
+    }
+
+    #[test]
+    fn normalize_root_path_adds_leading_slash_and_trims_trailing_slashes() {
+        assert_eq!(
+            normalize_root_path(Some("jinaai/jina-embeddings-v2-base-en")).unwrap(),
+            Some("/jinaai/jina-embeddings-v2-base-en".to_string())
+        );
+        assert_eq!(
+            normalize_root_path(Some("/jinaai/jina-embeddings-v2-base-en/")).unwrap(),
+            Some("/jinaai/jina-embeddings-v2-base-en".to_string())
+        );
+    }
+
+    #[test]
+    fn normalize_root_path_rejects_query_and_fragment() {
+        assert!(normalize_root_path(Some("/foo?bar")).is_err());
+        assert!(normalize_root_path(Some("/foo#bar")).is_err());
+    }
+}

router/src/lib.rs

@@ -62,6 +62,7 @@ pub async fn run(
     hf_token: Option<String>,
     hostname: Option<String>,
     port: u16,
+    root_path: Option<String>,
     uds_path: Option<String>,
     huggingface_hub_cache: Option<String>,
     payload_limit: usize,
@@ -390,15 +391,17 @@ pub async fn run(
             payload_limit,
             api_key,
             cors_allow_origin,
+            root_path,
         )
         .await
     }
 
     #[cfg(feature = "grpc")]
     {
-        // cors_allow_origin and payload_limit are not used for gRPC servers
+        // cors_allow_origin, payload_limit and root_path are not used for gRPC servers
         let _ = cors_allow_origin;
         let _ = payload_limit;
+        let _ = root_path;
         grpc::server::run(infer, info, addr, prom_builder, api_key).await
     }
 }

router/src/main.rs

@@ -151,6 +151,15 @@ struct Args {
     #[clap(default_value = "3000", long, short, env)]
     port: u16,
 
+    /// Root path to mount all HTTP routes under.
+    ///
+    /// For example, `--root-path jinaai/jina-embeddings-v2-base-en` serves the
+    /// OpenAI-compatible embeddings route at `/jinaai/jina-embeddings-v2-base-en/v1/embeddings`.
+    ///
+    /// Unused for gRPC servers.
+    #[clap(long, env)]
+    root_path: Option<String>,
+
     /// The name of the unix socket some text-embeddings-inference backends will use as they
     /// communicate internally with gRPC.
     #[clap(default_value = "/tmp/text-embeddings-inference-server", long, env)]
@@ -258,6 +267,7 @@ async fn main() -> Result<()> {
         token,
         Some(args.hostname),
         args.port,
+        args.root_path,
         Some(args.uds_path),
         args.huggingface_hub_cache,
         args.payload_limit,