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, ChatOptions, ChatRequest, ChatStreamEvent};
10use genai::Client;
11use std::sync::Arc;
12use std::time::Instant;
13use tokio::sync::{mpsc, RwLock};
14
15/// End of transmission signal
16pub const EOT_SIGNAL: &str = "<|EOT|>";
17
18/// Shared state for rate limiting and token counting
19#[derive(Default)]
20struct SharedState {
21    total_tokens_used: usize,
22    request_count: usize,
23}
24
25/// Thread-safe LLM provider implementation using Arc<RwLock<>>.
26///
27/// This provider can be cheaply cloned and shared across multiple agents.
28/// Each clone shares the same underlying client and rate limiting state.
29#[derive(Clone)]
30pub struct GenAIProvider {
31    /// The underlying genai client
32    client: Arc<Client>,
33    /// Shared state for rate limiting and metrics
34    shared: Arc<RwLock<SharedState>>,
35}
36
37impl GenAIProvider {
38    /// Creates a new GenAI provider with automatic configuration.
39    pub fn new() -> Result<Self> {
40        let client = Client::default();
41        Ok(Self {
42            client: Arc::new(client),
43            shared: Arc::new(RwLock::new(SharedState::default())),
44        })
45    }
46
47    /// Creates a new GenAI provider with explicit configuration.
48    pub fn new_with_config(provider_type: Option<&str>, api_key: Option<&str>) -> Result<Self> {
49        // Set environment variable if API key is provided
50        if let (Some(provider), Some(key)) = (provider_type, api_key) {
51            let env_var = match provider {
52                "openai" => "OPENAI_API_KEY",
53                "anthropic" => "ANTHROPIC_API_KEY",
54                "gemini" => "GEMINI_API_KEY",
55                "groq" => "GROQ_API_KEY",
56                "cohere" => "COHERE_API_KEY",
57                "xai" => "XAI_API_KEY",
58                "deepseek" => "DEEPSEEK_API_KEY",
59                "ollama" => {
60                    log::info!("Ollama provider detected - no API key required for local setup");
61                    return Self::new();
62                }
63                _ => {
64                    log::warn!("Unknown provider type for API key: {provider}");
65                    return Self::new();
66                }
67            };
68
69            log::info!("Setting {env_var} environment variable for genai client");
70            std::env::set_var(env_var, key);
71        }
72
73        Self::new()
74    }
75
76    /// Get total tokens used across all requests
77    pub async fn get_total_tokens_used(&self) -> usize {
78        self.shared.read().await.total_tokens_used
79    }
80
81    /// Get total request count
82    pub async fn get_request_count(&self) -> usize {
83        self.shared.read().await.request_count
84    }
85
86    /// Increment request counter (for metrics)
87    async fn increment_request(&self) {
88        let mut state = self.shared.write().await;
89        state.request_count += 1;
90    }
91
92    /// Add tokens to the total count
93    pub async fn add_tokens(&self, count: usize) {
94        let mut state = self.shared.write().await;
95        state.total_tokens_used += count;
96    }
97
98    /// Retrieves all available models for a specific provider.
99    pub async fn get_available_models(&self, provider: &str) -> Result<Vec<String>> {
100        let adapter_kind = str_to_adapter_kind(provider)?;
101
102        let models = self
103            .client
104            .all_model_names(adapter_kind)
105            .await
106            .context(format!("Failed to get models for provider: {provider}"))?;
107
108        Ok(models)
109    }
110
111    /// Generates a simple text response without streaming.
112    /// Includes exponential backoff retry for rate limits and transient errors.
113    pub async fn generate_response_simple(&self, model: &str, prompt: &str) -> Result<String> {
114        self.generate_response_with_retry(model, prompt, 3).await
115    }
116
117    /// Generates a response with configurable retry count and exponential backoff.
118    pub async fn generate_response_with_retry(
119        &self,
120        model: &str,
121        prompt: &str,
122        max_retries: usize,
123    ) -> Result<String> {
124        self.increment_request().await;
125
126        let chat_req = ChatRequest::default().append_message(ChatMessage::user(prompt));
127
128        log::debug!(
129            "Sending chat request to model: {model} with prompt length: {} chars",
130            prompt.len()
131        );
132
133        let start_time = Instant::now();
134        let mut last_error: Option<anyhow::Error> = None;
135        let mut retry_count = 0;
136
137        while retry_count <= max_retries {
138            if retry_count > 0 {
139                // Exponential backoff: 1s, 2s, 4s, 8s, ... (capped at 16s)
140                let delay_secs = std::cmp::min(1u64 << (retry_count - 1), 16);
141                log::warn!(
142                    "Retry {}/{} for model {} after {}s delay (previous error: {:?})",
143                    retry_count,
144                    max_retries,
145                    model,
146                    delay_secs,
147                    last_error.as_ref().map(|e| e.to_string())
148                );
149                println!(
150                    "   ⏳ Rate limited, retrying in {}s (attempt {}/{})",
151                    delay_secs, retry_count, max_retries
152                );
153                tokio::time::sleep(tokio::time::Duration::from_secs(delay_secs)).await;
154            }
155
156            match self.client.exec_chat(model, chat_req.clone(), None).await {
157                Ok(chat_res) => {
158                    let content = chat_res
159                        .first_text()
160                        .context("No text content in response")?;
161                    log::debug!(
162                        "Received response with {} characters in {}ms",
163                        content.len(),
164                        start_time.elapsed().as_millis()
165                    );
166
167                    return Ok(content.to_string());
168                }
169                Err(e) => {
170                    let err_str = e.to_string();
171
172                    // Check if it's a retryable error (rate limit, server error, network)
173                    let is_retryable = err_str.contains("429")
174                        || err_str.contains("rate limit")
175                        || err_str.contains("Rate limit")
176                        || err_str.contains("RESOURCE_EXHAUSTED")
177                        || err_str.contains("500")
178                        || err_str.contains("502")
179                        || err_str.contains("503")
180                        || err_str.contains("504")
181                        || err_str.contains("timeout")
182                        || err_str.contains("connection");
183
184                    if is_retryable && retry_count < max_retries {
185                        log::warn!("Retryable error for model {}: {}", model, err_str);
186                        last_error = Some(anyhow::anyhow!("{}", err_str));
187                        retry_count += 1;
188                        continue;
189                    } else {
190                        return Err(anyhow::anyhow!(
191                            "Failed to execute chat request for model {}: {}",
192                            model,
193                            err_str
194                        ));
195                    }
196                }
197            }
198        }
199
200        // Should not reach here, but handle gracefully
201        Err(last_error
202            .unwrap_or_else(|| anyhow::anyhow!("Unknown error after {} retries", max_retries)))
203    }
204
205    /// Generates a streaming response and sends chunks via mpsc channel.
206    pub async fn generate_response_stream_to_channel(
207        &self,
208        model: &str,
209        prompt: &str,
210        tx: mpsc::UnboundedSender<String>,
211    ) -> Result<()> {
212        self.increment_request().await;
213
214        let chat_req = ChatRequest::default().append_message(ChatMessage::user(prompt));
215
216        log::debug!("Sending streaming chat request to model: {model} with prompt: {prompt}");
217
218        let chat_res_stream = self
219            .client
220            .exec_chat_stream(model, chat_req, None)
221            .await
222            .context(format!(
223                "Failed to execute streaming chat request for model: {model}"
224            ))?;
225
226        let mut stream = chat_res_stream.stream;
227        let mut chunk_count = 0;
228        let mut total_content_length = 0;
229        let mut stream_ended_explicitly = false;
230        let start_time = Instant::now();
231
232        log::info!(
233            "=== STREAM START === Model: {}, Prompt length: {} chars",
234            model,
235            prompt.len()
236        );
237
238        while let Some(chunk_result) = stream.next().await {
239            let elapsed = start_time.elapsed();
240
241            match chunk_result {
242                Ok(ChatStreamEvent::Start) => {
243                    log::info!(">>> STREAM STARTED for model: {model} at {elapsed:?}");
244                }
245                Ok(ChatStreamEvent::Chunk(chunk)) => {
246                    chunk_count += 1;
247                    total_content_length += chunk.content.len();
248
249                    if chunk_count % 10 == 0 || chunk.content.len() > 100 {
250                        log::info!(
251                            "CHUNK #{}: {} chars, total: {} chars, elapsed: {:?}",
252                            chunk_count,
253                            chunk.content.len(),
254                            total_content_length,
255                            elapsed
256                        );
257                    }
258
259                    if !chunk.content.is_empty() && tx.send(chunk.content.clone()).is_err() {
260                        log::error!(
261                            "!!! CHANNEL SEND FAILED for chunk #{chunk_count} - STOPPING STREAM !!!"
262                        );
263                        break;
264                    }
265                }
266                Ok(ChatStreamEvent::ReasoningChunk(chunk)) => {
267                    log::info!(
268                        "REASONING CHUNK: {} chars at {:?}",
269                        chunk.content.len(),
270                        elapsed
271                    );
272                }
273                Ok(ChatStreamEvent::End(_)) => {
274                    log::info!(">>> STREAM ENDED EXPLICITLY for model: {model} after {chunk_count} chunks, {total_content_length} chars, {elapsed:?} elapsed");
275                    stream_ended_explicitly = true;
276                    break;
277                }
278                Ok(ChatStreamEvent::ToolCallChunk(_)) => {
279                    log::debug!("Tool call chunk received (ignored)");
280                }
281                Err(e) => {
282                    log::error!(
283                        "!!! STREAM ERROR after {chunk_count} chunks at {elapsed:?}: {e} !!!"
284                    );
285                    let error_msg = format!("Stream error: {e}");
286                    let _ = tx.send(error_msg);
287                    return Err(e.into());
288                }
289            }
290        }
291
292        let final_elapsed = start_time.elapsed();
293        if !stream_ended_explicitly {
294            log::warn!("!!! STREAM ENDED IMPLICITLY (exhausted) for model: {model} after {chunk_count} chunks, {total_content_length} chars, {final_elapsed:?} elapsed !!!");
295        }
296
297        log::info!(
298            "=== STREAM COMPLETE === Model: {model}, Final: {chunk_count} chunks, {total_content_length} chars, {final_elapsed:?} elapsed"
299        );
300
301        // Add approximate token count
302        self.add_tokens(total_content_length / 4).await; // Rough estimate
303
304        if tx.send(EOT_SIGNAL.to_string()).is_err() {
305            log::error!("!!! FAILED TO SEND EOT SIGNAL - channel may be closed !!!");
306            return Err(anyhow::anyhow!("Channel closed during EOT signal send"));
307        }
308
309        log::info!(">>> EOT SIGNAL SENT for model: {model} <<<");
310        Ok(())
311    }
312
313    /// Generate response with conversation history
314    pub async fn generate_response_with_history(
315        &self,
316        model: &str,
317        messages: Vec<ChatMessage>,
318    ) -> Result<String> {
319        self.increment_request().await;
320
321        let chat_req = ChatRequest::new(messages);
322
323        log::debug!("Sending chat request to model: {model} with conversation history");
324
325        let chat_res = self
326            .client
327            .exec_chat(model, chat_req, None)
328            .await
329            .context(format!("Failed to execute chat request for model: {model}"))?;
330
331        let content = chat_res
332            .first_text()
333            .context("No text content in response")?;
334
335        log::debug!("Received response with {} characters", content.len());
336        Ok(content.to_string())
337    }
338
339    /// Generate response with custom chat options
340    pub async fn generate_response_with_options(
341        &self,
342        model: &str,
343        prompt: &str,
344        options: ChatOptions,
345    ) -> Result<String> {
346        self.increment_request().await;
347
348        let chat_req = ChatRequest::default().append_message(ChatMessage::user(prompt));
349
350        log::debug!("Sending chat request to model: {model} with custom options");
351
352        let chat_res = self
353            .client
354            .exec_chat(model, chat_req, Some(&options))
355            .await
356            .context(format!("Failed to execute chat request for model: {model}"))?;
357
358        let content = chat_res
359            .first_text()
360            .context("No text content in response")?;
361
362        log::debug!("Received response with {} characters", content.len());
363        Ok(content.to_string())
364    }
365
366    /// Get a list of supported providers
367    pub fn get_supported_providers() -> Vec<&'static str> {
368        vec![
369            "openai",
370            "anthropic",
371            "gemini",
372            "groq",
373            "cohere",
374            "ollama",
375            "xai",
376            "deepseek",
377        ]
378    }
379
380    /// Get all available providers
381    pub async fn get_available_providers(&self) -> Result<Vec<String>> {
382        Ok(Self::get_supported_providers()
383            .iter()
384            .map(|s| s.to_string())
385            .collect())
386    }
387
388    /// Test if a model is available and working
389    pub async fn test_model(&self, model: &str) -> Result<bool> {
390        match self.generate_response_simple(model, "Hello").await {
391            Ok(_) => {
392                log::info!("Model {model} is available and working");
393                Ok(true)
394            }
395            Err(e) => {
396                log::warn!("Model {model} test failed: {e}");
397                Ok(false)
398            }
399        }
400    }
401
402    /// Validate and get the best available model for a provider
403    pub async fn validate_model(&self, model: &str, provider_type: Option<&str>) -> Result<String> {
404        if self.test_model(model).await? {
405            return Ok(model.to_string());
406        }
407
408        if let Some(provider) = provider_type {
409            if let Ok(models) = self.get_available_models(provider).await {
410                if !models.is_empty() {
411                    log::info!("Model {} not available, using {} instead", model, models[0]);
412                    return Ok(models[0].clone());
413                }
414            }
415        }
416
417        log::warn!("Could not validate model {model}, proceeding anyway");
418        Ok(model.to_string())
419    }
420}
421
422/// Convert a provider string to genai AdapterKind
423fn str_to_adapter_kind(provider: &str) -> Result<AdapterKind> {
424    match provider.to_lowercase().as_str() {
425        "openai" => Ok(AdapterKind::OpenAI),
426        "anthropic" => Ok(AdapterKind::Anthropic),
427        "gemini" | "google" => Ok(AdapterKind::Gemini),
428        "groq" => Ok(AdapterKind::Groq),
429        "cohere" => Ok(AdapterKind::Cohere),
430        "ollama" => Ok(AdapterKind::Ollama),
431        "xai" => Ok(AdapterKind::Xai),
432        "deepseek" => Ok(AdapterKind::DeepSeek),
433        _ => Err(anyhow::anyhow!("Unsupported provider: {}", provider)),
434    }
435}
436
437#[cfg(test)]
438mod tests {
439    use super::*;
440
441    #[test]
442    fn test_str_to_adapter_kind() {
443        assert!(str_to_adapter_kind("openai").is_ok());
444        assert!(str_to_adapter_kind("anthropic").is_ok());
445        assert!(str_to_adapter_kind("gemini").is_ok());
446        assert!(str_to_adapter_kind("google").is_ok());
447        assert!(str_to_adapter_kind("groq").is_ok());
448        assert!(str_to_adapter_kind("cohere").is_ok());
449        assert!(str_to_adapter_kind("ollama").is_ok());
450        assert!(str_to_adapter_kind("xai").is_ok());
451        assert!(str_to_adapter_kind("deepseek").is_ok());
452        assert!(str_to_adapter_kind("invalid").is_err());
453    }
454
455    #[tokio::test]
456    async fn test_provider_creation() {
457        let provider = GenAIProvider::new();
458        assert!(provider.is_ok());
459    }
460
461    #[tokio::test]
462    async fn test_provider_is_clonable() {
463        let provider = GenAIProvider::new().unwrap();
464        let _clone1 = provider.clone();
465        let _clone2 = provider.clone();
466        // All clones share the same underlying state
467    }
468}
perspt_core/llm_provider.rs

perspt_core/
llm_provider.rs
