Issue each account its own LiteLLM virtual key instead of the shared master key
All checks were successful
build-and-release / build (cron) (push) Successful in 5m1s
build-and-release / build (catering-services) (push) Successful in 9m1s
build-and-release / build (customers) (push) Successful in 9m3s
build-and-release / build (developers) (push) Successful in 9m23s
build-and-release / build (employees) (push) Successful in 10m36s
build-and-release / build (companies) (push) Successful in 10m59s
build-and-release / build (gateway) (push) Successful in 3m26s
build-and-release / build (fitness-trainers) (push) Successful in 8m52s
build-and-release / build (jobs) (push) Successful in 4m46s
build-and-release / build (graphic-designers) (push) Successful in 8m44s
build-and-release / build (job-seekers) (push) Successful in 9m22s
build-and-release / build (makeup-artists) (push) Successful in 8m33s
build-and-release / build (leads) (push) Successful in 10m17s
build-and-release / build (payments) (push) Successful in 8m40s
build-and-release / build (photographers) (push) Successful in 9m36s
build-and-release / build (social-media-managers) (push) Successful in 8m38s
build-and-release / build (tutors) (push) Successful in 8m42s
build-and-release / build (ugc-content-creators) (push) Successful in 7m27s
build-and-release / build (video-editors) (push) Successful in 7m46s
build-and-release / build (users) (push) Successful in 10m0s
All checks were successful
build-and-release / build (cron) (push) Successful in 5m1s
build-and-release / build (catering-services) (push) Successful in 9m1s
build-and-release / build (customers) (push) Successful in 9m3s
build-and-release / build (developers) (push) Successful in 9m23s
build-and-release / build (employees) (push) Successful in 10m36s
build-and-release / build (companies) (push) Successful in 10m59s
build-and-release / build (gateway) (push) Successful in 3m26s
build-and-release / build (fitness-trainers) (push) Successful in 8m52s
build-and-release / build (jobs) (push) Successful in 4m46s
build-and-release / build (graphic-designers) (push) Successful in 8m44s
build-and-release / build (job-seekers) (push) Successful in 9m22s
build-and-release / build (makeup-artists) (push) Successful in 8m33s
build-and-release / build (leads) (push) Successful in 10m17s
build-and-release / build (payments) (push) Successful in 8m40s
build-and-release / build (photographers) (push) Successful in 9m36s
build-and-release / build (social-media-managers) (push) Successful in 8m38s
build-and-release / build (tutors) (push) Successful in 8m42s
build-and-release / build (ugc-content-creators) (push) Successful in 7m27s
build-and-release / build (video-editors) (push) Successful in 7m46s
build-and-release / build (users) (push) Successful in 10m0s
register() now generates a per-account LiteLLM key (best-effort, non-blocking)
and stores it on the user. New internal endpoint GET /internal/users/{id}/llm-key
lets other services fetch (or lazily backfill) an account's key, authenticated
via the existing X-AI-Service-Key shared secret.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
parent
74664176da
commit
f5201965d8
8 changed files with 177 additions and 5 deletions
|
|
@ -181,6 +181,36 @@ impl LiteLlmClient {
|
||||||
Ok((content, usage, request_id))
|
Ok((content, usage, request_id))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Issues a per-account LiteLLM virtual key via `/key/generate`, scoped to `user_id`.
|
||||||
|
/// Requires `self.api_key` to be the LiteLLM master key.
|
||||||
|
pub async fn generate_key(&self, user_id: uuid::Uuid) -> Result<String, LiteLlmError> {
|
||||||
|
let url = format!("{}/key/generate", self.base_url);
|
||||||
|
|
||||||
|
let mut req = self.client.post(&url).json(&serde_json::json!({
|
||||||
|
"user_id": user_id.to_string(),
|
||||||
|
"key_alias": format!("user-{}", user_id),
|
||||||
|
}));
|
||||||
|
if let Some(key) = &self.api_key {
|
||||||
|
req = req.header("Authorization", format!("Bearer {}", key));
|
||||||
|
}
|
||||||
|
|
||||||
|
let response = req.send().await?;
|
||||||
|
let status = response.status();
|
||||||
|
if !status.is_success() {
|
||||||
|
let body = response.text().await.unwrap_or_default();
|
||||||
|
return Err(LiteLlmError::Api {
|
||||||
|
status: status.as_u16(),
|
||||||
|
body,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
let body: Value = response.json().await?;
|
||||||
|
body.get("key")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.map(|s| s.to_string())
|
||||||
|
.ok_or(LiteLlmError::NoCompletion)
|
||||||
|
}
|
||||||
|
|
||||||
/// Direct pass-through for callers that want the raw JSON response.
|
/// Direct pass-through for callers that want the raw JSON response.
|
||||||
#[allow(dead_code)]
|
#[allow(dead_code)]
|
||||||
pub async fn raw_chat_completion(&self, body: Value) -> Result<Value, LiteLlmError> {
|
pub async fn raw_chat_completion(&self, body: Value) -> Result<Value, LiteLlmError> {
|
||||||
|
|
|
||||||
|
|
@ -308,6 +308,21 @@ async fn register(
|
||||||
}
|
}
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
|
// Issue this account its own LiteLLM virtual key (best-effort — a LiteLLM
|
||||||
|
// outage must not block registration; the internal /llm-key endpoint will
|
||||||
|
// lazily generate one on first use if this fails).
|
||||||
|
match crate::ai::litellm::LiteLlmClient::new() {
|
||||||
|
Ok(client) => match client.generate_key(user.id).await {
|
||||||
|
Ok(key) => {
|
||||||
|
if let Err(e) = db::models::user::UserRepository::set_litellm_key(&state.pool, user.id, &key).await {
|
||||||
|
tracing::error!(user_id = %user.id, error = %e, "Failed to store LiteLLM key after registration");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => tracing::warn!(user_id = %user.id, error = %e, "Failed to generate LiteLLM key at registration"),
|
||||||
|
},
|
||||||
|
Err(e) => tracing::warn!(error = %e, "LiteLLM client unavailable at registration"),
|
||||||
|
}
|
||||||
|
|
||||||
// Check if this is a demo account (payment gateway integration)
|
// Check if this is a demo account (payment gateway integration)
|
||||||
let is_demo_account = is_dummy_account_email(&email);
|
let is_demo_account = is_dummy_account_email(&email);
|
||||||
|
|
||||||
|
|
|
||||||
90
apps/users/src/handlers/internal.rs
Normal file
90
apps/users/src/handlers/internal.rs
Normal file
|
|
@ -0,0 +1,90 @@
|
||||||
|
use crate::AppState;
|
||||||
|
use axum::{
|
||||||
|
extract::{Path, State},
|
||||||
|
http::{HeaderMap, StatusCode},
|
||||||
|
response::IntoResponse,
|
||||||
|
routing::get,
|
||||||
|
Json, Router,
|
||||||
|
};
|
||||||
|
use db::models::user::UserRepository;
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
/// Internal, service-to-service routes (not exposed on the public API gateway).
|
||||||
|
/// Callers authenticate with the shared `X-AI-Service-Key` header, same as the
|
||||||
|
/// existing `/api/support/tickets/ai/create` endpoint.
|
||||||
|
pub fn router() -> Router<AppState> {
|
||||||
|
Router::new().route("/{id}/llm-key", get(get_llm_key))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn check_service_key(headers: &HeaderMap) -> bool {
|
||||||
|
let expected_key = std::env::var("AI_SERVICE_KEY").unwrap_or_default();
|
||||||
|
let provided_key = headers
|
||||||
|
.get("X-AI-Service-Key")
|
||||||
|
.and_then(|v| v.to_str().ok())
|
||||||
|
.unwrap_or_default();
|
||||||
|
!expected_key.is_empty() && provided_key == expected_key
|
||||||
|
}
|
||||||
|
|
||||||
|
/// GET /internal/users/{id}/llm-key
|
||||||
|
///
|
||||||
|
/// Returns this account's per-user LiteLLM virtual key, generating and
|
||||||
|
/// persisting one on the fly if it doesn't have one yet (covers accounts
|
||||||
|
/// created before this feature existed, or where generation failed at
|
||||||
|
/// registration time because LiteLLM was unreachable).
|
||||||
|
async fn get_llm_key(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Path(id): Path<Uuid>,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
if !check_service_key(&headers) {
|
||||||
|
return (
|
||||||
|
StatusCode::UNAUTHORIZED,
|
||||||
|
Json(serde_json::json!({ "error": "Invalid or missing AI service key" })),
|
||||||
|
)
|
||||||
|
.into_response();
|
||||||
|
}
|
||||||
|
|
||||||
|
match UserRepository::get_litellm_key(&state.pool, id).await {
|
||||||
|
Ok(Some(key)) => (StatusCode::OK, Json(serde_json::json!({ "key": key }))).into_response(),
|
||||||
|
Ok(None) => generate_and_store(&state, id).await,
|
||||||
|
Err(e) => {
|
||||||
|
tracing::error!(user_id = %id, error = %e, "Failed to look up LiteLLM key");
|
||||||
|
(
|
||||||
|
StatusCode::INTERNAL_SERVER_ERROR,
|
||||||
|
Json(serde_json::json!({ "error": "Failed to look up LiteLLM key" })),
|
||||||
|
)
|
||||||
|
.into_response()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn generate_and_store(state: &AppState, user_id: Uuid) -> axum::response::Response {
|
||||||
|
let client = match crate::ai::litellm::LiteLlmClient::new() {
|
||||||
|
Ok(c) => c,
|
||||||
|
Err(e) => {
|
||||||
|
tracing::error!(error = %e, "LiteLLM client unavailable");
|
||||||
|
return (
|
||||||
|
StatusCode::BAD_GATEWAY,
|
||||||
|
Json(serde_json::json!({ "error": "LiteLLM unavailable" })),
|
||||||
|
)
|
||||||
|
.into_response();
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
match client.generate_key(user_id).await {
|
||||||
|
Ok(key) => {
|
||||||
|
if let Err(e) = UserRepository::set_litellm_key(&state.pool, user_id, &key).await {
|
||||||
|
tracing::error!(user_id = %user_id, error = %e, "Failed to store LiteLLM key");
|
||||||
|
}
|
||||||
|
(StatusCode::OK, Json(serde_json::json!({ "key": key }))).into_response()
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::error!(user_id = %user_id, error = %e, "Failed to generate LiteLLM key");
|
||||||
|
(
|
||||||
|
StatusCode::BAD_GATEWAY,
|
||||||
|
Json(serde_json::json!({ "error": "Failed to generate LiteLLM key" })),
|
||||||
|
)
|
||||||
|
.into_response()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -12,6 +12,7 @@ pub mod auth;
|
||||||
pub mod config;
|
pub mod config;
|
||||||
pub mod coupons;
|
pub mod coupons;
|
||||||
pub mod dashboard;
|
pub mod dashboard;
|
||||||
|
pub mod internal;
|
||||||
pub mod kb;
|
pub mod kb;
|
||||||
pub mod modules;
|
pub mod modules;
|
||||||
pub mod notifications;
|
pub mod notifications;
|
||||||
|
|
|
||||||
|
|
@ -128,6 +128,8 @@ async fn main() {
|
||||||
.nest("/api/ai", handlers::ai::ai_router())
|
.nest("/api/ai", handlers::ai::ai_router())
|
||||||
// ── Ask Ash AI Credits (Phase 1) ────────────────────────────────────
|
// ── Ask Ash AI Credits (Phase 1) ────────────────────────────────────
|
||||||
.nest("/api/ai-credits", handlers::ai_credits::router())
|
.nest("/api/ai-credits", handlers::ai_credits::router())
|
||||||
|
|
||||||
|
.nest("/internal/users", handlers::internal::router())
|
||||||
.nest("/api/admin/ai-credits", handlers::ai_credits::admin_router())
|
.nest("/api/admin/ai-credits", handlers::ai_credits::admin_router())
|
||||||
.route("/health", get(|| async { "Users OK" }))
|
.route("/health", get(|| async { "Users OK" }))
|
||||||
.with_state(state);
|
.with_state(state);
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,5 @@
|
||||||
|
BEGIN;
|
||||||
|
|
||||||
|
ALTER TABLE users DROP COLUMN litellm_key;
|
||||||
|
|
||||||
|
COMMIT;
|
||||||
|
|
@ -0,0 +1,8 @@
|
||||||
|
-- Per-account LiteLLM virtual key, so AI features are billed/tracked per user
|
||||||
|
-- instead of every account sharing the LiteLLM master key.
|
||||||
|
|
||||||
|
BEGIN;
|
||||||
|
|
||||||
|
ALTER TABLE users ADD COLUMN litellm_key TEXT;
|
||||||
|
|
||||||
|
COMMIT;
|
||||||
|
|
@ -22,6 +22,7 @@ pub struct User {
|
||||||
pub created_at: DateTime<Utc>,
|
pub created_at: DateTime<Utc>,
|
||||||
pub updated_at: DateTime<Utc>,
|
pub updated_at: DateTime<Utc>,
|
||||||
pub deleted_at: Option<DateTime<Utc>>,
|
pub deleted_at: Option<DateTime<Utc>>,
|
||||||
|
pub litellm_key: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -58,7 +59,7 @@ impl UserRepository {
|
||||||
email_verified, phone_verified, status,
|
email_verified, phone_verified, status,
|
||||||
email_verification_token, email_verification_expires_at,
|
email_verification_token, email_verification_expires_at,
|
||||||
reset_password_token, reset_password_expires_at,
|
reset_password_token, reset_password_expires_at,
|
||||||
created_at, updated_at, deleted_at
|
created_at, updated_at, deleted_at, litellm_key
|
||||||
"#,
|
"#,
|
||||||
)
|
)
|
||||||
.bind(&payload.first_name)
|
.bind(&payload.first_name)
|
||||||
|
|
@ -78,7 +79,7 @@ impl UserRepository {
|
||||||
email_verified, phone_verified, status,
|
email_verified, phone_verified, status,
|
||||||
email_verification_token, email_verification_expires_at,
|
email_verification_token, email_verification_expires_at,
|
||||||
reset_password_token, reset_password_expires_at,
|
reset_password_token, reset_password_expires_at,
|
||||||
created_at, updated_at, deleted_at
|
created_at, updated_at, deleted_at, litellm_key
|
||||||
FROM users
|
FROM users
|
||||||
WHERE email = $1 AND deleted_at IS NULL
|
WHERE email = $1 AND deleted_at IS NULL
|
||||||
"#,
|
"#,
|
||||||
|
|
@ -95,7 +96,7 @@ impl UserRepository {
|
||||||
email_verified, phone_verified, status,
|
email_verified, phone_verified, status,
|
||||||
email_verification_token, email_verification_expires_at,
|
email_verification_token, email_verification_expires_at,
|
||||||
reset_password_token, reset_password_expires_at,
|
reset_password_token, reset_password_expires_at,
|
||||||
created_at, updated_at, deleted_at
|
created_at, updated_at, deleted_at, litellm_key
|
||||||
FROM users
|
FROM users
|
||||||
WHERE id = $1 AND deleted_at IS NULL
|
WHERE id = $1 AND deleted_at IS NULL
|
||||||
"#,
|
"#,
|
||||||
|
|
@ -152,7 +153,7 @@ impl UserRepository {
|
||||||
email_verified, phone_verified, status,
|
email_verified, phone_verified, status,
|
||||||
email_verification_token, email_verification_expires_at,
|
email_verification_token, email_verification_expires_at,
|
||||||
reset_password_token, reset_password_expires_at,
|
reset_password_token, reset_password_expires_at,
|
||||||
created_at, updated_at, deleted_at
|
created_at, updated_at, deleted_at, litellm_key
|
||||||
FROM users
|
FROM users
|
||||||
WHERE email_verification_token = $1 AND deleted_at IS NULL
|
WHERE email_verification_token = $1 AND deleted_at IS NULL
|
||||||
"#,
|
"#,
|
||||||
|
|
@ -200,7 +201,7 @@ impl UserRepository {
|
||||||
email_verified, phone_verified, status,
|
email_verified, phone_verified, status,
|
||||||
email_verification_token, email_verification_expires_at,
|
email_verification_token, email_verification_expires_at,
|
||||||
reset_password_token, reset_password_expires_at,
|
reset_password_token, reset_password_expires_at,
|
||||||
created_at, updated_at, deleted_at
|
created_at, updated_at, deleted_at, litellm_key
|
||||||
FROM users
|
FROM users
|
||||||
WHERE reset_password_token = $1 AND deleted_at IS NULL
|
WHERE reset_password_token = $1 AND deleted_at IS NULL
|
||||||
"#,
|
"#,
|
||||||
|
|
@ -242,6 +243,26 @@ impl UserRepository {
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub async fn set_litellm_key(pool: &PgPool, user_id: Uuid, key: &str) -> Result<(), sqlx::Error> {
|
||||||
|
sqlx::query(
|
||||||
|
"UPDATE users SET litellm_key = $1, updated_at = NOW() WHERE id = $2",
|
||||||
|
)
|
||||||
|
.bind(key)
|
||||||
|
.bind(user_id)
|
||||||
|
.execute(pool)
|
||||||
|
.await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_litellm_key(pool: &PgPool, user_id: Uuid) -> Result<Option<String>, sqlx::Error> {
|
||||||
|
sqlx::query_scalar::<_, Option<String>>(
|
||||||
|
"SELECT litellm_key FROM users WHERE id = $1 AND deleted_at IS NULL",
|
||||||
|
)
|
||||||
|
.bind(user_id)
|
||||||
|
.fetch_one(pool)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
pub async fn store_refresh_token(
|
pub async fn store_refresh_token(
|
||||||
pool: &PgPool,
|
pool: &PgPool,
|
||||||
user_id: Uuid,
|
user_id: Uuid,
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue