1use 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
23pub const EOT_SIGNAL: &str = "<|EOT|>";
25
26#[derive(Debug, Clone)]
28pub struct ResolvedProvider {
29 pub provider: String,
31 pub model: String,
33}
34
35pub fn detect_provider_from_env() -> (&'static str, &'static str) {
40 if vertex_project_from_env().is_some() {
41 ("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 ("ollama", "llama3.2")
62 }
63}
64
65#[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#[derive(Default)]
75struct SharedState {
76 total_tokens_used: usize,
77 request_count: usize,
78}
79
80#[derive(Clone)]
85pub struct GenAIProvider {
86 client: Arc<Client>,
88 shared: Arc<RwLock<SharedState>>,
90}
91
92impl GenAIProvider {
93 pub(crate) fn client(&self) -> &Client {
95 &self.client
96 }
97
98 pub fn new() -> Result<Self> {
100 let client = Client::default();
101 Ok(Self::from_client(client))
102 }
103
104 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 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 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 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 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 pub async fn get_total_tokens_used(&self) -> usize {
206 self.shared.read().await.total_tokens_used
207 }
208
209 pub async fn get_request_count(&self) -> usize {
211 self.shared.read().await.request_count
212 }
213
214 async fn increment_request(&self) {
216 let mut state = self.shared.write().await;
217 state.request_count += 1;
218 }
219
220 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 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 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 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 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 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 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 Err(last_error
349 .unwrap_or_else(|| anyhow::anyhow!("Unknown error after {} retries", max_retries)))
350 }
351
352 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 self.add_tokens(total_content_length / 4).await; 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 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 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 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 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
528fn 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 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 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 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 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
721fn 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 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
835pub(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 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 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")); }
979
980 #[test]
981 fn test_valid_vertex_segment_rejects_host_redirection() {
982 assert_eq!(valid_vertex_segment("evil.com/"), None);
984 assert_eq!(valid_vertex_segment("evil.com"), None); 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 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 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 }
1091}