From fb5a5971d32c3ba83f148dfb1ce4ceb12158c061 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E8=83=A1=E9=94=8B?= <fengmi99@gmail.com>
Date: Thu, 21 Sep 2023 15:06:51 +0800
Subject: [PATCH] feat: proxy server address mapping to the model server (#461)

* feat: proxy server address mapping to the model server

* fix: add swagger in Config

* refactor: add_proxy_server

* fix: missing semicolo
---
 crates/tabby-common/src/config.rs             |  8 +++++++
 .../tabby-scheduler/tests/integration_test.rs |  3 ++-
 crates/tabby/src/serve/mod.rs                 | 23 +++++++++++++++++++
 3 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/crates/tabby-common/src/config.rs b/crates/tabby-common/src/config.rs
index 4c1a406..281ce22 100644
--- a/crates/tabby-common/src/config.rs
+++ b/crates/tabby-common/src/config.rs
@@ -15,6 +15,9 @@ pub struct Config {
 
     #[serde(default)]
     pub experimental: Experimental,
+
+    #[serde(default)]
+    pub swagger: SwaggerConfig,
 }
 
 #[derive(Serialize, Deserialize, Default)]
@@ -23,6 +26,11 @@ pub struct Experimental {
     pub enable_prompt_rewrite: bool,
 }
 
+#[derive(Serialize, Deserialize, Default)]
+pub struct SwaggerConfig {
+    pub server_url: Option<String>,
+}
+
 impl Config {
     pub fn load() -> Result<Self, Error> {
         let file = serdeconv::from_toml_file(crate::path::config_file().as_path());
diff --git a/crates/tabby-scheduler/tests/integration_test.rs b/crates/tabby-scheduler/tests/integration_test.rs
index bb80a1b..2fcde8d 100644
--- a/crates/tabby-scheduler/tests/integration_test.rs
+++ b/crates/tabby-scheduler/tests/integration_test.rs
@@ -3,7 +3,7 @@ mod tests {
     use std::fs::create_dir_all;
 
     use tabby_common::{
-        config::{Config, Experimental, Repository},
+        config::{Config, Experimental, Repository, SwaggerConfig},
         path::set_tabby_root,
     };
     use temp_testdir::*;
@@ -20,6 +20,7 @@ mod tests {
             repositories: vec![Repository {
                 git_url: "https://github.com/TabbyML/interview-questions".to_owned(),
             }],
+            swagger: SwaggerConfig { server_url: None },
             experimental: Experimental::default(),
         };
 
diff --git a/crates/tabby/src/serve/mod.rs b/crates/tabby/src/serve/mod.rs
index 3ee40d5..943ca7a 100644
--- a/crates/tabby/src/serve/mod.rs
+++ b/crates/tabby/src/serve/mod.rs
@@ -162,6 +162,7 @@ pub async fn main(config: &Config, args: &ServeArgs) {
     info!("Starting server, this might takes a few minutes...");
 
     let doc = add_localhost_server(ApiDoc::openapi(), args.port);
+    let doc = add_proxy_server(doc, config.swagger.server_url.clone());
     let app = Router::new()
         .merge(SwaggerUi::new("/swagger-ui").url("/api-docs/openapi.json", doc))
         .nest("/v1", api_router(args, config))
@@ -238,3 +239,25 @@ fn add_localhost_server(doc: utoipa::openapi::OpenApi, port: u16) -> utoipa::ope
 
     doc
 }
+
+fn add_proxy_server(
+    doc: utoipa::openapi::OpenApi,
+    server_url: Option<String>,
+) -> utoipa::openapi::OpenApi {
+    if server_url.is_none() {
+        return doc;
+    }
+
+    let server_url: String = server_url.unwrap();
+    let mut doc = doc;
+    if let Some(servers) = doc.servers.as_mut() {
+        servers.push(
+            ServerBuilder::new()
+                .url(server_url)
+                .description(Some("Swagger Server"))
+                .build(),
+        );
+    }
+
+    doc
+}