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) + } +}