From 526e12ac12029cc6abbc57bd803d9054e3b17f11 Mon Sep 17 00:00:00 2001 From: Ashwin Kumar Sivakumar Date: Thu, 16 Jul 2026 20:29:11 +0530 Subject: [PATCH] Resolve per-account LiteLLM key instead of always using the shared master key Adds UserKeyClient, which fetches (and caches) an account's LiteLLM virtual key from the users service's internal endpoint. AiProvider gains complete_as_user(), and LiteLLMProvider resolves the caller's key when available, falling back to the master key otherwise so existing behavior is unchanged where no per-account key is wired up yet. Co-Authored-By: Claude Sonnet 5 --- src/main.rs | 23 +++++--- src/providers/llm/ai_provider.rs | 15 ++++++ src/providers/llm/litellm_provider.rs | 44 ++++++++++++++-- src/providers/llm/mod.rs | 1 + src/providers/llm/user_key_client.rs | 75 +++++++++++++++++++++++++++ 5 files changed, 146 insertions(+), 12 deletions(-) create mode 100644 src/providers/llm/user_key_client.rs diff --git a/src/main.rs b/src/main.rs index 214a543..0dc8386 100644 --- a/src/main.rs +++ b/src/main.rs @@ -28,6 +28,7 @@ use providers::help_center::nxtgauge_help_center_provider::NxtgaugeHelpCenterPro use providers::llm::fake_provider::FakeAiProvider; use providers::llm::litellm_provider::LiteLLMProvider; use providers::llm::ollama_provider::OllamaAiProvider; +use providers::llm::user_key_client::UserKeyClient; use providers::tickets::nxtgauge_ticket_provider::NxtgaugeTicketProvider; use providers::llm::ai_provider::AiProvider; use retrieval::embeddings::ollama_embedding_provider::OllamaEmbeddingProvider; @@ -48,6 +49,8 @@ async fn main() { } }; + let http_client = reqwest::Client::new(); + let ai_provider: Arc = match cfg.llm_provider.as_str() { "fake" => { info!("Using Fake AI provider for testing"); @@ -55,11 +58,19 @@ async fn main() { } "litellm" => { info!("Using LiteLLM provider with model {}", cfg.litellm_model); - Arc::new(LiteLLMProvider::new( - cfg.litellm_base_url.clone(), - cfg.litellm_api_key.clone(), - cfg.litellm_model.clone(), - )) as Arc + let user_key_client = Arc::new(UserKeyClient::new( + http_client.clone(), + cfg.nxtgauge_users_url.clone(), + cfg.ai_service_key.clone(), + )); + Arc::new( + LiteLLMProvider::new( + cfg.litellm_base_url.clone(), + cfg.litellm_api_key.clone(), + cfg.litellm_model.clone(), + ) + .with_user_key_client(user_key_client), + ) as Arc } _ => { info!("Using Ollama provider with model {}", cfg.ollama_chat_model); @@ -75,8 +86,6 @@ async fn main() { cfg.ollama_embed_model.clone(), )); - let http_client = reqwest::Client::new(); - let help_center_provider = Arc::new(NxtgaugeHelpCenterProvider::new( http_client.clone(), cfg.nxtgauge_users_url.clone(), diff --git a/src/providers/llm/ai_provider.rs b/src/providers/llm/ai_provider.rs index 1165475..bfa5b1a 100644 --- a/src/providers/llm/ai_provider.rs +++ b/src/providers/llm/ai_provider.rs @@ -29,4 +29,19 @@ pub trait AiProvider: Send + Sync { ) -> Result { self.complete(system_prompt, user_prompt).await } + + /// Same as `complete_as`, but scoped to a specific account's LiteLLM key + /// when the provider supports per-account credentials (`user_id` is the + /// nxtgauge account id). Providers that don't support this (Ollama, the + /// fake test provider) or that have no key for this user fall back to + /// `complete_as`, which uses the shared/master credential. + async fn complete_as_user( + &self, + _user_id: Option<&str>, + model: &str, + system_prompt: &str, + user_prompt: &str, + ) -> Result { + self.complete_as(model, system_prompt, user_prompt).await + } } diff --git a/src/providers/llm/litellm_provider.rs b/src/providers/llm/litellm_provider.rs index 91efef8..f9d71bb 100644 --- a/src/providers/llm/litellm_provider.rs +++ b/src/providers/llm/litellm_provider.rs @@ -1,8 +1,13 @@ +use std::sync::Arc; + use async_trait::async_trait; use reqwest::Client; use serde::{Deserialize, Serialize}; -use crate::{error::AppError, providers::llm::ai_provider::AiProvider}; +use crate::{ + error::AppError, + providers::llm::{ai_provider::AiProvider, user_key_client::UserKeyClient}, +}; #[derive(Clone)] pub struct LiteLLMProvider { @@ -10,6 +15,7 @@ pub struct LiteLLMProvider { base_url: String, api_key: String, model: String, + user_key_client: Option>, } impl LiteLLMProvider { @@ -19,9 +25,16 @@ impl LiteLLMProvider { base_url, api_key, model, + user_key_client: None, } } + /// Enables per-account key resolution for `complete_as_user` calls. + pub fn with_user_key_client(mut self, user_key_client: Arc) -> Self { + self.user_key_client = Some(user_key_client); + self + } + fn fallback_response(user_prompt: &str) -> String { format!( "LiteLLM model is unavailable right now. I captured your request so workflow can continue: {}", @@ -29,7 +42,13 @@ impl LiteLLMProvider { ) } - async fn chat(&self, model: &str, system_prompt: &str, user_prompt: &str) -> Result { + async fn chat( + &self, + api_key: &str, + model: &str, + system_prompt: &str, + user_prompt: &str, + ) -> Result { let url = format!("{}/chat/completions", self.base_url.trim_end_matches('/')); let payload = ChatCompletionRequest { @@ -51,7 +70,7 @@ impl LiteLLMProvider { let res = self .client .post(url) - .header("Authorization", format!("Bearer {}", self.api_key)) + .header("Authorization", format!("Bearer {}", api_key)) .header("Content-Type", "application/json") .json(&payload) .send() @@ -118,7 +137,7 @@ struct Message { impl AiProvider for LiteLLMProvider { async fn complete(&self, system_prompt: &str, user_prompt: &str) -> Result { let model = self.model.clone(); - self.chat(&model, system_prompt, user_prompt).await + self.chat(&self.api_key, &model, system_prompt, user_prompt).await } async fn complete_as( @@ -127,6 +146,21 @@ impl AiProvider for LiteLLMProvider { system_prompt: &str, user_prompt: &str, ) -> Result { - self.chat(model, system_prompt, user_prompt).await + self.chat(&self.api_key, model, system_prompt, user_prompt).await + } + + async fn complete_as_user( + &self, + user_id: Option<&str>, + model: &str, + system_prompt: &str, + user_prompt: &str, + ) -> Result { + let key = match (user_id, &self.user_key_client) { + (Some(uid), Some(client)) => client.get_key(uid).await, + _ => None, + }; + self.chat(key.as_deref().unwrap_or(&self.api_key), model, system_prompt, user_prompt) + .await } } diff --git a/src/providers/llm/mod.rs b/src/providers/llm/mod.rs index 04f258f..2c7b510 100644 --- a/src/providers/llm/mod.rs +++ b/src/providers/llm/mod.rs @@ -2,3 +2,4 @@ pub mod ai_provider; pub mod fake_provider; pub mod litellm_provider; pub mod ollama_provider; +pub mod user_key_client; diff --git a/src/providers/llm/user_key_client.rs b/src/providers/llm/user_key_client.rs new file mode 100644 index 0000000..b56c395 --- /dev/null +++ b/src/providers/llm/user_key_client.rs @@ -0,0 +1,75 @@ +use std::collections::HashMap; + +use reqwest::Client; +use serde::Deserialize; +use tokio::sync::RwLock; + +/// Fetches (and caches) a per-account LiteLLM virtual key from the users +/// service's internal `/internal/users/{id}/llm-key` endpoint, so LLM calls +/// can be attributed/billed to the account making them instead of always +/// using the shared LiteLLM master key. +pub struct UserKeyClient { + client: Client, + base_url: String, + ai_service_key: String, + cache: RwLock>, +} + +#[derive(Debug, Deserialize)] +struct LlmKeyResponse { + key: String, +} + +impl UserKeyClient { + pub fn new(client: Client, base_url: String, ai_service_key: String) -> Self { + Self { + client, + base_url, + ai_service_key, + cache: RwLock::new(HashMap::new()), + } + } + + /// Returns the account's LiteLLM key, or `None` if it can't be fetched + /// (users service unreachable, key not configured, etc). Callers should + /// fall back to the shared master key on `None`. + pub async fn get_key(&self, user_id: &str) -> Option { + if let Some(key) = self.cache.read().await.get(user_id) { + return Some(key.clone()); + } + + if self.ai_service_key.is_empty() { + return None; + } + + let url = format!( + "{}/internal/users/{}/llm-key", + self.base_url.trim_end_matches('/'), + user_id + ); + + let response = self + .client + .get(&url) + .header("X-AI-Service-Key", &self.ai_service_key) + .send() + .await + .ok()?; + + if !response.status().is_success() { + tracing::warn!( + user_id, + status = %response.status(), + "Failed to fetch per-user LiteLLM key; falling back to master key" + ); + return None; + } + + let body: LlmKeyResponse = response.json().await.ok()?; + self.cache + .write() + .await + .insert(user_id.to_string(), body.key.clone()); + Some(body.key) + } +}