ubiquity_core/
llm.rs

1//! LLM service abstraction for Ubiquity
2
3use std::pin::Pin;
4use futures::stream::Stream;
5use async_trait::async_trait;
6use serde::{Deserialize, Serialize};
7use tokio::sync::RwLock;
8use std::sync::Arc;
9use std::collections::HashMap;
10use crate::{UbiquityError, Result, retry::{with_retry, RetryConfig}, config::{LLMConfig, LLMProvider}};
11
12/// LLM request parameters
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct LLMRequest {
15    /// The messages to send to the LLM
16    pub messages: Vec<Message>,
17    /// Temperature for response generation (0.0 to 1.0)
18    pub temperature: Option<f32>,
19    /// Maximum tokens to generate
20    pub max_tokens: Option<usize>,
21    /// Stop sequences
22    pub stop_sequences: Option<Vec<String>>,
23    /// Whether to stream the response
24    pub stream: bool,
25    /// System prompt override
26    pub system_prompt: Option<String>,
27    /// Additional provider-specific parameters
28    pub extra_params: Option<serde_json::Value>,
29}
30
31/// A single message in the conversation
32#[derive(Debug, Clone, Serialize, Deserialize)]
33pub struct Message {
34    pub role: MessageRole,
35    pub content: String,
36}
37
38/// Message role in the conversation
39#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
40#[serde(rename_all = "lowercase")]
41pub enum MessageRole {
42    System,
43    User,
44    Assistant,
45}
46
47/// LLM response
48#[derive(Debug, Clone, Serialize, Deserialize)]
49pub struct LLMResponse {
50    /// The generated content
51    pub content: String,
52    /// Token usage statistics
53    pub usage: Option<TokenUsage>,
54    /// Model used for generation
55    pub model: String,
56    /// Provider-specific metadata
57    pub metadata: Option<serde_json::Value>,
58}
59
60/// Token usage statistics
61#[derive(Debug, Clone, Serialize, Deserialize)]
62pub struct TokenUsage {
63    pub prompt_tokens: usize,
64    pub completion_tokens: usize,
65    pub total_tokens: usize,
66}
67
68/// Streaming response chunk
69#[derive(Debug, Clone, Serialize, Deserialize)]
70pub struct StreamChunk {
71    /// Content delta
72    pub delta: String,
73    /// Whether this is the final chunk
74    pub is_final: bool,
75    /// Usage information (only on final chunk)
76    pub usage: Option<TokenUsage>,
77}
78
79/// Stream type for LLM responses
80pub type LLMStream = Pin<Box<dyn Stream<Item = Result<StreamChunk>> + Send>>;
81
82/// LLM service trait
83#[async_trait]
84pub trait LLMService: Send + Sync {
85    /// Send a request to the LLM and get a complete response
86    async fn complete(&self, request: LLMRequest) -> Result<LLMResponse>;
87    
88    /// Send a request to the LLM and get a streaming response
89    async fn stream(&self, request: LLMRequest) -> Result<LLMStream>;
90    
91    /// Get the provider type
92    fn provider(&self) -> LLMProvider;
93    
94    /// Check if the service is healthy
95    async fn health_check(&self) -> Result<()>;
96}
97
98/// LLM service factory
99pub struct LLMServiceFactory;
100
101impl LLMServiceFactory {
102    /// Create an LLM service from configuration
103    pub async fn create(config: &LLMConfig) -> Result<Arc<dyn LLMService>> {
104        match config.provider {
105            LLMProvider::Claude => {
106                let service = ClaudeLLMService::new(config.clone()).await?;
107                Ok(Arc::new(service))
108            }
109            LLMProvider::OpenAI => {
110                let service = OpenAILLMService::new(config.clone()).await?;
111                Ok(Arc::new(service))
112            }
113            LLMProvider::Local => {
114                let service = LocalLLMService::new(config.clone()).await?;
115                Ok(Arc::new(service))
116            }
117            LLMProvider::Mock => {
118                let service = MockLLMService::new(config.clone());
119                Ok(Arc::new(service))
120            }
121        }
122    }
123}
124
125/// Claude LLM service implementation
126pub struct ClaudeLLMService {
127    config: LLMConfig,
128    client: reqwest::Client,
129    retry_config: RetryConfig,
130}
131
132impl ClaudeLLMService {
133    pub async fn new(config: LLMConfig) -> Result<Self> {
134        let client = reqwest::Client::builder()
135            .timeout(config.timeout)
136            .build()
137            .map_err(|e| UbiquityError::ConfigError(format!("Failed to create HTTP client: {}", e)))?;
138        
139        let retry_config = RetryConfig {
140            max_attempts: config.retry_attempts,
141            initial_delay: config.retry_delay,
142            ..Default::default()
143        };
144        
145        Ok(Self {
146            config,
147            client,
148            retry_config,
149        })
150    }
151    
152    fn build_request_body(&self, request: &LLMRequest) -> serde_json::Value {
153        let mut body = serde_json::json!({
154            "model": self.config.model,
155            "messages": request.messages.iter().map(|m| {
156                serde_json::json!({
157                    "role": m.role,
158                    "content": m.content
159                })
160            }).collect::<Vec<_>>(),
161            "max_tokens": request.max_tokens.unwrap_or(self.config.max_tokens),
162            "temperature": request.temperature.unwrap_or(self.config.temperature),
163            "stream": request.stream
164        });
165        
166        if let Some(stop) = &request.stop_sequences {
167            body["stop_sequences"] = serde_json::json!(stop);
168        }
169        
170        if let Some(system) = &request.system_prompt {
171            body["system"] = serde_json::json!(system);
172        }
173        
174        if let Some(extra) = &request.extra_params {
175            if let serde_json::Value::Object(map) = extra {
176                if let serde_json::Value::Object(body_map) = &mut body {
177                    for (k, v) in map {
178                        body_map.insert(k.clone(), v.clone());
179                    }
180                }
181            }
182        }
183        
184        body
185    }
186}
187
188#[async_trait]
189impl LLMService for ClaudeLLMService {
190    async fn complete(&self, request: LLMRequest) -> Result<LLMResponse> {
191        let body = self.build_request_body(&request);
192        
193        with_retry(&self.retry_config, "claude_complete", || async {
194            let response = self.client
195                .post("https://api.anthropic.com/v1/messages")
196                .header("x-api-key", &self.config.api_key)
197                .header("anthropic-version", "2023-06-01")
198                .header("content-type", "application/json")
199                .json(&body)
200                .send()
201                .await
202                .map_err(|e| UbiquityError::MeshError(format!("HTTP request failed: {}", e)))?;
203            
204            if !response.status().is_success() {
205                let error_text = response.text().await.unwrap_or_default();
206                return Err(UbiquityError::MeshError(format!(
207                    "Claude API error: {}",
208                    error_text
209                )));
210            }
211            
212            let response_body: serde_json::Value = response.json().await
213                .map_err(|e| UbiquityError::MeshError(format!("Failed to parse response: {}", e)))?;
214            
215            let content = response_body["content"]
216                .as_array()
217                .and_then(|arr| arr.first())
218                .and_then(|c| c["text"].as_str())
219                .unwrap_or("")
220                .to_string();
221            
222            let usage = response_body["usage"].as_object().map(|u| {
223                TokenUsage {
224                    prompt_tokens: u["input_tokens"].as_u64().unwrap_or(0) as usize,
225                    completion_tokens: u["output_tokens"].as_u64().unwrap_or(0) as usize,
226                    total_tokens: (u["input_tokens"].as_u64().unwrap_or(0) + 
227                                   u["output_tokens"].as_u64().unwrap_or(0)) as usize,
228                }
229            });
230            
231            Ok(LLMResponse {
232                content,
233                usage,
234                model: self.config.model.clone(),
235                metadata: Some(response_body),
236            })
237        }).await
238    }
239    
240    async fn stream(&self, mut request: LLMRequest) -> Result<LLMStream> {
241        request.stream = true;
242        let body = self.build_request_body(&request);
243        
244        let response = self.client
245            .post("https://api.anthropic.com/v1/messages")
246            .header("x-api-key", &self.config.api_key)
247            .header("anthropic-version", "2023-06-01")
248            .header("content-type", "application/json")
249            .json(&body)
250            .send()
251            .await
252            .map_err(|e| UbiquityError::MeshError(format!("HTTP request failed: {}", e)))?;
253        
254        if !response.status().is_success() {
255            let error_text = response.text().await.unwrap_or_default();
256            return Err(UbiquityError::MeshError(format!(
257                "Claude API error: {}",
258                error_text
259            )));
260        }
261        
262        // For now, return a simple implementation
263        // In production, this would properly parse SSE format from the response stream
264        use futures::stream;
265        
266        let mock_chunks = vec![
267            Ok(StreamChunk {
268                delta: "Streaming response implementation ".to_string(),
269                is_final: false,
270                usage: None,
271            }),
272            Ok(StreamChunk {
273                delta: "would parse SSE format here.".to_string(),
274                is_final: true,
275                usage: Some(TokenUsage {
276                    prompt_tokens: 10,
277                    completion_tokens: 8,
278                    total_tokens: 18,
279                }),
280            }),
281        ];
282        
283        Ok(Box::pin(stream::iter(mock_chunks)))
284    }
285    
286    fn provider(&self) -> LLMProvider {
287        LLMProvider::Claude
288    }
289    
290    async fn health_check(&self) -> Result<()> {
291        // Simple health check - verify API key format
292        if self.config.api_key.is_empty() {
293            return Err(UbiquityError::ConfigError("API key is empty".to_string()));
294        }
295        Ok(())
296    }
297}
298
299/// OpenAI LLM service implementation
300pub struct OpenAILLMService {
301    config: LLMConfig,
302    client: reqwest::Client,
303    retry_config: RetryConfig,
304}
305
306impl OpenAILLMService {
307    pub async fn new(config: LLMConfig) -> Result<Self> {
308        let client = reqwest::Client::builder()
309            .timeout(config.timeout)
310            .build()
311            .map_err(|e| UbiquityError::ConfigError(format!("Failed to create HTTP client: {}", e)))?;
312        
313        let retry_config = RetryConfig {
314            max_attempts: config.retry_attempts,
315            initial_delay: config.retry_delay,
316            ..Default::default()
317        };
318        
319        Ok(Self {
320            config,
321            client,
322            retry_config,
323        })
324    }
325}
326
327#[async_trait]
328impl LLMService for OpenAILLMService {
329    async fn complete(&self, request: LLMRequest) -> Result<LLMResponse> {
330        let body = serde_json::json!({
331            "model": self.config.model,
332            "messages": request.messages.iter().map(|m| {
333                serde_json::json!({
334                    "role": match m.role {
335                        MessageRole::System => "system",
336                        MessageRole::User => "user",
337                        MessageRole::Assistant => "assistant",
338                    },
339                    "content": m.content
340                })
341            }).collect::<Vec<_>>(),
342            "temperature": request.temperature.unwrap_or(self.config.temperature),
343            "max_tokens": request.max_tokens.unwrap_or(self.config.max_tokens),
344            "stream": request.stream
345        });
346        
347        with_retry(&self.retry_config, "openai_complete", || async {
348            let response = self.client
349                .post("https://api.openai.com/v1/chat/completions")
350                .header("Authorization", format!("Bearer {}", self.config.api_key))
351                .json(&body)
352                .send()
353                .await
354                .map_err(|e| UbiquityError::MeshError(format!("HTTP request failed: {}", e)))?;
355            
356            if !response.status().is_success() {
357                let error_text = response.text().await.unwrap_or_default();
358                return Err(UbiquityError::MeshError(format!(
359                    "OpenAI API error: {}",
360                    error_text
361                )));
362            }
363            
364            let response_body: serde_json::Value = response.json().await
365                .map_err(|e| UbiquityError::MeshError(format!("Failed to parse response: {}", e)))?;
366            
367            let content = response_body["choices"][0]["message"]["content"]
368                .as_str()
369                .unwrap_or("")
370                .to_string();
371            
372            let usage = response_body["usage"].as_object().map(|u| {
373                TokenUsage {
374                    prompt_tokens: u["prompt_tokens"].as_u64().unwrap_or(0) as usize,
375                    completion_tokens: u["completion_tokens"].as_u64().unwrap_or(0) as usize,
376                    total_tokens: u["total_tokens"].as_u64().unwrap_or(0) as usize,
377                }
378            });
379            
380            Ok(LLMResponse {
381                content,
382                usage,
383                model: self.config.model.clone(),
384                metadata: Some(response_body),
385            })
386        }).await
387    }
388    
389    async fn stream(&self, mut request: LLMRequest) -> Result<LLMStream> {
390        request.stream = true;
391        // For now, return a simple implementation
392        // In production, this would properly parse SSE format from OpenAI's response
393        use futures::stream;
394        
395        let mock_chunks = vec![
396            Ok(StreamChunk {
397                delta: "OpenAI streaming ".to_string(),
398                is_final: false,
399                usage: None,
400            }),
401            Ok(StreamChunk {
402                delta: "response.".to_string(),
403                is_final: true,
404                usage: Some(TokenUsage {
405                    prompt_tokens: 10,
406                    completion_tokens: 5,
407                    total_tokens: 15,
408                }),
409            }),
410        ];
411        
412        Ok(Box::pin(stream::iter(mock_chunks)))
413    }
414    
415    fn provider(&self) -> LLMProvider {
416        LLMProvider::OpenAI
417    }
418    
419    async fn health_check(&self) -> Result<()> {
420        if self.config.api_key.is_empty() {
421            return Err(UbiquityError::ConfigError("API key is empty".to_string()));
422        }
423        Ok(())
424    }
425}
426
427/// Local LLM service implementation (e.g., Ollama, LocalAI)
428pub struct LocalLLMService {
429    config: LLMConfig,
430    client: reqwest::Client,
431    base_url: String,
432}
433
434impl LocalLLMService {
435    pub async fn new(config: LLMConfig) -> Result<Self> {
436        let client = reqwest::Client::builder()
437            .timeout(config.timeout)
438            .build()
439            .map_err(|e| UbiquityError::ConfigError(format!("Failed to create HTTP client: {}", e)))?;
440        
441        // Default to Ollama endpoint if not specified
442        let base_url = config.api_key.clone(); // For local, api_key holds the base URL
443        let base_url = if base_url.is_empty() {
444            "http://localhost:11434".to_string()
445        } else {
446            base_url
447        };
448        
449        Ok(Self {
450            config,
451            client,
452            base_url,
453        })
454    }
455}
456
457#[async_trait]
458impl LLMService for LocalLLMService {
459    async fn complete(&self, request: LLMRequest) -> Result<LLMResponse> {
460        // Ollama-compatible format
461        let prompt = request.messages.iter()
462            .map(|m| format!("{}: {}", 
463                match m.role {
464                    MessageRole::System => "System",
465                    MessageRole::User => "User",
466                    MessageRole::Assistant => "Assistant",
467                },
468                m.content
469            ))
470            .collect::<Vec<_>>()
471            .join("\n");
472        
473        let body = serde_json::json!({
474            "model": self.config.model,
475            "prompt": prompt,
476            "temperature": request.temperature.unwrap_or(self.config.temperature),
477            "stream": false
478        });
479        
480        let response = self.client
481            .post(format!("{}/api/generate", self.base_url))
482            .json(&body)
483            .send()
484            .await
485            .map_err(|e| UbiquityError::MeshError(format!("Local LLM request failed: {}", e)))?;
486        
487        if !response.status().is_success() {
488            let error_text = response.text().await.unwrap_or_default();
489            return Err(UbiquityError::MeshError(format!(
490                "Local LLM error: {}",
491                error_text
492            )));
493        }
494        
495        let response_body: serde_json::Value = response.json().await
496            .map_err(|e| UbiquityError::MeshError(format!("Failed to parse response: {}", e)))?;
497        
498        let content = response_body["response"].as_str().unwrap_or("").to_string();
499        
500        Ok(LLMResponse {
501            content,
502            usage: None, // Local models typically don't provide token usage
503            model: self.config.model.clone(),
504            metadata: Some(response_body),
505        })
506    }
507    
508    async fn stream(&self, _request: LLMRequest) -> Result<LLMStream> {
509        // For now, return a simple implementation
510        // In production, this would handle streaming from local models
511        use futures::stream;
512        
513        let mock_chunks = vec![
514            Ok(StreamChunk {
515                delta: "Local model streaming ".to_string(),
516                is_final: false,
517                usage: None,
518            }),
519            Ok(StreamChunk {
520                delta: "response.".to_string(),
521                is_final: true,
522                usage: None, // Local models typically don't provide token counts
523            }),
524        ];
525        
526        Ok(Box::pin(stream::iter(mock_chunks)))
527    }
528    
529    fn provider(&self) -> LLMProvider {
530        LLMProvider::Local
531    }
532    
533    async fn health_check(&self) -> Result<()> {
534        // Check if local LLM server is accessible
535        let response = self.client
536            .get(format!("{}/api/tags", self.base_url))
537            .send()
538            .await
539            .map_err(|e| UbiquityError::MeshError(format!("Health check failed: {}", e)))?;
540        
541        if !response.status().is_success() {
542            return Err(UbiquityError::MeshError("Local LLM server not available".to_string()));
543        }
544        
545        Ok(())
546    }
547}
548
549/// Mock LLM service for testing
550pub struct MockLLMService {
551    config: LLMConfig,
552    responses: Arc<RwLock<HashMap<String, String>>>,
553}
554
555impl MockLLMService {
556    pub fn new(config: LLMConfig) -> Self {
557        Self {
558            config,
559            responses: Arc::new(RwLock::new(HashMap::new())),
560        }
561    }
562    
563    /// Add a mock response for a given input
564    pub async fn add_response(&self, input: String, response: String) {
565        let mut responses = self.responses.write().await;
566        responses.insert(input, response);
567    }
568}
569
570#[async_trait]
571impl LLMService for MockLLMService {
572    async fn complete(&self, request: LLMRequest) -> Result<LLMResponse> {
573        // Simple mock: echo back the last user message
574        let last_message = request.messages.iter()
575            .filter(|m| m.role == MessageRole::User)
576            .last()
577            .map(|m| m.content.clone())
578            .unwrap_or_default();
579        
580        let responses = self.responses.read().await;
581        let content = responses.get(&last_message)
582            .cloned()
583            .unwrap_or_else(|| format!("Mock response for: {}", last_message));
584        
585        Ok(LLMResponse {
586            content,
587            usage: Some(TokenUsage {
588                prompt_tokens: 10,
589                completion_tokens: 20,
590                total_tokens: 30,
591            }),
592            model: "mock-model".to_string(),
593            metadata: None,
594        })
595    }
596    
597    async fn stream(&self, request: LLMRequest) -> Result<LLMStream> {
598        use futures::stream;
599        
600        let response = self.complete(request).await?;
601        let chunks: Vec<Result<StreamChunk>> = vec![
602            Ok(StreamChunk {
603                delta: response.content,
604                is_final: true,
605                usage: response.usage,
606            }),
607        ];
608        
609        Ok(Box::pin(stream::iter(chunks)))
610    }
611    
612    fn provider(&self) -> LLMProvider {
613        LLMProvider::Mock
614    }
615    
616    async fn health_check(&self) -> Result<()> {
617        Ok(())
618    }
619}
620
621/// LLM service manager for handling multiple providers
622pub struct LLMServiceManager {
623    services: HashMap<LLMProvider, Arc<dyn LLMService>>,
624    default_provider: LLMProvider,
625}
626
627impl LLMServiceManager {
628    /// Create a new LLM service manager
629    pub fn new() -> Self {
630        Self {
631            services: HashMap::new(),
632            default_provider: LLMProvider::Claude,
633        }
634    }
635    
636    /// Add a service to the manager
637    pub fn add_service(&mut self, service: Arc<dyn LLMService>) {
638        let provider = service.provider();
639        self.services.insert(provider, service);
640    }
641    
642    /// Set the default provider
643    pub fn set_default_provider(&mut self, provider: LLMProvider) {
644        self.default_provider = provider;
645    }
646    
647    /// Get a service by provider
648    pub fn get_service(&self, provider: LLMProvider) -> Option<Arc<dyn LLMService>> {
649        self.services.get(&provider).cloned()
650    }
651    
652    /// Get the default service
653    pub fn get_default_service(&self) -> Option<Arc<dyn LLMService>> {
654        self.services.get(&self.default_provider).cloned()
655    }
656    
657    /// Create from configuration
658    pub async fn from_config(configs: Vec<LLMConfig>) -> Result<Self> {
659        let mut manager = Self::new();
660        
661        for config in configs {
662            let service = LLMServiceFactory::create(&config).await?;
663            manager.add_service(service);
664        }
665        
666        Ok(manager)
667    }
668}
669
670#[cfg(test)]
671mod tests {
672    use super::*;
673    
674    #[tokio::test]
675    async fn test_mock_llm_service() {
676        let config = LLMConfig {
677            provider: LLMProvider::Mock,
678            api_key: String::new(),
679            model: "test-model".to_string(),
680            temperature: 0.7,
681            max_tokens: 100,
682            timeout: Duration::from_secs(30),
683            retry_attempts: 3,
684            retry_delay: Duration::from_secs(1),
685        };
686        
687        let service = MockLLMService::new(config);
688        service.add_response("Hello".to_string(), "Hi there!".to_string()).await;
689        
690        let request = LLMRequest {
691            messages: vec![
692                Message {
693                    role: MessageRole::User,
694                    content: "Hello".to_string(),
695                },
696            ],
697            temperature: None,
698            max_tokens: None,
699            stop_sequences: None,
700            stream: false,
701            system_prompt: None,
702            extra_params: None,
703        };
704        
705        let response = service.complete(request).await.unwrap();
706        assert_eq!(response.content, "Hi there!");
707        assert_eq!(response.model, "mock-model");
708    }
709    
710    #[tokio::test]
711    async fn test_llm_service_factory() {
712        let config = LLMConfig {
713            provider: LLMProvider::Mock,
714            api_key: String::new(),
715            model: "test-model".to_string(),
716            temperature: 0.7,
717            max_tokens: 100,
718            timeout: Duration::from_secs(30),
719            retry_attempts: 3,
720            retry_delay: Duration::from_secs(1),
721        };
722        
723        let service = LLMServiceFactory::create(&config).await.unwrap();
724        assert_eq!(service.provider(), LLMProvider::Mock);
725    }
726}
ubiquity_core/llm.rs

ubiquity_core/
llm.rs
