Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 10 additions & 4 deletions crates/nilai-api/src/routes/pricing.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ pub fn router() -> Router<AppState> {
Router::new()
.route("/v1/pricing", get(get_all_prices))
.route(
"/v1/pricing/*model_name",
"/v1/pricing/{*model_name}",
get(get_price).put(set_price).delete(delete_price),
)
}
Expand Down Expand Up @@ -41,6 +41,8 @@ async fn get_price(
State(state): State<AppState>,
Path(model_name): Path<String>,
) -> Result<Json<LLMPriceConfig>, (StatusCode, Json<serde_json::Value>)> {
let model_name = model_name.strip_prefix('/').unwrap_or(&model_name);

let store = state.pricing_store.as_ref().ok_or_else(|| {
(
StatusCode::SERVICE_UNAVAILABLE,
Expand All @@ -49,7 +51,7 @@ async fn get_price(
})?;

let price = store
.get_price(&ModelName::new(&model_name))
.get_price(&ModelName::new(model_name))
.await
.map_err(|e| {
(
Expand All @@ -67,6 +69,8 @@ async fn set_price(
headers: HeaderMap,
Json(config): Json<LLMPriceConfig>,
) -> Result<Json<LLMPriceConfig>, (StatusCode, Json<serde_json::Value>)> {
let model_name = model_name.strip_prefix('/').unwrap_or(&model_name);

// Verify admin token
verify_admin_token(&state, &headers)?;

Expand All @@ -89,7 +93,7 @@ async fn set_price(
})?;

store
.set_price(&ModelName::new(&model_name), &config)
.set_price(&ModelName::new(model_name), &config)
.await
.map_err(|e| {
(
Expand All @@ -106,6 +110,8 @@ async fn delete_price(
Path(model_name): Path<String>,
headers: HeaderMap,
) -> Result<StatusCode, (StatusCode, Json<serde_json::Value>)> {
let model_name = model_name.strip_prefix('/').unwrap_or(&model_name);

verify_admin_token(&state, &headers)?;

if model_name == "default" {
Expand All @@ -123,7 +129,7 @@ async fn delete_price(
})?;

let existed = store
.delete_price(&ModelName::new(&model_name))
.delete_price(&ModelName::new(model_name))
.await
.map_err(|e| {
(
Expand Down
2 changes: 1 addition & 1 deletion docker/nilai-api.Dockerfile
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
# Stage 1: Build
FROM rust:1.80-bookworm AS builder
FROM rust:1.88-bookworm AS builder

WORKDIR /app

Expand Down
Loading