[router] Refactor: decouple select and send stage (#2440)

This commit is contained in:
Byron Hsu
2024-12-11 00:51:21 -08:00
committed by GitHub
parent 7310aede97
commit d4de9a6235
2 changed files with 107 additions and 101 deletions

View File

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