[router] Refactor: decouple select and send stage (#2440)
This commit is contained in:
@@ -29,84 +29,41 @@ impl AppState {
|
||||
}
|
||||
}
|
||||
|
||||
async fn forward_request(
|
||||
client: &reqwest::Client,
|
||||
worker_url: String,
|
||||
route: String,
|
||||
) -> HttpResponse {
|
||||
match client.get(format!("{}{}", worker_url, route)).send().await {
|
||||
Ok(res) => {
|
||||
let status = actix_web::http::StatusCode::from_u16(res.status().as_u16())
|
||||
.unwrap_or(actix_web::http::StatusCode::INTERNAL_SERVER_ERROR);
|
||||
|
||||
// print the status
|
||||
println!(
|
||||
"Forwarding Request Worker URL: {}, Route: {}, Status: {}",
|
||||
worker_url, route, status
|
||||
);
|
||||
match res.bytes().await {
|
||||
Ok(body) => HttpResponse::build(status).body(body.to_vec()),
|
||||
Err(_) => HttpResponse::InternalServerError().finish(),
|
||||
}
|
||||
}
|
||||
Err(_) => HttpResponse::InternalServerError().finish(),
|
||||
}
|
||||
}
|
||||
|
||||
#[get("/health")]
|
||||
async fn health(data: web::Data<AppState>) -> impl Responder {
|
||||
let worker_url = match data.router.get_first() {
|
||||
Some(url) => url,
|
||||
None => return HttpResponse::InternalServerError().finish(),
|
||||
};
|
||||
|
||||
forward_request(&data.client, worker_url, "/health".to_string()).await
|
||||
data.router.route_to_first(&data.client, "/health").await
|
||||
}
|
||||
|
||||
#[get("/health_generate")]
|
||||
async fn health_generate(data: web::Data<AppState>) -> impl Responder {
|
||||
let worker_url = match data.router.get_first() {
|
||||
Some(url) => url,
|
||||
None => return HttpResponse::InternalServerError().finish(),
|
||||
};
|
||||
|
||||
forward_request(&data.client, worker_url, "/health_generate".to_string()).await
|
||||
data.router
|
||||
.route_to_first(&data.client, "/health_generate")
|
||||
.await
|
||||
}
|
||||
|
||||
#[get("/get_server_info")]
|
||||
async fn get_server_info(data: web::Data<AppState>) -> impl Responder {
|
||||
let worker_url = match data.router.get_first() {
|
||||
Some(url) => url,
|
||||
None => return HttpResponse::InternalServerError().finish(),
|
||||
};
|
||||
|
||||
forward_request(&data.client, worker_url, "/get_server_info".to_string()).await
|
||||
data.router
|
||||
.route_to_first(&data.client, "/get_server_info")
|
||||
.await
|
||||
}
|
||||
|
||||
#[get("/v1/models")]
|
||||
async fn v1_models(data: web::Data<AppState>) -> impl Responder {
|
||||
let worker_url = match data.router.get_first() {
|
||||
Some(url) => url,
|
||||
None => return HttpResponse::InternalServerError().finish(),
|
||||
};
|
||||
|
||||
forward_request(&data.client, worker_url, "/v1/models".to_string()).await
|
||||
data.router.route_to_first(&data.client, "/v1/models").await
|
||||
}
|
||||
|
||||
#[get("/get_model_info")]
|
||||
async fn get_model_info(data: web::Data<AppState>) -> impl Responder {
|
||||
let worker_url = match data.router.get_first() {
|
||||
Some(url) => url,
|
||||
None => return HttpResponse::InternalServerError().finish(),
|
||||
};
|
||||
|
||||
forward_request(&data.client, worker_url, "/get_model_info".to_string()).await
|
||||
data.router
|
||||
.route_to_first(&data.client, "/get_model_info")
|
||||
.await
|
||||
}
|
||||
|
||||
#[post("/generate")]
|
||||
async fn generate(req: HttpRequest, body: Bytes, data: web::Data<AppState>) -> impl Responder {
|
||||
data.router
|
||||
.dispatch(&data.client, req, body, "generate")
|
||||
.route_generate_request(&data.client, req, body, "/generate")
|
||||
.await
|
||||
}
|
||||
|
||||
@@ -117,7 +74,7 @@ async fn v1_chat_completions(
|
||||
data: web::Data<AppState>,
|
||||
) -> impl Responder {
|
||||
data.router
|
||||
.dispatch(&data.client, req, body, "v1/chat/completions")
|
||||
.route_generate_request(&data.client, req, body, "/v1/chat/completions")
|
||||
.await
|
||||
}
|
||||
|
||||
@@ -128,7 +85,7 @@ async fn v1_completions(
|
||||
data: web::Data<AppState>,
|
||||
) -> impl Responder {
|
||||
data.router
|
||||
.dispatch(&data.client, req, body, "v1/completions")
|
||||
.route_generate_request(&data.client, req, body, "/v1/completions")
|
||||
.await
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user