Skip to main content

perspt_core/
llm_provider.rs

1//! # LLM Provider Module
2//!
3//! Thread-safe LLM provider abstraction for multi-agent use.
4//! Wraps genai::Client with Arc<RwLock<>> for shared state.
5
6use anyhow::{Context, Result};
7use futures::StreamExt;
8use genai::adapter::AdapterKind;
9use genai::chat::{ChatMessage, ChatRequest, ChatStreamEvent};
10use genai::resolver::{AuthData, AuthResolver, Endpoint, ProviderConfig, ServiceTargetResolver};
11use genai::Client;
12use genai::ModelIden;
13use genai::ServiceTarget;
14use std::future::Future;
15use std::path::PathBuf;
16use std::pin::Pin;
17use std::sync::Arc;
18use std::time::Instant;
19use tokio::sync::{mpsc, RwLock};
20
21use crate::config::Config;
22
23/// End of transmission signal
24pub const EOT_SIGNAL: &str = "<|EOT|>";
25
26/// Effective provider id and model after merging config, CLI flags, and env.
27#[derive(Debug, Clone)]
28pub struct ResolvedProvider {
29    /// Provider id, e.g. `openai`, `ollama`.
30    pub provider: String,
31    /// Model name to use (passed to genai verbatim so namespacing works).
32    pub model: String,
33}
34
35/// Detect the provider id and a sensible default model from environment keys.
36///
37/// Used as the fallback when no provider is configured. Falls back to a local
38/// Ollama setup when no API keys are present.
39pub fn detect_provider_from_env() -> (&'static str, &'static str) {
40    if vertex_project_from_env().is_some() {
41        // Google Vertex AI (Agent/AI Platform). Models are namespace-routed as
42        // `vertex::<model>`; auth is an OAuth2 Bearer token from ADC or
43        // VERTEX_API_KEY.
44        ("vertex", "vertex::gemini-2.5-flash")
45    } else if std::env::var("GEMINI_API_KEY").is_ok() {
46        ("gemini", "gemini-3.1-flash-lite-preview")
47    } else if std::env::var("OPENAI_API_KEY").is_ok() {
48        ("openai", "gpt-4o-mini")
49    } else if std::env::var("ANTHROPIC_API_KEY").is_ok() {
50        ("anthropic", "claude-3-5-sonnet-20241022")
51    } else if std::env::var("GROQ_API_KEY").is_ok() {
52        ("groq", "llama-3.1-8b-instant")
53    } else if std::env::var("COHERE_API_KEY").is_ok() {
54        ("cohere", "command-r-plus")
55    } else if std::env::var("XAI_API_KEY").is_ok() {
56        ("xai", "grok-beta")
57    } else if std::env::var("DEEPSEEK_API_KEY").is_ok() {
58        ("deepseek", "deepseek-chat")
59    } else {
60        // Default to Ollama for local usage
61        ("ollama", "llama3.2")
62    }
63}
64
65/// Response from a non-streaming LLM call, carrying text and token usage.
66#[derive(Debug, Clone)]
67pub struct LlmResponse {
68    pub text: String,
69    pub tokens_in: Option<i32>,
70    pub tokens_out: Option<i32>,
71}
72
73/// Shared state for rate limiting and token counting
74#[derive(Default)]
75struct SharedState {
76    total_tokens_used: usize,
77    request_count: usize,
78}
79
80/// Thread-safe LLM provider implementation using Arc<RwLock<>>.
81///
82/// This provider can be cheaply cloned and shared across multiple agents.
83/// Each clone shares the same underlying client and rate limiting state.
84#[derive(Clone)]
85pub struct GenAIProvider {
86    /// The underlying genai client
87    client: Arc<Client>,
88    /// Shared state for rate limiting and metrics
89    shared: Arc<RwLock<SharedState>>,
90}
91
92impl GenAIProvider {
93    /// The underlying bound client, for the sibling tool-calling driver.
94    pub(crate) fn client(&self) -> &Client {
95        &self.client
96    }
97
98    /// Creates a new GenAI provider with automatic configuration.
99    pub fn new() -> Result<Self> {
100        let client = Client::default();
101        Ok(Self::from_client(client))
102    }
103
104    /// Creates a new GenAI provider with explicit configuration.
105    pub fn new_with_config(provider_type: Option<&str>, api_key: Option<&str>) -> Result<Self> {
106        let adapter_kind = provider_type.and_then(|provider| match str_to_adapter_kind(provider) {
107            Ok(adapter_kind) => Some(adapter_kind),
108            Err(_) => {
109                log::warn!("Unknown provider type for genai client: {provider}");
110                None
111            }
112        });
113
114        // Set environment variable if API key is provided
115        if let (Some(provider), Some(key)) = (provider_type, api_key) {
116            if let Some(env_var) = provider_api_key_env_var(provider) {
117                log::info!("Setting {env_var} environment variable for genai client");
118                std::env::set_var(env_var, key);
119            } else if provider.eq_ignore_ascii_case("ollama") {
120                log::info!("Ollama provider detected - no API key required for local setup");
121            } else {
122                log::warn!("Unknown provider type for API key: {provider}");
123            }
124        }
125
126        let is_vertex = provider_type
127            .map(|p| p.eq_ignore_ascii_case("vertex"))
128            .unwrap_or(false);
129
130        let client = if is_vertex {
131            // Vertex AI authenticates with an OAuth2 Bearer token from ADC; no
132            // static API key is required when ADC is configured.
133            build_vertex_client()
134        } else {
135            match adapter_kind {
136                Some(adapter_kind) => build_bound_client(adapter_kind, provider_type),
137                None => Client::default(),
138            }
139        };
140
141        Ok(Self::from_client(client))
142    }
143
144    /// Build a provider from a `Config`, merging in environment detection and an
145    /// optional CLI model override, and return the effective provider/model.
146    ///
147    /// Precedence:
148    ///   - provider: `config.provider` > environment detection
149    ///   - model:    `cli_model` > `config.model` > provider default
150    ///   - api_key:  `config.api_key` > ambient environment
151    ///   - base_url: `config.base_url` > ambient environment
152    ///
153    /// The returned client is bound to the resolved adapter, so custom/local
154    /// OpenAI-compatible model names (e.g. `phi-4-npu-ov`) route correctly while
155    /// recognized names still resolve by prefix. Model names are passed through
156    /// verbatim so genai namespacing (`openai::model`) keeps working.
157    pub fn from_config(
158        config: &Config,
159        cli_model: Option<&str>,
160    ) -> Result<(Self, ResolvedProvider)> {
161        let (env_provider, env_model) = detect_provider_from_env();
162
163        let env_model_override = std::env::var("OPENAI_MODEL")
164            .or_else(|_| std::env::var("MODEL"))
165            .ok();
166
167        let model = cli_model
168            .map(str::to_string)
169            .or_else(|| config.model.clone())
170            .or(env_model_override)
171            .unwrap_or_else(|| env_model.to_string());
172
173        let provider = config
174            .provider
175            .clone()
176            .or_else(|| provider_from_model_namespace(&model).map(str::to_string))
177            .unwrap_or_else(|| env_provider.to_string());
178
179        // Propagate a configured base URL into the env var that build_bound_client
180        // reads, without clobbering an explicit ambient override.
181        if let Some(base_url) = config.base_url.as_deref() {
182            if let Some(env_var) = provider_base_url_env_var(&provider) {
183                if std::env::var(env_var).is_err() {
184                    std::env::set_var(env_var, base_url);
185                }
186            }
187        }
188
189        if provider.eq_ignore_ascii_case("vertex") {
190            configure_vertex_environment(config);
191        }
192
193        let provider_obj = Self::new_with_config(Some(&provider), config.api_key.as_deref())?;
194        Ok((provider_obj, ResolvedProvider { provider, model }))
195    }
196
197    fn from_client(client: Client) -> Self {
198        Self {
199            client: Arc::new(client),
200            shared: Arc::new(RwLock::new(SharedState::default())),
201        }
202    }
203
204    /// Get total tokens used across all requests
205    pub async fn get_total_tokens_used(&self) -> usize {
206        self.shared.read().await.total_tokens_used
207    }
208
209    /// Get total request count
210    pub async fn get_request_count(&self) -> usize {
211        self.shared.read().await.request_count
212    }
213
214    /// Increment request counter (for metrics)
215    async fn increment_request(&self) {
216        let mut state = self.shared.write().await;
217        state.request_count += 1;
218    }
219
220    /// Add tokens to the total count
221    pub async fn add_tokens(&self, count: usize) {
222        let mut state = self.shared.write().await;
223        state.total_tokens_used += count;
224    }
225
226    /// Retrieves all available models for a specific provider.
227    pub async fn get_available_models(&self, provider: &str) -> Result<Vec<String>> {
228        let adapter_kind = str_to_adapter_kind(provider)?;
229        let provider_config = provider_base_url_from_env(provider)
230            .map(|base_url| {
231                ProviderConfig::from_endpoint(Endpoint::from_owned(normalize_base_url(&base_url)))
232            })
233            .unwrap_or_default();
234
235        let models = self
236            .client
237            .all_model_names(adapter_kind, provider_config)
238            .await
239            .context(format!("Failed to get models for provider: {provider}"))?;
240
241        Ok(models)
242    }
243
244    /// Generates a simple text response without streaming.
245    /// Includes exponential backoff retry for rate limits and transient errors.
246    pub async fn generate_response_simple(&self, model: &str, prompt: &str) -> Result<LlmResponse> {
247        self.generate_response_with_retry(model, prompt, 3).await
248    }
249
250    /// Generates a response with configurable retry count and exponential backoff.
251    pub async fn generate_response_with_retry(
252        &self,
253        model: &str,
254        prompt: &str,
255        max_retries: usize,
256    ) -> Result<LlmResponse> {
257        self.increment_request().await;
258
259        let chat_req = ChatRequest::default().append_message(ChatMessage::user(prompt));
260
261        log::debug!(
262            "Sending chat request to model: {model} with prompt length: {} chars",
263            prompt.len()
264        );
265
266        let start_time = Instant::now();
267        let mut last_error: Option<anyhow::Error> = None;
268        let mut retry_count = 0;
269
270        while retry_count <= max_retries {
271            if retry_count > 0 {
272                // Exponential backoff: 1s, 2s, 4s, 8s, ... (capped at 16s)
273                let delay_secs = std::cmp::min(1u64 << (retry_count - 1), 16);
274                log::warn!(
275                    "Retry {}/{} for model {} after {}s delay (previous error: {:?})",
276                    retry_count,
277                    max_retries,
278                    model,
279                    delay_secs,
280                    last_error.as_ref().map(|e| e.to_string())
281                );
282                println!(
283                    "   ⏳ Rate limited, retrying in {}s (attempt {}/{})",
284                    delay_secs, retry_count, max_retries
285                );
286                tokio::time::sleep(tokio::time::Duration::from_secs(delay_secs)).await;
287            }
288
289            match self.client.exec_chat(model, chat_req.clone(), None).await {
290                Ok(chat_res) => {
291                    let tokens_in = chat_res.usage.prompt_tokens;
292                    let tokens_out = chat_res.usage.completion_tokens;
293                    let content = chat_res
294                        .first_text()
295                        .context("No text content in response")?;
296                    log::debug!(
297                        "Received response with {} characters in {}ms (tokens: in={:?}, out={:?})",
298                        content.len(),
299                        start_time.elapsed().as_millis(),
300                        tokens_in,
301                        tokens_out,
302                    );
303
304                    // Update shared token counter with real values when available
305                    let total = tokens_in.unwrap_or(0) + tokens_out.unwrap_or(0);
306                    if total > 0 {
307                        self.add_tokens(total as usize).await;
308                    }
309
310                    return Ok(LlmResponse {
311                        text: content.to_string(),
312                        tokens_in,
313                        tokens_out,
314                    });
315                }
316                Err(e) => {
317                    let err_str = e.to_string();
318
319                    // Check if it's a retryable error (rate limit, server error, network)
320                    let is_retryable = err_str.contains("429")
321                        || err_str.contains("rate limit")
322                        || err_str.contains("Rate limit")
323                        || err_str.contains("RESOURCE_EXHAUSTED")
324                        || err_str.contains("500")
325                        || err_str.contains("502")
326                        || err_str.contains("503")
327                        || err_str.contains("504")
328                        || err_str.contains("timeout")
329                        || err_str.contains("connection");
330
331                    if is_retryable && retry_count < max_retries {
332                        log::warn!("Retryable error for model {}: {}", model, err_str);
333                        last_error = Some(anyhow::anyhow!("{}", err_str));
334                        retry_count += 1;
335                        continue;
336                    } else {
337                        return Err(anyhow::anyhow!(
338                            "Failed to execute chat request for model {}: {}",
339                            model,
340                            err_str
341                        ));
342                    }
343                }
344            }
345        }
346
347        // Should not reach here, but handle gracefully
348        Err(last_error
349            .unwrap_or_else(|| anyhow::anyhow!("Unknown error after {} retries", max_retries)))
350    }
351
352    /// Generates a streaming response and sends chunks via mpsc channel.
353    pub async fn generate_response_stream_to_channel(
354        &self,
355        model: &str,
356        prompt: &str,
357        tx: mpsc::UnboundedSender<String>,
358    ) -> Result<()> {
359        self.increment_request().await;
360
361        let chat_req = ChatRequest::default().append_message(ChatMessage::user(prompt));
362
363        log::debug!("Sending streaming chat request to model: {model} with prompt: {prompt}");
364
365        let chat_res_stream = self
366            .client
367            .exec_chat_stream(model, chat_req, None)
368            .await
369            .context(format!(
370                "Failed to execute streaming chat request for model: {model}"
371            ))?;
372
373        let mut stream = chat_res_stream.stream;
374        let mut chunk_count = 0;
375        let mut total_content_length = 0;
376        let mut stream_ended_explicitly = false;
377        let start_time = Instant::now();
378
379        log::info!(
380            "=== STREAM START === Model: {}, Prompt length: {} chars",
381            model,
382            prompt.len()
383        );
384
385        while let Some(chunk_result) = stream.next().await {
386            let elapsed = start_time.elapsed();
387
388            match chunk_result {
389                Ok(ChatStreamEvent::Start) => {
390                    log::info!(">>> STREAM STARTED for model: {model} at {elapsed:?}");
391                }
392                Ok(ChatStreamEvent::Chunk(chunk)) => {
393                    chunk_count += 1;
394                    total_content_length += chunk.content.len();
395
396                    if chunk_count % 10 == 0 || chunk.content.len() > 100 {
397                        log::info!(
398                            "CHUNK #{}: {} chars, total: {} chars, elapsed: {:?}",
399                            chunk_count,
400                            chunk.content.len(),
401                            total_content_length,
402                            elapsed
403                        );
404                    }
405
406                    if !chunk.content.is_empty() && tx.send(chunk.content.clone()).is_err() {
407                        log::error!(
408                            "!!! CHANNEL SEND FAILED for chunk #{chunk_count} - STOPPING STREAM !!!"
409                        );
410                        break;
411                    }
412                }
413                Ok(ChatStreamEvent::ReasoningChunk(chunk)) => {
414                    log::info!(
415                        "REASONING CHUNK: {} chars at {:?}",
416                        chunk.content.len(),
417                        elapsed
418                    );
419                    if !chunk.content.is_empty() {
420                        let _ = tx.send(format!("__PERSPT_REASONING__:{}", chunk.content));
421                    }
422                }
423                Ok(ChatStreamEvent::End(_)) => {
424                    log::info!(">>> STREAM ENDED EXPLICITLY for model: {model} after {chunk_count} chunks, \
425                        {total_content_length} chars, {elapsed:?} elapsed");
426                    stream_ended_explicitly = true;
427                    break;
428                }
429                Ok(ChatStreamEvent::ToolCallChunk(_)) => {
430                    log::debug!("Tool call chunk received (ignored)");
431                }
432                Ok(ChatStreamEvent::ThoughtSignatureChunk(_)) => {
433                    log::debug!("Thought signature chunk received (ignored)");
434                }
435                Err(e) => {
436                    log::error!(
437                        "!!! STREAM ERROR after {chunk_count} chunks at {elapsed:?}: {e} !!!"
438                    );
439                    let error_msg = format!("Stream error: {e}");
440                    let _ = tx.send(error_msg);
441                    return Err(e.into());
442                }
443            }
444        }
445
446        let final_elapsed = start_time.elapsed();
447        if !stream_ended_explicitly {
448            log::warn!(
449                "!!! STREAM ENDED IMPLICITLY (exhausted) for model: {model} after {chunk_count} \
450                chunks, {total_content_length} chars, {final_elapsed:?} elapsed !!!"
451            );
452        }
453
454        log::info!(
455            "=== STREAM COMPLETE === Model: {model}, Final: {chunk_count} chunks, {total_content_length} \
456                chars, {final_elapsed:?} elapsed"
457        );
458
459        // Add approximate token count
460        self.add_tokens(total_content_length / 4).await; // Rough estimate
461
462        if tx.send(EOT_SIGNAL.to_string()).is_err() {
463            log::error!("!!! FAILED TO SEND EOT SIGNAL - channel may be closed !!!");
464            return Err(anyhow::anyhow!("Channel closed during EOT signal send"));
465        }
466
467        log::info!(">>> EOT SIGNAL SENT for model: {model} <<<");
468        Ok(())
469    }
470
471    /// Get a list of supported providers
472    pub fn get_supported_providers() -> Vec<&'static str> {
473        vec![
474            "openai",
475            "anthropic",
476            "gemini",
477            "groq",
478            "cohere",
479            "ollama",
480            "vertex",
481            "xai",
482            "deepseek",
483        ]
484    }
485
486    /// Get all available providers
487    pub async fn get_available_providers(&self) -> Result<Vec<String>> {
488        Ok(Self::get_supported_providers()
489            .iter()
490            .map(|s| s.to_string())
491            .collect())
492    }
493
494    /// Test if a model is available and working
495    pub async fn test_model(&self, model: &str) -> Result<bool> {
496        match self.generate_response_simple(model, "Hello").await {
497            Ok(_) => {
498                log::info!("Model {model} is available and working");
499                Ok(true)
500            }
501            Err(e) => {
502                log::warn!("Model {model} test failed: {e}");
503                Ok(false)
504            }
505        }
506    }
507
508    /// Validate and get the best available model for a provider
509    pub async fn validate_model(&self, model: &str, provider_type: Option<&str>) -> Result<String> {
510        if self.test_model(model).await? {
511            return Ok(model.to_string());
512        }
513
514        if let Some(provider) = provider_type {
515            if let Ok(models) = self.get_available_models(provider).await {
516                if !models.is_empty() {
517                    log::info!("Model {} not available, using {} instead", model, models[0]);
518                    return Ok(models[0].clone());
519                }
520            }
521        }
522
523        log::warn!("Could not validate model {model}, proceeding anyway");
524        Ok(model.to_string())
525    }
526}
527
528/// Build a genai client for Google Vertex AI authenticated via Application
529/// Default Credentials (ADC).
530///
531/// genai's Vertex adapter reads `VERTEX_PROJECT_ID` (required) and
532/// `VERTEX_LOCATION` to construct the request
533/// URL, and expects an OAuth2 Bearer token from an [`AuthResolver`]. This
534/// resolver fetches that token from ADC (gcloud login, a service account, or the
535/// metadata server) on each request, so no static API key is needed. If a
536/// `VERTEX_API_KEY` bearer token is explicitly set, it is used as an override.
537fn build_vertex_client() -> Client {
538    let resolver = AuthResolver::from_resolver_async_fn(
539        |_model: ModelIden| -> Pin<
540            Box<dyn Future<Output = genai::resolver::Result<Option<AuthData>>> + Send>,
541        > {
542            Box::pin(async move {
543                // Explicit bearer-token override wins when present.
544                if let Ok(token) = std::env::var("VERTEX_API_KEY") {
545                    if !token.trim().is_empty() {
546                        return Ok(Some(AuthData::from_single(token)));
547                    }
548                }
549                // Otherwise resolve an access token from ADC.
550                let provider = gcp_auth::provider().await.map_err(|e| {
551                    genai::resolver::Error::Custom(format!(
552                        "Vertex ADC provider init failed (run `gcloud auth application-default login`): {e}"
553                    ))
554                })?;
555                let scopes = ["https://www.googleapis.com/auth/cloud-platform"];
556                let token = provider.token(&scopes).await.map_err(|e| {
557                    genai::resolver::Error::Custom(format!("Vertex ADC token fetch failed: {e}"))
558                })?;
559                Ok(Some(AuthData::from_single(token.as_str())))
560            })
561        },
562    );
563
564    let mut builder = Client::builder()
565        .with_adapter_kind(AdapterKind::Vertex)
566        .with_auth_resolver(resolver);
567
568    if let Some(endpoint) = resolved_vertex_endpoint() {
569        builder = builder.with_service_target_resolver_fn(move |mut target: ServiceTarget| {
570            target.endpoint = Endpoint::from_owned(endpoint.clone());
571            Ok(target)
572        });
573    }
574
575    builder.build()
576}
577
578fn build_bound_client(adapter_kind: AdapterKind, provider_type: Option<&str>) -> Client {
579    let mut builder = Client::builder().with_adapter_kind(adapter_kind);
580
581    if let Some(base_url) = provider_type.and_then(provider_base_url_from_env) {
582        let endpoint = normalize_base_url(&base_url);
583        let target_resolver = ServiceTargetResolver::from_resolver_fn(
584            move |mut service_target: ServiceTarget| -> genai::resolver::Result<ServiceTarget> {
585                if service_target.model.adapter_kind == adapter_kind {
586                    service_target.endpoint = Endpoint::from_owned(endpoint.clone());
587                }
588                Ok(service_target)
589            },
590        );
591        builder = builder.with_service_target_resolver(target_resolver);
592    }
593
594    builder.build()
595}
596
597fn provider_from_model_namespace(model: &str) -> Option<&'static str> {
598    let lower = model.to_ascii_lowercase();
599    lower.split_once("::").and_then(|(prefix, _)| match prefix {
600        "openai" => Some("openai"),
601        "anthropic" => Some("anthropic"),
602        "gemini" | "google" => Some("gemini"),
603        "vertex" => Some("vertex"),
604        "groq" => Some("groq"),
605        "cohere" => Some("cohere"),
606        "ollama" => Some("ollama"),
607        "xai" => Some("xai"),
608        "deepseek" => Some("deepseek"),
609        _ => None,
610    })
611}
612
613pub(crate) fn configure_vertex_environment(config: &Config) {
614    if std::env::var("VERTEX_PROJECT_ID").is_err() {
615        if let Some(project) = config
616            .vertex_project_id
617            .as_deref()
618            .map(str::trim)
619            .filter(|v| !v.is_empty())
620            .map(str::to_string)
621            .or_else(vertex_project_from_env)
622            .or_else(read_gcloud_project)
623        {
624            // Only propagate a syntactically valid project into the process env;
625            // genai reads VERTEX_PROJECT_ID to build the request URL, so an
626            // invalid value should not be planted there from config discovery.
627            match valid_vertex_segment(&project) {
628                Some(valid) => std::env::set_var("VERTEX_PROJECT_ID", valid),
629                None => log::warn!(
630                    "Ignoring discovered Vertex project ID (must contain only ASCII letters, \
631                     digits, and hyphens)"
632                ),
633            }
634        }
635    }
636
637    if std::env::var("VERTEX_LOCATION").is_err() {
638        if let Some(location) = config
639            .vertex_location
640            .as_deref()
641            .map(str::trim)
642            .filter(|v| !v.is_empty())
643        {
644            // The location is interpolated into the endpoint host, so never
645            // plant an invalid value into the env from a config file.
646            match valid_vertex_segment(location) {
647                Some(valid) => std::env::set_var("VERTEX_LOCATION", valid),
648                None => log::warn!(
649                    "Ignoring invalid vertex_location from config (must contain only ASCII \
650                     letters, digits, and hyphens)"
651                ),
652            }
653        }
654    }
655}
656
657fn vertex_project_from_env() -> Option<String> {
658    [
659        "VERTEX_PROJECT_ID",
660        "GOOGLE_CLOUD_PROJECT",
661        "GCLOUD_PROJECT",
662        "CLOUDSDK_CORE_PROJECT",
663    ]
664    .into_iter()
665    .filter_map(|key| std::env::var(key).ok())
666    .map(|value| value.trim().to_string())
667    .find(|value| !value.is_empty())
668}
669
670fn gcloud_config_dir() -> Option<PathBuf> {
671    if let Ok(dir) = std::env::var("CLOUDSDK_CONFIG") {
672        let trimmed = dir.trim();
673        if !trimmed.is_empty() {
674            return Some(PathBuf::from(trimmed));
675        }
676    }
677    dirs::home_dir().map(|home| home.join(".config").join("gcloud"))
678}
679
680fn read_gcloud_project() -> Option<String> {
681    let config_dir = gcloud_config_dir()?;
682    let active_config = std::fs::read_to_string(config_dir.join("active_config"))
683        .ok()
684        .map(|s| s.trim().to_string())
685        .filter(|s| !s.is_empty())
686        .unwrap_or_else(|| "default".to_string());
687    let config_path = config_dir
688        .join("configurations")
689        .join(format!("config_{active_config}"));
690    let content = std::fs::read_to_string(config_path).ok()?;
691    parse_gcloud_project(&content)
692}
693
694fn parse_gcloud_project(content: &str) -> Option<String> {
695    let mut in_core = false;
696    for raw in content.lines() {
697        let line = raw.trim();
698        if line.is_empty() || line.starts_with('#') {
699            continue;
700        }
701        if line.starts_with('[') && line.ends_with(']') {
702            in_core = line.eq_ignore_ascii_case("[core]");
703            continue;
704        }
705        if !in_core {
706            continue;
707        }
708        let Some((key, value)) = line.split_once('=') else {
709            continue;
710        };
711        if key.trim() == "project" {
712            let project = value.trim();
713            if !project.is_empty() {
714                return Some(project.to_string());
715            }
716        }
717    }
718    None
719}
720
721/// Validate a Vertex project/location segment before it is interpolated into the
722/// request endpoint URL.
723///
724/// GCP project IDs and Vertex locations are limited to ASCII letters, digits,
725/// and hyphens. Rejecting anything else prevents a crafted value (e.g. one
726/// containing `/`, `.`, `:`, or `@`) from altering the endpoint *host* — the
727/// location is interpolated into `{location}-aiplatform.googleapis.com`, so an
728/// unvalidated value could otherwise redirect the OAuth bearer token to an
729/// attacker-controlled host. Returns the trimmed value when valid.
730fn valid_vertex_segment(value: &str) -> Option<&str> {
731    let trimmed = value.trim();
732    if trimmed.is_empty() {
733        return None;
734    }
735    trimmed
736        .chars()
737        .all(|c| c.is_ascii_alphanumeric() || c == '-')
738        .then_some(trimmed)
739}
740
741fn resolved_vertex_endpoint() -> Option<String> {
742    let project_raw = std::env::var("VERTEX_PROJECT_ID").ok()?;
743    let project = match valid_vertex_segment(&project_raw) {
744        Some(p) => p.to_string(),
745        None => {
746            log::warn!(
747                "Ignoring VERTEX_PROJECT_ID for endpoint construction: must be non-empty and \
748                 contain only ASCII letters, digits, and hyphens"
749            );
750            return None;
751        }
752    };
753    let location = match std::env::var("VERTEX_LOCATION") {
754        Ok(raw) if !raw.trim().is_empty() => match valid_vertex_segment(&raw) {
755            Some(l) => l.to_string(),
756            None => {
757                log::warn!(
758                    "Ignoring invalid VERTEX_LOCATION (must contain only ASCII letters, digits, \
759                     and hyphens); falling back to 'global'"
760                );
761                "global".to_string()
762            }
763        },
764        _ => "global".to_string(),
765    };
766    Some(vertex_endpoint_base(&project, &location))
767}
768
769fn vertex_endpoint_base(project: &str, location: &str) -> String {
770    let project = project.trim();
771    let location = location.trim();
772    // Default to v1beta1: current Gemini releases (e.g. gemini-3.7-flash,
773    // gemini-3.5-flash-lite) are served there, and it is a superset of v1.
774    // Override with VERTEX_API_VERSION when an older surface is required.
775    let version = std::env::var("VERTEX_API_VERSION")
776        .ok()
777        .map(|v| v.trim().to_string())
778        .filter(|v| !v.is_empty() && valid_vertex_segment(v).is_some())
779        .unwrap_or_else(|| "v1beta1".to_string());
780    if location.eq_ignore_ascii_case("global") {
781        format!("https://aiplatform.googleapis.com/{version}/projects/{project}/locations/global/")
782    } else {
783        format!(
784            "https://{location}-aiplatform.googleapis.com/{version}/\
785             projects/{project}/locations/{location}/"
786        )
787    }
788}
789
790pub(crate) fn provider_base_url_env_var(provider: &str) -> Option<&'static str> {
791    match provider.to_lowercase().as_str() {
792        "openai" => Some("OPENAI_BASE_URL"),
793        "anthropic" => Some("ANTHROPIC_BASE_URL"),
794        "gemini" | "google" => Some("GEMINI_BASE_URL"),
795        "groq" => Some("GROQ_BASE_URL"),
796        "cohere" => Some("COHERE_BASE_URL"),
797        "ollama" => Some("OLLAMA_BASE_URL"),
798        "xai" => Some("XAI_BASE_URL"),
799        "deepseek" => Some("DEEPSEEK_BASE_URL"),
800        _ => None,
801    }
802}
803
804fn provider_base_url_from_env(provider: &str) -> Option<String> {
805    let env_var = provider_base_url_env_var(provider)?;
806
807    std::env::var(env_var)
808        .ok()
809        .map(|value| value.trim().to_string())
810        .filter(|value| !value.is_empty())
811}
812
813pub(crate) fn provider_api_key_env_var(provider: &str) -> Option<&'static str> {
814    match provider.to_lowercase().as_str() {
815        "openai" => Some("OPENAI_API_KEY"),
816        "anthropic" => Some("ANTHROPIC_API_KEY"),
817        "gemini" | "google" => Some("GEMINI_API_KEY"),
818        "vertex" => Some("VERTEX_API_KEY"),
819        "groq" => Some("GROQ_API_KEY"),
820        "cohere" => Some("COHERE_API_KEY"),
821        "xai" => Some("XAI_API_KEY"),
822        "deepseek" => Some("DEEPSEEK_API_KEY"),
823        _ => None,
824    }
825}
826
827fn normalize_base_url(base_url: &str) -> String {
828    if base_url.ends_with('/') {
829        base_url.to_string()
830    } else {
831        format!("{base_url}/")
832    }
833}
834
835/// Convert a provider string to genai AdapterKind
836pub(crate) fn str_to_adapter_kind(provider: &str) -> Result<AdapterKind> {
837    match provider.to_lowercase().as_str() {
838        "openai" => Ok(AdapterKind::OpenAI),
839        "anthropic" => Ok(AdapterKind::Anthropic),
840        "gemini" | "google" => Ok(AdapterKind::Gemini),
841        "vertex" => Ok(AdapterKind::Vertex),
842        "groq" => Ok(AdapterKind::Groq),
843        "cohere" => Ok(AdapterKind::Cohere),
844        "ollama" => Ok(AdapterKind::Ollama),
845        "xai" => Ok(AdapterKind::Xai),
846        "deepseek" => Ok(AdapterKind::DeepSeek),
847        _ => Err(anyhow::anyhow!("Unsupported provider: {}", provider)),
848    }
849}
850
851#[cfg(test)]
852mod tests {
853    use super::*;
854
855    #[test]
856    fn test_str_to_adapter_kind() {
857        assert!(str_to_adapter_kind("openai").is_ok());
858        assert!(str_to_adapter_kind("anthropic").is_ok());
859        assert!(str_to_adapter_kind("gemini").is_ok());
860        assert!(str_to_adapter_kind("google").is_ok());
861        assert!(str_to_adapter_kind("groq").is_ok());
862        assert!(str_to_adapter_kind("cohere").is_ok());
863        assert!(str_to_adapter_kind("ollama").is_ok());
864        assert!(str_to_adapter_kind("vertex").is_ok());
865        assert!(str_to_adapter_kind("xai").is_ok());
866        assert!(str_to_adapter_kind("deepseek").is_ok());
867        assert!(str_to_adapter_kind("invalid").is_err());
868    }
869
870    #[tokio::test]
871    async fn test_provider_creation() {
872        let provider = GenAIProvider::new();
873        assert!(provider.is_ok());
874    }
875
876    #[tokio::test]
877    async fn test_configured_provider_binds_adapter_for_custom_model_names() {
878        let provider = GenAIProvider::new_with_config(Some("openai"), None).unwrap();
879        let target = provider
880            .client
881            .resolve_service_target("gemma4-32b-it")
882            .await
883            .unwrap();
884
885        assert_eq!(target.model.adapter_kind, AdapterKind::OpenAI);
886    }
887
888    #[tokio::test]
889    async fn test_namespaced_model_resolves_on_unbound_client() {
890        // genai-native namespacing must work without a bound client.
891        let provider = GenAIProvider::new().unwrap();
892        let target = provider
893            .client
894            .resolve_service_target("openai::phi-4-npu-ov")
895            .await
896            .unwrap();
897
898        assert_eq!(target.model.adapter_kind, AdapterKind::OpenAI);
899    }
900
901    #[tokio::test]
902    async fn test_from_config_binds_adapter_for_custom_model() {
903        let config = Config {
904            provider: Some("openai".to_string()),
905            model: Some("phi-4-npu-ov".to_string()),
906            ..Default::default()
907        };
908        let (provider, resolved) = GenAIProvider::from_config(&config, None).unwrap();
909        assert_eq!(resolved.provider, "openai");
910        assert_eq!(resolved.model, "phi-4-npu-ov");
911
912        let target = provider
913            .client
914            .resolve_service_target(&resolved.model)
915            .await
916            .unwrap();
917        assert_eq!(target.model.adapter_kind, AdapterKind::OpenAI);
918    }
919
920    #[test]
921    fn test_from_config_model_precedence() {
922        let config = Config {
923            provider: Some("openai".to_string()),
924            model: Some("config-model".to_string()),
925            ..Default::default()
926        };
927        // CLI override wins over config model.
928        let (_p, resolved) = GenAIProvider::from_config(&config, Some("cli-model")).unwrap();
929        assert_eq!(resolved.model, "cli-model");
930    }
931
932    #[test]
933    fn test_provider_from_model_namespace_detects_vertex() {
934        assert_eq!(
935            provider_from_model_namespace("vertex::gemini-2.5-flash"),
936            Some("vertex")
937        );
938        assert_eq!(provider_from_model_namespace("gemini-2.5-flash"), None);
939    }
940
941    #[tokio::test]
942    async fn test_from_config_uses_namespaced_vertex_model_when_provider_absent() {
943        let previous_project = std::env::var("VERTEX_PROJECT_ID").ok();
944        let previous_location = std::env::var("VERTEX_LOCATION").ok();
945        std::env::set_var("VERTEX_PROJECT_ID", "unit-test-project");
946        std::env::remove_var("VERTEX_LOCATION");
947
948        let config = Config::default();
949        let (_provider, resolved) =
950            GenAIProvider::from_config(&config, Some("vertex::gemini-2.5-flash")).unwrap();
951        assert_eq!(resolved.provider, "vertex");
952        assert_eq!(resolved.model, "vertex::gemini-2.5-flash");
953        assert!(std::env::var("VERTEX_LOCATION").is_err());
954        assert_eq!(
955            resolved_vertex_endpoint().as_deref(),
956            Some(
957                "https://aiplatform.googleapis.com/v1beta1/projects/unit-test-project/locations/global/"
958            )
959        );
960
961        match previous_project {
962            Some(value) => std::env::set_var("VERTEX_PROJECT_ID", value),
963            None => std::env::remove_var("VERTEX_PROJECT_ID"),
964        }
965        match previous_location {
966            Some(value) => std::env::set_var("VERTEX_LOCATION", value),
967            None => std::env::remove_var("VERTEX_LOCATION"),
968        }
969    }
970
971    #[test]
972    fn test_valid_vertex_segment_accepts_real_values() {
973        assert_eq!(valid_vertex_segment("perspt"), Some("perspt"));
974        assert_eq!(valid_vertex_segment("us-central1"), Some("us-central1"));
975        assert_eq!(valid_vertex_segment("global"), Some("global"));
976        assert_eq!(valid_vertex_segment("europe-west4"), Some("europe-west4"));
977        assert_eq!(valid_vertex_segment("  perspt  "), Some("perspt")); // trimmed
978    }
979
980    #[test]
981    fn test_valid_vertex_segment_rejects_host_redirection() {
982        // Values that could alter the endpoint host or path must be rejected.
983        assert_eq!(valid_vertex_segment("evil.com/"), None);
984        assert_eq!(valid_vertex_segment("evil.com"), None); // '.'
985        assert_eq!(valid_vertex_segment("a/b"), None);
986        assert_eq!(valid_vertex_segment("a:b"), None);
987        assert_eq!(valid_vertex_segment("a@b"), None);
988        assert_eq!(valid_vertex_segment("a b"), None);
989        assert_eq!(valid_vertex_segment(""), None);
990        assert_eq!(valid_vertex_segment("   "), None);
991    }
992
993    #[test]
994    fn test_resolved_vertex_endpoint_rejects_malicious_location() {
995        let prev_project = std::env::var("VERTEX_PROJECT_ID").ok();
996        let prev_location = std::env::var("VERTEX_LOCATION").ok();
997
998        // A crafted location must not be interpolated into the host; it falls
999        // back to the safe global endpoint instead.
1000        std::env::set_var("VERTEX_PROJECT_ID", "perspt");
1001        std::env::set_var("VERTEX_LOCATION", "evil.com/");
1002        assert_eq!(
1003            resolved_vertex_endpoint().as_deref(),
1004            Some("https://aiplatform.googleapis.com/v1beta1/projects/perspt/locations/global/"),
1005            "malicious location must fall back to global, never redirect the host"
1006        );
1007
1008        // An invalid project yields no endpoint override at all.
1009        std::env::set_var("VERTEX_PROJECT_ID", "bad/project");
1010        std::env::set_var("VERTEX_LOCATION", "us-central1");
1011        assert_eq!(resolved_vertex_endpoint(), None);
1012
1013        match prev_project {
1014            Some(v) => std::env::set_var("VERTEX_PROJECT_ID", v),
1015            None => std::env::remove_var("VERTEX_PROJECT_ID"),
1016        }
1017        match prev_location {
1018            Some(v) => std::env::set_var("VERTEX_LOCATION", v),
1019            None => std::env::remove_var("VERTEX_LOCATION"),
1020        }
1021    }
1022
1023    #[test]
1024    fn test_vertex_endpoint_base_matches_genai_vertex_shape() {
1025        assert_eq!(
1026            vertex_endpoint_base("test-project", "global"),
1027            "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/"
1028        );
1029        assert_eq!(
1030            vertex_endpoint_base("test-project", "test-location"),
1031            "https://test-location-aiplatform.googleapis.com/v1beta1/\
1032             projects/test-project/locations/test-location/"
1033        );
1034    }
1035
1036    #[test]
1037    fn test_parse_gcloud_project_reads_core_project() {
1038        let content = r#"
1039        [compute]
1040        region = ignored-location
1041
1042        [core]
1043        account = user@example.com
1044        project = test-project
1045        "#;
1046        assert_eq!(
1047            parse_gcloud_project(content).as_deref(),
1048            Some("test-project")
1049        );
1050    }
1051
1052    #[tokio::test]
1053    async fn test_openai_base_url_overrides_bound_provider_endpoint() {
1054        let previous = std::env::var("OPENAI_BASE_URL").ok();
1055        std::env::set_var("OPENAI_BASE_URL", "https://custom.example/v1");
1056
1057        let provider = GenAIProvider::new_with_config(Some("openai"), None).unwrap();
1058        let target = provider
1059            .client
1060            .resolve_service_target("gemma4-32b-it")
1061            .await
1062            .unwrap();
1063
1064        assert_eq!(target.endpoint.base_url(), "https://custom.example/v1/");
1065
1066        match previous {
1067            Some(value) => std::env::set_var("OPENAI_BASE_URL", value),
1068            None => std::env::remove_var("OPENAI_BASE_URL"),
1069        }
1070    }
1071
1072    #[test]
1073    fn test_normalize_base_url() {
1074        assert_eq!(
1075            normalize_base_url("https://custom.example/v1"),
1076            "https://custom.example/v1/"
1077        );
1078        assert_eq!(
1079            normalize_base_url("https://custom.example/v1/"),
1080            "https://custom.example/v1/"
1081        );
1082    }
1083
1084    #[tokio::test]
1085    async fn test_provider_is_clonable() {
1086        let provider = GenAIProvider::new().unwrap();
1087        let _clone1 = provider.clone();
1088        let _clone2 = provider.clone();
1089        // All clones share the same underlying state
1090    }
1091}