git_iris/agents/
provider.rs

1//! Dynamic provider abstraction for rig-core 0.27+
2//!
3//! This module provides runtime provider selection using enum dispatch,
4//! allowing git-iris to work with any supported provider based on config.
5
6use anyhow::Result;
7use rig::{
8    agent::{Agent, AgentBuilder, PromptResponse},
9    client::{CompletionClient, ProviderClient},
10    completion::{Prompt, PromptError},
11    providers::{anthropic, gemini, openai},
12};
13
14use crate::providers::Provider;
15
16/// Completion model types for each provider
17pub type OpenAIModel = openai::completion::CompletionModel;
18pub type AnthropicModel = anthropic::completion::CompletionModel;
19pub type GeminiModel = gemini::completion::CompletionModel;
20
21/// Agent builder types for each provider
22pub type OpenAIBuilder = AgentBuilder<OpenAIModel>;
23pub type AnthropicBuilder = AgentBuilder<AnthropicModel>;
24pub type GeminiBuilder = AgentBuilder<GeminiModel>;
25
26/// Dynamic agent that can be any provider's agent type
27pub enum DynAgent {
28    OpenAI(Agent<OpenAIModel>),
29    Anthropic(Agent<AnthropicModel>),
30    Gemini(Agent<GeminiModel>),
31}
32
33impl DynAgent {
34    /// Simple prompt - returns response string
35    pub async fn prompt(&self, msg: &str) -> Result<String, PromptError> {
36        match self {
37            Self::OpenAI(a) => a.prompt(msg).await,
38            Self::Anthropic(a) => a.prompt(msg).await,
39            Self::Gemini(a) => a.prompt(msg).await,
40        }
41    }
42
43    /// Multi-turn prompt with specified depth for tool calling
44    pub async fn prompt_multi_turn(&self, msg: &str, depth: usize) -> Result<String, PromptError> {
45        match self {
46            Self::OpenAI(a) => a.prompt(msg).multi_turn(depth).await,
47            Self::Anthropic(a) => a.prompt(msg).multi_turn(depth).await,
48            Self::Gemini(a) => a.prompt(msg).multi_turn(depth).await,
49        }
50    }
51
52    /// Multi-turn prompt with extended details (token usage, etc.)
53    pub async fn prompt_extended(
54        &self,
55        msg: &str,
56        depth: usize,
57    ) -> Result<PromptResponse, PromptError> {
58        match self {
59            Self::OpenAI(a) => a.prompt(msg).multi_turn(depth).extended_details().await,
60            Self::Anthropic(a) => a.prompt(msg).multi_turn(depth).extended_details().await,
61            Self::Gemini(a) => a.prompt(msg).multi_turn(depth).extended_details().await,
62        }
63    }
64}
65
66/// Source of the resolved API key (for logging/debugging)
67#[derive(Debug, Clone, Copy, PartialEq, Eq)]
68pub enum ApiKeySource {
69    Config,
70    Environment,
71    ClientDefault,
72}
73
74/// Validate API key format and log warnings for suspicious keys
75fn validate_and_warn(key: &str, provider: Provider, source: &str) {
76    if let Err(warning) = provider.validate_api_key_format(key) {
77        tracing::warn!(
78            provider = %provider,
79            source = source,
80            "API key format warning: {}",
81            warning
82        );
83    }
84}
85
86/// Resolve API key from config or environment variable.
87///
88/// Resolution order:
89/// 1. If `api_key` is `Some` and non-empty, use it (from config)
90/// 2. Otherwise, check the provider's environment variable
91/// 3. If neither has a key, returns `None` (caller will use `from_env()`)
92///
93/// Note: An empty string in config is treated as "not configured" and falls
94/// back to the environment variable. This allows users to override env vars
95/// in config while still supporting env-only setups.
96pub fn resolve_api_key(api_key: Option<&str>, provider: Provider) -> (Option<String>, ApiKeySource) {
97    // If explicit key provided and non-empty, use it
98    if let Some(key) = api_key
99        && !key.is_empty()
100    {
101        tracing::trace!(
102            provider = %provider,
103            source = "config",
104            "Using API key from configuration"
105        );
106        validate_and_warn(key, provider, "config");
107        return (Some(key.to_string()), ApiKeySource::Config);
108    }
109
110    // Fall back to environment variable
111    if let Ok(key) = std::env::var(provider.api_key_env()) {
112        tracing::trace!(
113            provider = %provider,
114            env_var = %provider.api_key_env(),
115            source = "environment",
116            "Using API key from environment variable"
117        );
118        validate_and_warn(&key, provider, "environment");
119        return (Some(key), ApiKeySource::Environment);
120    }
121
122    tracing::trace!(
123        provider = %provider,
124        source = "client_default",
125        "No API key found, will use client's from_env()"
126    );
127    (None, ApiKeySource::ClientDefault)
128}
129
130/// Create an `OpenAI` agent builder
131///
132/// # Arguments
133/// * `model` - The model name to use
134/// * `api_key` - Optional API key from config. Resolution order:
135///   1. Non-empty `api_key` parameter (from config)
136///   2. `OPENAI_API_KEY` environment variable
137///   3. Client's `from_env()` (requires env var to be set)
138///
139/// # Errors
140/// Returns an error if client creation fails (invalid credentials or missing env var).
141///
142/// # Security
143/// Error messages are sanitized to prevent potential API key exposure.
144pub fn openai_builder(model: &str, api_key: Option<&str>) -> Result<OpenAIBuilder> {
145    let (resolved_key, _source) = resolve_api_key(api_key, Provider::OpenAI);
146    let client = match resolved_key {
147        Some(key) => openai::Client::new(&key)
148            // Sanitize error to prevent potential key exposure in error messages
149            .map_err(|_| anyhow::anyhow!(
150                "Failed to create OpenAI client: authentication or configuration error"
151            ))?,
152        None => openai::Client::from_env(),
153    };
154    Ok(client.completions_api().agent(model))
155}
156
157/// Create an Anthropic agent builder
158///
159/// # Arguments
160/// * `model` - The model name to use
161/// * `api_key` - Optional API key from config. Resolution order:
162///   1. Non-empty `api_key` parameter (from config)
163///   2. `ANTHROPIC_API_KEY` environment variable
164///   3. Client's `from_env()` (requires env var to be set)
165///
166/// # Errors
167/// Returns an error if client creation fails (invalid credentials or missing env var).
168///
169/// # Security
170/// Error messages are sanitized to prevent potential API key exposure.
171pub fn anthropic_builder(model: &str, api_key: Option<&str>) -> Result<AnthropicBuilder> {
172    let (resolved_key, _source) = resolve_api_key(api_key, Provider::Anthropic);
173    let client = match resolved_key {
174        Some(key) => anthropic::Client::new(&key)
175            // Sanitize error to prevent potential key exposure in error messages
176            .map_err(|_| anyhow::anyhow!(
177                "Failed to create Anthropic client: authentication or configuration error"
178            ))?,
179        None => anthropic::Client::from_env(),
180    };
181    Ok(client.agent(model))
182}
183
184/// Create a Gemini agent builder
185///
186/// # Arguments
187/// * `model` - The model name to use
188/// * `api_key` - Optional API key from config. Resolution order:
189///   1. Non-empty `api_key` parameter (from config)
190///   2. `GOOGLE_API_KEY` environment variable
191///   3. Client's `from_env()` (requires env var to be set)
192///
193/// # Errors
194/// Returns an error if client creation fails (invalid credentials or missing env var).
195///
196/// # Security
197/// Error messages are sanitized to prevent potential API key exposure.
198pub fn gemini_builder(model: &str, api_key: Option<&str>) -> Result<GeminiBuilder> {
199    let (resolved_key, _source) = resolve_api_key(api_key, Provider::Google);
200    let client = match resolved_key {
201        Some(key) => gemini::Client::new(&key)
202            // Sanitize error to prevent potential key exposure in error messages
203            .map_err(|_| anyhow::anyhow!(
204                "Failed to create Gemini client: authentication or configuration error"
205            ))?,
206        None => gemini::Client::from_env(),
207    };
208    Ok(client.agent(model))
209}
210
211#[cfg(test)]
212mod tests {
213    use super::*;
214
215    #[test]
216    fn test_resolve_api_key_uses_config_when_provided() {
217        // Config key takes precedence
218        let (key, source) =
219            resolve_api_key(Some("sk-config-key-1234567890"), Provider::OpenAI);
220        assert_eq!(key, Some("sk-config-key-1234567890".to_string()));
221        assert_eq!(source, ApiKeySource::Config);
222    }
223
224    #[test]
225    fn test_resolve_api_key_empty_config_not_used() {
226        // Empty config should NOT be treated as a valid key
227        // It should fall through to env var or client default
228        let empty_config: Option<&str> = Some("");
229        let (_key, source) = resolve_api_key(empty_config, Provider::OpenAI);
230
231        // Empty config should NOT return Config source
232        // This test verifies the empty string is treated as "not configured"
233        assert_ne!(source, ApiKeySource::Config);
234    }
235
236    #[test]
237    fn test_resolve_api_key_none_config_checks_env() {
238        // When config is None, should check env var
239        let (key, source) = resolve_api_key(None, Provider::OpenAI);
240
241        // Result depends on whether OPENAI_API_KEY is set in the environment
242        // We just verify the function doesn't panic and returns appropriate source
243        match source {
244            ApiKeySource::Environment => {
245                assert!(key.is_some());
246            }
247            ApiKeySource::ClientDefault => {
248                assert!(key.is_none());
249            }
250            ApiKeySource::Config => {
251                panic!("Should not return Config source when config is None");
252            }
253        }
254    }
255
256    #[test]
257    fn test_api_key_source_enum_equality() {
258        assert_eq!(ApiKeySource::Config, ApiKeySource::Config);
259        assert_eq!(ApiKeySource::Environment, ApiKeySource::Environment);
260        assert_eq!(ApiKeySource::ClientDefault, ApiKeySource::ClientDefault);
261        assert_ne!(ApiKeySource::Config, ApiKeySource::Environment);
262    }
263
264    #[test]
265    fn test_resolve_api_key_all_providers() {
266        // Test that resolve_api_key works for all supported providers
267        for provider in Provider::ALL {
268            let (key, source) =
269                resolve_api_key(Some("test-key-123456789012345"), *provider);
270            assert_eq!(key, Some("test-key-123456789012345".to_string()));
271            assert_eq!(source, ApiKeySource::Config);
272        }
273    }
274
275    #[test]
276    fn test_resolve_api_key_config_precedence() {
277        // Even if env var is set, config should take precedence
278        // We can't easily mock env vars in unit tests, but we can verify
279        // that a provided config key is always used regardless of env state
280        let config_key = "sk-from-config-abcdef1234567890";
281        let (key, source) = resolve_api_key(Some(config_key), Provider::OpenAI);
282
283        assert_eq!(key.as_deref(), Some(config_key));
284        assert_eq!(source, ApiKeySource::Config);
285    }
286
287    #[test]
288    fn test_api_key_source_debug_impl() {
289        // Verify Debug is implemented for logging purposes
290        let source = ApiKeySource::Config;
291        let debug_str = format!("{:?}", source);
292        assert!(debug_str.contains("Config"));
293    }
294}
git_iris/agents/provider.rs

git_iris/agents/
provider.rs
