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 <noreply@anthropic.com>

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<dyn AiProvider> = 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(
+            let user_key_client = Arc::new(UserKeyClient::new(
-                cfg.litellm_base_url.clone(),
+                http_client.clone(),
-                cfg.litellm_api_key.clone(),
+                cfg.nxtgauge_users_url.clone(),
-                cfg.litellm_model.clone(),
+                cfg.ai_service_key.clone(),
-            )) as Arc<dyn AiProvider>
+            ));
+            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<dyn AiProvider>
         }
         _ => {
             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(),

src/providers/llm/ai_provider.rs

@@ -29,4 +29,19 @@ pub trait AiProvider: Send + Sync {
     ) -> Result<String, AppError> {
         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<String, AppError> {
+        self.complete_as(model, system_prompt, user_prompt).await
+    }
 }

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<Arc<UserKeyClient>>,
 }
 
 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<UserKeyClient>) -> 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<String, AppError> {
+    async fn chat(
+        &self,
+        api_key: &str,
+        model: &str,
+        system_prompt: &str,
+        user_prompt: &str,
+    ) -> Result<String, AppError> {
         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<String, AppError> {
         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<String, AppError> {
-        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<String, AppError> {
+        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
     }
 }

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;

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<HashMap<String, String>>,
+}
+
+#[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<String> {
+        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)
+    }
+}