liter_llm 1.0.0.pre.rc.6

data/vendor/liter-llm/src/provider/custom.rs

@@ -0,0 +1,404 @@
+//! Runtime registration of custom LLM providers.
+//!
+//! Allows users to register providers that are not part of the built-in
+//! `providers.json` registry.  Custom providers are checked **first** during
+//! model detection, so they can override built-in routing.
+
+use std::borrow::Cow;
+use std::sync::RwLock;
+
+use serde::{Deserialize, Serialize};
+
+use super::Provider;
+use crate::error::{LiterLlmError, Result};
+
+// ── Global custom-provider registry ──────────────────────────────────────────
+
+/// Thread-safe registry of runtime-registered custom providers.
+///
+/// Uses `RwLock` so that reads (the hot path inside `detect_provider`) only
+/// take a shared lock, while mutations (`register` / `unregister`) take an
+/// exclusive lock.
+static CUSTOM_PROVIDERS: RwLock<Vec<CustomProviderConfig>> = RwLock::new(Vec::new());
+
+/// Configuration for registering a custom LLM provider at runtime.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CustomProviderConfig {
+    /// Unique name for this provider (e.g., "my-provider").
+    pub name: String,
+    /// Base URL for the provider's API (e.g., "https://api.my-provider.com/v1").
+    pub base_url: String,
+    /// Authentication header format.
+    pub auth_header: AuthHeaderFormat,
+    /// Model name prefixes that route to this provider (e.g., ["my-"]).
+    pub model_prefixes: Vec<String>,
+}
+
+/// How the API key is sent in the HTTP request.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum AuthHeaderFormat {
+    /// Bearer token: `Authorization: Bearer <key>`
+    Bearer,
+    /// Custom header: e.g., `X-Api-Key: <key>`
+    ApiKey(String),
+    /// No authentication required.
+    None,
+}
+
+// ── Public API ───────────────────────────────────────────────────────────────
+
+/// Register a custom provider in the global runtime registry.
+///
+/// The provider will be checked **before** all built-in providers during model
+/// detection.  If a provider with the same `name` already exists it is replaced.
+///
+/// # Errors
+///
+/// Returns an error if the config is invalid (empty name, empty base_url, or
+/// no model prefixes).
+pub fn register_custom_provider(config: CustomProviderConfig) -> Result<()> {
+    validate_config(&config)?;
+
+    let mut providers = CUSTOM_PROVIDERS.write().map_err(|e| LiterLlmError::ServerError {
+        message: format!("custom provider registry lock poisoned: {e}"),
+    })?;
+
+    // Replace existing entry with the same name, or append.
+    if let Some(existing) = providers.iter_mut().find(|p| p.name == config.name) {
+        *existing = config;
+    } else {
+        providers.push(config);
+    }
+
+    Ok(())
+}
+
+/// Remove a previously registered custom provider by name.
+///
+/// Returns `true` if a provider with the given name was found and removed,
+/// `false` if no such provider existed.
+///
+/// # Errors
+///
+/// Returns an error only if the internal lock is poisoned.
+pub fn unregister_custom_provider(name: &str) -> Result<bool> {
+    let mut providers = CUSTOM_PROVIDERS.write().map_err(|e| LiterLlmError::ServerError {
+        message: format!("custom provider registry lock poisoned: {e}"),
+    })?;
+
+    let before = providers.len();
+    providers.retain(|p| p.name != name);
+    Ok(providers.len() < before)
+}
+
+/// Try to match a model name against the custom-provider registry.
+///
+/// Returns a boxed [`Provider`] if a custom provider claims the model,
+/// `None` otherwise.  This is called at the **top** of `detect_provider`
+/// so custom providers always take priority over built-in ones.
+pub(crate) fn detect_custom_provider(model: &str) -> Option<Box<dyn Provider>> {
+    let providers = CUSTOM_PROVIDERS.read().ok()?;
+
+    for cfg in providers.iter() {
+        let matches = cfg
+            .model_prefixes
+            .iter()
+            .any(|prefix| model.starts_with(prefix.as_str()));
+
+        if matches {
+            return Some(Box::new(CustomProvider { config: cfg.clone() }));
+        }
+    }
+
+    None
+}
+
+/// Clear all custom providers.  Intended for test isolation only.
+#[cfg(test)]
+pub(crate) fn clear_custom_providers() {
+    if let Ok(mut providers) = CUSTOM_PROVIDERS.write() {
+        providers.clear();
+    }
+}
+
+// ── Validation ───────────────────────────────────────────────────────────────
+
+fn validate_config(config: &CustomProviderConfig) -> Result<()> {
+    if config.name.trim().is_empty() {
+        return Err(LiterLlmError::BadRequest {
+            message: "custom provider name must not be empty or whitespace-only".into(),
+        });
+    }
+    if config.base_url.trim().is_empty() {
+        return Err(LiterLlmError::BadRequest {
+            message: "custom provider base_url must not be empty or whitespace-only".into(),
+        });
+    }
+    if config.model_prefixes.is_empty() {
+        return Err(LiterLlmError::BadRequest {
+            message: "custom provider must have at least one model prefix".into(),
+        });
+    }
+    for prefix in &config.model_prefixes {
+        if prefix.is_empty() {
+            return Err(LiterLlmError::BadRequest {
+                message: "custom provider model prefix must not be empty (would match all models)".into(),
+            });
+        }
+    }
+    Ok(())
+}
+
+// ── Provider implementation ──────────────────────────────────────────────────
+
+/// A runtime-registered custom provider.
+///
+/// Wraps a [`CustomProviderConfig`] and implements the [`Provider`] trait so
+/// the client can use it exactly like a built-in provider.
+struct CustomProvider {
+    config: CustomProviderConfig,
+}
+
+impl Provider for CustomProvider {
+    fn name(&self) -> &str {
+        &self.config.name
+    }
+
+    fn base_url(&self) -> &str {
+        &self.config.base_url
+    }
+
+    fn auth_header<'a>(&'a self, api_key: &'a str) -> Option<(Cow<'static, str>, Cow<'a, str>)> {
+        match &self.config.auth_header {
+            AuthHeaderFormat::Bearer => Some((Cow::Borrowed("Authorization"), Cow::Owned(format!("Bearer {api_key}")))),
+            AuthHeaderFormat::ApiKey(header_name) => Some((Cow::Owned(header_name.clone()), Cow::Borrowed(api_key))),
+            AuthHeaderFormat::None => None,
+        }
+    }
+
+    fn matches_model(&self, model: &str) -> bool {
+        self.config
+            .model_prefixes
+            .iter()
+            .any(|prefix| model.starts_with(prefix.as_str()))
+    }
+}
+
+// ── Tests ────────────────────────────────────────────────────────────────────
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Mutex to serialize tests that share the global custom-provider registry.
+    static TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
+
+    /// Acquire the test lock and clear the registry.
+    fn setup() -> std::sync::MutexGuard<'static, ()> {
+        let guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
+        clear_custom_providers();
+        guard
+    }
+
+    #[test]
+    fn register_and_detect_by_model_prefix() {
+        let _guard = setup();
+
+        let config = CustomProviderConfig {
+            name: "my-provider".into(),
+            base_url: "https://api.my-provider.com/v1".into(),
+            auth_header: AuthHeaderFormat::Bearer,
+            model_prefixes: vec!["my-".into(), "my-provider/".into()],
+        };
+
+        register_custom_provider(config).expect("registration should succeed");
+
+        let provider = detect_custom_provider("my-model-7b");
+        assert!(provider.is_some(), "should detect custom provider by prefix 'my-'");
+        let provider = provider.unwrap();
+        assert_eq!(provider.name(), "my-provider");
+        assert_eq!(provider.base_url(), "https://api.my-provider.com/v1");
+
+        // Also detect via slash-prefix routing.
+        let provider2 = detect_custom_provider("my-provider/llama-70b");
+        assert!(provider2.is_some(), "should detect custom provider by slash prefix");
+
+        // Non-matching model should not detect.
+        let none = detect_custom_provider("gpt-4");
+        assert!(none.is_none(), "should not match unrelated model");
+    }
+
+    #[test]
+    fn unregister_removes_provider() {
+        let _guard = setup();
+
+        let config = CustomProviderConfig {
+            name: "ephemeral".into(),
+            base_url: "https://api.ephemeral.com/v1".into(),
+            auth_header: AuthHeaderFormat::Bearer,
+            model_prefixes: vec!["eph-".into()],
+        };
+
+        register_custom_provider(config).expect("registration should succeed");
+        assert!(detect_custom_provider("eph-model").is_some());
+
+        let removed = unregister_custom_provider("ephemeral").expect("unregister should succeed");
+        assert!(removed, "should return true when provider was found");
+
+        assert!(
+            detect_custom_provider("eph-model").is_none(),
+            "should no longer detect after unregister"
+        );
+
+        // Unregistering again returns false.
+        let removed_again = unregister_custom_provider("ephemeral").expect("unregister should succeed");
+        assert!(!removed_again, "should return false when provider not found");
+    }
+
+    #[test]
+    fn custom_provider_with_api_key_auth() {
+        let _guard = setup();
+
+        let config = CustomProviderConfig {
+            name: "secure-provider".into(),
+            base_url: "https://api.secure.com/v1".into(),
+            auth_header: AuthHeaderFormat::ApiKey("X-Custom-Auth".into()),
+            model_prefixes: vec!["secure/".into()],
+        };
+
+        register_custom_provider(config).expect("registration should succeed");
+
+        let provider = detect_custom_provider("secure/model-1").expect("should detect provider");
+        let (header_name, header_value) = provider
+            .auth_header("my-secret-key")
+            .expect("should return auth header");
+        assert_eq!(header_name.as_ref(), "X-Custom-Auth");
+        assert_eq!(header_value.as_ref(), "my-secret-key");
+    }
+
+    #[test]
+    fn custom_provider_with_no_auth() {
+        let _guard = setup();
+
+        let config = CustomProviderConfig {
+            name: "local-provider".into(),
+            base_url: "http://localhost:8080/v1".into(),
+            auth_header: AuthHeaderFormat::None,
+            model_prefixes: vec!["local/".into()],
+        };
+
+        register_custom_provider(config).expect("registration should succeed");
+
+        let provider = detect_custom_provider("local/model").expect("should detect provider");
+        assert!(
+            provider.auth_header("unused").is_none(),
+            "no-auth provider should return None"
+        );
+    }
+
+    #[test]
+    fn custom_provider_bearer_auth() {
+        let _guard = setup();
+
+        let config = CustomProviderConfig {
+            name: "bearer-provider".into(),
+            base_url: "https://api.bearer.com/v1".into(),
+            auth_header: AuthHeaderFormat::Bearer,
+            model_prefixes: vec!["bearer/".into()],
+        };
+
+        register_custom_provider(config).expect("registration should succeed");
+
+        let provider = detect_custom_provider("bearer/model").expect("should detect provider");
+        let (header_name, header_value) = provider.auth_header("my-token").expect("should return auth header");
+        assert_eq!(header_name.as_ref(), "Authorization");
+        assert_eq!(header_value.as_ref(), "Bearer my-token");
+    }
+
+    #[test]
+    fn register_replaces_existing_provider() {
+        let _guard = setup();
+
+        let config1 = CustomProviderConfig {
+            name: "updatable".into(),
+            base_url: "https://old.example.com/v1".into(),
+            auth_header: AuthHeaderFormat::Bearer,
+            model_prefixes: vec!["upd/".into()],
+        };
+        register_custom_provider(config1).expect("first registration should succeed");
+
+        let config2 = CustomProviderConfig {
+            name: "updatable".into(),
+            base_url: "https://new.example.com/v1".into(),
+            auth_header: AuthHeaderFormat::Bearer,
+            model_prefixes: vec!["upd/".into()],
+        };
+        register_custom_provider(config2).expect("second registration should succeed");
+
+        let provider = detect_custom_provider("upd/model").expect("should detect provider");
+        assert_eq!(
+            provider.base_url(),
+            "https://new.example.com/v1",
+            "should use the updated config"
+        );
+    }
+
+    #[test]
+    fn validation_rejects_empty_name() {
+        let _guard = setup();
+
+        let config = CustomProviderConfig {
+            name: String::new(),
+            base_url: "https://example.com".into(),
+            auth_header: AuthHeaderFormat::Bearer,
+            model_prefixes: vec!["x/".into()],
+        };
+        let result = register_custom_provider(config);
+        assert!(result.is_err(), "should reject empty name");
+    }
+
+    #[test]
+    fn validation_rejects_empty_base_url() {
+        let _guard = setup();
+
+        let config = CustomProviderConfig {
+            name: "valid-name".into(),
+            base_url: String::new(),
+            auth_header: AuthHeaderFormat::Bearer,
+            model_prefixes: vec!["x/".into()],
+        };
+        let result = register_custom_provider(config);
+        assert!(result.is_err(), "should reject empty base_url");
+    }
+
+    #[test]
+    fn validation_rejects_no_prefixes() {
+        let _guard = setup();
+
+        let config = CustomProviderConfig {
+            name: "valid-name".into(),
+            base_url: "https://example.com".into(),
+            auth_header: AuthHeaderFormat::Bearer,
+            model_prefixes: vec![],
+        };
+        let result = register_custom_provider(config);
+        assert!(result.is_err(), "should reject empty model_prefixes");
+    }
+
+    #[test]
+    fn config_serde_round_trip() {
+        let config = CustomProviderConfig {
+            name: "serde-test".into(),
+            base_url: "https://example.com/v1".into(),
+            auth_header: AuthHeaderFormat::ApiKey("X-Api-Key".into()),
+            model_prefixes: vec!["serde/".into()],
+        };
+
+        let json = serde_json::to_string(&config).expect("should serialize");
+        let parsed: CustomProviderConfig = serde_json::from_str(&json).expect("should deserialize");
+
+        assert_eq!(parsed.name, "serde-test");
+        assert_eq!(parsed.base_url, "https://example.com/v1");
+        assert_eq!(parsed.model_prefixes, vec!["serde/"]);
+    }
+}

data/vendor/liter-llm/src/provider/google_ai.rs

@@ -0,0 +1,281 @@
+use std::borrow::Cow;
+
+use crate::error::Result;
+use crate::provider::Provider;
+use crate::types::ChatCompletionChunk;
+
+use super::vertex::{parse_gemini_stream_event, transform_gemini_request, transform_gemini_response};
+
+/// Google AI Studio base URL (Gemini API via API key).
+const BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta";
+
+/// Google AI Studio (Gemini) provider.
+///
+/// Uses the same Gemini `generateContent` format as Vertex AI but with:
+/// - Different base URL: `https://generativelanguage.googleapis.com/v1beta`
+/// - API key authentication via `x-goog-api-key` header (not OAuth2).
+/// - Model routing via `gemini/` or `google_ai/` prefix.
+///
+/// Request/response translation is shared with [`super::vertex::VertexAiProvider`]
+/// via the `pub(crate)` Gemini transform functions.
+///
+/// # Configuration
+///
+/// ```rust,ignore
+/// let config = ClientConfigBuilder::new("AIza...your-api-key").build();
+/// let client = DefaultClient::new(config, Some("gemini/gemini-2.0-flash"))?;
+///
+/// // Or using the google_ai/ prefix:
+/// let client = DefaultClient::new(config, Some("google_ai/gemini-2.0-flash"))?;
+/// ```
+pub struct GoogleAiProvider;
+
+impl Provider for GoogleAiProvider {
+    fn name(&self) -> &str {
+        "google_ai"
+    }
+
+    fn base_url(&self) -> &str {
+        BASE_URL
+    }
+
+    /// Google AI Studio uses `x-goog-api-key` header for authentication.
+    fn auth_header<'a>(&'a self, api_key: &'a str) -> Option<(Cow<'static, str>, Cow<'a, str>)> {
+        Some((Cow::Borrowed("x-goog-api-key"), Cow::Borrowed(api_key)))
+    }
+
+    fn matches_model(&self, model: &str) -> bool {
+        model.starts_with("gemini/") || model.starts_with("google_ai/")
+    }
+
+    fn strip_model_prefix<'m>(&self, model: &'m str) -> &'m str {
+        model
+            .strip_prefix("gemini/")
+            .or_else(|| model.strip_prefix("google_ai/"))
+            .unwrap_or(model)
+    }
+
+    /// Build the full URL for a Google AI Studio request.
+    ///
+    /// Chat completions → `{base}/models/{model}:generateContent`
+    /// Embeddings       → `{base}/models/{model}:embedContent`
+    /// Other paths      → `{base}{endpoint_path}`
+    fn build_url(&self, endpoint_path: &str, model: &str) -> String {
+        let base = self.base_url().trim_end_matches('/');
+        if endpoint_path.contains("chat/completions") {
+            format!("{base}/models/{model}:generateContent")
+        } else if endpoint_path.contains("embeddings") {
+            format!("{base}/models/{model}:embedContent")
+        } else {
+            format!("{base}{endpoint_path}")
+        }
+    }
+
+    fn transform_request(&self, body: &mut serde_json::Value) -> Result<()> {
+        transform_gemini_request(body)
+    }
+
+    fn transform_response(&self, body: &mut serde_json::Value) -> Result<()> {
+        transform_gemini_response(body)
+    }
+
+    /// Build the streaming URL: appends `?alt=sse` to enable SSE streaming.
+    fn build_stream_url(&self, endpoint_path: &str, model: &str) -> String {
+        let url = self.build_url(endpoint_path, model);
+        format!("{url}?alt=sse")
+    }
+
+    fn parse_stream_event(&self, event_data: &str) -> Result<Option<ChatCompletionChunk>> {
+        parse_gemini_stream_event(event_data)
+    }
+}
+
+// ── Unit tests ────────────────────────────────────────────────────────────────
+
+#[cfg(test)]
+mod tests {
+    use serde_json::json;
+
+    use super::*;
+    use crate::provider::Provider;
+
+    fn provider() -> GoogleAiProvider {
+        GoogleAiProvider
+    }
+
+    // ── base_url ──────────────────────────────────────────────────────────────
+
+    #[test]
+    fn base_url_is_generative_language_api() {
+        let p = provider();
+        assert_eq!(p.base_url(), "https://generativelanguage.googleapis.com/v1beta");
+    }
+
+    // ── auth_header ──────────────────────────────────────────────────────────
+
+    #[test]
+    fn auth_header_uses_x_goog_api_key() {
+        let p = provider();
+        let (name, value) = p.auth_header("test-key").unwrap();
+        assert_eq!(name.as_ref(), "x-goog-api-key");
+        assert_eq!(value.as_ref(), "test-key");
+    }
+
+    // ── model matching ───────────────────────────────────────────────────────
+
+    #[test]
+    fn matches_gemini_prefix() {
+        let p = provider();
+        assert!(p.matches_model("gemini/gemini-2.0-flash"));
+        assert!(p.matches_model("google_ai/gemini-2.0-flash"));
+        assert!(!p.matches_model("vertex_ai/gemini-2.0-flash"));
+        assert!(!p.matches_model("gpt-4"));
+    }
+
+    #[test]
+    fn strip_model_prefix_gemini() {
+        let p = provider();
+        assert_eq!(p.strip_model_prefix("gemini/gemini-2.0-flash"), "gemini-2.0-flash");
+        assert_eq!(p.strip_model_prefix("google_ai/gemini-pro"), "gemini-pro");
+        assert_eq!(p.strip_model_prefix("gemini-2.0-flash"), "gemini-2.0-flash");
+    }
+
+    // ── build_url ─────────────────────────────────────────────────────────────
+
+    #[test]
+    fn build_url_chat_completions() {
+        let p = provider();
+        let url = p.build_url("/chat/completions", "gemini-2.0-flash");
+        assert_eq!(
+            url,
+            "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent"
+        );
+    }
+
+    #[test]
+    fn build_url_embeddings() {
+        let p = provider();
+        let url = p.build_url("/embeddings", "text-embedding-004");
+        assert_eq!(
+            url,
+            "https://generativelanguage.googleapis.com/v1beta/models/text-embedding-004:embedContent"
+        );
+    }
+
+    #[test]
+    fn build_stream_url_appends_alt_sse() {
+        let p = provider();
+        let url = p.build_stream_url("/chat/completions", "gemini-2.0-flash");
+        assert!(url.ends_with("?alt=sse"));
+        assert!(url.contains(":generateContent"));
+    }
+
+    // ── transform_request (delegates to shared Gemini logic) ─────────────────
+
+    #[test]
+    fn transform_request_basic_chat() {
+        let p = provider();
+        let mut body = json!({
+            "messages": [
+                {"role": "system", "content": "You are helpful."},
+                {"role": "user", "content": "Hello!"}
+            ],
+            "max_tokens": 100
+        });
+
+        p.transform_request(&mut body).unwrap();
+
+        assert_eq!(body["systemInstruction"]["parts"][0]["text"], "You are helpful.");
+        assert_eq!(body["contents"][0]["role"], "user");
+        assert_eq!(body["contents"][0]["parts"][0]["text"], "Hello!");
+        assert_eq!(body["generationConfig"]["maxOutputTokens"], 100);
+    }
+
+    #[test]
+    fn transform_request_safety_settings() {
+        let p = provider();
+        let mut body = json!({
+            "messages": [{"role": "user", "content": "hi"}],
+            "extra_body": {
+                "safety_settings": [
+                    {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"}
+                ]
+            }
+        });
+
+        p.transform_request(&mut body).unwrap();
+
+        let settings = body["safetySettings"].as_array().unwrap();
+        assert_eq!(settings.len(), 1);
+        assert_eq!(settings[0]["category"], "HARM_CATEGORY_HATE_SPEECH");
+    }
+
+    #[test]
+    fn transform_request_cached_content() {
+        let p = provider();
+        let mut body = json!({
+            "messages": [{"role": "user", "content": "hi"}],
+            "extra_body": {
+                "cached_content": "cachedContents/abc123"
+            }
+        });
+
+        p.transform_request(&mut body).unwrap();
+
+        assert_eq!(body["cachedContent"], "cachedContents/abc123");
+    }
+
+    // ── transform_response (delegates to shared Gemini logic) ────────────────
+
+    #[test]
+    fn transform_response_basic() {
+        let p = provider();
+        let mut body = json!({
+            "candidates": [{
+                "content": {
+                    "role": "model",
+                    "parts": [{"text": "Hello from Google AI!"}]
+                },
+                "finishReason": "STOP"
+            }],
+            "usageMetadata": {
+                "promptTokenCount": 5,
+                "candidatesTokenCount": 4
+            }
+        });
+
+        p.transform_response(&mut body).unwrap();
+
+        assert_eq!(body["object"], "chat.completion");
+        assert_eq!(body["choices"][0]["message"]["content"], "Hello from Google AI!");
+        assert_eq!(body["choices"][0]["finish_reason"], "stop");
+        assert_eq!(body["usage"]["prompt_tokens"], 5);
+        assert_eq!(body["usage"]["completion_tokens"], 4);
+    }
+
+    // ── parse_stream_event (delegates to shared Gemini logic) ────────────────
+
+    #[test]
+    fn parse_stream_event_empty_returns_none() {
+        let p = provider();
+        let result = p.parse_stream_event("").unwrap();
+        assert!(result.is_none());
+    }
+
+    #[test]
+    fn parse_stream_event_basic_chunk() {
+        let p = provider();
+        let event_data = r#"{
+            "candidates": [{
+                "content": {"role": "model", "parts": [{"text": "Hi"}]},
+                "finishReason": "STOP"
+            }],
+            "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 1}
+        }"#;
+
+        let chunk = p.parse_stream_event(event_data).unwrap().unwrap();
+
+        assert_eq!(chunk.object, "chat.completion.chunk");
+        assert_eq!(chunk.choices[0].delta.content.as_deref(), Some("Hi"));
+    }
+}