Skip to main content

minillmlib/generator/
info.rs

1//! Generator information - configuration for LLM providers
2
3use secrecy::SecretString;
4
5/// Configuration for an LLM provider/generator
6#[derive(Debug, Clone)]
7pub struct GeneratorInfo {
8    /// Display name for this generator
9    pub name: String,
10
11    /// Base URL for the API (e.g., `https://openrouter.ai/api/v1`)
12    pub base_url: String,
13
14    /// Model identifier (e.g., "anthropic/claude-3.5-sonnet")
15    pub model: String,
16
17    /// API key (stored securely)
18    pub api_key: Option<SecretString>,
19
20    /// Optional organization ID
21    pub organization_id: Option<String>,
22
23    /// Custom headers to include in requests
24    pub custom_headers: Vec<(String, String)>,
25
26    /// Whether this provider supports streaming
27    pub supports_streaming: bool,
28
29    /// Whether this provider supports vision/images
30    pub supports_vision: bool,
31
32    /// Whether this provider supports audio input
33    pub supports_audio: bool,
34
35    /// Maximum context length (tokens)
36    pub max_context_length: Option<usize>,
37
38    /// Default completion parameters for this generator
39    pub default_params: super::CompletionParameters,
40}
41
42impl GeneratorInfo {
43    /// Create a new GeneratorInfo with minimal configuration
44    pub fn new(
45        name: impl Into<String>,
46        base_url: impl Into<String>,
47        model: impl Into<String>,
48    ) -> Self {
49        Self {
50            name: name.into(),
51            base_url: base_url.into(),
52            model: model.into(),
53            api_key: None,
54            organization_id: None,
55            custom_headers: Vec::new(),
56            supports_streaming: true,
57            supports_vision: false,
58            supports_audio: false,
59            max_context_length: None,
60            default_params: super::CompletionParameters::default(),
61        }
62    }
63
64    /// Set the API key
65    pub fn with_api_key(mut self, key: impl Into<String>) -> Self {
66        self.api_key = Some(SecretString::from(key.into()));
67        self
68    }
69
70    /// Set API key from environment variable
71    pub fn with_api_key_from_env(mut self, env_var: &str) -> Self {
72        if let Ok(key) = std::env::var(env_var) {
73            self.api_key = Some(SecretString::from(key));
74        }
75        self
76    }
77
78    /// Add a custom header
79    pub fn with_header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
80        self.custom_headers.push((name.into(), value.into()));
81        self
82    }
83
84    /// Enable vision support
85    pub fn with_vision(mut self) -> Self {
86        self.supports_vision = true;
87        self
88    }
89
90    /// Enable audio support
91    pub fn with_audio(mut self) -> Self {
92        self.supports_audio = true;
93        self
94    }
95
96    /// Set max context length
97    pub fn with_max_context(mut self, length: usize) -> Self {
98        self.max_context_length = Some(length);
99        self
100    }
101
102    /// Set default completion parameters
103    pub fn with_default_params(mut self, params: super::CompletionParameters) -> Self {
104        self.default_params = params;
105        self
106    }
107
108    /// Get the full completions endpoint URL
109    pub fn completions_url(&self) -> String {
110        let base = self.base_url.trim_end_matches('/');
111        format!("{}/chat/completions", base)
112    }
113}
114
115// Pre-configured generators for common providers
116impl GeneratorInfo {
117    /// Create an OpenRouter generator
118    pub fn openrouter(model: impl Into<String>) -> Self {
119        Self::new("OpenRouter", "https://openrouter.ai/api/v1", model)
120            .with_api_key_from_env("OPENROUTER_API_KEY")
121            .with_header("HTTP-Referer", "https://github.com/minillmlib")
122            .with_header("X-Title", "MiniLLMLib")
123    }
124
125    /// Create an OpenAI generator
126    pub fn openai(model: impl Into<String>) -> Self {
127        Self::new("OpenAI", "https://api.openai.com/v1", model)
128            .with_api_key_from_env("OPENAI_API_KEY")
129    }
130
131    /// Create an Anthropic generator (via OpenRouter format)
132    pub fn anthropic(model: impl Into<String>) -> Self {
133        Self::new("Anthropic", "https://api.anthropic.com/v1", model)
134            .with_api_key_from_env("ANTHROPIC_API_KEY")
135    }
136
137    /// Create a custom URL-based generator
138    pub fn custom(
139        name: impl Into<String>,
140        base_url: impl Into<String>,
141        model: impl Into<String>,
142    ) -> Self {
143        Self::new(name, base_url, model)
144    }
145}