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
release-0.2
胡锋 2023-09-21 15:06:51 +08:00 committed by GitHub
parent 2dbfb0f35b
commit fb5a5971d3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 33 additions and 1 deletions

View File

@ -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());

View File

@ -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(),
};

View File

@ -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
}